diff --git a/docs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.md b/docs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.md index 91274ffc5459..e964e1149bc7 100644 --- a/docs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.md +++ b/docs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.md @@ -220,7 +220,7 @@ Figure 4 illustrates the prediction implementation within TensorRT LLM. To suppo **Auxiliary memory management.** Managing the paged KT cache presented another challenge. `RocketKVCacheManager` inherits from `KVCacheManager` and extends it with a dedicated `BlockManager` for the auxiliary KT cache at the Python level. The main KV cache and the KT cache share block IDs for each request, so that the lifecycle of KT cache blocks is automatically tied to the corresponding KV cache blocks. The `BlockManager` handles slot allocation and deallocation for the KT cache independently, while `RocketKVCacheManager` overrides methods such as `get_cache_bytes_per_token` and `prepare_resources` to ensure that memory sizing accounts for the extra KT cache footprint and that the correct KT cache pointers are passed to prediction kernels at each step. This design keeps the integration lightweight and easy to iterate on, though it inherits the limitations of Python-level management—namely, no automatic support for KV cache reuse or disaggregated serving. -The concrete implementation can be found in `tensorrt_llm/_torch/attention_backend/sparse/rocket.py`. +The concrete implementation can be found in `tensorrt_llm/_torch/attention_backend/sparse/rocket/`. ### DeepSeek Sparse Attention (DSA) @@ -261,7 +261,7 @@ As with RocketKV, a dedicated metadata class `DSATrtllmAttentionMetadata` is def **Auxiliary memory management.** DSA requires an auxiliary **indexer K cache** to store the low-rank K projections for reuse across decoding steps. `DSAKVCacheManager` inherits from `KVCacheManager`, but unlike RocketKV's Python-level KT cache management, DSA's indexer K cache is integrated directly into the C++ `KVCacheManager`. This design enables compatibility with advanced features such as KV cache reuse, chunked prefill, and disaggregated serving—features that would be difficult to support with a Python-level manager. -The concrete implementation can be found in `tensorrt_llm/_torch/attention_backend/sparse/dsa.py`. +The concrete implementation can be found in `tensorrt_llm/_torch/attention_backend/sparse/dsa/`. For a comprehensive description of DSA kernel optimizations, precision strategies, feature support (MTP, disaggregated serving, Wide-EP), and benchmark results, please refer to the dedicated blog post: [Optimizing DeepSeek-V3.2 on NVIDIA Blackwell GPUs](blog15_Optimizing_DeepSeek_V32_on_NVIDIA_Blackwell_GPUs.md). diff --git a/docs/source/developer-guide/sparse-attention-development-guide.md b/docs/source/developer-guide/sparse-attention-development-guide.md index c645d4ec4514..0829acab1474 100644 --- a/docs/source/developer-guide/sparse-attention-development-guide.md +++ b/docs/source/developer-guide/sparse-attention-development-guide.md @@ -137,19 +137,26 @@ KV eviction are tracked as future work. subclasses override: ```python -sparse_kv_indices, sparse_kv_offsets = self.sparse_kv_predict(q, k, metadata, **kwargs) -sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict(q, k, metadata, **kwargs) +sparse_kv_indices, sparse_kv_offsets = self.sparse_kv_predict(q, k, metadata, forward_args) +sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict(q, k, metadata, forward_args) ``` +`hooks.py` writes these results to `SparseRuntimeParams`. SkipSoftmax writes +its thresholds to the same runtime interface consumed by `AttentionOp`. + Different KV heads are allowed to emit different sparse index sets; Q heads that map to the same KV head share the KV head's sparse pattern. Algorithm implementations live under `tensorrt_llm/_torch/attention_backend/sparse/`: -- `rocket.py`, `dsa.py` — concrete algorithms. -- `kernel.py` — custom Triton kernels (importance scoring, Top-K). -- `utils.py` — dispatch helpers. +- `rocket/` — RocketKV backend, metadata, cache manager, parameters, and kernels. +- `dsa/` — DSA backend, indexer, metadata, cache manager, parameters, custom ops, and kernels. +- `deepseek_v4/` — DeepSeek-V4 backend, indexer, metadata, cache manager, + parameters, module hooks, and index conversion kernels. +- `skip_softmax/` — SkipSoftmax parameter parsing and runtime scheduler. +- `hooks.py` — module hooks and common backend prediction orchestration. +- `registry.py` — backend, metadata, and cache-manager dispatch helpers. ### AttentionOp behavior @@ -231,7 +238,7 @@ Create a new backend class inheriting from `TrtllmAttention` (or `tensorrt_llm/_torch/attention_backend/sparse/`. Override one or both prediction methods. -**`sparse_kv_predict(self, q, k, metadata, **kwargs)`** +**`sparse_kv_predict(self, q, k, metadata, forward_args)`** - **Behavior**: return the indices of tokens to retain in the KV cache. - **Outputs**: @@ -245,7 +252,7 @@ prediction methods. in-place gather (`updateSparseKvCacheAfterFmha`) is safe. The sort cost buys compatibility with chunked prefill and similar features. -**`sparse_attn_predict(self, q, k, metadata, **kwargs)`** +**`sparse_attn_predict(self, q, k, metadata, forward_args)`** - **Behavior**: return the sparse indices used by the generation-phase attention computation. @@ -285,7 +292,7 @@ If the algorithm needs extra tensors beyond the main KV cache: ### 4. Registration and dispatch - Register the new config + backend in - `tensorrt_llm/_torch/attention_backend/sparse/utils.py` and + `tensorrt_llm/_torch/attention_backend/sparse/registry.py` and `tensorrt_llm/_torch/pyexecutor/_util.py` so the runtime routes requests to your backend when the config is present. - If your algorithm exposes new C++ parameters, plumb them through @@ -318,5 +325,3 @@ for the kernel-side specifics. eviction as a compromise that keeps KV cache flexibility manageable. - **Unified auxiliary memory management** — let custom auxiliary pools inherit KV-cache features (reuse, offloading) by default. -- **Code refactoring** — as more algorithms land, unify the - framework-level scaffolding for maintainability. diff --git a/docs/source/torch/adding_custom_kernels.md b/docs/source/torch/adding_custom_kernels.md index 7a3bd12104ec..568ee683cf94 100644 --- a/docs/source/torch/adding_custom_kernels.md +++ b/docs/source/torch/adding_custom_kernels.md @@ -290,13 +290,13 @@ The adapter and env helpers live in [`tensorrt_llm/_torch/custom_ops/cutedsl_mat - **The kernel.** [`cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h) declares `invokeIndexerKCacheScatter(...)`; [`cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu) defines the `__global__` kernel and the host launcher. Picked up automatically by the `*.cu` glob in `cpp/tensorrt_llm/kernels/CMakeLists.txt`. - **The binding.** [`cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp) validates inputs, fetches the stream, and calls the launcher; registers the op via `TORCH_LIBRARY_FRAGMENT(trtllm, m)` and `TORCH_LIBRARY_IMPL(trtllm, CUDA, m)`. A one-line addition to [`cpp/tensorrt_llm/thop/CMakeLists.txt`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/CMakeLists.txt) wires it into `th_common` (this directory is not globbed). -- **The integration.** In [`tensorrt_llm/_torch/attention_backend/sparse/dsa.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/attention_backend/sparse/dsa.py), `_update_k_cache` calls `torch.ops.trtllm.indexer_k_cache_scatter_op(...)` in place of the prior Python scatter loop. +- **The integration.** In [`tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py), `_update_k_cache` calls `torch.ops.trtllm.indexer_k_cache_scatter_op(...)` in place of the prior Python scatter loop. - **The tests.** [`tests/unittest/_torch/attention/sparse/test_dsa_indexer.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py) — `test_indexer_k_cache_scatter_custom_op` runs the CUDA op against a PyTorch reference on two layers of the same cache pool and asserts the outputs match. The full call chain is three lines: ```py -# tensorrt_llm/_torch/attention_backend/sparse/dsa.py +# tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py torch.ops.trtllm.indexer_k_cache_scatter_op( k_fp8, k_scale, k_cache, metadata.slot_mapping_fp8, metadata.slot_mapping_scale, diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py index 9c081bf19dfc..2d67536c2603 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py @@ -32,7 +32,7 @@ # transitively reads; entries here are exempt. _THOP_EXCLUDED_FIELDS: frozenset = frozenset( { - "topk_indices", # DSA-only + "sparse_backend_args", # consumed by sparse prediction before the attention op "attention_mask_data", # custom-mask code path "out_scale_sf", # promoted into ``out_scale`` in ``TrtllmAttention.forward`` for NVFP4 path "skip_mla_rope_generation", # handled in ``TrtllmAttention.forward`` for the test-only MLA path @@ -60,7 +60,7 @@ def forward( attn = self.attn # Every kwarg sources from ``attn`` / ``metadata`` / ``forward_args`` - # (with ``forward_args.sparse_prediction`` for sparse-attn inputs), + # (with ``forward_args.sparse_runtime_params`` for sparse inputs), # or a literal allowlisted in ``_THOP_LITERALS``. # ``test_attention_op_sync.py`` enforces this statically. thop.attention( @@ -141,8 +141,6 @@ def forward( cross_kv=forward_args.cross_kv, relative_attention_bias=forward_args.relative_attention_bias, relative_attention_max_distance=forward_args.relative_attention_max_distance, - skip_softmax_threshold_scale_factor_prefill=forward_args.skip_softmax_kernel_params.threshold_scale_factor_prefill, - skip_softmax_threshold_scale_factor_decode=forward_args.skip_softmax_kernel_params.threshold_scale_factor_decode, # --- Module config (TrtllmAttention) --- rotary_inv_freq=attn.rotary_inv_freq, rotary_cos_sin=attn.rotary_cos_sin, @@ -171,12 +169,22 @@ def forward( rope_append=attn.rope_append, attention_chunk_size=attn.attention_chunk_size, skip_softmax_stat=attn.skip_softmax_stat, - # --- Sparse-specific (AttentionForwardArgs.sparse_prediction) --- - sparse_kv_indices=forward_args.sparse_prediction.sparse_kv_indices, - sparse_kv_offsets=forward_args.sparse_prediction.sparse_kv_offsets, - sparse_attn_indices=forward_args.sparse_prediction.sparse_attn_indices, - sparse_attn_offsets=forward_args.sparse_prediction.sparse_attn_offsets, - sparse_attn_indices_block_size=forward_args.sparse_prediction.sparse_attn_indices_block_size, - sparse_mla_topk_lens=forward_args.sparse_prediction.sparse_mla_topk_lens, - compressed_kv_cache_pool_ptr=forward_args.sparse_prediction.compressed_kv_cache_pool_ptr, + # --- Sparse runtime parameters --- + sparse_kv_indices=forward_args.sparse_runtime_params.sparse_kv_indices, + sparse_kv_offsets=forward_args.sparse_runtime_params.sparse_kv_offsets, + sparse_attn_indices=forward_args.sparse_runtime_params.sparse_attn_indices, + sparse_attn_offsets=forward_args.sparse_runtime_params.sparse_attn_offsets, + sparse_attn_indices_block_size=( + forward_args.sparse_runtime_params.sparse_attn_indices_block_size + ), + sparse_mla_topk_lens=forward_args.sparse_runtime_params.sparse_mla_topk_lens, + compressed_kv_cache_pool_ptr=( + forward_args.sparse_runtime_params.compressed_kv_cache_pool_ptr + ), + skip_softmax_threshold_scale_factor_prefill=( + forward_args.sparse_runtime_params.threshold_scale_factor_prefill + ), + skip_softmax_threshold_scale_factor_decode=( + forward_args.sparse_runtime_params.threshold_scale_factor_decode + ), ) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index ff2002c36d57..c5d1b6431733 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -620,8 +620,8 @@ def _is_supported_with_reason( return False, "sage attention." if meta.helix_position_offsets is not None: return False, "helix parallelism." - sparse_kv_indices = fwd.sparse_prediction.sparse_kv_indices - sparse_attn_indices = fwd.sparse_prediction.sparse_attn_indices + sparse_kv_indices = fwd.sparse_runtime_params.sparse_kv_indices + sparse_attn_indices = fwd.sparse_runtime_params.sparse_attn_indices if ( (sparse_kv_indices is not None and sparse_kv_indices.numel() > 0) or (sparse_attn_indices is not None and sparse_attn_indices.numel() > 0) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py index 0bbb1ad1509a..5e2f5415bdef 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -5,7 +5,7 @@ MsaSparseGqaFmha wraps the fmha_sm100 paged sparse GQA kernel and participates in the standard TrtllmAttention.forward dispatch loop. The owning MiniMax-M3 MSA attention layer runs an MsaIndexer to select the -per-query KV blocks and publishes them on forward_args.sparse_prediction; +per-query KV blocks and publishes them on forward_args.sparse_runtime_params; this class attends over them. """ @@ -173,7 +173,7 @@ class MsaSparseGqaFmha(Fmha): """SM100 paged GQA FMHA powered by MSA's fmha_sm100 kernel. Handles every MiniMax-M3 MSA layer. Sparse layers pass the indexer's - selected KV block indices on forward_args.sparse_prediction.sparse_attn_indices + selected KV block indices on forward_args.sparse_runtime_params.sparse_attn_indices and attend those blocks; dense layers leave the indices None and attend the full page table. @@ -216,7 +216,7 @@ def forward( # Sparse layers attend the per-query top-k blocks with the sparse plan; # dense layers leave the indices None and attend the full page table # with the dense plan. - kv_block_indexes = forward_args.sparse_prediction.sparse_attn_indices + kv_block_indexes = forward_args.sparse_runtime_params.sparse_attn_indices plan = ( metadata.msa_decode_gqa_plan if kv_block_indexes is not None diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index fec6ba4e6b92..2eff51528fc2 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import copy import weakref from collections import namedtuple @@ -27,7 +30,8 @@ from ..pyexecutor.resource_manager import KVCacheManager from ..pyexecutor.trace_log_utils import log_tensor_size from ..utils import get_model_extra_attrs -from .sparse.params import SkipSoftmaxKernelParams, SparseMetadataParams +from .sparse.params import (SparseBackendForwardArgs, SparseMetadataParams, + SparseRuntimeParams) try: # Transformers v5 @@ -886,27 +890,6 @@ class CustomAttentionMask(str, Enum): AttentionMask = Union[PredefinedAttentionMask, CustomAttentionMask] -@dataclass(kw_only=True, slots=True) -class SparsePrediction: - """Sparse KV / attention indices predicted by the framework backends. - - RocketKV and DSA produce these from ``sparse_kv_predict`` / - ``sparse_attn_predict``, telling the attention op which KV tokens to keep - and which blocks to attend to. Backends that don't predict leave - ``AttentionForwardArgs.sparse_prediction`` at its default-constructed value - (all-``None`` / ``0`` fields). - """ - sparse_kv_indices: Optional[torch.Tensor] = None - sparse_kv_offsets: Optional[torch.Tensor] = None - sparse_attn_indices: Optional[torch.Tensor] = None - sparse_attn_offsets: Optional[torch.Tensor] = None - sparse_attn_indices_block_size: int = 0 - # DeepSeek-V4 sparse-MLA only: per-token compressed top-k lengths and the - # base pointer of the compressed KV cache pool (compress_ratio > 1). - sparse_mla_topk_lens: Optional[torch.Tensor] = None - compressed_kv_cache_pool_ptr: Optional[int] = None - - @dataclass(kw_only=True, slots=True) class AttentionForwardArgs: """Per-forward optional arguments for attention backends.""" @@ -960,17 +943,14 @@ class AttentionForwardArgs: sage_attn_num_elts_per_blk_v: int = 0 sage_attn_qk_int8: bool = False - topk_indices: Optional[torch.Tensor] = None - is_fused_qkv: bool = False update_kv_cache: bool = True # Optional normalized diffusion timestep for timestep-varying sparse attention. timestep: Optional[torch.Tensor] = None - sparse_prediction: SparsePrediction = field( - default_factory=SparsePrediction) - skip_softmax_kernel_params: SkipSoftmaxKernelParams = field( - default_factory=SkipSoftmaxKernelParams) + sparse_backend_args: Optional[SparseBackendForwardArgs] = None + sparse_runtime_params: SparseRuntimeParams = field( + default_factory=SparseRuntimeParams) @property def mask_type(self) -> int: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/__init__.py index f293f9547506..7a93344644ff 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/__init__.py @@ -1,7 +1,10 @@ -from .utils import (get_flashinfer_sparse_attn_attention_backend, - get_sparse_attn_kv_cache_manager, - get_trtllm_sparse_attn_attention_backend, - get_vanilla_sparse_attn_attention_backend) +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .registry import (get_flashinfer_sparse_attn_attention_backend, + get_sparse_attn_kv_cache_manager, + get_trtllm_sparse_attn_attention_backend, + get_vanilla_sparse_attn_attention_backend) __all__ = [ "get_sparse_attn_kv_cache_manager", diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py index ee94ac46dc89..a822bd90bd40 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py @@ -13,22 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .backend import DeepseekV4TrtllmAttention from .cache_manager import DeepseekV4CacheManager -from .deepseek_v4 import ( - DeepseekV4AttentionType, - DeepSeekV4MetadataParams, - DeepSeekV4Params, - DeepseekV4TrtllmAttention, - make_deepseek_v4_sparse_metadata_params, - make_deepseek_v4_sparse_params, -) +from .indexer import DeepseekV4Indexer +from .metadata import DeepseekV4TrtllmAttentionMetadata +from .params import DeepseekV4AttentionType, DeepSeekV4MetadataParams, DeepSeekV4Params __all__ = [ "DeepSeekV4MetadataParams", "DeepSeekV4Params", "DeepseekV4AttentionType", "DeepseekV4CacheManager", + "DeepseekV4Indexer", "DeepseekV4TrtllmAttention", - "make_deepseek_v4_sparse_metadata_params", - "make_deepseek_v4_sparse_params", + "DeepseekV4TrtllmAttentionMetadata", ] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.py new file mode 100644 index 000000000000..d7699cbd07a7 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import replace +from typing import TYPE_CHECKING, Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionForwardArgs, + AttentionInputType, + MLAParams, + PositionalEmbeddingParams, + merge_attention_forward_args, +) +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm.models.modeling_utils import QuantConfig + +from .cache_manager import get_token_bytes +from .compressor import Compressor +from .indexer import DeepseekV4Indexer +from .kernels import deepseek_v4_local_to_global_indices +from .metadata import DeepseekV4TrtllmAttentionMetadata +from .params import DeepseekV4AttentionType, DeepSeekV4Params + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig + + +class DeepseekV4TrtllmAttention(TrtllmAttention): + Metadata = DeepseekV4TrtllmAttentionMetadata + + def __init__( + self, + layer_idx: int, + num_heads: int, + head_dim: int, + num_kv_heads: Optional[int] = None, + quant_config: Optional[QuantConfig] = None, + q_scaling: Optional[float] = None, + pos_embd_params: Optional[PositionalEmbeddingParams] = None, + mla_params: Optional[MLAParams] = None, + skip_create_weights_in_init: bool = False, + attention_chunk_size: Optional[int] = None, + sparse_attention_config: Optional["SparseAttentionConfig"] = None, + sparse_params: Optional[DeepSeekV4Params] = None, + dtype: Optional[torch.dtype] = None, + aux_stream: Optional[torch.cuda.Stream] = None, + **kwargs, + ): + if sparse_attention_config is None: + sparse_attention_config = sparse_params + assert sparse_attention_config is not None, ( + "sparse_attention_config is required for DeepseekV4TrtllmAttention and cannot be None" + ) + if sparse_params is None: + sparse_params = sparse_attention_config.to_sparse_params() + assert sparse_params is not None, ( + "sparse_params is required for DeepseekV4TrtllmAttention and cannot be None" + ) + assert mla_params is not None, "DeepSeek-V4 attention requires MLA parameters" + mla_params = replace( + mla_params, + v_head_dim=head_dim, + rope_append=False, + ) + TrtllmAttention.__init__( + self, + layer_idx, + num_heads, + head_dim, + sparse_params=sparse_params, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + q_scaling=q_scaling, + pos_embd_params=pos_embd_params, + mla_params=mla_params, + skip_create_weights_in_init=skip_create_weights_in_init, + attention_chunk_size=attention_chunk_size, + **kwargs, + ) + + self.sparse_attention_config = sparse_attention_config + self.compress_ratio = sparse_attention_config.compress_ratios[layer_idx] + + if self.compress_ratio == 4: + self.indexer = DeepseekV4Indexer( + quant_config, + pos_embd_params, + mla_params, + skip_create_weights_in_init, + sparse_attention_config, + dtype, + self.compress_ratio, + layer_idx, + aux_stream, + ) + + if self.compress_ratio > 1: + rms_norm_eps = 1e-6 + has_fp8_kv_cache = False + if quant_config is not None: + has_fp8_kv_cache = quant_config.layer_quant_mode.has_fp8_kv_cache() + kv_cache_dtype = "fp8_pertensor" if has_fp8_kv_cache else "default" + self.compressor = Compressor( + mla_params, + layer_idx, + self.compress_ratio, + rms_norm_eps, + skip_create_weights_in_init, + pos_embd_params, + kv_cache_dtype=kv_cache_dtype, + dtype=dtype, + rotate_activation=False, + ) + + def _prepare_sparse_forward_args( + self, + metadata: DeepseekV4TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> None: + attention_input_type = forward_args.attention_input_type + if attention_input_type == AttentionInputType.context_only: + start_idx = 0 + end_idx = metadata.num_ctx_tokens + elif attention_input_type == AttentionInputType.generation_only: + start_idx = metadata.num_ctx_tokens + end_idx = metadata.num_tokens + else: + start_idx = 0 + end_idx = metadata.num_tokens + + sparse_args = forward_args.sparse_runtime_params + sparse_args.sparse_mla_topk_lens = metadata.sparse_mla_topk_lens[self.compress_ratio][ + start_idx:end_idx + ] + if self.compress_ratio > 1: + sparse_args.compressed_kv_cache_pool_ptr = metadata.sparse_mla_base_ptrs[ + self.compress_ratio + ] + else: + sparse_args.compressed_kv_cache_pool_ptr = None + + metadata.num_sparse_topk = ( + self.sparse_attention_config.window_size + + metadata.max_compressed_indices[self.compress_ratio] + ) + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: DeepseekV4TrtllmAttentionMetadata, + forward_args: Optional[AttentionForwardArgs] = None, + **kwargs, + ): + forward_args = merge_attention_forward_args(forward_args, kwargs) + attn_sink = getattr(self, "attn_sink", None) + if attn_sink is not None: + if forward_args.attention_sinks is None: + forward_args = replace(forward_args, attention_sinks=attn_sink.data) + + self._prepare_sparse_forward_args(metadata, forward_args) + return super().forward(q, k, v, metadata, forward_args=forward_args) + + def sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: DeepseekV4TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Convert local indices (SWA + compressed) to global pool indices.""" + 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] + + # Get cached buffer pointers + swa_buffer_ptr = metadata.swa_buffer_ptrs[layer_idx] + + # Token stride + index_head_dim = self.sparse_attention_config.index_head_dim + has_fp8_kv_cache = False + if self.quant_config is not None: + has_fp8_kv_cache = self.quant_config.layer_quant_mode.has_fp8_kv_cache() + token_stride = get_token_bytes( + self.head_dim, + index_head_dim, + self.compress_ratio, + DeepseekV4AttentionType.SWA, + has_fp8_kv_cache, + ) + + # Select token range based on phase + if attention_input_type == AttentionInputType.context_only: + start_idx = 0 + end_idx = metadata.num_ctx_tokens + elif attention_input_type == AttentionInputType.generation_only: + start_idx = metadata.num_ctx_tokens + end_idx = metadata.num_tokens + else: + start_idx = 0 + 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] + local_layer_idx = kv_cache_manager.layer_offsets[layer_idx] + block_table_swa = metadata.sliding_block_tables[ + local_layer_idx, DeepseekV4AttentionType.SWA.value + ] + + 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] + if self.compress_ratio == 4: + sparse_backend_args = forward_args.sparse_backend_args + assert sparse_backend_args is not None, ( + "sparse_backend_args is required when compress_ratio=4" + ) + topk_indices = sparse_backend_args.topk_indices + assert topk_indices is not None, "topk_indices is required when compress_ratio=4" + compressed_local_indices = topk_indices + else: + compressed_local_indices = metadata.compressed_local_indices_cuda[start_idx:end_idx] + else: + compressed_buffer_ptr = 0 + compress_pool_base_ptr = 0 + block_table_compressed = None + compressed_local_indices = None + + global_indices = deepseek_v4_local_to_global_indices( + req_id=req_id, + block_table_swa=block_table_swa, + swa_local_indices=swa_local_indices, + swa_pool_base_ptr=swa_pool_base_ptr, + swa_buffer_ptr=swa_buffer_ptr, + tokens_per_block=kv_cache_manager.tokens_per_block, + token_stride=token_stride, + block_table_compressed=block_table_compressed, + compressed_local_indices=compressed_local_indices, + 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], + ) + + return global_indices, None + + def sparse_kv_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: DeepseekV4TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + return None, None diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py index 6eb9c366ff7c..7c22dff239e7 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -46,18 +46,77 @@ from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX from .compressor import KVCacheDtype -from .deepseek_v4 import ( +from .params import ( DEEPSEEK_V4_NON_SLIDING_ATTENTION, DEEPSEEK_V4_SLIDING_ATTENTION, DEEPSEEK_V4_SPARSE_RATIO, DeepseekV4AttentionType, compress_ratio_has_attention, - get_attn_dim, - get_token_bytes, is_overlap_compressor, ) +def get_attn_dim( + head_dim: int, index_head_dim: int, compress_ratio: int, attn_type: DeepseekV4AttentionType +) -> int: + state_factor = 2 if is_overlap_compressor(compress_ratio) else 1 + if attn_type == DeepseekV4AttentionType.SWA: + return head_dim + if attn_type == DeepseekV4AttentionType.COMPRESS: + return head_dim + if attn_type == DeepseekV4AttentionType.COMPRESSOR_KV: + return state_factor * head_dim + if attn_type == DeepseekV4AttentionType.COMPRESSOR_SCORE: + return state_factor * head_dim + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + return index_head_dim + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_KV: + return state_factor * index_head_dim + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE: + return state_factor * index_head_dim + raise ValueError(f"Unsupported DeepSeek-V4 attention type: {attn_type}") + + +def get_token_bytes( + head_dim: int, + index_head_dim: int, + compress_ratio: int, + attn_type: DeepseekV4AttentionType, + has_fp8_kv_cache: bool, + indexer_k_dtype: str = "fp8", +) -> int: + if not compress_ratio_has_attention(compress_ratio, attn_type): + raise ValueError( + f"Layer with compress ratio {compress_ratio} does not have attention type {attn_type}" + ) + + attn_dim = get_attn_dim(head_dim, index_head_dim, compress_ratio, attn_type) + + dtype_bytes = 1 if has_fp8_kv_cache else 2 + # (indexer) compressor kv and score always use float32 + if attn_type in [ + DeepseekV4AttentionType.COMPRESSOR_KV, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_KV, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + ]: + dtype_bytes = 4 # (indexer) compressor kv and score use float32 + # Indexer cache always packs data + per-block scales into one row. Only + # the two indexer presets ("fp8" blockwise / "fp4" mxfp4) are valid + # here — bf16 / fp8_pertensor are reserved for the main-attention + # compressor. + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + if indexer_k_dtype == "fp8": + return attn_dim + index_head_dim // 128 * 4 + if indexer_k_dtype == "fp4": + return index_head_dim // 2 + index_head_dim // 32 + raise ValueError( + f"Unsupported indexer_k_dtype {indexer_k_dtype!r}; expected 'fp8' or 'fp4'." + ) + + return attn_dim * dtype_bytes + + def _estimate_non_sliding_attn_size_per_token( head_dim: int, index_head_dim: int, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py index 6dd402369d9f..4377117ed4d4 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py @@ -13,8 +13,10 @@ from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding +from .params import DeepseekV4AttentionType + if TYPE_CHECKING: - from .deepseek_v4 import DeepseekV4TrtllmAttentionMetadata + from .metadata import DeepseekV4TrtllmAttentionMetadata class KVCacheDtype(IntEnum): @@ -144,9 +146,6 @@ def forward( - mxfp4 indexer: (fp4_output, ue8m0 scale) - no compressed tokens: (None, None) """ - # Import at runtime to avoid circular dependency - from .deepseek_v4 import DeepseekV4AttentionType - # Extract metadata num_contexts = metadata.num_contexts num_generations = metadata.num_generations diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py new file mode 100644 index 000000000000..6fb1a1a1c0e6 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py @@ -0,0 +1,394 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-V4 indexer implementation.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention_backend.interface import MLAParams, PositionalEmbeddingParams +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream +from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding +from tensorrt_llm._utils import is_sm_100f +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.utils import fp8_utils + +from ..dsa.indexer import HAS_FAST_HADAMARD, Indexer, rotate_activation +from .compressor import Compressor, KVCacheDtype, resolve_kv_cache_dtype + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig + + from ..dsa.metadata import DSAtrtllmAttentionMetadata + from .metadata import DeepseekV4TrtllmAttentionMetadata + + +class DeepseekV4Indexer(Indexer): + def __init__( + self, + quant_config: Optional[QuantConfig], + pos_embd_params: Optional[PositionalEmbeddingParams], + mla_params: Optional[MLAParams], + skip_create_weights_in_init: bool, + sparse_attention_config: "SparseAttentionConfig", + dtype: Optional[torch.dtype], + compress_ratio: int = 1, + layer_idx: int = 0, + aux_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__( + quant_config, + pos_embd_params, + mla_params, + skip_create_weights_in_init, + sparse_attention_config, + dtype, + compress_ratio, + layer_idx, + aux_stream, + ) + # Keep the checkpoint's FP8 quantization while deriving the scale view + # consumed by the fused CuTe DSL indexer-Q projection. + self.wq_b.use_indexer_q_cutedsl_fusion = True + # Override base Indexer.weights_proj to bf16 (matches V4 checkpoint). + self.weights_proj = Linear( + self.hidden_size, + self.n_heads, + bias=False, + dtype=dtype, + quant_config=None, + skip_create_weights_in_init=skip_create_weights_in_init, + use_custom_cublas_mm=True, + ) + self.rotary_emb = RotaryEmbedding( + pos_embd_params.rope, + head_dim=self.rope_dim, + is_neox=False, + ) + rms_norm_eps = 1e-6 + index_head_dim = sparse_attention_config.index_head_dim + indexer_mla_params = MLAParams( + hidden_size=mla_params.hidden_size, + qk_rope_head_dim=mla_params.qk_rope_head_dim, + qk_nope_head_dim=index_head_dim - mla_params.qk_rope_head_dim, + ) + # Map the user-facing FP4 knob ("fp8" / "fp4") onto the Compressor's + # cache layout preset string. The Compressor's preset namespace also + # covers main-attention layouts (bf16 / fp8_pertensor), which the + # indexer doesn't use, so the translation lives here at the + # boundary instead of leaking through the user-facing config. + self.indexer_k_dtype = sparse_attention_config.indexer_k_dtype + compressor_preset = "mxfp4" if self.indexer_k_dtype == "fp4" else "fp8_blockwise" + self.indexer_cache_dtype = resolve_kv_cache_dtype(compressor_preset) + self.compressor = Compressor( + indexer_mla_params, + layer_idx, + compress_ratio, + rms_norm_eps, + skip_create_weights_in_init, + pos_embd_params, + dtype=dtype, + kv_cache_dtype=compressor_preset, + is_indexer=True, + rotate_activation=HAS_FAST_HADAMARD, + ) + self.indexer_start_event = torch.cuda.Event() + self.weights_proj_event = torch.cuda.Event() + self.k_cache_update_event = torch.cuda.Event() + + def post_load_weights(self): + # V4 does not use the V3 fused fp32 wk+weights_proj GEMM, and the + # base concat would now hit an fp32/bf16 dtype mismatch. + return + + def _qk_projection_and_rope(self, qr: torch.Tensor, position_ids: torch.Tensor): + """Project Q and apply RoPE. + + Returns q with layout [num_tokens, n_heads, head_dim] where + head_dim = nope_dim + rope_dim, RoPE already applied in-place. + """ + q = self.wq_b(qr) + q = q.view(-1, self.n_heads, self.head_dim) + return self._apply_q_rope(q, position_ids) + + def _apply_q_rope(self, q: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + """Apply RoPE in-place to a projected indexer Q tensor.""" + # Fused in-place RoPE on the rope portion of each head + nope_dim = self.head_dim - self.rope_dim + torch.ops.trtllm.mla_rope_inplace( + q, + position_ids.view(-1), + self.rotary_emb.rotary_cos_sin, + self.n_heads, + nope_dim, + self.rope_dim, + False, + self.rotary_emb.is_neox, + ) + return q + + def _project_and_quantize_q( + self, qr: torch.Tensor, position_ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Project and quantize Q, using the fused MXFP4 path when supported.""" + use_fused_project_mxfp4 = ( + self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE + and not HAS_FAST_HADAMARD + and not self.rotary_emb.is_neox + and self.head_dim == 128 + and self.rope_dim == 64 + and self.wq_b.has_fp8_block_scales + and hasattr(self.wq_b, "indexer_q_weight_scale_cutedsl") + and hasattr( + torch.ops.trtllm, + "cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell", + ) + and qr.dtype == torch.bfloat16 + and is_sm_100f() + ) + if use_fused_project_mxfp4: + q_fp4, q_scale = torch.ops.trtllm.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell( + qr, + self.wq_b.weight, + self.wq_b.indexer_q_weight_scale_cutedsl, + position_ids.view(-1), + self.rotary_emb.rotary_cos_sin.view(-1, self.rope_dim), + self.wq_b.indexer_q_alpha_cutedsl, + use_tvm_ffi=True, + ) + return q_fp4.view(-1, self.n_heads, self.head_dim // 2), q_scale.view( + -1, self.n_heads, 1 + ) + + q = self.wq_b(qr).view(-1, self.n_heads, self.head_dim) + q = self._apply_q_rope(q, position_ids) + return self._quantize_q(q) + + def _quantize_q(self, q: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + # Rotate + quantize (layout matches compressor K: [nope|pe]). After + # rotate_activation (Hadamard) the nope/rope split becomes a linear + # mix; treating the row as a (head_dim - rope_dim, rope_dim) split + # for fused_cat_fp4 is equivalent to a single per-token FP4 quantize + # because the kernel only cares about the concatenated row. K goes + # through the same rotation in compressor_postprocess_scatter so + # layouts match. + q = rotate_activation(q) + q = q.view(-1, self.head_dim) + if self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE: + nope_dim = self.head_dim - self.rope_dim + q_nope, q_pe = q.split([nope_dim, self.rope_dim], dim=-1) + q_fp8, q_scale = torch.ops.trtllm.fused_cat_fp4(q_nope, q_pe) + # Two FP4 codes pack into one byte: trailing dim is head_dim // 2. + q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim // 2) + q_scale = q_scale.view(-1, self.n_heads, 1) + else: + q_fp8, q_scale = fp8_utils.fp8_quantize_1x128_sf_transpose( + q, use_ue8m0=self.scale_fmt == "ue8m0" + ) + q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim) + q_scale = q_scale.view(-1, self.n_heads, 1) + return q_fp8, q_scale + + def _apply_weight_scale(self, weights: torch.Tensor, q_scale: torch.Tensor) -> torch.Tensor: + # The DeepGEMM FP4 kernel applies per-block q_scale internally, so + # weights only carry softmax_scale * n_heads^-0.5. `weights_proj` is + # bf16 to match the V4 checkpoint, and the FP4 branch's only + # post-projection multiplier is a Python float (`bf16 * float` stays + # bf16). DeepGEMM's fp8_fp4_(paged_)mqa_logits asserts + # `weights.scalar_type() == kFloat`, so cast explicitly here. The FP8 + # branch gets upcast for free via the fp32 `q_scale` multiply in + # `_weight_scale`, so no cast is needed there. + if self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE: + return weights.float() * self.weight_scale_factor + return self._weight_scale(weights, q_scale) + + def _update_k_cache_if_needed( + self, + k_fp8: Optional[torch.Tensor], + k_scale: Optional[torch.Tensor], + metadata: DeepseekV4TrtllmAttentionMetadata, + ) -> None: + if k_fp8 is None: + return + + # The MXFP4 compressor's postprocess_scatter performs rotation + + # quantization + cache scatter in one fused op, so the cache row is + # already up to date by the time we get here. + if self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE: + return + + assert k_scale is not None, "FP8 blockwise indexer cache update requires scale tensor" + self._update_k_cache(k_fp8, k_scale, metadata) + + def precompute_aux( + self, + hidden_states: torch.Tensor, + metadata: DeepseekV4TrtllmAttentionMetadata, + ) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]]: + """Pre-launch the qr-independent half of the indexer prepare phase. + + Runs weights_proj + internal compressor + k_cache_update on + ``self.aux_stream`` and records ``weights_proj_event`` / + ``k_cache_update_event``. The caller hands the returned tuple back to + ``forward()`` via the ``pre_aux`` kwarg, which makes the overlapped + prepare path skip its own aux-stream launch and consume these results + directly. + + Returns ``None`` when multi-stream mode is off (caller should fall + back to the normal ``forward()`` call without ``pre_aux``). + """ + if not (do_multi_stream() and self.aux_stream is not None): + return None + self.indexer_start_event.record() + with torch.cuda.stream(self.aux_stream): + self.indexer_start_event.wait() + weights = self.weights_proj(hidden_states) + self.weights_proj_event.record() + k_fp8, k_scale = self.compressor(hidden_states, metadata) + self._update_k_cache_if_needed(k_fp8, k_scale, metadata) + self.k_cache_update_event.record() + return (weights, k_fp8, k_scale) + + def _run_overlapped_indexer_prepare( + self, + qr: torch.Tensor, + hidden_states: torch.Tensor, + metadata: DeepseekV4TrtllmAttentionMetadata, + position_ids: torch.Tensor, + pre_aux: Optional[ + Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] + ] = None, + ) -> Tuple[ + torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], torch.Tensor + ]: + """Prepare indexer inputs by splitting independent work across two streams. + + The current stream owns the Q path. The auxiliary stream starts from + the recorded launch point and owns weights projection, compressor, and + the K-cache update. + + When ``pre_aux`` is provided the aux-stream work has already been + launched (and ``weights_proj_event`` / ``k_cache_update_event`` + recorded) by ``precompute_aux``; the aux-stream block here is skipped + and the precomputed tensors are consumed directly. + + Timeline (pre_aux=None): + current stream: + record indexer_start_event + q_proj + RoPE -> quant_q + wait weights_proj_event -> weight_scale + wait k_cache_update_event -> return + + aux_stream: + wait indexer_start_event + weights_proj -> record weights_proj_event + compressor -> update_k_cache -> record k_cache_update_event + + Dependency graph: + q_proj + RoPE -> quant_q -- q_scale --. + v + weights_proj --------------------> weight_scale + compressor -> update_k_cache ----> final wait + """ + if pre_aux is None: + self.indexer_start_event.record() + + q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) + + with torch.cuda.stream(self.aux_stream): + self.indexer_start_event.wait() + + weights = self.weights_proj(hidden_states) + self.weights_proj_event.record() + + k_fp8, k_scale = self.compressor(hidden_states, metadata) + self._update_k_cache_if_needed(k_fp8, k_scale, metadata) + self.k_cache_update_event.record() + else: + weights, k_fp8, k_scale = pre_aux + # pre_aux tensors were allocated on aux_stream; record on the + # consuming stream so the caching allocator can't recycle them mid-use. + cur_stream = torch.cuda.current_stream() + weights.record_stream(cur_stream) + if k_fp8 is not None: + k_fp8.record_stream(cur_stream) + if k_scale is not None: + k_scale.record_stream(cur_stream) + q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) + + self.weights_proj_event.wait() + weights = self._apply_weight_scale(weights, q_scale) + + self.k_cache_update_event.wait() + return q_fp8, q_scale, k_fp8, k_scale, weights + + def _run_serial_indexer_prepare( + self, + qr: torch.Tensor, + hidden_states: torch.Tensor, + metadata: DeepseekV4TrtllmAttentionMetadata, + position_ids: torch.Tensor, + ) -> Tuple[ + torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], torch.Tensor + ]: + q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) + + weights = self.weights_proj(hidden_states) + + weights = self._apply_weight_scale(weights, q_scale) + + k_fp8, k_scale = self.compressor(hidden_states, metadata) + self._update_k_cache_if_needed(k_fp8, k_scale, metadata) + return q_fp8, q_scale, k_fp8, k_scale, weights + + def _update_k_cache( + self, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + metadata: DSAtrtllmAttentionMetadata, + ) -> None: + # DSV4's indexer compressor already scatters INDEXER_COMPRESS. The + # shared DSA scatter would duplicate that write. + return + + def forward( + self, + qr: torch.Tensor, + hidden_states: torch.Tensor, + metadata: DeepseekV4TrtllmAttentionMetadata, + position_ids: torch.Tensor, + pre_aux: Optional[ + Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] + ] = None, + ): + if do_multi_stream() and self.aux_stream is not None: + q_fp8, q_scale, k_fp8, k_scale, weights = self._run_overlapped_indexer_prepare( + qr, + hidden_states, + metadata, + position_ids, + pre_aux=pre_aux, + ) + else: + assert pre_aux is None, "pre_aux requires multi-stream mode" + q_fp8, q_scale, k_fp8, k_scale, weights = self._run_serial_indexer_prepare( + qr, hidden_states, metadata, position_ids + ) + + # 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]] + else: + topk_indices = self.sparse_attn_indexer( + metadata, + hidden_states, + q_fp8, + k_fp8, + k_scale, + weights, + q_scale=q_scale, + ) + return topk_indices diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.py new file mode 100644 index 000000000000..ac10fea430fa --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.py @@ -0,0 +1,287 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +######################################################## +# Index gather kernel +######################################################## + + +@triton.jit +def _deepseek_v4_local_to_global_kernel( + req_id_ptr, + block_table_swa_ptr, + block_table_compressed_ptr, + swa_local_indices_ptr, + compressed_local_indices_ptr, + out_ptr, + swa_buffer_offset_in_tokens, + compressed_buffer_offset_in_tokens, + tokens_per_block_swa: tl.constexpr, + tokens_per_block_compressed: tl.constexpr, + max_blocks_swa, + max_blocks_compressed, + num_swa_indices: tl.constexpr, + num_compressed_indices: tl.constexpr, + total_output_indices, + has_compressed: tl.constexpr, + bt_swa_stride0, + bt_swa_stride1, + bt_compressed_stride0, + bt_compressed_stride1, + swa_indices_stride0, + swa_indices_stride1, + compressed_indices_stride0, + compressed_indices_stride1, + out_stride0, + out_stride1, + LAUNCH_WITH_PDL: tl.constexpr, +): + """ + Triton kernel for converting local indices to global KV cache pool indices. + + Dual-pool output layout: + - SWA region [0, num_swa_indices): indices relative to swa_pool_base_ptr + - Compress region [num_swa_indices, num_swa_indices + num_compressed_indices): + indices relative to compress_pool_base_ptr + - Invalid positions padded with -1 at their fixed positions. + + This enables the FMHA kernel to determine which TMA descriptor to use based + solely on tile index (tile 0 = SWA via tmaKSecondary_, rest = compress via tmaK_). + """ + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + token_id = tl.program_id(0) + + # Load request ID for this token + req = tl.load(req_id_ptr + token_id) + + # Load all SWA local indices for this token + swa_ids = tl.arange(0, num_swa_indices) + swa_ptr = swa_local_indices_ptr + token_id * swa_indices_stride0 + swa_ids * swa_indices_stride1 + swa_local_idx = tl.load(swa_ptr) + + # Compute global indices for all SWA positions + swa_valid_mask = swa_local_idx >= 0 + swa_block_ordinal = swa_local_idx // tokens_per_block_swa + swa_token_in_block = swa_local_idx % tokens_per_block_swa + swa_valid_block = swa_block_ordinal < max_blocks_swa + swa_full_mask = swa_valid_mask & swa_valid_block + + swa_bt_ptr = block_table_swa_ptr + req * bt_swa_stride0 + swa_block_ordinal * bt_swa_stride1 + swa_page_index = tl.load(swa_bt_ptr, mask=swa_full_mask, other=0) + + swa_global_index = ( + swa_buffer_offset_in_tokens + swa_page_index * tokens_per_block_swa + swa_token_in_block + ) + swa_global_index = tl.where(swa_full_mask, swa_global_index, -1) + + # Store SWA results at fixed positions [0, num_swa_indices) + swa_out_ptr = out_ptr + token_id * out_stride0 + swa_ids * out_stride1 + tl.store(swa_out_ptr, swa_global_index) + + if has_compressed: + # Load all compressed local indices for this token + compressed_ids = tl.arange(0, num_compressed_indices) + compressed_ptr = ( + compressed_local_indices_ptr + + token_id * compressed_indices_stride0 + + compressed_ids * compressed_indices_stride1 + ) + compressed_local_idx = tl.load(compressed_ptr) + + # Compute global indices for all compressed positions + compressed_valid_mask = compressed_local_idx >= 0 + compressed_block_ordinal = compressed_local_idx // tokens_per_block_compressed + compressed_token_in_block = compressed_local_idx % tokens_per_block_compressed + compressed_valid_block = compressed_block_ordinal < max_blocks_compressed + compressed_full_mask = compressed_valid_mask & compressed_valid_block + + compressed_bt_ptr = ( + block_table_compressed_ptr + + req * bt_compressed_stride0 + + compressed_block_ordinal * bt_compressed_stride1 + ) + compressed_page_index = tl.load(compressed_bt_ptr, mask=compressed_full_mask, other=0) + + compressed_global_index = ( + compressed_buffer_offset_in_tokens + + compressed_page_index * tokens_per_block_compressed + + compressed_token_in_block + ) + compressed_global_index = tl.where(compressed_full_mask, compressed_global_index, -1) + + # Store compressed results at fixed positions [num_swa_indices, total) + compressed_write_pos = num_swa_indices + compressed_ids + compressed_out_ptr = out_ptr + token_id * out_stride0 + compressed_write_pos * out_stride1 + tl.store(compressed_out_ptr, compressed_global_index) + + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +def deepseek_v4_local_to_global_indices( + req_id: torch.Tensor, # int32 [num_tokens] + block_table_swa: torch.Tensor, # int32 [num_requests, max_blocks_swa] + swa_local_indices: torch.Tensor, # int32 [num_tokens, num_swa_indices] + swa_pool_base_ptr: int, # int64: base address of SWA pool + swa_buffer_ptr: int, # int64: base address of SWA buffer + tokens_per_block: int, # tokens per block for SWA + token_stride: int, # bytes per token (use SWA token stride) + # Optional compressed arguments (for compress_ratio > 1) + block_table_compressed: torch.Tensor | None = None, + compressed_local_indices: torch.Tensor | None = None, + compress_pool_base_ptr: int = 0, # int64: base address of compress pool + compressed_buffer_ptr: int = 0, + compress_ratio: int = 1, + num_compressed_indices: int = 0, # max number of compressed indices +) -> torch.Tensor: + """ + Convert local token indices to global KV cache pool indices. + + Dual-pool output layout: + - SWA region [0, window_size): indices relative to swa_pool_base_ptr + - Compress region [window_size, window_size + num_compressed_indices): + indices relative to compress_pool_base_ptr + - Invalid positions padded with -1 at their fixed positions. + + For compress_ratio=1: Only SWA region, all indices relative to swa_pool_base_ptr. + For compress_ratio>1: SWA region + compress region with separate base pointers. + + Args: + req_id: Request ID per token [num_tokens], int32 + block_table_swa: SWA block table [num_requests, max_blocks_swa], int32 + swa_local_indices: Local indices for SWA cache [num_tokens, num_swa_indices], int32 + Use -1 for invalid/padding indices. + swa_pool_base_ptr: Base address of SWA pool + swa_buffer_ptr: Base address of SWA buffer (base_pool_ptr + buffer_offset_in_slot) + tokens_per_block: Number of tokens per block for SWA cache + token_stride: Bytes per token (use SWA token stride) + block_table_compressed: Compressed block table [num_requests, max_blocks_compressed], int32 (optional) + compressed_local_indices: Local indices for compressed cache + [num_tokens, num_compressed_indices], int32 (optional) + Use -1 for invalid/padding indices. + compress_pool_base_ptr: Base address of compress pool + compressed_buffer_ptr: Base address of compressed buffer (optional) + compress_ratio: Compression ratio (1: no compression, >1: with compression) + num_compressed_indices: Max number of compressed indices for CUDA graph compatibility + Output width = num_swa_indices + num_compressed_indices. + + Returns: + global_indices: int32 [num_tokens, num_swa_indices + num_compressed_indices] + """ + assert req_id.dtype == torch.int32, f"req_id must be int32, got {req_id.dtype}" + assert block_table_swa.dtype == torch.int32, ( + f"block_table_swa must be int32, got {block_table_swa.dtype}" + ) + assert swa_local_indices.dtype == torch.int32, ( + f"swa_local_indices must be int32, got {swa_local_indices.dtype}" + ) + + num_tokens = req_id.shape[0] + num_swa_indices = swa_local_indices.shape[1] + + assert swa_local_indices.shape[0] == num_tokens + + has_compressed = compress_ratio > 1 + + # Compute SWA buffer offset relative to swa_pool_base_ptr in tokens + swa_buffer_offset_in_tokens = (swa_buffer_ptr - swa_pool_base_ptr) // token_stride + + if has_compressed: + assert block_table_compressed is not None, ( + "block_table_compressed required when compress_ratio > 1" + ) + assert compressed_local_indices is not None, ( + "compressed_local_indices required when compress_ratio > 1" + ) + assert block_table_compressed.dtype == torch.int32, ( + f"block_table_compressed must be int32, got {block_table_compressed.dtype}" + ) + assert compressed_local_indices.dtype == torch.int32, ( + f"compressed_local_indices must be int32, got {compressed_local_indices.dtype}" + ) + assert compressed_local_indices.shape[0] == num_tokens + + tokens_per_block_compressed = tokens_per_block // compress_ratio + # Compute compressed buffer offset relative to compress_pool_base_ptr in tokens + assert (compressed_buffer_ptr - compress_pool_base_ptr) % token_stride == 0, ( + "compressed_buffer_ptr must be aligned to token_stride" + ) + compressed_buffer_offset_in_tokens = ( + compressed_buffer_ptr - compress_pool_base_ptr + ) // token_stride + _, max_blocks_compressed = block_table_compressed.shape + block_table_compressed_c = block_table_compressed.contiguous() + compressed_local_indices_c = compressed_local_indices.contiguous() + else: + # Dummy values + tokens_per_block_compressed = tokens_per_block + compressed_buffer_offset_in_tokens = 0 + max_blocks_compressed = 1 + block_table_compressed_c = torch.zeros((1, 1), dtype=torch.int32, device=req_id.device) + compressed_local_indices_c = torch.zeros((1, 1), dtype=torch.int32, device=req_id.device) + + total_output_indices = num_swa_indices + num_compressed_indices + _, max_blocks_swa = block_table_swa.shape + + # Ensure contiguous tensors + req_id_c = req_id.contiguous() + block_table_swa_c = block_table_swa.contiguous() + swa_local_indices_c = swa_local_indices.contiguous() + + # Create output tensor + out = torch.empty((num_tokens, total_output_indices), dtype=torch.int32, device=req_id.device) + + # Grid: one program per token + grid = (num_tokens,) + + # Get strides + bt_swa_stride0, bt_swa_stride1 = block_table_swa_c.stride() + bt_compressed_stride0, bt_compressed_stride1 = block_table_compressed_c.stride() + swa_indices_stride0, swa_indices_stride1 = swa_local_indices_c.stride() + compressed_indices_stride0, compressed_indices_stride1 = compressed_local_indices_c.stride() + out_stride0, out_stride1 = out.stride() + launch_with_pdl = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + # Launch kernel + _deepseek_v4_local_to_global_kernel[grid]( + req_id_c, + block_table_swa_c, + block_table_compressed_c, + swa_local_indices_c, + compressed_local_indices_c, + out, + swa_buffer_offset_in_tokens, + compressed_buffer_offset_in_tokens, + tokens_per_block, + tokens_per_block_compressed, + max_blocks_swa, + max_blocks_compressed, + num_swa_indices, + num_compressed_indices, + total_output_indices, + has_compressed, + bt_swa_stride0, + bt_swa_stride1, + bt_compressed_stride0, + bt_compressed_stride1, + swa_indices_stride0, + swa_indices_stride1, + compressed_indices_stride0, + compressed_indices_stride1, + out_stride0, + out_stride1, + LAUNCH_WITH_PDL=launch_with_pdl, + launch_pdl=launch_with_pdl, + ) + + return out diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py similarity index 51% rename from tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py rename to tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py index 1f76e5a73e51..a48679f93c44 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py @@ -1,263 +1,29 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +"""Runtime metadata for DeepSeek-V4 sparse attention.""" + +from __future__ import annotations import math -from dataclasses import dataclass, field, replace -from enum import Enum -from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Set, Tuple +from typing import Dict, Optional, Set, Tuple import torch -from tensorrt_llm._torch.attention_backend.interface import ( - AttentionForwardArgs, - AttentionInputType, - MLAParams, - PositionalEmbeddingParams, - merge_attention_forward_args, -) -from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata -from tensorrt_llm._torch.modules.linear import Linear # noqa: E402 (avoid cycle) -from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream -from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.utils import maybe_compile -from tensorrt_llm._utils import is_sm_100f, prefer_pinned -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization.utils import fp8_utils -from tensorrt_llm.runtime.kv_cache_manager_v2 import DataRole - -from ..dsa import ( - HAS_FAST_HADAMARD, - DSAMetadataParams, - DSAParams, - DSAtrtllmAttentionMetadata, - Indexer, - rotate_activation, -) -from ..kernel import deepseek_v4_local_to_global_indices -from .compressor import Compressor, KVCacheDtype, resolve_kv_cache_dtype - -if TYPE_CHECKING: - from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig - -DEEPSEEK_V4_SPARSE_RATIO = 4 -DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO = 4 - - -class DeepseekV4AttentionType(Enum): - # attentions managed in sliding-window mode are put in the front on purpose - SWA = 0 - COMPRESSOR_KV = 1 - COMPRESSOR_SCORE = 2 - INDEXER_COMPRESSOR_KV = 3 - INDEXER_COMPRESSOR_SCORE = 4 - - # attentions not managed in sliding-window mode - COMPRESS = 5 - INDEXER_COMPRESS = 6 - - @property - def role(self) -> DataRole: - return DataRole(f"deepseek_v4_{self.name.lower()}") - - -DEEPSEEK_V4_SLIDING_ATTENTION = ( - DeepseekV4AttentionType.SWA, - DeepseekV4AttentionType.COMPRESSOR_KV, - DeepseekV4AttentionType.COMPRESSOR_SCORE, - DeepseekV4AttentionType.INDEXER_COMPRESSOR_KV, - DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, -) -assert tuple(attn_type.value for attn_type in DEEPSEEK_V4_SLIDING_ATTENTION) == tuple( - range(len(DEEPSEEK_V4_SLIDING_ATTENTION)) -) - -DEEPSEEK_V4_NON_SLIDING_ATTENTION = ( - DeepseekV4AttentionType.COMPRESS, - DeepseekV4AttentionType.INDEXER_COMPRESS, +from tensorrt_llm._utils import prefer_pinned + +from ..dsa.metadata import DSAtrtllmAttentionMetadata +from .indexer import DeepseekV4Indexer +from .params import ( + DEEPSEEK_V4_SLIDING_ATTENTION, + DEEPSEEK_V4_SPARSE_RATIO, + DeepseekV4AttentionType, + DeepSeekV4MetadataParams, + is_compress_layer, ) -def is_overlap_compressor(compress_ratio: int) -> bool: - return compress_ratio == DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO - - -def is_sparse_layer(compress_ratio: int) -> bool: - return compress_ratio == DEEPSEEK_V4_SPARSE_RATIO - - -def is_compress_layer(compress_ratio: int) -> bool: - return compress_ratio > 1 - - -def compress_ratio_has_attention(compress_ratio: int, attn_type: DeepseekV4AttentionType) -> bool: - is_sparse = is_sparse_layer(compress_ratio) - is_compress = is_compress_layer(compress_ratio) - - if attn_type == DeepseekV4AttentionType.SWA: - return True - if attn_type == DeepseekV4AttentionType.COMPRESS: - return is_compress - if attn_type == DeepseekV4AttentionType.COMPRESSOR_KV: - return is_compress - if attn_type == DeepseekV4AttentionType.COMPRESSOR_SCORE: - return is_compress - if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: - return is_sparse - if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_KV: - return is_sparse - if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE: - return is_sparse - raise ValueError(f"Unsupported DeepSeek-V4 attention type: {attn_type}") - - -def get_attn_dim( - head_dim: int, index_head_dim: int, compress_ratio: int, attn_type: DeepseekV4AttentionType -) -> int: - state_factor = 2 if is_overlap_compressor(compress_ratio) else 1 - if attn_type == DeepseekV4AttentionType.SWA: - return head_dim - if attn_type == DeepseekV4AttentionType.COMPRESS: - return head_dim - if attn_type == DeepseekV4AttentionType.COMPRESSOR_KV: - return state_factor * head_dim - if attn_type == DeepseekV4AttentionType.COMPRESSOR_SCORE: - return state_factor * head_dim - if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: - return index_head_dim - if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_KV: - return state_factor * index_head_dim - if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE: - return state_factor * index_head_dim - raise ValueError(f"Unsupported DeepSeek-V4 attention type: {attn_type}") - - -def get_token_bytes( - head_dim: int, - index_head_dim: int, - compress_ratio: int, - attn_type: DeepseekV4AttentionType, - has_fp8_kv_cache: bool, - indexer_k_dtype: str = "fp8", -) -> int: - if not compress_ratio_has_attention(compress_ratio, attn_type): - raise ValueError( - f"Layer with compress ratio {compress_ratio} does not have attention type {attn_type}" - ) - - attn_dim = get_attn_dim(head_dim, index_head_dim, compress_ratio, attn_type) - - dtype_bytes = 1 if has_fp8_kv_cache else 2 - # (indexer) compressor kv and score always use float32 - if attn_type in [ - DeepseekV4AttentionType.COMPRESSOR_KV, - DeepseekV4AttentionType.COMPRESSOR_SCORE, - DeepseekV4AttentionType.INDEXER_COMPRESSOR_KV, - DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, - ]: - dtype_bytes = 4 # (indexer) compressor kv and score use float32 - # Indexer cache always packs data + per-block scales into one row. Only - # the two indexer presets ("fp8" blockwise / "fp4" mxfp4) are valid - # here — bf16 / fp8_pertensor are reserved for the main-attention - # compressor. - if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: - if indexer_k_dtype == "fp8": - return attn_dim + index_head_dim // 128 * 4 - if indexer_k_dtype == "fp4": - return index_head_dim // 2 + index_head_dim // 32 - raise ValueError( - f"Unsupported indexer_k_dtype {indexer_k_dtype!r}; expected 'fp8' or 'fp4'." - ) - - return attn_dim * dtype_bytes - - -@dataclass(frozen=True) -class DeepSeekV4Params(DSAParams): - """DeepSeek-V4 sparse attention runtime parameters.""" - - algorithm: Literal["deepseek_v4"] = field(init=False, default="deepseek_v4") - compress_ratios: List[int] = field(default_factory=list) - window_size: int = 128 - - -@dataclass(frozen=True) -class DeepSeekV4MetadataParams(DSAMetadataParams): - """DeepSeek-V4 sparse attention metadata settings.""" - - compress_ratios: List[int] = field(default_factory=list) - window_size: int = 128 - - -def _sparse_config_value( - sparse_attention_config: "SparseAttentionConfig", pretrained_config, name: str, default=None -): - value = getattr(sparse_attention_config, name) - if value is not None: - return value - if pretrained_config is not None: - return getattr(pretrained_config, name, default) - return default - - -def make_deepseek_v4_sparse_params( - sparse_attention_config: "SparseAttentionConfig", - pretrained_config=None, -) -> DeepSeekV4Params: - return DeepSeekV4Params( - index_n_heads=_sparse_config_value( - sparse_attention_config, pretrained_config, "index_n_heads" - ), - index_head_dim=_sparse_config_value( - sparse_attention_config, pretrained_config, "index_head_dim" - ), - index_topk=_sparse_config_value(sparse_attention_config, pretrained_config, "index_topk"), - indexer_max_chunk_size=sparse_attention_config.indexer_max_chunk_size, - skip_indexer_for_short_seqs=(sparse_attention_config.skip_indexer_for_short_seqs), - use_cute_dsl_topk=sparse_attention_config.use_cute_dsl_topk, - use_cute_dsl_paged_mqa_logits=(sparse_attention_config.use_cute_dsl_paged_mqa_logits), - q_split_threshold=sparse_attention_config.q_split_threshold, - indexer_rope_interleave=(sparse_attention_config.indexer_rope_interleave), - enable_heuristic_topk=sparse_attention_config.enable_heuristic_topk, - indexer_k_dtype=sparse_attention_config.indexer_k_dtype, - compress_ratios=sparse_attention_config.compress_ratios, - window_size=sparse_attention_config.window_size, - ) - - -def make_deepseek_v4_sparse_metadata_params( - sparse_attention_config: "SparseAttentionConfig", - pretrained_config=None, -) -> DeepSeekV4MetadataParams: - return DeepSeekV4MetadataParams( - indexer_max_chunk_size=(sparse_attention_config.indexer_max_chunk_size or 32768), - max_sparse_topk=_sparse_config_value( - sparse_attention_config, pretrained_config, "index_topk" - ), - index_head_dim=_sparse_config_value( - sparse_attention_config, pretrained_config, "index_head_dim", 128 - ), - enable_indexer_skip=sparse_attention_config.skip_indexer_for_short_seqs, - enable_heuristic_topk=sparse_attention_config.enable_heuristic_topk, - use_cute_dsl_topk=sparse_attention_config.use_cute_dsl_topk, - use_cute_dsl_paged_mqa_logits=(sparse_attention_config.use_cute_dsl_paged_mqa_logits), - q_split_threshold=sparse_attention_config.q_split_threshold, - compress_ratios=sparse_attention_config.compress_ratios, - window_size=sparse_attention_config.window_size, - ) - - class DeepseekV4TrtllmAttentionMetadata(DSAtrtllmAttentionMetadata): # The set of compress ratios for the layers compress_ratio_set: Set[int] @@ -1053,604 +819,3 @@ def _compute_ctx_compressed_position_ids( compressed_position_ids_bufs[compress_ratio][:total_ctx_comp] = ( (past_kv[:num_contexts][ctx_req] + ctx_offset) * compress_ratio ).to(torch.int) - - -class DeepseekV4Indexer(Indexer): - def __init__( - self, - quant_config: Optional[QuantConfig], - pos_embd_params: Optional[PositionalEmbeddingParams], - mla_params: Optional[MLAParams], - skip_create_weights_in_init: bool, - sparse_attention_config: "SparseAttentionConfig", - dtype: Optional[torch.dtype], - compress_ratio: int = 1, - layer_idx: int = 0, - aux_stream: Optional[torch.cuda.Stream] = None, - ): - super().__init__( - quant_config, - pos_embd_params, - mla_params, - skip_create_weights_in_init, - sparse_attention_config, - dtype, - compress_ratio, - layer_idx, - aux_stream, - ) - # Preserve the checkpoint's 128x128 FP8 quantization while deriving a - # native-MXF8 (sf_vec=32) scale view for the TRT-LLM CuTe DSL fused - # indexer-Q projection. FP8BlockScalesLinearMethod materializes the - # derived, swizzled scale once after weight loading. - self.wq_b.use_indexer_q_cutedsl_fusion = True - # Override base Indexer.weights_proj to bf16 (matches V4 checkpoint). - self.weights_proj = Linear( - self.hidden_size, - self.n_heads, - bias=False, - dtype=dtype, - quant_config=None, - skip_create_weights_in_init=skip_create_weights_in_init, - use_custom_cublas_mm=True, - ) - self.rotary_emb = RotaryEmbedding( - pos_embd_params.rope, - head_dim=self.rope_dim, - is_neox=False, - ) - rms_norm_eps = 1e-6 - index_head_dim = sparse_attention_config.index_head_dim - indexer_mla_params = MLAParams( - hidden_size=mla_params.hidden_size, - qk_rope_head_dim=mla_params.qk_rope_head_dim, - qk_nope_head_dim=index_head_dim - mla_params.qk_rope_head_dim, - ) - # Map the user-facing FP4 knob ("fp8" / "fp4") onto the Compressor's - # cache layout preset string. The Compressor's preset namespace also - # covers main-attention layouts (bf16 / fp8_pertensor), which the - # indexer doesn't use, so the translation lives here at the - # boundary instead of leaking through the user-facing config. - self.indexer_k_dtype = sparse_attention_config.indexer_k_dtype - compressor_preset = "mxfp4" if self.indexer_k_dtype == "fp4" else "fp8_blockwise" - self.indexer_cache_dtype = resolve_kv_cache_dtype(compressor_preset) - self.compressor = Compressor( - indexer_mla_params, - layer_idx, - compress_ratio, - rms_norm_eps, - skip_create_weights_in_init, - pos_embd_params, - dtype=dtype, - kv_cache_dtype=compressor_preset, - is_indexer=True, - rotate_activation=HAS_FAST_HADAMARD, - ) - self.indexer_start_event = torch.cuda.Event() - self.weights_proj_event = torch.cuda.Event() - self.k_cache_update_event = torch.cuda.Event() - - def post_load_weights(self): - # V4 does not use the V3 fused fp32 wk+weights_proj GEMM, and the - # base concat would now hit an fp32/bf16 dtype mismatch. - return - - def _qk_projection_and_rope(self, qr: torch.Tensor, position_ids: torch.Tensor): - """Project Q and apply RoPE. - - Returns q with layout [num_tokens, n_heads, head_dim] where - head_dim = nope_dim + rope_dim, RoPE already applied in-place. - """ - q = self.wq_b(qr) - q = q.view(-1, self.n_heads, self.head_dim) - return self._apply_q_rope(q, position_ids) - - def _apply_q_rope(self, q: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: - """Apply RoPE in-place to a projected indexer Q tensor.""" - # Fused in-place RoPE on the rope portion of each head - nope_dim = self.head_dim - self.rope_dim - torch.ops.trtllm.mla_rope_inplace( - q, - position_ids.view(-1), - self.rotary_emb.rotary_cos_sin, - self.n_heads, - nope_dim, - self.rope_dim, - False, - self.rotary_emb.is_neox, - ) - return q - - def _project_and_quantize_q( - self, qr: torch.Tensor, position_ids: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Project Q and produce the cache precision consumed by the indexer. - - The DSv4 MXFP4 configuration can fuse projection, interleaved RoPE, - and FP4 quantization. If the optional Hadamard transform is active, - retain the legacy path because that transform changes all 128 values - in a head and is not part of the CuTe DSL kernel. - """ - use_fused_project_mxfp4 = ( - self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE - and not HAS_FAST_HADAMARD - and not self.rotary_emb.is_neox - and self.head_dim == 128 - and self.rope_dim == 64 - and self.wq_b.has_fp8_block_scales - and hasattr(self.wq_b, "indexer_q_weight_scale_cutedsl") - and hasattr( - torch.ops.trtllm, - "cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell", - ) - and qr.dtype == torch.bfloat16 - and is_sm_100f() - ) - if use_fused_project_mxfp4: - q_fp4, q_scale = torch.ops.trtllm.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell( - qr, - self.wq_b.weight, - self.wq_b.indexer_q_weight_scale_cutedsl, - position_ids.view(-1), - self.rotary_emb.rotary_cos_sin.view(-1, self.rope_dim), - self.wq_b.indexer_q_alpha_cutedsl, - use_tvm_ffi=True, - ) - return q_fp4.view(-1, self.n_heads, self.head_dim // 2), q_scale.view( - -1, self.n_heads, 1 - ) - - q = self.wq_b(qr).view(-1, self.n_heads, self.head_dim) - q = self._apply_q_rope(q, position_ids) - return self._quantize_q(q) - - def _quantize_q(self, q: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - # Rotate + quantize (layout matches compressor K: [nope|pe]). After - # rotate_activation (Hadamard) the nope/rope split becomes a linear - # mix; treating the row as a (head_dim - rope_dim, rope_dim) split - # for fused_cat_fp4 is equivalent to a single per-token FP4 quantize - # because the kernel only cares about the concatenated row. K goes - # through the same rotation in compressor_postprocess_scatter so - # layouts match. - q = rotate_activation(q) - q = q.view(-1, self.head_dim) - if self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE: - nope_dim = self.head_dim - self.rope_dim - q_nope, q_pe = q.split([nope_dim, self.rope_dim], dim=-1) - q_fp8, q_scale = torch.ops.trtllm.fused_cat_fp4(q_nope, q_pe) - # Two FP4 codes pack into one byte: trailing dim is head_dim // 2. - q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim // 2) - q_scale = q_scale.view(-1, self.n_heads, 1) - else: - q_fp8, q_scale = fp8_utils.fp8_quantize_1x128_sf_transpose( - q, use_ue8m0=self.scale_fmt == "ue8m0" - ) - q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim) - q_scale = q_scale.view(-1, self.n_heads, 1) - return q_fp8, q_scale - - def _apply_weight_scale(self, weights: torch.Tensor, q_scale: torch.Tensor) -> torch.Tensor: - # The DeepGEMM FP4 kernel applies per-block q_scale internally, so - # weights only carry softmax_scale * n_heads^-0.5. `weights_proj` is - # bf16 to match the V4 checkpoint, and the FP4 branch's only - # post-projection multiplier is a Python float (`bf16 * float` stays - # bf16). DeepGEMM's fp8_fp4_(paged_)mqa_logits asserts - # `weights.scalar_type() == kFloat`, so cast explicitly here. The FP8 - # branch gets upcast for free via the fp32 `q_scale` multiply in - # `_weight_scale`, so no cast is needed there. - if self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE: - return weights.float() * self.weight_scale_factor - return self._weight_scale(weights, q_scale) - - def _update_k_cache_if_needed( - self, - k_fp8: Optional[torch.Tensor], - k_scale: Optional[torch.Tensor], - metadata: DeepseekV4TrtllmAttentionMetadata, - ) -> None: - if k_fp8 is None: - return - - # The MXFP4 compressor's postprocess_scatter performs rotation + - # quantization + cache scatter in one fused op, so the cache row is - # already up to date by the time we get here. - if self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE: - return - - assert k_scale is not None, "FP8 blockwise indexer cache update requires scale tensor" - self._update_k_cache(k_fp8, k_scale, metadata) - - def precompute_aux( - self, - hidden_states: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, - ) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]]: - """Pre-launch the qr-independent half of the indexer prepare phase. - - Runs weights_proj + internal compressor + k_cache_update on - ``self.aux_stream`` and records ``weights_proj_event`` / - ``k_cache_update_event``. The caller hands the returned tuple back to - ``forward()`` via the ``pre_aux`` kwarg, which makes the overlapped - prepare path skip its own aux-stream launch and consume these results - directly. - - Returns ``None`` when multi-stream mode is off (caller should fall - back to the normal ``forward()`` call without ``pre_aux``). - """ - if not (do_multi_stream() and self.aux_stream is not None): - return None - self.indexer_start_event.record() - with torch.cuda.stream(self.aux_stream): - self.indexer_start_event.wait() - weights = self.weights_proj(hidden_states) - self.weights_proj_event.record() - k_fp8, k_scale = self.compressor(hidden_states, metadata) - self._update_k_cache_if_needed(k_fp8, k_scale, metadata) - self.k_cache_update_event.record() - return (weights, k_fp8, k_scale) - - def _run_overlapped_indexer_prepare( - self, - qr: torch.Tensor, - hidden_states: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, - position_ids: torch.Tensor, - pre_aux: Optional[ - Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] - ] = None, - ) -> Tuple[ - torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], torch.Tensor - ]: - """Prepare indexer inputs by splitting independent work across two streams. - - The current stream owns the Q path. The auxiliary stream starts from - the recorded launch point and owns weights projection, compressor, and - the K-cache update. - - When ``pre_aux`` is provided the aux-stream work has already been - launched (and ``weights_proj_event`` / ``k_cache_update_event`` - recorded) by ``precompute_aux``; the aux-stream block here is skipped - and the precomputed tensors are consumed directly. - - Timeline (pre_aux=None): - current stream: - record indexer_start_event - q_proj + RoPE -> quant_q - wait weights_proj_event -> weight_scale - wait k_cache_update_event -> return - - aux_stream: - wait indexer_start_event - weights_proj -> record weights_proj_event - compressor -> update_k_cache -> record k_cache_update_event - - Dependency graph: - q_proj + RoPE -> quant_q -- q_scale --. - v - weights_proj --------------------> weight_scale - compressor -> update_k_cache ----> final wait - """ - if pre_aux is None: - self.indexer_start_event.record() - - q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) - - with torch.cuda.stream(self.aux_stream): - self.indexer_start_event.wait() - - weights = self.weights_proj(hidden_states) - self.weights_proj_event.record() - - k_fp8, k_scale = self.compressor(hidden_states, metadata) - self._update_k_cache_if_needed(k_fp8, k_scale, metadata) - self.k_cache_update_event.record() - else: - weights, k_fp8, k_scale = pre_aux - # pre_aux tensors were allocated on aux_stream; record on the - # consuming stream so the caching allocator can't recycle them mid-use. - cur_stream = torch.cuda.current_stream() - weights.record_stream(cur_stream) - if k_fp8 is not None: - k_fp8.record_stream(cur_stream) - if k_scale is not None: - k_scale.record_stream(cur_stream) - q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) - - self.weights_proj_event.wait() - weights = self._apply_weight_scale(weights, q_scale) - - self.k_cache_update_event.wait() - return q_fp8, q_scale, k_fp8, k_scale, weights - - def _run_serial_indexer_prepare( - self, - qr: torch.Tensor, - hidden_states: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, - position_ids: torch.Tensor, - ) -> Tuple[ - torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], torch.Tensor - ]: - q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) - - weights = self.weights_proj(hidden_states) - - weights = self._apply_weight_scale(weights, q_scale) - - k_fp8, k_scale = self.compressor(hidden_states, metadata) - self._update_k_cache_if_needed(k_fp8, k_scale, metadata) - return q_fp8, q_scale, k_fp8, k_scale, weights - - def _update_k_cache( - self, - k_fp8: torch.Tensor, - k_scale: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata, - ) -> None: - # DSV4's indexer compressor already scatters INDEXER_COMPRESS. The - # shared DSA scatter would duplicate that write. - return - - def forward( - self, - qr: torch.Tensor, - hidden_states: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, - position_ids: torch.Tensor, - pre_aux: Optional[ - Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] - ] = None, - ): - if do_multi_stream() and self.aux_stream is not None: - q_fp8, q_scale, k_fp8, k_scale, weights = self._run_overlapped_indexer_prepare( - qr, - hidden_states, - metadata, - position_ids, - pre_aux=pre_aux, - ) - else: - assert pre_aux is None, "pre_aux requires multi-stream mode" - q_fp8, q_scale, k_fp8, k_scale, weights = self._run_serial_indexer_prepare( - qr, hidden_states, metadata, position_ids - ) - - # 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]] - else: - topk_indices = self.sparse_attn_indexer( - metadata, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - q_scale=q_scale, - update_k_cache=False, - ) - return topk_indices - - -class DeepseekV4TrtllmAttention(TrtllmAttention): - Metadata = DeepseekV4TrtllmAttentionMetadata - - def __init__( - self, - layer_idx: int, - num_heads: int, - head_dim: int, - num_kv_heads: Optional[int] = None, - quant_config: Optional[QuantConfig] = None, - q_scaling: Optional[float] = None, - pos_embd_params: Optional[PositionalEmbeddingParams] = None, - mla_params: Optional[MLAParams] = None, - skip_create_weights_in_init: bool = False, - attention_chunk_size: Optional[int] = None, - sparse_attention_config: Optional["SparseAttentionConfig"] = None, - sparse_params: Optional[DeepSeekV4Params] = None, - dtype: Optional[torch.dtype] = None, - aux_stream: Optional[torch.cuda.Stream] = None, - **kwargs, - ): - if sparse_attention_config is None: - sparse_attention_config = sparse_params - assert sparse_attention_config is not None, ( - "sparse_attention_config is required for DeepseekV4TrtllmAttention and cannot be None" - ) - if sparse_params is None: - sparse_params = make_deepseek_v4_sparse_params(sparse_attention_config) - assert sparse_params is not None, ( - "sparse_params is required for DeepseekV4TrtllmAttention and cannot be None" - ) - TrtllmAttention.__init__( - self, - layer_idx, - num_heads, - head_dim, - sparse_params=sparse_params, - num_kv_heads=num_kv_heads, - quant_config=quant_config, - q_scaling=q_scaling, - pos_embd_params=pos_embd_params, - mla_params=mla_params, - skip_create_weights_in_init=skip_create_weights_in_init, - attention_chunk_size=attention_chunk_size, - **kwargs, - ) - - self.sparse_attention_config = sparse_attention_config - self.compress_ratio = sparse_attention_config.compress_ratios[layer_idx] - - if self.compress_ratio == 4: - self.indexer = DeepseekV4Indexer( - quant_config, - pos_embd_params, - mla_params, - skip_create_weights_in_init, - sparse_attention_config, - dtype, - self.compress_ratio, - layer_idx, - aux_stream, - ) - - if self.compress_ratio > 1: - rms_norm_eps = 1e-6 - has_fp8_kv_cache = False - if quant_config is not None: - has_fp8_kv_cache = quant_config.layer_quant_mode.has_fp8_kv_cache() - kv_cache_dtype = "fp8_pertensor" if has_fp8_kv_cache else "default" - self.compressor = Compressor( - mla_params, - layer_idx, - self.compress_ratio, - rms_norm_eps, - skip_create_weights_in_init, - pos_embd_params, - kv_cache_dtype=kv_cache_dtype, - dtype=dtype, - rotate_activation=False, - ) - - def _prepare_sparse_forward_args( - self, - metadata: DeepseekV4TrtllmAttentionMetadata, - forward_args: AttentionForwardArgs, - ) -> None: - attention_input_type = forward_args.attention_input_type - if attention_input_type == AttentionInputType.context_only: - start_idx = 0 - end_idx = metadata.num_ctx_tokens - elif attention_input_type == AttentionInputType.generation_only: - start_idx = metadata.num_ctx_tokens - end_idx = metadata.num_tokens - else: - start_idx = 0 - end_idx = metadata.num_tokens - - sparse_args = forward_args.sparse_prediction - sparse_args.sparse_mla_topk_lens = metadata.sparse_mla_topk_lens[self.compress_ratio][ - start_idx:end_idx - ] - if self.compress_ratio > 1: - sparse_args.compressed_kv_cache_pool_ptr = metadata.sparse_mla_base_ptrs[ - self.compress_ratio - ] - else: - sparse_args.compressed_kv_cache_pool_ptr = None - - metadata.num_sparse_topk = ( - self.sparse_attention_config.window_size - + metadata.max_compressed_indices[self.compress_ratio] - ) - - def forward( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - metadata: DeepseekV4TrtllmAttentionMetadata, - forward_args: Optional[AttentionForwardArgs] = None, - **kwargs, - ): - forward_args = merge_attention_forward_args(forward_args, kwargs) - attn_sink = getattr(self, "attn_sink", None) - if attn_sink is not None: - if forward_args.attention_sinks is None: - forward_args = replace(forward_args, attention_sinks=attn_sink.data) - - self._prepare_sparse_forward_args(metadata, forward_args) - return super().forward(q, k, v, metadata, forward_args=forward_args) - - def sparse_attn_predict( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - metadata: DeepseekV4TrtllmAttentionMetadata, - forward_args: AttentionForwardArgs, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """Convert local indices (SWA + compressed) to global pool indices.""" - 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] - - # Get cached buffer pointers - swa_buffer_ptr = metadata.swa_buffer_ptrs[layer_idx] - - # Token stride - index_head_dim = self.sparse_attention_config.index_head_dim - has_fp8_kv_cache = False - if self.quant_config is not None: - has_fp8_kv_cache = self.quant_config.layer_quant_mode.has_fp8_kv_cache() - token_stride = get_token_bytes( - self.head_dim, - index_head_dim, - self.compress_ratio, - DeepseekV4AttentionType.SWA, - has_fp8_kv_cache, - ) - - # Select token range based on phase - if attention_input_type == AttentionInputType.context_only: - start_idx = 0 - end_idx = metadata.num_ctx_tokens - elif attention_input_type == AttentionInputType.generation_only: - start_idx = metadata.num_ctx_tokens - end_idx = metadata.num_tokens - else: - start_idx = 0 - 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] - local_layer_idx = kv_cache_manager.layer_offsets[layer_idx] - block_table_swa = metadata.sliding_block_tables[ - local_layer_idx, DeepseekV4AttentionType.SWA.value - ] - - 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] - if self.compress_ratio == 4: - topk_indices = forward_args.topk_indices - assert topk_indices is not None, "topk_indices is required when compress_ratio=4" - compressed_local_indices = topk_indices - else: - compressed_local_indices = metadata.compressed_local_indices_cuda[start_idx:end_idx] - else: - compressed_buffer_ptr = 0 - compress_pool_base_ptr = 0 - block_table_compressed = None - compressed_local_indices = None - - global_indices = deepseek_v4_local_to_global_indices( - req_id=req_id, - block_table_swa=block_table_swa, - swa_local_indices=swa_local_indices, - swa_pool_base_ptr=swa_pool_base_ptr, - swa_buffer_ptr=swa_buffer_ptr, - tokens_per_block=kv_cache_manager.tokens_per_block, - token_stride=token_stride, - block_table_compressed=block_table_compressed, - compressed_local_indices=compressed_local_indices, - 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], - ) - - return global_indices, None - - def sparse_kv_predict( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - metadata: DeepseekV4TrtllmAttentionMetadata, - forward_args: AttentionForwardArgs, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - return None, None diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py new file mode 100644 index 000000000000..b91daa6604a2 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py @@ -0,0 +1,817 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-V4 integration for the shared MLA module.""" + +import os +from typing import TYPE_CHECKING, Optional + +import torch +from torch import nn + +from tensorrt_llm._torch.attention_backend.interface import AttentionInputType, AttentionMetadata +from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode +from tensorrt_llm._torch.modules.multi_stream_utils import ( + do_multi_stream, + maybe_execute_in_parallel, +) +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding +from tensorrt_llm._utils import get_sm_version, is_sm_100f + +from ..params import SparseBackendForwardArgs + +if TYPE_CHECKING: + from tensorrt_llm._torch.distributed import AllReduceParams + +_q_b_proj_cute_dsl_import_ok: Optional[bool] = None + + +# Module initialization and weight lifecycle for DeepSeek-V4 MLA. + + +def initialize_sparse_attn( + self, + *, + config, + mapping, + mapping_o, + rms_norm_eps: float, + quant_config, + q_scaling: float, + bias: bool, + dtype: torch.dtype, + reduce_output: bool, + aux_stream: Optional[torch.cuda.Stream], +) -> None: + """Initialize DeepSeek-V4 module state and remove unused dense modules.""" + del bias, q_scaling + tp_size = mapping.tp_size + if self.num_groups % tp_size != 0: + raise ValueError( + f"DeepSeek-V4 num_groups ({self.num_groups}) must be divisible by tp_size ({tp_size})." + ) + if self.num_heads % self.num_groups != 0: + raise ValueError( + f"DeepSeek-V4 num_heads ({self.num_heads}) must be divisible by " + f"num_groups ({self.num_groups})." + ) + if self.is_lite: + raise ValueError("DeepSeek-V4 does not support lite MLA") + + del self.kv_b_proj + del self.v_b_proj + del self.o_proj + self.mha = None + self.indexer = getattr(self.mqa, "indexer", None) + self.compressor = getattr(self.mqa, "compressor", None) + + self.n_local_groups = self.num_groups // tp_size + self.q_b_layernorm = RMSNorm( + hidden_size=self.qk_head_dim, + eps=rms_norm_eps, + dtype=dtype, + has_weights=False, + ) + self.kv_a_layernorm = RMSNorm( + hidden_size=self.kv_lora_rank + self.qk_rope_head_dim, + dtype=dtype, + eps=rms_norm_eps, + ) + self.o_a_proj = nn.Parameter( + torch.empty( + ( + self.n_local_groups, + self.o_lora_rank, + self.num_heads * self.qk_head_dim // self.num_groups, + ), + dtype=dtype, + ), + requires_grad=False, + ) + self.o_b_proj = Linear( + self.num_groups * self.o_lora_rank, + self.hidden_size, + bias=False, + dtype=dtype, + mapping=mapping_o, + tensor_parallel_mode=TensorParallelMode.ROW, + quant_config=quant_config, + skip_create_weights_in_init=config.skip_create_weights_in_init, + reduce_output=reduce_output, + allreduce_strategy=config.allreduce_strategy, + force_dynamic_quantization=config.force_dynamic_quantization, + use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, + use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, + ) + + self.has_dsv4_indexer = ( + self.layer_idx is not None and self.sparse_params.compress_ratios[self.layer_idx] == 4 + ) + self.indexer_stream = None + self.indexer_aux_stream = None + self.compressor_stream = None + if self.has_dsv4_indexer and aux_stream is not None: + self.indexer_stream = torch.cuda.Stream(device=aux_stream.device) + self.indexer_aux_stream = torch.cuda.Stream(device=aux_stream.device) + self.compressor_stream = torch.cuda.Stream(device=aux_stream.device) + if self.indexer_aux_stream is not None: + assert self.indexer is not None + self.indexer.aux_stream = self.indexer_aux_stream + + self.inverse_rotary_emb = RotaryEmbedding( + self.pos_embd_params.rope, + head_dim=self.qk_rope_head_dim, + is_neox=self.pos_embd_params.is_neox, + inverse=True, + ) + self._disable_dsv4_epilogue_fusion = os.environ.get( + "TRTLLM_DSV4_DISABLE_FMHA_EPILOGUE_FUSION", "" + ).strip().lower() in ("1", "true", "on") + self.dsv4_overlap_start_event = torch.cuda.Event() + self.dsv4_compressor_start_event = torch.cuda.Event() + self.dsv4_compressor_event = torch.cuda.Event() + self.dsv4_indexer_event = torch.cuda.Event() + self.attention_output_hidden_size = self.num_heads_tp_cp * self.v_head_dim + + +def create_sparse_attn_weights(self) -> None: + has_fp8_block_scales = bool( + self.o_b_proj.quant_config and self.o_b_proj.quant_config.quant_mode.has_fp8_block_scales() + ) + self.o_a_proj_dequant = None + if has_fp8_block_scales: + self.o_a_proj_scale = nn.Parameter( + torch.empty( + ( + self.n_local_groups, + self.o_lora_rank // 128, + self.num_heads * self.qk_head_dim // self.num_groups // 128, + ), + dtype=torch.float32, + ), + requires_grad=False, + ) + if is_sm_100f(): + self.o_a_proj = nn.Parameter( + torch.empty( + ( + self.n_local_groups, + self.o_lora_rank, + self.num_heads * self.qk_head_dim // self.num_groups, + ), + dtype=torch.float8_e4m3fn, + ), + requires_grad=False, + ) + else: + self.o_a_proj_scale = None + + +def transform_sparse_attn_weights(self) -> None: + """Skip the dense MLA weight transformation for DeepSeek-V4.""" + return None + + +# Fused epilogue buffer management and output projection. + + +def _validate_dsv4_epilogue_buffers( + self, + num_tokens: int, + dsv4_epilogue_output: tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + fp8_o, output_sf = dsv4_epilogue_output + scale_buf_m = (num_tokens + 3) // 4 * 4 + if fp8_o.shape[1] != num_tokens or output_sf.shape[2] != scale_buf_m: + raise RuntimeError("Invalid DSv4 fused epilogue buffers for current token count.") + return fp8_o, output_sf + + +def prepare_sparse_attn_outputs( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata +) -> list[torch.Tensor]: + def _should_use_dsv4_epilogue_fusion() -> bool: + num_contexts = attn_metadata.num_contexts + num_generations = attn_metadata.num_generations + if self._disable_dsv4_epilogue_fusion: + return False + if num_contexts == 0 and num_generations == 0: + return False + if num_contexts > 0 and num_generations > 0: + # The fused buffers do not carry token offsets for a mixed batch. + return False + if self.mapping.has_cp_helix() or not is_sm_100f(): + return False + if not getattr(self.mapping, "enable_attention_dp", False): + return False + if self.num_heads != 128 or self.num_heads_tp != 128: + return False + if getattr(self.mqa, "sparse_params", None) is None: + return False + if not getattr(self.mqa, "has_fp8_kv_cache", False): + return False + if self.o_a_proj.dtype != torch.float8_e4m3fn: + return False + if self.kv_lora_rank != 448 or self.qk_rope_head_dim != 64: + return False + if self.qk_head_dim != 512 or self.v_head_dim != 512: + return False + if self.n_local_groups <= 0 or self.num_heads_tp % self.n_local_groups != 0: + return False + return not self.inverse_rotary_emb.is_neox + + def _create_dsv4_epilogue_buffers() -> tuple[torch.Tensor, torch.Tensor]: + if self.n_local_groups <= 0 or self.num_heads_tp % self.n_local_groups != 0: + raise ValueError( + "DSv4 fused epilogue requires num_heads_tp to be divisible by n_local_groups." + ) + heads_per_group = self.num_heads_tp // self.n_local_groups + num_tokens = attn_metadata.num_tokens + scale_buf_m = (num_tokens + 3) // 4 * 4 + fp8_o = hidden_states.new_empty( + (self.n_local_groups, num_tokens, heads_per_group * self.v_head_dim), + dtype=torch.float8_e4m3fn, + ) + output_sf = hidden_states.new_empty( + ( + self.n_local_groups, + heads_per_group * (self.v_head_dim // 128), + scale_buf_m, + ), + dtype=torch.float32, + ) + return fp8_o, output_sf + + if _should_use_dsv4_epilogue_fusion(): + attn_output = [self.create_output(hidden_states[:0], attn_metadata.num_contexts)] + attn_output.extend(_create_dsv4_epilogue_buffers()) + return attn_output + return [self.create_output(hidden_states, attn_metadata.num_contexts)] + + +def project_sparse_attn_output( + self, + attn_output: list[torch.Tensor], + position_ids: Optional[torch.Tensor] = None, + attn_metadata: Optional[AttentionMetadata] = None, + all_reduce_params: Optional["AllReduceParams"] = None, +) -> torch.Tensor: + del attn_metadata, all_reduce_params + if len(attn_output) > 1: + attn_fp8, attn_scale = attn_output[1:] + num_tokens = attn_fp8.shape[1] + o_lora = torch.empty( + [num_tokens, self.n_local_groups, self.o_lora_rank], + device=attn_fp8.device, + dtype=self.dtype, + ) + torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( + attn_fp8, + self.o_a_proj, + attn_scale, + self.o_a_proj_scale, + o_lora.transpose(0, 1), + ) + return self.o_b_proj(o_lora.flatten(1)) + + attn_output_tensor = attn_output[0] + assert position_ids is not None + num_tokens = attn_output_tensor.shape[0] + attn_output_tensor = attn_output_tensor.view(num_tokens, self.num_heads_tp, -1) + + # Fuse inverse RoPE with FP8 quantization to avoid a BF16 latent read/write. + # This is independent of the K/V absorption BMM implementation. + fused_inv_rope_fp8 = self.o_a_proj.dtype == torch.float8_e4m3fn and is_sm_100f() + if fused_inv_rope_fp8: + heads_per_group = self.num_heads_tp // self.n_local_groups + attn_fp8, attn_scale = torch.ops.trtllm.fused_inv_rope_fp8_quant_vllm_port( + attn_output_tensor, + position_ids.view(-1), + self.inverse_rotary_emb.rotary_cos_sin, + self.n_local_groups, + heads_per_group, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + 128, + self.inverse_rotary_emb.is_neox, + ) + o_lora = torch.empty( + [num_tokens, self.n_local_groups, self.o_lora_rank], + device=attn_output_tensor.device, + dtype=self.dtype, + ) + torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( + attn_fp8, + self.o_a_proj, + attn_scale, + self.o_a_proj_scale, + o_lora.transpose(0, 1), + ) + o_lora = o_lora.flatten(1) + return self.o_b_proj(o_lora) + + # Restore the RoPE portion before output projection. + torch.ops.trtllm.mla_rope_inplace( + attn_output_tensor, + position_ids.view(-1), + self.inverse_rotary_emb.rotary_cos_sin, + self.num_heads_tp, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + True, + self.inverse_rotary_emb.is_neox, + ) + + o_lora = torch.empty( + [num_tokens, self.n_local_groups, self.o_lora_rank], + device=attn_output_tensor.device, + dtype=attn_output_tensor.dtype, + ) + if self.o_a_proj.dtype == torch.bfloat16: + # [groups, tokens, dim] @ [groups, dim, rank] -> [groups, tokens, rank] + torch.ops.trtllm.bmm_out( + attn_output_tensor.view(num_tokens, self.n_local_groups, -1).transpose(0, 1), + self.o_a_proj.transpose(1, 2), + o_lora.transpose(0, 1), + ) + elif self.o_a_proj.dtype == torch.float8_e4m3fn: + from tensorrt_llm._torch.modules.mla import fp8_block_scaling_bmm_out + + fp8_block_scaling_bmm_out( + attn_output_tensor.view(num_tokens, self.n_local_groups, -1), + self.o_a_proj, + self.o_a_proj_scale, + o_lora.transpose(0, 1), + self.o_a_proj_dequant, + self.use_cute_dsl_blockscaling_bmm, + ) + else: + raise NotImplementedError(f"Missing bmm impl for dtype: {self.o_a_proj.dtype}.") + o_lora = o_lora.flatten(1) + output = self.o_b_proj(o_lora) + return output + + +# Context and generation attention execution paths. + + +def forward_generation_sparse_attn( + self, + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + attn_metadata: AttentionMetadata, + output: torch.Tensor, + position_ids: Optional[torch.Tensor] = None, + latent_cache: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, + sparse_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run the DeepSeek-V4 generation absorption path.""" + if get_sm_version() < 100: + raise RuntimeError("DeepSeek-V4 is not supported on pre-Blackwell GPUs") + del compressed_kv, k_pe + num_tokens = q.shape[0] + q_pe = q.view(-1, self.num_heads_tp, self.qk_head_dim)[..., self.qk_nope_head_dim :] + + num_seqs = attn_metadata.num_seqs + cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=q.device) + cu_kv_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=q.device) + fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=q.device) + + has_fp8_kv_cache = bool(getattr(self.mqa, "has_fp8_kv_cache", False)) + mla_bmm1_scale = None + mla_bmm2_scale = None + quant_q_buffer = None + if has_fp8_kv_cache: + mla_bmm1_scale = torch.empty(2, dtype=torch.float32, device=q.device) + mla_bmm2_scale = torch.empty(1, dtype=torch.float32, device=q.device) + quant_q_buffer = torch.empty( + num_tokens, + self.num_heads_tp, + self.kv_lora_rank + self.qk_rope_head_dim, + dtype=torch.uint8, + device=q.device, + ) + + self.mqa.mla_rope_generation( + q, + q_pe, + latent_cache, + attn_metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + ) + + attention_output = output + output_sf = None + inverse_rope_cos_sin = None + if sparse_epilogue_output is not None: + attention_output, output_sf = _validate_dsv4_epilogue_buffers( + self, num_tokens, sparse_epilogue_output + ) + inverse_rope_cos_sin = self.inverse_rotary_emb.rotary_cos_sin + + attn_out_latent = self._attn_forward_gen( + self.mqa, + q, + None, + None, + position_ids, + attn_metadata, + attention_input_type=AttentionInputType.generation_only, + out_scale=self.out_scale, + output=attention_output, + output_sf=output_sf, + latent_cache=latent_cache, + q_pe=q_pe, + sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk_indices), + cu_q_seqlens=cu_q_seqlens, + cu_kv_seqlens=cu_kv_seqlens, + fmha_scheduler_counter=fmha_scheduler_counter, + mla_bmm1_scale=mla_bmm1_scale, + mla_bmm2_scale=mla_bmm2_scale, + quant_q_buffer=quant_q_buffer, + dsv4_inv_rope_cos_sin_cache=inverse_rope_cos_sin, + enable_dsv4_epilogue_fusion=sparse_epilogue_output is not None, + ) + if sparse_epilogue_output is not None: + return attn_out_latent + + if self.mapping.has_cp_helix(): + raise RuntimeError( + "DeepSeek-V4 + CP Helix is not supported because the post-process " + "does not preserve the pre-allocated output buffer." + ) + assert attn_out_latent.data_ptr() == output.data_ptr(), ( + "Attention backend did not write into the provided output buffer." + ) + return output + + +def forward_context_sparse_attn( + self, + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + attn_metadata: AttentionMetadata, + output: torch.Tensor, + latent_cache: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + sparse_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run the DeepSeek-V4 context absorption path.""" + if get_sm_version() < 100: + raise RuntimeError("DeepSeek-V4 is not supported on pre-Blackwell GPUs") + del compressed_kv, k_pe + num_tokens = q.shape[0] + q_pe = q.view(-1, self.num_heads_tp, self.qk_head_dim)[..., self.qk_nope_head_dim :] + + quant_q_buffer = getattr(self, "_fused_quant_q_buffer", None) + fused_q_pe = getattr(self, "_fused_q_pe", None) + quant_scale_qkv = getattr(self, "_quant_scale_qkv", None) + use_fused_q_fp8 = ( + quant_q_buffer is not None and fused_q_pe is not None and quant_scale_qkv is not None + ) + if use_fused_q_fp8: + q_pe = fused_q_pe[:num_tokens] + quant_q_buffer = quant_q_buffer[:num_tokens].view( + num_tokens, + self.num_heads_tp, + self.kv_lora_rank + self.qk_rope_head_dim, + ) + else: + quant_q_buffer = None + quant_scale_qkv = None + + attention_output = output + output_sf = None + inverse_rope_cos_sin = None + if sparse_epilogue_output is not None: + attention_output, output_sf = _validate_dsv4_epilogue_buffers( + self, num_tokens, sparse_epilogue_output + ) + inverse_rope_cos_sin = self.inverse_rotary_emb.rotary_cos_sin + + attn_out_latent = self._attn_forward_gen( + self.mqa, + q, + None, + None, + position_ids, + attn_metadata, + attention_input_type=AttentionInputType.context_only, + out_scale=self.out_scale, + output=attention_output, + output_sf=output_sf, + latent_cache=latent_cache, + q_pe=q_pe, + quant_q_buffer=quant_q_buffer, + quant_scale_qkv=quant_scale_qkv, + sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk_indices), + dsv4_inv_rope_cos_sin_cache=inverse_rope_cos_sin, + enable_dsv4_epilogue_fusion=sparse_epilogue_output is not None, + ) + self._fused_quant_q_buffer = None + self._fused_q_pe = None + + if sparse_epilogue_output is not None: + return attn_out_latent + if self.mapping.has_cp_helix(): + raise RuntimeError( + "DeepSeek-V4 + CP Helix is not supported because the post-process " + "does not preserve the pre-allocated output buffer." + ) + assert attn_out_latent.data_ptr() == output.data_ptr(), ( + "Attention backend did not write into the provided output buffer." + ) + return output + + +# End-to-end DeepSeek-V4 forward scheduling hook. + + +def forward_sparse_attn( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + attn_output: list[torch.Tensor], +) -> None: + """Run DeepSeek-V4 MLA and write into the algorithm-defined output buffers.""" + assert self.mha is None and self.mqa is not None, "DeepSeek-V4 is only supported in MQA mode" + output = attn_output[0] + sparse_epilogue_output = (attn_output[1], attn_output[2]) if len(attn_output) > 1 else None + num_contexts = attn_metadata.num_contexts + num_generations = attn_metadata.num_generations + num_ctx_tokens = attn_metadata.num_ctx_tokens + num_tokens = attn_metadata.num_tokens + if sparse_epilogue_output is not None and ((num_contexts > 0) == (num_generations > 0)): + raise RuntimeError("DSv4 epilogue fusion requires a context-only or generation-only batch.") + + hidden_states = hidden_states[:num_tokens, ...] + if position_ids is not None: + position_ids = position_ids[..., :num_tokens] + + # TRTLLM_MLA_EXTRA_OVERLAP overlaps the compressor and ratio-4 indexer with + # Q projection on dedicated streams. + _v4_extra_overlap = ( + os.environ.get("TRTLLM_MLA_EXTRA_OVERLAP", "1") == "1" + and self.compressor is not None + and self.aux_stream is not None + ) + _use_indexer_overlap = ( + _v4_extra_overlap + and do_multi_stream() + and self.indexer is not None + and self.indexer_stream is not None + ) + + # The compressor depends only on hidden states and metadata, so start it before + # KV projection. Q work later shares this stream and its completion event. + if _use_indexer_overlap: + self.dsv4_compressor_start_event.record() + with torch.cuda.stream(self.compressor_stream): + self.dsv4_compressor_start_event.wait() + self.compressor(hidden_states, attn_metadata) + + # Precompute QR-independent indexer work while the caller stream prepares KV. + # Passing pre_aux later prevents the indexer from launching this work again. + _indexer_pre_aux = None + if _use_indexer_overlap: + _indexer_pre_aux = self.indexer.precompute_aux(hidden_states, attn_metadata) + + q, kv = self.kv_a_proj_with_mqa(hidden_states).split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], -1 + ) + + q, kv = maybe_execute_in_parallel( + lambda: self.q_a_layernorm(q), + lambda: self.kv_a_layernorm(kv), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + compressed_kv, k_pe = kv.split([self.kv_lora_rank, self.qk_rope_head_dim], -1) + qr = q + latent_cache = torch.concat([compressed_kv, k_pe], dim=-1) + + # Use the CuTe BF16 Q projection only for ratio-4 CSA layers with unquantized, + # bias-free weights. TRTLLM_MLA_Q_B_PROJ_USE_CUTE_DSL disables this path. + _use_q_b_cute = ( + self.has_dsv4_indexer + and os.environ.get("TRTLLM_MLA_Q_B_PROJ_USE_CUTE_DSL", "1") == "1" + and self.q_b_proj.bias is None + and self.q_b_proj.weight.dtype == torch.bfloat16 + ) + + def _q_b_proj_cute_dsl_bf16(q: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + global _q_b_proj_cute_dsl_import_ok + if _q_b_proj_cute_dsl_import_ok is None: + try: + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + _q_b_proj_cute_dsl_import_ok = IS_CUTLASS_DSL_AVAILABLE + except ImportError: + _q_b_proj_cute_dsl_import_ok = False + if not _q_b_proj_cute_dsl_import_ok or not is_sm_100f(): + return torch.nn.functional.linear(q, weight) + + assert q.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( + "q_b_proj cute_dsl path requires bfloat16 inputs" + ) + q = q.contiguous() + weight = weight.contiguous() + m, n = q.shape[0], weight.shape[0] + out = q.new_empty((m, n), dtype=torch.bfloat16) + torch.ops.trtllm.cute_dsl_bf16_gemm_blackwell(q, weight, out) + return out + + def _is_fused_q_fp8_quant_enabled() -> bool: + # Mixed/generation batches must use q itself in the generation path. + if os.environ.get("TRTLLM_DISABLE_FUSED_Q_FP8_QUANT", "0") == "1": + return False + if self.qk_head_dim != 512 or self.kv_lora_rank != 448: + return False + if num_generations > 0: + return False + return bool(getattr(self.mqa, "has_fp8_kv_cache", False)) + + def _q_b_layernorm(q: torch.Tensor) -> torch.Tensor: + assert q.dim() == 2 and q.shape[1] == self.num_heads_tp * self.qk_head_dim + return torch.ops.trtllm.deepseek_v4_q_norm( + q, + self.num_heads_tp, + self.qk_head_dim, + float(self.q_b_layernorm.variance_epsilon), + ) + + def _q_b_layernorm_fused_fp8( + q_proj: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + assert q_proj.dim() == 2 + assert q_proj.shape[1] == self.num_heads_tp * self.qk_head_dim + if getattr(self, "_quant_scale_qkv", None) is None: + self._quant_scale_qkv = torch.tensor([1.0], dtype=torch.float32, device=q_proj.device) + + # The 3D q_pe layout is required by the sparse-MLA context op. + num_q_tokens = q_proj.shape[0] + rope_dim = self.qk_head_dim - self.kv_lora_rank + quant_q_buffer = q_proj.new_empty( + (num_q_tokens, self.num_heads_tp * self.qk_head_dim), + dtype=torch.float8_e4m3fn, + ) + q_pe = q_proj.new_empty((num_q_tokens, self.num_heads_tp, rope_dim)) + torch.ops.trtllm.deepseek_v4_q_norm_fused_fp8( + q_proj, + quant_q_buffer, + q_pe.view(num_q_tokens, self.num_heads_tp * rope_dim), + self.num_heads_tp, + self.qk_head_dim, + self.kv_lora_rank, + float(self.q_b_layernorm.variance_epsilon), + self._quant_scale_qkv, + ) + assert self._quant_scale_qkv is not None + return q_proj, quant_q_buffer, q_pe, self._quant_scale_qkv + + def _q_branch(): + # CuTe BF16 projection is incompatible with fused FP8 Q quantization. + if _use_q_b_cute: + assert not _is_fused_q_fp8_quant_enabled(), ( + "CuTe DSL q_b_proj path is incompatible with the fused FP8 q-quant branch" + ) + q_proj = _q_b_proj_cute_dsl_bf16(q, self.q_b_proj.weight) + # The context path detects fusion from these buffers, so clear stale state. + self._fused_quant_q_buffer = None + self._fused_q_pe = None + return _q_b_layernorm(q_proj) + q_proj = self.q_b_proj(q) + if _is_fused_q_fp8_quant_enabled(): + placeholder_q, quant_q_buffer, q_pe, quant_scale_qkv = _q_b_layernorm_fused_fp8(q_proj) + self._fused_quant_q_buffer = quant_q_buffer + self._fused_q_pe = q_pe + self._quant_scale_qkv = quant_scale_qkv + return placeholder_q + self._fused_quant_q_buffer = None + self._fused_q_pe = None + return _q_b_layernorm(q_proj) + + def _compressor_branch(): + self.compressor(hidden_states, attn_metadata) + return None + + def _indexer_branch(): + return self.indexer( + qr, + hidden_states, + attn_metadata, + position_ids, + pre_aux=_indexer_pre_aux, + ) + + topk_indices = None + indexer_ran = False + if _v4_extra_overlap: + if _use_indexer_overlap: + # Compressor and indexer prework are already in flight. Run Q after the + # compressor on the same stream, then synchronize both streams below. + self.dsv4_overlap_start_event.record() + + with torch.cuda.stream(self.indexer_stream): + self.dsv4_overlap_start_event.wait() + topk_indices = _indexer_branch() + indexer_ran = True + self.dsv4_indexer_event.record() + + with torch.cuda.stream(self.compressor_stream): + self.dsv4_overlap_start_event.wait() + q = _q_branch() + self.dsv4_compressor_event.record() + + self.dsv4_compressor_event.wait() + self.dsv4_indexer_event.wait() + + # Keep cross-stream outputs alive on the consuming stream. + cur_stream = torch.cuda.current_stream() + if q is not None: + q.record_stream(cur_stream) + if topk_indices is not None: + topk_indices.record_stream(cur_stream) + else: + q, _ = maybe_execute_in_parallel( + _q_branch, + _compressor_branch, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + else: + q = _q_branch() + if self.compressor is not None: + self.compressor(hidden_states, attn_metadata) + + if self.indexer is not None: + if not indexer_ran: + topk_indices = _indexer_branch() + + assert q.shape[0] == num_tokens, f"Expect q.shape[0] to be {num_tokens}, but got {q.shape[0]}" + + assert output is not None, "output must be provided" + + if num_contexts > 0: + q_ctx = q[:num_ctx_tokens, ...] + topk_indices_ctx = topk_indices[:num_ctx_tokens, :] if topk_indices is not None else None + compressed_kv_ctx = compressed_kv[:num_ctx_tokens, ...] + k_pe_ctx = k_pe[:num_ctx_tokens, ...] + latent_cache_ctx = latent_cache[:num_ctx_tokens, ...] + ctx_position_ids = position_ids[..., :num_ctx_tokens] if position_ids is not None else None + if self.apply_rotary_emb: + assert ctx_position_ids is not None + k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids) + + forward_context_sparse_attn( + self, + q_ctx, + compressed_kv_ctx, + k_pe_ctx, + attn_metadata, + output[:num_ctx_tokens, :], + position_ids=ctx_position_ids, + latent_cache=latent_cache_ctx, + topk_indices=topk_indices_ctx, + sparse_epilogue_output=sparse_epilogue_output, + ) + + if num_generations > 0: + q_gen = q[num_ctx_tokens:, ...] + topk_indices_gen = ( + topk_indices[num_ctx_tokens:num_tokens, :] if topk_indices is not None else None + ) + compressed_kv_gen = compressed_kv[num_ctx_tokens:, ...] + k_pe_gen = k_pe[num_ctx_tokens:, ...] + latent_cache_gen = latent_cache[num_ctx_tokens:, ...] + gen_position_ids = ( + position_ids[..., num_ctx_tokens:num_tokens] if position_ids is not None else None + ) + if self.apply_rotary_emb: + assert gen_position_ids is not None + k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids) + + forward_generation_sparse_attn( + self, + q_gen, + compressed_kv_gen, + k_pe_gen, + attn_metadata, + output[num_ctx_tokens:num_tokens, :], + position_ids=gen_position_ids, + latent_cache=latent_cache_gen, + topk_indices=topk_indices_gen, + sparse_epilogue_output=sparse_epilogue_output, + ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.py new file mode 100644 index 000000000000..7a1b7e25138e --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DeepSeek-V4 parameter and cache-role types.""" + +from dataclasses import dataclass, field +from enum import Enum +from typing import List, Literal + +from tensorrt_llm.runtime.kv_cache_manager_v2 import DataRole + +from ..dsa.params import DSAMetadataParams, DSAParams + +DEEPSEEK_V4_SPARSE_RATIO = 4 +DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO = 4 + + +class DeepseekV4AttentionType(Enum): + """DeepSeek-V4 cache roles.""" + + # Sliding-window roles must remain contiguous. + SWA = 0 + COMPRESSOR_KV = 1 + COMPRESSOR_SCORE = 2 + INDEXER_COMPRESSOR_KV = 3 + INDEXER_COMPRESSOR_SCORE = 4 + + # Non-sliding cache roles. + COMPRESS = 5 + INDEXER_COMPRESS = 6 + + @property + def role(self) -> DataRole: + return DataRole(f"deepseek_v4_{self.name.lower()}") + + +DEEPSEEK_V4_SLIDING_ATTENTION = ( + DeepseekV4AttentionType.SWA, + DeepseekV4AttentionType.COMPRESSOR_KV, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_KV, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, +) +assert tuple(attn_type.value for attn_type in DEEPSEEK_V4_SLIDING_ATTENTION) == tuple( + range(len(DEEPSEEK_V4_SLIDING_ATTENTION)) +) + +DEEPSEEK_V4_NON_SLIDING_ATTENTION = ( + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, +) + + +def is_overlap_compressor(compress_ratio: int) -> bool: + return compress_ratio == DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO + + +def is_sparse_layer(compress_ratio: int) -> bool: + return compress_ratio == DEEPSEEK_V4_SPARSE_RATIO + + +def is_compress_layer(compress_ratio: int) -> bool: + return compress_ratio > 1 + + +def compress_ratio_has_attention(compress_ratio: int, attn_type: DeepseekV4AttentionType) -> bool: + is_sparse = is_sparse_layer(compress_ratio) + is_compress = is_compress_layer(compress_ratio) + + if attn_type == DeepseekV4AttentionType.SWA: + return True + if attn_type == DeepseekV4AttentionType.COMPRESS: + return is_compress + if attn_type == DeepseekV4AttentionType.COMPRESSOR_KV: + return is_compress + if attn_type == DeepseekV4AttentionType.COMPRESSOR_SCORE: + return is_compress + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + return is_sparse + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_KV: + return is_sparse + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE: + return is_sparse + raise ValueError(f"Unsupported DeepSeek-V4 attention type: {attn_type}") + + +@dataclass(frozen=True) +class DeepSeekV4Params(DSAParams): + """DeepSeek-V4 backend parameters.""" + + algorithm: Literal["deepseek_v4"] = field(init=False, default="deepseek_v4") + compress_ratios: List[int] = field(default_factory=list) + window_size: int = 128 + + +@dataclass(frozen=True) +class DeepSeekV4MetadataParams(DSAMetadataParams): + """DeepSeek-V4 metadata parameters.""" + + compress_ratios: List[int] = field(default_factory=list) + window_size: int = 128 diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py deleted file mode 100644 index 8466bbf59900..000000000000 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ /dev/null @@ -1,3356 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Dense Sparse Attention (DSA) backend for TRT-LLM with indexer-based TopK selection.""" -import math -import os -import threading -from contextlib import contextmanager -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, List, Literal, Optional, Set, Tuple, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F - -import tensorrt_llm -import tensorrt_llm.bindings -from tensorrt_llm._torch.attention_backend.interface import ( - AttentionForwardArgs, AttentionInputType, MLAParams, - PositionalEmbeddingParams) -from tensorrt_llm._torch.attention_backend.trtllm import ( - TrtllmAttention, TrtllmAttentionMetadata) -from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from tensorrt_llm._torch.distributed.ops import allgather -from tensorrt_llm._torch.modules.layer_norm import LayerNorm -from tensorrt_llm._torch.modules.linear import Linear -from tensorrt_llm._torch.modules.multi_stream_utils import \ - maybe_execute_in_parallel -from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding -from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager -from tensorrt_llm._torch.utils import Fp4QuantizedTensor, maybe_compile -from tensorrt_llm._utils import (get_size_in_bytes, get_sm_version, - maybe_pin_memory, prefer_pinned) -from tensorrt_llm.bindings import DataType -from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.bindings.internal.batch_manager import \ - CacheType as CacheTypeCpp -from tensorrt_llm.deep_gemm import (fp8_fp4_mqa_logits, - fp8_fp4_paged_mqa_logits, fp8_mqa_logits, - fp8_paged_mqa_logits, - get_paged_mqa_logits_metadata) -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantConfig - -from .params import SparseMetadataParams, SparseParams - -ModelConfig = tensorrt_llm.bindings.ModelConfig - -# Cap the per-call indexer MQA-logits transient (in elements). fp8_mqa_logits -# allocates its [q x kv] logits output via torch.empty; the KV dimension is the -# full (compressed) context and is unbounded, so for a large query chunk on a -# long-context prefill this single allocation can reach tens of GB. Under -# PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True such an allocation can stall -# indefinitely in cuMemCreate on the longest-context (attention_dp laggard) rank -# -> GPU idle -> peers block at the next MoE all-to-all -> watchdog hang. Tiling -# the query dimension caps the transient to q_tile x kv with identical results -# (each query row's logits/top-k are independent). Override via env if needed. -_INDEXER_MQA_LOGITS_ELEM_BUDGET = int( - os.environ.get("TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET", 1 << 31)) - -if TYPE_CHECKING: - from tensorrt_llm._torch.speculative.interface import SpecMetadata - from tensorrt_llm._torch.speculative.spec_tree_manager import \ - SpecTreeManager - from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, - SparseAttentionConfig) - -# Optional import: fast-hadamard-transform causes CI build issues (requires wheel+torch pre-installed) -try: - from fast_hadamard_transform import hadamard_transform - HAS_FAST_HADAMARD = True -except ImportError: - hadamard_transform = None - HAS_FAST_HADAMARD = False - -# Idempotency guard for warmup_heuristic_topk_decode — keyed by -# (device_index, top_k, hint_size, num_cols). Prevents repeated allocations -# and synchronizations when multiple Indexer modules invoke the warmup with -# the same parameters during model construction. -_HEURISTIC_TOPK_WARMUP_DONE: Set[Tuple[int, int, int, int]] = set() -_HEURISTIC_TOPK_WARMUP_LOCK = threading.Lock() - - -@dataclass(frozen=True) -class DSAMetadataParams(SparseMetadataParams): - """DSA metadata settings derived from user config.""" - - indexer_max_chunk_size: int - max_sparse_topk: Optional[int] - index_head_dim: int - enable_indexer_skip: bool - enable_heuristic_topk: bool - use_cute_dsl_topk: bool - use_cute_dsl_paged_mqa_logits: bool - q_split_threshold: int - - -@dataclass(frozen=True) -class DSAParams(SparseParams): - """DSA sparse attention backend parameters.""" - - algorithm: Literal["dsa"] = field(init=False, default="dsa") - index_n_heads: Optional[int] = None - index_head_dim: Optional[int] = None - index_topk: Optional[int] = None - indexer_max_chunk_size: Optional[int] = None - skip_indexer_for_short_seqs: bool = True - use_cute_dsl_topk: bool = False - use_cute_dsl_paged_mqa_logits: bool = False - q_split_threshold: int = 8192 - indexer_rope_interleave: bool = False - enable_heuristic_topk: bool = False - indexer_k_dtype: Literal["fp8", "fp4"] = "fp8" - # Cross-layer indexer sharing: whether this layer runs its own indexer - # ("full") or reuses the previous full layer's top-k ("shared"). Always - # True for a dense per-layer indexer (e.g. DeepSeek-V3.2). - is_full_indexer_layer: bool = True - - @property - def indices_block_size(self) -> int: - return 1 - - -def warmup_heuristic_topk_decode(top_k: int = 2048, - hint_size: int = 2048, - num_cols: int = 4096) -> None: - """Pre-initialize cached hardware attributes in the C++ Scheme X dispatcher. - - The dispatcher inside ``invokeIndexerTopKDecode`` lazily queries - ``cudaDeviceGetAttribute`` for ``MultiProcessorCount`` and - ``L2CacheSize`` on its first call. Those host-side queries must not - be issued during ``cudaStreamBeginCapture / EndCapture``: the values - captured there become frozen into the graph and cannot be refreshed - across replays on a different device. - - This warmup issues one small heuristic decode call so the static - caches are populated before any CUDA Graph capture begins. Must be - called from the Indexer setup hook (``layer_idx == 0``) when - ``enable_heuristic_topk`` is true. - - Repeated invocations with the same ``(device, top_k, hint_size, - num_cols)`` key are short-circuited so that constructing many Indexer - modules in the same process does not re-allocate scratch tensors or - issue redundant synchronizations. - """ - key = (torch.cuda.current_device(), top_k, hint_size, num_cols) - with _HEURISTIC_TOPK_WARMUP_LOCK: - if key in _HEURISTIC_TOPK_WARMUP_DONE: - return - _HEURISTIC_TOPK_WARMUP_DONE.add(key) - - device = torch.device("cuda") - logits = torch.zeros((1, num_cols), dtype=torch.float32, device=device) - seq_lens = torch.tensor([num_cols], dtype=torch.int32, device=device) - indices = torch.empty((1, top_k), dtype=torch.int32, device=device) - pre_idx = torch.zeros((1, hint_size), dtype=torch.int32, device=device) - scratch = torch.empty((top_k, ), dtype=torch.float32, device=device) - # The default warmup geometry (num_cols=4096) falls below kSeqSmall=12288 - # and routes to the Radix path with blocks_per_row=2 (num_rows=1 sweeps - # bp ∈ [2, maxByCols=2]). The cpp op rejects blocks_per_row > 1 without - # caller-owned radix aux scratch, so supply worst-case (kMaxBlocksPerRowDecode=10) - # buffers here. Cost is negligible (~80 KB) and the warmup is a one-shot. - _radix_max_bp = 10 - radix_aux_indices = torch.empty((1, _radix_max_bp, top_k), - dtype=torch.int32, - device=device) - radix_aux_logits = torch.empty((1, _radix_max_bp, top_k), - dtype=torch.float32, - device=device) - torch.ops.trtllm.indexer_topk_decode(logits, - seq_lens, - indices, - 1, - top_k, - pre_idx=pre_idx, - heuristic_scratch=scratch, - radix_aux_indices=radix_aux_indices, - radix_aux_logits=radix_aux_logits) - torch.cuda.synchronize() - - -# `block_kv` arg passed to DeepGEMM's `get_paged_mqa_logits_metadata`. This is -# a SCHEDULE-granularity parameter, not the cache page size. The DG metadata -# kernel computes `SPLIT_KV = block_kv * 4` (the multiplier 4 is hardcoded; -# DeepGEMM commit 7f2a703 dropped the SM100-aware `arch == 10 ? 2 : 4` from -# nv_dev #fc97232 during the Public release 26/04 sync, leaving the formula -# uniform across SM90/SM100). Both the DG compute kernel (which hardcodes -# `split_kv = 256` at csrc/apis/attention.hpp:353) and our DSL paged-MQA-logits -# kernel (compute tile = 128 × kNumMathWarpGroups = 2 = 256) require -# `SPLIT_KV = 256`. So we must pass `block_kv = 256 / 4 = 64` here. Independent -# of the indexer K cache's physical page size (`tokens_per_block`), which only -# affects cache reads inside the compute kernel — the metadata wrapper does -# not read the cache. Passing the cache page size directly (the previous -# behavior) only works by accident when `tokens_per_block == 64`; for -# `tokens_per_block == 32` it produces a SPLIT_KV=128 schedule that the -# compute kernels misinterpret. TODO(remove once DeepGEMM restores the -# SM100-aware num_math_warpgroups in the metadata JIT impl). -_DG_SCHEDULE_BLOCK_KV = 64 - - -def _pick_dsl_expand( - next_n: int, - num_sms: int, - batch_size: int = 0, - max_ctx: int = 0, - kernel_atoms: Tuple[int, ...] = (1, 2, 3)) -> Tuple[int, int]: - """Pick (expand_factor, effective_next_n) for the DSL paged kernel - using a wave-aware strategy. Used by both FP4 and FP8 DSL paths. - - The DSL kernel natively supports ``effective_next_n ∈ kernel_atoms`` - (FP4: ``(1, 2, 3)``; FP8: ``(1, 2, 3, 4)``). For ``next_n`` not natively - supported or when SM utilization can be improved, reshape - ``[B, next_n, ...]`` -> ``[B * expand_factor, effective_next_n, ...]`` - caller-side. - - Strategy: enumerate ``(expand_factor, effective_next_n)`` pairs with - ``expand_factor * effective_next_n == next_n`` and ``effective_next_n - in kernel_atoms``. Score each by ``(waves, -expand_factor)`` where - ``waves = ceil(B * expand_factor * ceil(max_ctx/256) / num_sms)``. - Pick min waves; on tie, prefer LARGER expand_factor (more SMs busy per - wave; pays HBM cost of expand_factor x KV re-reads). - - When ``batch_size == 0`` or ``max_ctx == 0`` (workload unknown), fall - back to the legacy HBM-minimizing heuristic: largest effective_next_n - that divides next_n cleanly (still constrained to ``kernel_atoms``). - - Examples (wave-aware, num_sms=148 [B200], SPLIT_KV=256 tokens): - FP4, next_n=4, B=1, ctx=4096 -> (4, 1): ntask=64<148, 1 wave, max factor - FP4, next_n=4, B=32, ctx=4096 -> (2, 2): ntask=1024>148, multi-wave, min factor - FP4, next_n=2, B=1, ctx=4096 -> (2, 1): wave-tie, larger factor - FP8, next_n=4, B=1, ctx=4096 -> (4, 1): kernel_atoms incl. 4 doesn't change small-B pick - """ - # Legacy fallback when workload is unknown. - if batch_size <= 0 or max_ctx <= 0: - for eff in sorted(kernel_atoms, reverse=True): - if next_n % eff == 0: - return next_n // eff, eff - return next_n, 1 - - SPLIT_KV_TOKENS = 256 - cands = [] - for eff in kernel_atoms: - if next_n % eff == 0: - factor = next_n // eff - ntask = batch_size * factor * ( - (max_ctx + SPLIT_KV_TOKENS - 1) // SPLIT_KV_TOKENS) - waves = (ntask + num_sms - 1) // num_sms - cands.append((waves, factor, eff)) - if not cands: - return next_n, 1 - cands.sort(key=lambda x: (x[0], -x[1])) # min waves, max factor - _, factor, eff = cands[0] - return factor, eff - - -def _compute_slot_mappings( - global_positions: torch.Tensor, - block_offsets: torch.Tensor, - req_indices: torch.Tensor, - head_dim: int, - tokens_per_block: int, - quant_block_size: int, - data_bytes_per_token: Optional[int] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Compute flat byte indices for FP8/FP4 data and scales from global token positions. - - Shared by Indexer.prepare() (CPU) and on_update_kv_lens() (GPU) to avoid - duplicating the slot mapping arithmetic. - - Args: - global_positions: Per-token absolute position in the KV sequence. - block_offsets: [num_seqs, max_blocks_per_seq] block offset table. - req_indices: Per-token request index. - head_dim: Indexer head dimension (logical element count). - tokens_per_block: Tokens stored per cache block. - quant_block_size: Quantization block size. - data_bytes_per_token: Bytes of quantized data per token in the cache - (default: head_dim, i.e. one byte per FP8 value). For MXFP4 this - should be ``head_dim // 2`` since two E2M1 codes pack into one byte. - The scale layout is identical at head_dim=128: 4 bytes per token - (one float32 for FP8, four packed UE8M0 exponents for FP4). - - Returns: - (fp8_indices, scale_indices): Flat byte offsets into the cache pool. - """ - if data_bytes_per_token is None: - data_bytes_per_token = head_dim - scale_size = head_dim // quant_block_size * 4 # float32 = 4 bytes - block_stride = tokens_per_block * (data_bytes_per_token + scale_size) - scale_base_offset = tokens_per_block * data_bytes_per_token - - block_indices_in_seq = global_positions // tokens_per_block - pos_in_blocks = global_positions % tokens_per_block - - max_blocks = block_offsets.shape[1] - if block_indices_in_seq.is_cuda: - # Clamp to prevent OOB from stale token-to-seq mappings during - # CUDA graph capture/replay with MTP + DSA. - block_indices_in_seq = block_indices_in_seq.clamp(0, max_blocks - 1) - else: - 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()}" - - block_ids = block_offsets[req_indices, block_indices_in_seq].to(torch.int64) - - fp8_indices = block_ids * block_stride + pos_in_blocks * data_bytes_per_token - scale_indices = (block_ids * block_stride + scale_base_offset + - pos_in_blocks * scale_size) - return fp8_indices, scale_indices - - -def _unravel_indices(flat_indices: torch.Tensor, - shape: Tuple[int, ...]) -> Tuple[torch.Tensor, ...]: - """ - Unravel indices into multiple dimensions. - """ - d3 = shape[3] - i3 = flat_indices % d3 - flat_indices = flat_indices // d3 - d2 = shape[2] - i2 = flat_indices % d2 - flat_indices = flat_indices // d2 - d1 = shape[1] - i1 = flat_indices % d1 - flat_indices = flat_indices // d1 - i0 = flat_indices - return i0, i1, i2, i3 - - -def rotate_activation(x: torch.Tensor) -> torch.Tensor: - """Apply Hadamard rotation to activation tensor for DSA sparse attention.""" - assert x.dtype == torch.bfloat16 - - if not HAS_FAST_HADAMARD: - # Fallback: skip transformation (acceptable for test/dev) - logger.warning_once( - "fast-hadamard-transform not available. Sparse MLA will skip " - "hadamard transformation. Install with: " - "pip install git+https://github.com/Dao-AILab/fast-hadamard-transform.git", - key="fast_hadamard_import_missing") - return x - - hidden_size = x.size(-1) - assert (hidden_size & (hidden_size - 1) - ) == 0, "Hidden size must be a power of 2 for Hadamard transform." - return hadamard_transform(x, scale=hidden_size**-0.5) - - -def transform_local_topk_and_prepare_pool_view( - topk_indices: torch.Tensor, - attn_metadata: "DSAtrtllmAttentionMetadata", - layer_idx: int, - is_generation: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Convert local topk indices to global pool indices and prepare KV pool. - - Uses cached values from attn_metadata._ensure_pool_view_cached() - to avoid redundant Python/CUDA overhead across layers. - """ - assert topk_indices.dtype == torch.int32 - - attn_metadata._ensure_pool_view_cached() - - if is_generation: - block_table = attn_metadata._cached_block_table_gen - req_idx = attn_metadata._cached_req_idx_gen - else: - block_table = attn_metadata._cached_block_table_ctx - req_idx = attn_metadata._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, - topk_indices.shape[1], - attn_metadata._cached_stride_factor, - layer_idx, - ) - - return global_indices, attn_metadata._cached_pool_view - - -def split_prefill_chunks( - seq_lens: torch.Tensor, - max_chunk_size: int, - start_idx: int = 0, -) -> List[List[Tuple[int, int, int, int]]]: - """ - Split prefill requests into chunks based on max_chunk_size. - Supports two-level chunking: - 1. Request-boundary chunking: group multiple small requests into one chunk - 2. Intra-request chunking: split large requests into multiple Q-block chunks - - Args: - seq_lens: Sequence lengths for all requests - max_chunk_size: Maximum number of tokens per chunk - start_idx: Starting index for prefill requests - - Returns: - List of chunk groups, where each group is a list of chunk specs. - Each chunk spec is (req_idx, token_start_in_req, token_end_in_req, req_cum_start) - - - For multi-request chunks: group contains multiple specs (one per request) - - For intra-request chunks: each Q-block is a separate group with single spec - """ - chunk_groups = [] - num_reqs = len(seq_lens) - - current_req = start_idx - # Compute cumulative token positions - query_start_loc_cpu = torch.cat([ - torch.zeros(1, dtype=torch.int32, device='cpu'), - seq_lens.cumsum(dim=0).to(torch.int32) - ]) - - while current_req < num_reqs: - seq_len = seq_lens[current_req].item() - req_cum_start = query_start_loc_cpu[current_req].item() - - if seq_len <= max_chunk_size: - # This request fits in one chunk - try to pack with others - current_size = seq_len - chunk_specs = [(current_req, 0, seq_len, req_cum_start)] - next_req = current_req + 1 - - # Try to add more requests to this chunk - while next_req < num_reqs: - next_seq_len = seq_lens[next_req].item() - if next_seq_len > max_chunk_size: - # Next request is large, stop packing - break - if current_size + next_seq_len <= max_chunk_size: - next_cum_start = query_start_loc_cpu[next_req].item() - chunk_specs.append( - (next_req, 0, next_seq_len, next_cum_start)) - current_size += next_seq_len - next_req += 1 - else: - break - - # Add as one multi-request chunk group - chunk_groups.append(chunk_specs) - current_req = next_req - else: - # Large request - split into Q-blocks - # Each Q-block is a separate chunk group (processed in separate iteration) - num_q_blocks = (seq_len + max_chunk_size - 1) // max_chunk_size - for q_block_idx in range(num_q_blocks): - token_start = q_block_idx * max_chunk_size - token_end = min(token_start + max_chunk_size, seq_len) - q_block_spec = [(current_req, token_start, token_end, - req_cum_start)] - chunk_groups.append(q_block_spec) - - current_req += 1 - - return chunk_groups - - -# Shrink the indexer prefill chunk size for very long requests to bound the -# fp8(_fp4)_mqa_logits activation memory (~ chunk_size * K_compressed), keyed on -# the largest compressed KV length in the batch. Entries are -# (k_compressed_lower_bound_exclusive, chunk_size); >512K -> 8K, [256K, 512K] -# -> 16K, otherwise unchanged. -_INDEXER_CHUNK_SIZE_HEURISTIC = ( - (512 * 1024, 8 * 1024), - (256 * 1024 - 1, 16 * 1024), -) - - -def select_indexer_chunk_size(configured_chunk_size: int, - max_k_compressed: int) -> int: - """Pick the indexer prefill chunk size from the batch's largest K_compressed. - - Only reduces ``configured_chunk_size`` (never increases it). - """ - for threshold, chunk_size in _INDEXER_CHUNK_SIZE_HEURISTIC: - if max_k_compressed > threshold: - return min(configured_chunk_size, chunk_size) - return configured_chunk_size - - -def _select_indexer_compress_ratio(compress_ratios: List[int]) -> int: - if 4 in compress_ratios: - return 4 - if 1 in compress_ratios: - return 1 - return 0 - - -def _effective_compress_ratio_divisor(compress_ratio: int) -> int: - return compress_ratio if compress_ratio > 1 else 1 - - -def compute_cu_seqlen_kv_bounds_with_cache( - seq_lens: torch.Tensor, - num_contexts: int, - num_ctx_tokens: int, - cached_token_lens: Optional[torch.Tensor] = None, - kv_lens: Optional[torch.Tensor] = None, - compress_ratio: int = 1, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Compute attention window bounds for batched sequences with causal attention, - accounting for cached KV tokens. - - Args: - seq_lens: current token lengths [num_contexts], dtype=torch.int32 - num_contexts: Number of sequences in the batch - num_ctx_tokens: Total number of context tokens across all sequences in current batch - cached_token_lens: Cached KV token lengths [num_contexts], dtype=torch.int32 (optional) - kv_lens: KV token lengths [num_contexts], dtype=torch.int32 (optional) - compress_ratio: Compression ratio for KV tokens - - Returns: - cu_seqlen_ks: Start index in KV for each Q token [num_ctx_tokens] - cu_seqlen_ke: End index (exclusive) in KV for each Q token [num_ctx_tokens] - """ - device = seq_lens.device - # Total KV lengths per request - if kv_lens is None: - kv_lens = seq_lens if cached_token_lens is None else cached_token_lens + seq_lens # [num_contexts] - kv_lens = kv_lens // compress_ratio - - # Cumulative KV offsets: where each request's KV sequence starts in global KV space - cu_kv_offsets = torch.cat([ - torch.zeros(1, device=device, dtype=torch.int32), - torch.cumsum(kv_lens, dim=0).to(torch.int32) - ]) # [num_contexts + 1] - - # Map each Q token to its request: [0,0,...,0, 1,1,...,1, ..., B-1,B-1,...,B-1] - batch_ids = torch.repeat_interleave( - torch.arange(num_contexts, device=device, dtype=torch.int32), - seq_lens) # [num_ctx_tokens] - - # Each Q token's KV window starts at its request's KV sequence start - cu_seqlen_ks = cu_kv_offsets[batch_ids] # [num_ctx_tokens] - - # Compute local Q position within each request (0-based, relative to current batch context tokens) - cu_q_offsets = torch.cat([ - torch.zeros(1, device=device, dtype=torch.int32), - torch.cumsum(seq_lens, dim=0).to(torch.int32) - ]) # [num_contexts + 1] - - global_q_positions = torch.arange(num_ctx_tokens, - device=device, - dtype=torch.int32) - local_q_positions = global_q_positions - torch.repeat_interleave( - cu_q_offsets[:-1], seq_lens) # [num_ctx_tokens] - - if cached_token_lens is not None: - cached_per_token = torch.repeat_interleave(cached_token_lens, - seq_lens) # [num_ctx_tokens] - cu_seqlen_ke = cu_seqlen_ks + (cached_per_token + local_q_positions + - 1) // compress_ratio # [num_ctx_tokens] - else: - cu_seqlen_ke = cu_seqlen_ks + (local_q_positions + - 1) // compress_ratio # [num_ctx_tokens] - - return cu_seqlen_ks, cu_seqlen_ke - - -@dataclass -class IndexerPrefillChunkMetadata: - """Metadata for a single prefill chunk in the indexer""" - cu_seqlen_ks: torch.Tensor # Attention window start for each token - cu_seqlen_ke: torch.Tensor # Attention window end for each token - token_start: int # Q token start index in batch - token_end: int # Q token end index in batch - k_token_start: int # K token start index in batch - k_token_end: int # K token end index in batch - - -@dataclass(init=False) -class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): - """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" - - sparse_metadata_params: Optional[DSAMetadataParams] = None - # 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 - # 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 - # max number of draft tokens - max_draft_tokens: int = 0 - # Indexer head dimension - indexer_head_dim: int = 128 - # Indexer quant block size - indexer_quant_block_size: int = 128 - # Enable indexer skip for short sequences - enable_indexer_skip: bool = False - # Cross-layer indexer sharing: previous full layer's top-k, reused by - # "shared" layers (None for a dense per-layer indexer). - shared_topk_indices: 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 - # 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 - - def __init__(self, *args, **kwargs): - """Initialize DSA metadata with SM count and indexer chunk size.""" - 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_stride_factor = 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 - 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) - - def __post_init__(self): - """Allocate indexer K-cache buffers and heuristic TopK metadata.""" - super().__post_init__() - if not isinstance(self.kv_cache_manager, DSACacheManager): - has_deepseek_v4_cache_interface = all( - hasattr(self.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)}") - - 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.num_sparse_topk = sparse_metadata_params.max_sparse_topk - self.sparse_mla_topk = self.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) - self.use_cute_dsl_topk = (sparse_metadata_params.use_cute_dsl_topk - and IS_CUTLASS_DSL_AVAILABLE) - self.kv_lens_row_reorder = None - capture_graph = self.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]) - - # Effective tokens-per-block for the indexer k-cache slot mapping. - # 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 - self._indexer_compress_ratio = _select_indexer_compress_ratio( - self.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 - - self.create_buffers_for_mla_rope_append(capture_graph=capture_graph) - self.create_buffers_for_indexer(capture_graph=capture_graph) - - def prepare(self): - super().prepare() - self._invalidate_pool_view_cache() - # Cross-layer indexer sharing is per-step state; clear it so a "shared" - # layer can never reuse a previous step's top-k before a full layer runs. - self.shared_topk_indices = None - - # Get kv lengths - assert self.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, - dtype=torch.int, - device='cpu', - ) - if self.enable_helix: - # 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] - kv_lens = cached_token_lens.clone() - kv_lens[active_rank] += self.seq_lens_kv[active_rank] - else: - kv_lens = cached_token_lens + self.seq_lens_kv - - # For mla_rope_append_paged_kv_assign_q - self.prepare_for_mla_rope_append(cached_token_lens, kv_lens) - - # Prepare to support skip indexer - self.prepare_for_skip_indexer(kv_lens) - - # For indices conversion - self.prepare_for_indices_conversion() - - # For indexer k cache - self.prepare_for_indexer_k_cache() - - # For spec decode - self.prepare_for_spec_decode(kv_lens) - - # Prepare metadata for indexer - Indexer.prepare(metadata=self) - - def get_indexer_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor: - if self._indexer_compress_ratio <= 1: - return kv_lens - return kv_lens // self._indexer_compress_ratio - - def get_indexer_max_seq_len(self) -> int: - 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) - - def on_update_kv_lens(self): - # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas. - # Especially for the changes in the _preprocess_inputs() of model_engine.py. - # - # NOTE: - # In overlap scheduler + speculative decoding, kv_lens_cuda can be corrected at runtime - # (inside _preprocess_inputs) to account for variable accepted tokens. The indexer - # 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. - - # _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 - # transform_local_topk_and_prepare_pool_view() call. - self._invalidate_pool_view_cache() - # Per-step state for cross-layer indexer sharing; clear at the step - # boundary so a "shared" layer never reuses a stale top-k. - self.shared_topk_indices = None - - if self.kv_cache_manager is not None and self.num_tokens > 0: - seq_lens = self.seq_lens_cuda[:self.num_seqs] - # Runtime cached lengths after overlap/spec-dec correction. - start_positions = self.kv_lens_cuda[:self.num_seqs] - seq_lens - - # Reuse request-per-token mapping prepared in metadata.prepare(). - # This avoids repeat_interleave in graph-capture mode. - req_indices = self.req_idx_per_token[:self.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) - seq_starts[req_indices] - - global_positions = start_positions[req_indices] + token_offsets - # 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) - data_bytes_per_token = index_head_dim // 2 if use_fp4 else index_head_dim - fp8_indices, scale_indices = _compute_slot_mappings( - global_positions, - self.indexer_k_cache_block_offsets, - req_indices, - index_head_dim, - self._tokens_per_block, - self.kv_cache_manager.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 - - if self.num_generations > 0: - torch.cumsum( - self.kv_lens_cuda[self.num_contexts:self. - num_seqs], # num_contexts should be 0 - dim=0, - dtype=torch.int64, - out=self.gen_kv_indptr[1:self.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]), - dim=0, - dtype=torch.int64, - out=self.gen_cached_token_indptr[1:self.num_generations + 1]) - gen_kv_lens = self.kv_lens_cuda[self.num_contexts:self.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 - next_n_cap = self.kv_lens_cuda_2d.shape[1] - self.kv_lens_cuda_2d[:self.num_generations, :next_n_cap].copy_( - gen_indexer_kv_lens.unsqueeze(-1).expand(-1, next_n_cap)) - scheduler_metadata_buffer = get_paged_mqa_logits_metadata( - gen_indexer_kv_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, - self.num_sms) - self.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, - non_blocking=True) - if (self.max_draft_tokens > 0 - and not self.use_expanded_buffers_for_mtp): - scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( - self.kv_lens_cuda_2d[:self.num_generations, :next_n_cap], - _DG_SCHEDULE_BLOCK_KV, self.num_sms) - self.scheduler_metadata_buffer_full_next_n.copy_( - scheduler_metadata_buffer_full_next_n, non_blocking=True) - if self.use_expanded_buffers_for_mtp: - num_draft_tokens = 1 + self.max_draft_tokens - num_tokens = self.num_generations * num_draft_tokens - kv_lens_expanded = torch.stack([gen_indexer_kv_lens] * - num_draft_tokens, - dim=0) - self.kv_lens_expanded_cuda[:num_tokens] = \ - kv_lens_expanded.transpose(0, 1).contiguous().flatten() - scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( - self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), - _DG_SCHEDULE_BLOCK_KV, self.num_sms) - self.scheduler_metadata_buffer_expanded.copy_( - scheduler_metadata_buffer_expanded, non_blocking=True) - if self.expand_for_dsl and self.dsl_expand_factor > 1: - expand_factor = self.dsl_expand_factor - num_tokens = self.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( - self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), - _DG_SCHEDULE_BLOCK_KV, self.num_sms) - self.scheduler_metadata_buffer_expanded.copy_( - scheduler_metadata_buffer_expanded, non_blocking=True) - self._compute_kv_lens_row_reorder() - self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True) - - def _compute_kv_lens_row_reorder(self): - """LJF (longest-job-first) row-reorder for the GVR DSL top-k path. - - Writes ``argsort(gen_kv_lens, descending)`` into the stable buffer when - the multi-wave threshold is met, otherwise leaves ``order_row`` None. - Called from ``on_update_kv_lens()`` (both base and DeepSeek-V4 via - super()) unconditionally every forward step so the GVR op sees a fresh - valid permutation and never a stale one from a prior step. Copies into - the stable buffer (not a fresh tensor) so the CUDA-Graph-captured op - reads a valid permutation on every replay. - """ - # Gate on row count (num_generations * next_n) rather than request count - # so the threshold aligns with the kernel-side tuning note that records - # the win region starting at num_rows >= 2 * num_sms. Using - # num_generations alone is only correct for next_n == 2; for next_n == 1 - # it engages inside the measured regression band, and for next_n == 4 it - # misses the win region between 2*num_sms and 4*num_sms rows. - next_n = 1 + self.max_draft_tokens - if (self.enable_heuristic_topk and self.use_cute_dsl_topk - and self.num_generations * next_n >= 2 * self.num_sms): - gen_kv_lens = self.kv_lens_cuda[self.num_contexts:self.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] - 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, ), - 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, - device='cpu', - pin_memory=prefer_pinned(), - ) - self.ctx_kv_indptr = self.get_empty( - self.cuda_graph_buffers, - (self.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, - 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, ), - cache_name="gen_cached_token_indptr", - dtype=torch.int64, - capture_graph=capture_graph, - ) - self.host_gen_cached_token_indptr = torch.zeros_like( - self.gen_cached_token_indptr, - device='cpu', - pin_memory=prefer_pinned(), - ) - self.gen_kv_indptr = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_requests + 1, ), - cache_name="gen_kv_indptr", - dtype=torch.int64, - capture_graph=capture_graph, - ) - self.host_gen_kv_indptr = torch.zeros_like( - self.gen_kv_indptr, - device='cpu', - pin_memory=prefer_pinned(), - ) - - def _create_radix_aux_buffers(self, capture_graph=False): - # Persistent scratch for Radix-split-work indexer path (blocks_per_row > 1). - # Mirrors the fix the Heuristic path applied: per-call th::empty inside - # indexer_topk_decode produces stale pointers under CUDA Graph replay when - # the caching allocator is perturbed by chunked prefill at high CONC. - # Sized to the worst case kMaxBlocksPerRowDecode=10 from - # cpp/tensorrt_llm/kernels/indexerTopK.cu, times the max number of - # generation rows (num_seqs * (1 + max_draft_tokens)); the cpp op aborts - # if this is smaller than num_rows*blocks_per_row*index_topk. Allocated - # unconditionally: even with enable_heuristic_topk=True the dispatcher can - # fall back to Radix when canUseHeuristic returns False (small numColumns, - # etc.). MUST be re-created whenever max_draft_tokens changes (see - # update_spec_dec_param) or it is left too small once MTP raises the - # generation-row count. - _radix_max_blocks_per_row = 10 - _radix_max_gen_tokens = self.max_num_sequences * (1 + - self.max_draft_tokens) - self.radix_aux_indices = self.get_empty( - self.cuda_graph_buffers, - (_radix_max_gen_tokens, _radix_max_blocks_per_row, - self.num_sparse_topk), - cache_name="radix_aux_indices", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.radix_aux_logits = self.get_empty( - self.cuda_graph_buffers, - (_radix_max_gen_tokens, _radix_max_blocks_per_row, - self.num_sparse_topk), - cache_name="radix_aux_logits", - dtype=torch.float32, - capture_graph=capture_graph, - ) - - def create_buffers_for_indexer(self, capture_graph=False): - 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_k_cache_block_offsets = self.get_empty( - self.cuda_graph_buffers, - [self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq], - cache_name="indexer_k_cache_block_offsets", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.host_indexer_k_cache_block_offsets = torch.zeros_like( - self.indexer_k_cache_block_offsets, - device='cpu', - pin_memory=prefer_pinned(), - ) - - # 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, ), - cache_name="slot_mapping_fp8", - dtype=torch.int64, - capture_graph=capture_graph, - ) - self.host_slot_mapping_fp8 = torch.zeros_like( - self.slot_mapping_fp8, - device='cpu', - pin_memory=prefer_pinned(), - ) - self.slot_mapping_scale = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens, ), - cache_name="slot_mapping_scale", - dtype=torch.int64, - capture_graph=capture_graph, - ) - self.host_slot_mapping_scale = torch.zeros_like( - self.slot_mapping_scale, - device='cpu', - pin_memory=prefer_pinned(), - ) - # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation. - # These buffers will be allocated dynamically in Indexer.prepare() based on actual total_kv_len to save memory. - if self.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, ), - cache_name="req_idx_per_token", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.host_req_idx_per_token = torch.empty_like( - 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], - cache_name="block_table", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.scheduler_metadata_buffer = self.get_empty( - self.cuda_graph_buffers, - (self.num_sms + 1, 2), - cache_name="scheduler_metadata_buffer", - dtype=torch.int32, - capture_graph=capture_graph, - ) - # When MTP runs without the expanded-tokens path, the same forward step - # alternates between full-window calls (next_n == 1 + max_draft_tokens) - # and per-token draft calls (next_n == 1). The 2D DeepGEMM metadata - # API encodes next_n into the schedule, so the precomputed schedule - # for one shape cannot be reused for the other. Maintain a second - # buffer holding the schedule for the full next_n window; the draft - # 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.num_sms + 1, 2), - cache_name="scheduler_metadata_buffer_full_next_n", - dtype=torch.int32, - capture_graph=capture_graph, - ) - # Pre-allocated 2D kv_lens buffer for the new DeepGEMM 2D context_lens - # API. Shape: (max_num_sequences, 1 + max_draft_tokens). Each row - # broadcasts the same kv_len across next_n positions; kernel reads a - # 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, ), - 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, ), - 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), - cache_name="topk_indices_buffer", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.host_topk_indices_buffer = torch.zeros_like( - self.topk_indices_buffer, - device='cpu', - pin_memory=prefer_pinned(), - ) - # Per-layer persistent buffers for heuristic TopK pre_idx. - # Indexed by [local_layer_idx, generation_position, :]. - # The graph captures reads/writes on these stable-address buffers; - # each replay's write becomes the next replay's read (feedback loop). - self.enable_heuristic_topk = ( - sparse_metadata_params.enable_heuristic_topk - and get_sm_version() >= 100) - if self.enable_heuristic_topk: - num_local_layers = self.kv_cache_manager.num_local_layers - self.heuristic_prev_topk = self.get_empty( - self.cuda_graph_buffers, - (num_local_layers, self.max_num_sequences, - self.num_sparse_topk), - cache_name="heuristic_prev_topk", - dtype=torch.int32, - capture_graph=capture_graph, - ) - # Zero-initialize so the first decode step's pre_idx (kernel - # adds +1 offset) points to index 1 — a valid but benign candidate. - # Without this, uninitialized memory produces random hint indices. - self.heuristic_prev_topk.zero_() - # Scratch buffer for heuristic TopK kernel output values. - # Pre-allocated with stable address for CUDA Graph compatibility - # (replaces cudaMallocAsync/cudaFreeAsync inside the kernel launcher). - # Shape: [max_gen_tokens, topK] where max_gen_tokens = max_batch * (1 + max_draft). - # Only the C++ indexer_topk_decode path consumes it; the GVR DSL - # path does not, so skip the allocation when use_cute_dsl_topk. - if not self.use_cute_dsl_topk: - max_gen_tokens = self.max_num_sequences * ( - 1 + self.max_draft_tokens) - self.heuristic_scratch_values = self.get_empty( - self.cuda_graph_buffers, - (max_gen_tokens, self.num_sparse_topk), - cache_name="heuristic_scratch_values", - dtype=torch.float32, - capture_graph=capture_graph, - ) - # Stable-address buffer for the GVR DSL LJF row-reorder - # (order_row = argsort(gen_kv_lens, descending)). Must not be - # fresh-allocated per step: under CUDA Graph the captured op reads - # a frozen address, so prepare() copies into this buffer instead. - if self.use_cute_dsl_topk: - self.kv_lens_row_reorder_buffer = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences, ), - cache_name="kv_lens_row_reorder_buffer", - dtype=torch.int32, - capture_graph=capture_graph, - ) - - # Persistent scratch for the Radix-split-work indexer path. Re-created - # in update_spec_dec_param when max_draft_tokens changes so it stays - # large enough for the MTP generation-row count. - self._create_radix_aux_buffers(capture_graph=capture_graph) - - # Create expanded buffers for MTP support - self.create_expanded_buffers(capture_graph=capture_graph) - - def _create_kv_lens_2d_buffer(self, capture_graph=False): - """Pre-allocated buffer for the DeepGEMM 2D context_lens API. - - Avoids per-forward .expand().contiguous() allocations that break CUDA - graphs. The buffer is written in-place via .copy_() inside - on_update_kv_lens so its address stays stable across replays. - """ - self.kv_lens_cuda_2d = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences, 1 + self.max_draft_tokens), - cache_name="kv_lens_cuda_2d", - dtype=torch.int32, - capture_graph=capture_graph, - ) - - # 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.max_draft_tokens), ), - cache_name="kv_lens_expanded_cuda", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.kv_lens_expanded_host = torch.zeros_like( - self.kv_lens_expanded_cuda, - device='cpu', - pin_memory=prefer_pinned(), - ) - self.block_table_expanded = self.get_empty( - self.cuda_graph_buffers, - [ - self.max_num_sequences * (1 + self.max_draft_tokens), - self.kv_cache_manager.max_blocks_per_seq - ], - cache_name="block_table_expanded", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.host_block_table_expanded = torch.zeros_like( - self.block_table_expanded, - device='cpu', - pin_memory=prefer_pinned(), - ) - self.scheduler_metadata_buffer_expanded = self.get_empty( - self.cuda_graph_buffers, - (self.num_sms + 1, 2), - cache_name="scheduler_metadata_buffer_expanded", - dtype=torch.int32, - capture_graph=capture_graph, - ) - - # 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( - self, - batch_size, - is_spec_decoding_enabled, - is_spec_dec_tree, - is_spec_dec_dynamic_tree, - max_draft_len, - max_total_draft_tokens, - model_is_wrapped: bool = False, - spec_metadata: Optional['SpecMetadata'] = None, - spec_tree_manager: Optional['SpecTreeManager'] = None, - num_contexts: int = 0, - ): - """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, - model_is_wrapped, - spec_metadata, - spec_tree_manager, - num_contexts=num_contexts) - self.max_draft_tokens = max_draft_len - capture_graph = self.is_cuda_graph - if self.kv_lens_cuda_2d.shape[1] != 1 + self.max_draft_tokens: - self._create_kv_lens_2d_buffer(capture_graph=capture_graph) - init_shape = self.kv_lens_expanded_host.shape[0] - if self.max_num_sequences * (1 + self.max_draft_tokens) != init_shape: - self.create_expanded_buffers(capture_graph=capture_graph) - # Resize heuristic scratch buffer for new max_draft_tokens. - # Skip when use_cute_dsl_topk (GVR path never consumes it), matching - # the allocation guard in create_buffers_for_indexer. - if self.enable_heuristic_topk and not self.use_cute_dsl_topk: - max_gen_tokens = self.max_num_sequences * ( - 1 + self.max_draft_tokens) - self.heuristic_scratch_values = self.get_empty( - self.cuda_graph_buffers, - (max_gen_tokens, self.num_sparse_topk), - cache_name="heuristic_scratch_values", - dtype=torch.float32, - capture_graph=capture_graph, - ) - # The Radix-split-work scratch (radix_aux_*) is sized the same way - # (num_seqs * (1 + max_draft_tokens) rows) and is allocated - # unconditionally, so it must be resized here too -- otherwise the - # cpp indexer_topk_decode op aborts once MTP raises max_draft_tokens - # ("radix_aux_* must hold at least num_rows*blocks_per_row*index_topk - # elements"). - self._create_radix_aux_buffers(capture_graph=capture_graph) - - def _get_pool_block_indices(self) -> torch.Tensor: - """Extract memory pool block indices from host_kv_cache_block_offsets. - - The C++ setOffsets() encodes offsets as: - encoded = memPoolBlockIndex * numLayers * kvFactor - For SELFKONLY (MLA/DSA), kvFactor=1, so: - memPoolBlockIndex = encoded // num_local_layers - - Returns a (num_seqs, max_blocks_per_seq) int32 CPU tensor with valid - pool indices clamped to [0, blocks_in_primary_pool - 1]. - """ - num_local_layers = self.kv_cache_manager.num_local_layers - max_pool_idx = self.kv_cache_manager.blocks_in_primary_pool - 1 - # DSA uses SELFKONLY mode where only key cache is stored (kv_factor=1). - # host_kv_cache_block_offsets shape: (num_pools, max_batch*beam, 2, max_blocks_per_seq) - # Note: dim=2 is always 2 in the tensor layout (K and V slots), but for - # SELFKONLY only the K slot (index 0) contains valid data. - assert self.kv_cache_manager.kv_factor == 1, \ - f"DSA requires SELFKONLY mode (kv_factor=1), got kv_factor={self.kv_cache_manager.kv_factor}" - # Pool 0, first num_seqs entries, field 0 (key offsets) - encoded = self.kv_cache_manager.host_kv_cache_block_offsets[ - 0, :self.num_seqs, 0, :] - pool_indices = encoded // num_local_layers - # Clamp for safety: handles garbage padding from torch.empty in uninitialized slots - pool_indices = pool_indices.clamp(min=0, - max=max_pool_idx).to(torch.int32) - return pool_indices - - def _invalidate_pool_view_cache(self): - """Invalidate the cached pool view and related step-invariant values. - - Must be called at the start of each forward step (in prepare()) so that - _ensure_pool_view_cached() recomputes them for the new batch. - """ - self._pool_cache_valid = False - - def _ensure_pool_view_cached(self): - """Compute and cache values used by - transform_local_topk_and_prepare_pool_view(). - - These values (pool view, stride factor, block table slices, request - index slices) are constant across all layers sharing the same KV pool - and batch dimensions within a forward pass. Caching them avoids - redundant Python/CUDA overhead per layer. - - Safety: _invalidate_pool_view_cache() is called unconditionally at the - 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): - return - - pool = self.kv_cache_manager.get_unique_primary_pool() - kv_cache_manager = self.kv_cache_manager - num_blocks, num_layers, _, _ = pool.shape - self._cached_tokens_per_block = kv_cache_manager.tokens_per_block - head_dim = kv_cache_manager.head_dim - self._cached_pool_view = pool.squeeze(2).view(-1, 1, head_dim) - self._cached_stride_factor = (num_layers * - self._cached_tokens_per_block) - 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_req_idx_gen = ( - self.req_idx_per_token[self.num_ctx_tokens:self.num_tokens] - - self.num_contexts) - self._cached_kv_mgr_id = id(kv_cache_manager) - self._pool_cache_valid = True - - @maybe_compile(dynamic=True) - def _get_dense_topk_indices(self, seq_lens, kv_lens, num_tokens): - device = kv_lens.device - past_kv_lens = kv_lens - seq_lens - # get position ids - seq_ends = torch.cumsum(seq_lens, dim=0) - seq_starts = seq_ends - seq_lens - per_seq_offsets = past_kv_lens - seq_starts # Shape: [batch_size] - global_indices = torch.arange(num_tokens, device=device) - batch_indices = torch.searchsorted(seq_ends, - global_indices, - side='right') - 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) - 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) - 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), - 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) - 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 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), - 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) - self.topk_indices_buffer[gen_range, :].copy_( - self.host_topk_indices_buffer[gen_range, :], - non_blocking=True) - - def prepare_for_spec_decode(self, kv_lens: torch.Tensor): - # Because the fp8_paged_mqa_logits only supports seq_len == 1/2/4 (i.e., max_draft_tokens == 0/1/3) on sm100, and - # seq_len == 1/2 (i.e., max_draft_tokens == 0/1) on sm90, for other cases, we need to flatten the q tensor and - # expand the kv_lens and block_table for MTP support. - # TODO: - # - No distinction between sm90 and sm100 is needed once MTP3 is supported on sm90. - # - Remove this once fp8_paged_mqa_logits supports an arbitrary number of MTP draft tokens. - use_dsl = self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits - self.use_expanded_buffers_for_mtp = (not use_dsl and ( - (self.max_draft_tokens > 1 and get_sm_version() == 90) or - ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3) - and get_sm_version() >= 100))) - if self.use_expanded_buffers_for_mtp: - # Expand kv_lens_cuda (only generation) - num_tokens = self.num_generations * (1 + self.max_draft_tokens) - gen_kv_lens = self.get_indexer_kv_lens( - kv_lens[self.num_contexts:self.num_seqs]) - gen_kv_lens_expanded = torch.stack([gen_kv_lens] * - (1 + self.max_draft_tokens), - dim=0) - gen_kv_lens_expanded = gen_kv_lens_expanded.transpose( - 0, 1).contiguous().flatten() - 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) - - # Expand indexer_k_cache_block_offsets (only generation) - # host_indexer_k_cache_block_offsets already contains correct pool - # indices from _get_pool_block_indices() in prepare_for_indexer_k_cache(). - if self.kv_cache_manager is not None and self.num_generations > 0: - 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] - expanded_blocks = gen_block_tensor.repeat_interleave( - 1 + self.max_draft_tokens, dim=0) - self.host_block_table_expanded[:num_tokens, :max_len].copy_( - expanded_blocks, non_blocking=True) - self.block_table_expanded[:num_tokens].copy_( - self.host_block_table_expanded[:num_tokens], - non_blocking=True) - self.block_table_expanded.clamp_(min=0) - - self.expand_for_dsl = (use_dsl and self.kv_cache_manager is not None - and self.max_draft_tokens >= 1) - if self.expand_for_dsl and self.num_generations > 0: - next_n = 1 + self.max_draft_tokens - 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]) - max_ctx = int( - gen_kv_lens.max().item()) if gen_kv_lens.numel() else 0 - expand_factor, atom = _pick_dsl_expand( - next_n, - batch_size=self.num_generations, - max_ctx=max_ctx, - num_sms=self.num_sms, - kernel_atoms=kernel_atoms, - ) - self.dsl_expand_factor = expand_factor - self.dsl_atom = atom - if expand_factor > 1: - num_tokens = self.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_( - self.kv_lens_expanded_host[:num_tokens], non_blocking=True) - 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] - expanded_blocks = gen_block_tensor.repeat_interleave( - expand_factor, dim=0) - self.host_block_table_expanded[:num_tokens, :max_len].copy_( - expanded_blocks, non_blocking=True) - self.block_table_expanded[:num_tokens].copy_( - self.host_block_table_expanded[:num_tokens], - non_blocking=True) - self.block_table_expanded.clamp_(min=0) - else: - self.dsl_expand_factor = 1 - self.dsl_atom = 1 + self.max_draft_tokens - - def prepare_for_indexer_k_cache(self): - # Build indexer_k_cache_block_offsets using pool block indices derived - # from host_kv_cache_block_offsets (populated by super().prepare()). - # This correctly resolves block IDs to memory pool indices, which is - # required when host cache offload is enabled (block IDs != pool indices - # for onboarded secondary blocks). - if self.kv_cache_manager is None: - return - pool_indices = self._get_pool_block_indices() - self.host_indexer_k_cache_block_offsets[:self.num_seqs].copy_( - pool_indices) - self.indexer_k_cache_block_offsets[:self.num_seqs].copy_( - self.host_indexer_k_cache_block_offsets[:self.num_seqs], - non_blocking=True) - # Safety clamp: prevent OOB from CUDA graph padding entries which - # may contain stale negative or out-of-range values after block - # eviction/onboarding with host cache offload. - self.indexer_k_cache_block_offsets.clamp_(min=0) - - # Build block_table for topk_indices conversion (actual block allocation) - cached_token_lens = torch.tensor( - self.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] - kv_lens = cached_token_lens.clone() - kv_lens[active_rank] += self.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 - # 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. - host_block_table = torch.empty((pool_indices.shape[0], max_blocks_used), - dtype=pool_indices.dtype, - pin_memory=prefer_pinned()) - 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) - # Copy to GPU - self.block_table[:self.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: - # 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 - else: - self.skip_indexer_for_ctx_reqs = False - - if self.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 - 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.max_ctx_kv_len = kv_lens[:self.num_contexts].max().item() - self.max_ctx_seq_len = self.seq_lens[:self.num_contexts].max().item( - ) - # context cached token indptr - torch.cumsum( - cached_token_lens[:self.num_contexts], - dim=0, - dtype=torch.int64, - out=self.host_ctx_cached_token_indptr[1:self.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) - # context kv indptr - torch.cumsum(kv_lens[:self.num_contexts], - dim=0, - dtype=torch.int64, - out=self.host_ctx_kv_indptr[1:self.num_contexts + 1]) - self.ctx_kv_indptr[:self.num_contexts + 1].copy_( - self.host_ctx_kv_indptr[:self.num_contexts + 1], - non_blocking=True) - else: - self.num_ctx_cached_tokens = 0 - self.max_ctx_kv_len = 0 - self.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() - # generation cached token indptr - torch.cumsum( - cached_token_lens[self.num_contexts:self.num_seqs], - dim=0, - dtype=torch.int64, - out=self.host_gen_cached_token_indptr[1:self.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) - # generation kv indptr - torch.cumsum(kv_lens[self.num_contexts:self.num_seqs], - dim=0, - dtype=torch.int64, - out=self.host_gen_kv_indptr[1:self.num_generations + - 1]) - self.gen_kv_indptr[:self.num_generations + 1].copy_( - self.host_gen_kv_indptr[:self.num_generations + 1], - non_blocking=True) - else: - self.max_gen_seq_len = 0 - - 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, - dim=0, - )) - self.req_idx_per_token[:self.num_tokens].copy_( - self.host_req_idx_per_token[:self.num_tokens], - non_blocking=True, - ) - - -@maybe_compile(dynamic=True) -def _scale(weights: torch.Tensor, q_scale: torch.Tensor, - s: float) -> torch.Tensor: - """Scale attention weights by quantization scale and constant factor.""" - return weights * q_scale.squeeze(-1) * s - - -@maybe_compile(dynamic=True) -def _to_float(hidden_states: torch.Tensor) -> torch.Tensor: - """Cast hidden states to float32 for TF32 GEMM computation.""" - return hidden_states.float() - - -@contextmanager -def _tf32_matmul_enabled(): - """Temporarily enable TF32 tensor cores for FP32 matmul in this scope. - - Forces PyTorch/cuBLASLt to use CUBLAS_COMPUTE_32F_FAST_TF32, which - guarantees TF32 tensor cores. Plain CUBLAS_COMPUTE_32F (used by - torch.ops.trtllm.cublas_mm) falls back to SIMT SGEMM on CUDA cores - based on cuBLASLt heuristics for small M. - """ - prev = torch.backends.cuda.matmul.allow_tf32 - torch.backends.cuda.matmul.allow_tf32 = True - try: - yield - finally: - torch.backends.cuda.matmul.allow_tf32 = prev - - -@dataclass -class IndexerParams: - """ - Parameters for indexer. - """ - num_contexts: int - num_generations: int - num_ctx_tokens: int - head_dim: int - quant_block_size: int - tokens_per_block: int - compress_ratio: int - request_ids: List[int] - num_past_tokens: List[int] - seq_lens: torch.Tensor - # Bytes of quantized data per token in the indexer K cache; defaults to - # head_dim (one FP8 byte per element). For MXFP4 use head_dim // 2. - data_bytes_per_token: Optional[int] = None - - def __post_init__(self): - # Pre-compute frequently used tensors once instead of on every property access - num_past_tokens_tensor = torch.tensor(self.num_past_tokens, - dtype=torch.int32) - self._num_past_tokens_tensor = num_past_tokens_tensor - compress_ratio = self.compress_ratio - self._cached_kv_tokens = num_past_tokens_tensor // compress_ratio - self._all_kv_tokens = (self.seq_lens + - num_past_tokens_tensor) // compress_ratio - self._new_kv_tokens = self._all_kv_tokens - self._cached_kv_tokens - self._kv_lens = self._all_kv_tokens - if self.data_bytes_per_token is None: - self.data_bytes_per_token = self.head_dim - self._scale_size = self.head_dim // self.quant_block_size * 4 - self._block_stride = self.tokens_per_block * ( - self.data_bytes_per_token + self._scale_size) - - @property - def batch_size(self): - return len(self.request_ids) - - @property - def kv_lens(self): - return self._kv_lens - - @property - def total_tokens(self): - return self.seq_lens.sum().item() - - @property - def cached_kv_tokens(self): - return self._cached_kv_tokens - - @property - def all_kv_tokens(self): - return self._all_kv_tokens - - @property - def new_kv_tokens(self): - return self._new_kv_tokens - - @property - def scale_size(self): - return self._scale_size - - @property - def block_stride(self): - return self._block_stride - - -class Indexer(nn.Module): - """DSA sparse attention indexer that selects top-K KV cache entries per token.""" - - def __init__(self, - quant_config: Optional[QuantConfig], - pos_embd_params: Optional[PositionalEmbeddingParams], - mla_params: Optional[MLAParams], - skip_create_weights_in_init: bool, - sparse_params: DSAParams, - dtype: Optional[torch.dtype], - compress_ratio: int = 1, - layer_idx: int = 0, - aux_stream: Optional[torch.cuda.Stream] = None): - """Initialize indexer with projection weights, norms, and TopK configuration.""" - super().__init__() - self.hidden_size = mla_params.hidden_size - self.q_lora_rank = mla_params.q_lora_rank - self.rope_dim = mla_params.qk_rope_head_dim - self.n_heads = sparse_params.index_n_heads # 64 - self.head_dim = sparse_params.index_head_dim # 128 - self.index_topk = sparse_params.index_topk # 2048 - self.layer_idx = layer_idx - self.compress_ratio = compress_ratio - - self.wq_b = Linear( - self.q_lora_rank, - self.n_heads * self.head_dim, - bias=False, - dtype=dtype, - quant_config=quant_config, - skip_create_weights_in_init=skip_create_weights_in_init, - use_custom_cublas_mm=True) - self.wk = Linear( - self.hidden_size, - self.head_dim, - bias=False, - dtype=torch.float32, - quant_config=None, - skip_create_weights_in_init=skip_create_weights_in_init, - use_custom_cublas_mm=True) - self.k_norm = LayerNorm(hidden_size=self.head_dim, eps=1e-6) - self.weights_proj = Linear( - self.hidden_size, - self.n_heads, - bias=False, - dtype=torch.float32, - quant_config=None, - skip_create_weights_in_init=skip_create_weights_in_init, - use_custom_cublas_mm=True) - - # Fused wk + weights_proj weight for single F.linear FP32 GEMM under allow_tf32. - # Maps to TF32 tensor cores on Ampere+. - self._fused_wk_wp_weight: Optional[torch.Tensor] = None - - indexer_rope_interleave = sparse_params.indexer_rope_interleave - self.rotary_emb = RotaryEmbedding( - pos_embd_params.rope, - head_dim=self.rope_dim, - is_neox=not indexer_rope_interleave, - ) - - self.softmax_scale = self.head_dim**-0.5 - # TODO: make it configurable from hf config - self.scale_fmt = "ue8m0" - # indexer_k_dtype controls both Q and K precision. DeepGEMM's - # fp8_fp4_mqa_logits / fp8_fp4_paged_mqa_logits kernels only dispatch - # to FP4xFP4 or FP8xFP8 (no mixed-precision variant). The DeepGEMM - # kernel asserts SM100 + head_dim=128 at launch time under FP4. - self.use_fp4 = sparse_params.indexer_k_dtype == "fp4" - self.aux_stream = aux_stream - self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] - self.use_cute_dsl_topk = (sparse_params.use_cute_dsl_topk - and IS_CUTLASS_DSL_AVAILABLE) - self.use_cute_dsl_paged_mqa_logits = ( - sparse_params.use_cute_dsl_paged_mqa_logits - and IS_CUTLASS_DSL_AVAILABLE) - self.weight_scale_factor = self.softmax_scale * self.n_heads**-0.5 - - self._enable_heuristic_topk = (sparse_params.enable_heuristic_topk - and get_sm_version() >= 100) - - if (self.use_cute_dsl_topk - or self.use_cute_dsl_paged_mqa_logits) and layer_idx == 0: - from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops - - if self.use_cute_dsl_topk and not self._enable_heuristic_topk: - # the dtype of topk input tensor, which is float32 now. - # Note, need to update it if the dtype of topk input tensor is changed. - cute_dsl_custom_ops.warmup_cute_dsl_indexer_topk( - dtype=torch.float32, top_k=self.index_topk) - - if self._enable_heuristic_topk and layer_idx == 0: - # Populate static caches (sm_count, L2 cache size) inside the C++ - # Scheme X dispatcher before any CUDA Graph capture so the host - # attribute queries do not end up frozen into a captured graph. - warmup_heuristic_topk_decode(top_k=self.index_topk) - - # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM - # (populated in cache_derived_state; maps to TF32 tensor cores on Ampere+) - self._fused_wk_wp_weight: Optional[torch.Tensor] = None - - def cache_derived_state(self) -> None: - """Fuse wk + weights_proj into single FP32 weight for F.linear GEMM under allow_tf32 (TF32 tensor cores on Ampere+).""" - # wk: [head_dim, hidden_size] + weights_proj: [n_heads, hidden_size] - # → fused: [head_dim + n_heads, hidden_size] - self._fused_wk_wp_weight = torch.cat( - [self.wk.weight.data, self.weights_proj.weight.data], dim=0) - - def post_load_weights(self) -> None: - self.cache_derived_state() - - @staticmethod - def prepare_one_prefill_chunk( - metadata: DSAtrtllmAttentionMetadata, - chunk_specs: List[Tuple[int, int, int, int]], - ) -> IndexerPrefillChunkMetadata: - """ - Build metadata for one prefill chunk for indexer forward pass. - Handles both multi-request chunks and intra-request Q-block chunks. - - Args: - metadata: Attention metadata - chunk_specs: List of (req_idx, token_start_in_req, token_end_in_req, req_cum_start) - - token_start_in_req, token_end_in_req are indices into current batch context tokens - - For multi-request: multiple specs from different requests (full requests) - - For intra-request: single spec from one request's Q-block - - Note: Cached token counts are derived from metadata.host_ctx_cached_token_indptr - """ - device = metadata.cu_seqlen_ks.device - compress_ratio = _effective_compress_ratio_divisor( - _select_indexer_compress_ratio(metadata.compress_ratios)) - if len(chunk_specs) == 1: - # Single request or intra-request Q-block - req_idx, token_start_in_req, token_end_in_req, req_cum_start = chunk_specs[ - 0] - num_q_tokens = token_end_in_req - token_start_in_req - - # Get cached token count for this request from metadata - num_cached = ( - metadata.host_ctx_cached_token_indptr[req_idx + 1] - - metadata.host_ctx_cached_token_indptr[req_idx]).item() - - # Total compressed KV tokens for this request - req_kv_len = (num_cached + token_end_in_req) // compress_ratio - - # For intra-request chunks: Q block attends to all previous K in the request - # Q tokens [token_start_in_req:token_end_in_req] within the request's current tokens - # K tokens [0:req_kv_len] in compressed KV space - cu_seqlen_ks = torch.zeros(num_q_tokens, - dtype=torch.int32, - device='cpu') - cu_seqlen_ke = (torch.arange(token_start_in_req + 1, - token_end_in_req + 1, - dtype=torch.int32, - device='cpu') + - num_cached) // compress_ratio - - # Q token range in batch (indices into context tokens in the current batch) - token_start = req_cum_start + token_start_in_req - token_end = req_cum_start + token_end_in_req - - # K token range: index into full KV slot mapping in compressed KV space - # For req_idx=0 the offset is 0; for other indices compute compressed cumulative offset - kv_offset_in_extended = sum( - (metadata.host_ctx_kv_indptr[j + 1] - - metadata.host_ctx_kv_indptr[j]).item() // compress_ratio - for j in range(req_idx)) - total_kv_for_req = req_kv_len - k_token_start = kv_offset_in_extended - k_token_end = kv_offset_in_extended + total_kv_for_req - - else: - # Multi-request chunk: batch multiple full requests together - # Extract sequence lengths for these requests - req_seq_lens = [] - req_num_past_tokens = [] - req_kv_lens = [] - first_req_idx = chunk_specs[0][0] - - for spec in chunk_specs: - req_idx, token_start_in_req, token_end_in_req, _ = spec - req_seq_lens.append(token_end_in_req - token_start_in_req) - # Get cached token count from metadata - num_past_tokens = ( - metadata.host_ctx_cached_token_indptr[req_idx + 1] - - metadata.host_ctx_cached_token_indptr[req_idx]).item() - req_num_past_tokens.append(num_past_tokens) - req_kv_lens.append( - (num_past_tokens + req_seq_lens[-1]) // compress_ratio) - - req_seq_lens_tensor = torch.tensor(req_seq_lens, - dtype=torch.int32, - device='cpu') - req_num_past_tokens_tensor = torch.tensor(req_num_past_tokens, - dtype=torch.int32, - device='cpu') - req_kv_lens_tensor = torch.tensor(req_kv_lens, - dtype=torch.int32, - device='cpu') - num_q_tokens = sum(req_seq_lens) - - # Compute causal attention bounds for batched requests - cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - req_seq_lens_tensor, len(chunk_specs), num_q_tokens, - req_num_past_tokens_tensor, req_kv_lens_tensor, compress_ratio) - - # Global Q token ranges (indices into ctx tokens in the current batch) - token_start = chunk_specs[0][3] # req_cum_start of first request - token_end = token_start + num_q_tokens - - # K token range: index into full kv slot mapping (cached + current ctx tokens within the batch) - # Must use compressed offsets - kv_offset_in_extended = sum( - (metadata.host_ctx_kv_indptr[j + 1] - - metadata.host_ctx_kv_indptr[j]).item() // compress_ratio - for j in range(first_req_idx)) - total_kv_len = sum(req_kv_lens) - k_token_start = kv_offset_in_extended - k_token_end = kv_offset_in_extended + total_kv_len - - assert cu_seqlen_ks.shape[0] == num_q_tokens == token_end - token_start, \ - f"Indexer.prepare_one_prefill_chunk - cu_seqlen_ks length mismatch: {cu_seqlen_ks.shape[0]} != {num_q_tokens}" - assert cu_seqlen_ke.shape[0] == num_q_tokens == token_end - token_start, \ - f"Indexer.prepare_one_prefill_chunk - cu_seqlen_ke length mismatch: {cu_seqlen_ke.shape[0]} != {num_q_tokens}" - - return IndexerPrefillChunkMetadata( - cu_seqlen_ks=maybe_pin_memory(cu_seqlen_ks).to(device, - non_blocking=True), - cu_seqlen_ke=maybe_pin_memory(cu_seqlen_ke).to(device, - non_blocking=True), - token_start=token_start, - token_end=token_end, - k_token_start=k_token_start, - k_token_end=k_token_end, - ) - - @staticmethod - def build_indexer_params( - metadata: DSAtrtllmAttentionMetadata) -> Optional[IndexerParams]: - 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 - 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)) - 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, - compress_ratio=compress_ratio, - request_ids=metadata.request_ids, - num_past_tokens=metadata.kv_cache_params.num_cached_tokens_per_seq, - seq_lens=metadata.seq_lens, - data_bytes_per_token=data_bytes_per_token, - ) - - @staticmethod - def recompute_slot_mappings(metadata: DSAtrtllmAttentionMetadata, - indexer_params: Optional[IndexerParams] = None): - """Recompute slot_mapping_fp8/scale from current block offsets. - - This is the subset of prepare_for_update_k_cache() that maps each new - compressed KV token to its flat cache position. It is safe to call in - isolation after a caller has swapped the active KV cache manager and - block-offset buffers, such as during draft KV-cache replay. - """ - if indexer_params is None: - indexer_params = Indexer.build_indexer_params(metadata) - if indexer_params is None: - return - - batch_size = indexer_params.batch_size - tokens_per_block = indexer_params.tokens_per_block - head_dim = indexer_params.head_dim - new_kv_tokens = indexer_params.new_kv_tokens - total_new_kv_tokens = new_kv_tokens.sum().item() - data_bytes_per_token = indexer_params.data_bytes_per_token - - # Compute global positions for all kv tokens in the batch (fully vectorized) - req_indices = torch.repeat_interleave( - torch.arange(batch_size, dtype=torch.int64, device='cpu'), - new_kv_tokens) - # Vectorized token_offsets: arange(total) - cumulative start per request - cu_new_kv = torch.zeros(batch_size + 1, dtype=torch.int64, device='cpu') - cu_new_kv[1:] = new_kv_tokens.to(torch.int64).cumsum(0) - token_offsets = ( - torch.arange(total_new_kv_tokens, dtype=torch.int64, device='cpu') - - cu_new_kv[:-1].repeat_interleave(new_kv_tokens)) - global_positions = indexer_params.cached_kv_tokens[ - req_indices] + token_offsets - - # 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] - 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, - req_indices, - head_dim, - tokens_per_block, - indexer_params.quant_block_size, - 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 - - metadata.slot_mapping_fp8[:total_new_kv_tokens].copy_( - metadata.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) - - @staticmethod - def prepare_for_update_k_cache(metadata: DSAtrtllmAttentionMetadata, - indexer_params: IndexerParams): - """ - Prepare indexer for the update_k_cache stage. - - Compute slot_mapping for all requests (both context and generation) - This maps each token to its flat cache position for vectorized KV cache updates - """ - Indexer.recompute_slot_mappings(metadata, indexer_params) - - @staticmethod - def prepare_for_chunked_prefill(metadata: DSAtrtllmAttentionMetadata, - indexer_params: IndexerParams): - """ - Prepare indexer for the chunked prefill. - """ - num_contexts = indexer_params.num_contexts - seq_lens = indexer_params.seq_lens - tokens_per_block = indexer_params.tokens_per_block - head_dim = indexer_params.head_dim - - # When MLA chunked prefill is active, it already handles chunking - # Indexer should just process the current MLA chunk as a single chunk - has_mla_chunked_prefill = (metadata.enable_context_mla_with_cached_kv - and - metadata.runtime_features.chunked_prefill) - if has_mla_chunked_prefill: - chunk_specs = [(i, 0, seq_lens[i].item(), - seq_lens[:i].sum().item() if i > 0 else 0) - for i in range(num_contexts)] - metadata.indexer_prefill_chunks = [ - Indexer.prepare_one_prefill_chunk( - metadata, - chunk_specs, - ) - ] - else: - # Use indexer's own chunking logic to prevent L^2 complexity of indexer MQA logits computation for long sequences. - # This is only used when MLA chunked prefill is not enabled. - # Adapt chunk size to the batch's largest K_compressed (see - # select_indexer_chunk_size). - max_k_compressed = 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) - chunk_groups = split_prefill_chunks( - seq_lens[:num_contexts], - effective_chunk_size, - start_idx=0, - ) - - if len(chunk_groups - ) > 1 or metadata.enable_context_mla_with_cached_kv: - metadata.indexer_prefill_chunks = [ - Indexer.prepare_one_prefill_chunk( - metadata, - chunk_specs, - ) for chunk_specs in chunk_groups - ] - else: - metadata.indexer_prefill_chunks = None - - # When chunked prefill or KVCache reuse is enabled, we need to gather the full KV for indexer's logit computation. - # Indexer's own chunking does not need full KV gathering, instead it gathers only the current chunk with loop-based gathering. - if metadata.enable_context_mla_with_cached_kv: - # Use kv_lens which correctly computes (raw_past + seq_lens) // compress_ratio. - total_kv_per_request = indexer_params.kv_lens[:num_contexts] - total_kv_len = total_kv_per_request.sum().item() - host_slot_mapping_fp8_fullkv = torch.empty( - total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned()) - host_slot_mapping_scale_fullkv = torch.empty( - total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned()) - - req_indices = torch.repeat_interleave( - torch.arange(num_contexts, dtype=torch.int64, device='cpu'), - total_kv_per_request) - - cu_kv = torch.zeros(num_contexts + 1, - dtype=torch.int64, - device='cpu') - cu_kv[1:] = total_kv_per_request.to(torch.int64).cumsum(0) - kv_positions = ( - torch.arange(total_kv_len, dtype=torch.int64, device='cpu') - - cu_kv[:-1].repeat_interleave(total_kv_per_request)) - - fp8_flat_indices, scale_flat_indices = _compute_slot_mappings( - kv_positions, - metadata.host_indexer_k_cache_block_offsets, - req_indices, - head_dim, - tokens_per_block, - indexer_params.quant_block_size, - data_bytes_per_token=head_dim // - 2 if metadata.kv_cache_manager.use_fp4 else head_dim, - ) - - host_slot_mapping_fp8_fullkv[:total_kv_len] = fp8_flat_indices - host_slot_mapping_scale_fullkv[:total_kv_len] = scale_flat_indices - - assert len(fp8_flat_indices) == total_kv_len, \ - f"host_slot_mapping_fp8_fullkv/host_slot_mapping_scale_fullkv length mismatch: {len(fp8_flat_indices)} != total_kv_len={total_kv_len}" - - # Store extended mappings for indexer full KV gathering - 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) - else: - metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 - metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale - - @staticmethod - def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): - """ - Prepare scheduler metadata for the DeepGEMM decode MQA kernel. - """ - num_contexts = metadata.num_contexts - num_generations = metadata.num_generations - if not metadata.use_expanded_buffers_for_mtp: - gen_seq_lens = metadata.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 - next_n_cap = metadata.kv_lens_cuda_2d.shape[1] - metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_( - gen_seq_lens.unsqueeze(-1).expand(-1, next_n_cap)) - scheduler_metadata_buffer = get_paged_mqa_logits_metadata( - gen_seq_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, - metadata.num_sms) - metadata.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, - non_blocking=True) - if metadata.max_draft_tokens > 0: - scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( - metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap], - _DG_SCHEDULE_BLOCK_KV, metadata.num_sms) - metadata.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 = metadata.num_generations * (1 + - metadata.max_draft_tokens) - scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( - metadata.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), - _DG_SCHEDULE_BLOCK_KV, metadata.num_sms) - metadata.scheduler_metadata_buffer_expanded.copy_( - scheduler_metadata_buffer_expanded, non_blocking=True) - - @staticmethod - def prepare(metadata: DSAtrtllmAttentionMetadata): - """ - Prepare indexer for the forward pass. - This should be called during metadata.prepare() stage. - - - Computes slot_mapping for KV cache updates - - Prepares schedule_metadata for fp8_paged_mqa_logits - - Stores generation request IDs for decode phase - """ - - # Skip indexer preparation if the kv_cache_manager doesn't have index_head_dim. - # 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. - indexer_params = Indexer.build_indexer_params(metadata) - if indexer_params is None: - return - - num_contexts = metadata.num_contexts - num_generations = metadata.num_generations - 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)) - # Store compressed KV token count for context requests - metadata.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) - - # Prepare for prefill phase if there are context requests - if num_contexts > 0: - # Compute attention window bounds for each query token in batched sequences - # cu_seqlen_ks[i]: start index in global KV for query token i - # cu_seqlen_ke[i]: end index (exclusive) in global KV for query token i - host_seq_lens = seq_lens[:num_contexts] - host_num_past_tokens = indexer_params._num_past_tokens_tensor[: - num_contexts] - host_kv_lens = indexer_params.kv_lens[:num_contexts] - host_cu_seqlen_ks, host_cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - host_seq_lens, num_contexts, num_ctx_tokens, - host_num_past_tokens, host_kv_lens, compress_ratio) - - metadata.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_( - maybe_pin_memory(host_cu_seqlen_ke), non_blocking=True) - Indexer.prepare_for_chunked_prefill(metadata, indexer_params) - - # Prepare for decode phase if there are generation requests - if num_generations > 0: - # Prepare schedule metadata for fp8_paged_mqa_logits - # This is a preprocessing step that computes scheduling information for the kernel - Indexer.prepare_scheduler_metadata(metadata) - - def _update_k_cache(self, k_fp8: torch.Tensor, k_scale: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata) -> None: - """ - Insert/append k values and scales into the indexer k cache using pre-computed slot mappings. - Uses flat byte indices with vectorized scatter. - - Args: - 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: - return - - k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers( - self.layer_idx) - - num_tokens = k_fp8.shape[0] - - # The C++ op reinterprets k_fp8 (FP8) and k_scale (float32) as raw - # bytes internally and only reads the first num_tokens entries from - # the slot mapping buffers, avoiding Python-side view/slice overhead. - if k_scale.element_size() == 1: - # The op expects 4 byte elements. - k_scale = k_scale.view(torch.int32) - torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8, k_scale, k_cache, - metadata.slot_mapping_fp8, - metadata.slot_mapping_scale, - num_tokens) - - def _gather_k_cache_for_chunk( - self, - metadata: DSAtrtllmAttentionMetadata, - chunk: IndexerPrefillChunkMetadata, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Gather K values from indexer cache for a specific chunk. - - Uses pre-computed extended slot mappings that cover cached + current batch context tokens. - chunk.k_token_start/k_token_end directly index into the extended slot mapping. - - Args: - metadata: Attention metadata - chunk: Chunk metadata with k_token_start/end as indices into extended slot mapping - - Returns: - 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, \ - "_gather_k_cache_for_chunk requires extended slot mappings (only available with cached tokens)" - - k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers( - self.layer_idx) - - head_dim = self.head_dim - scale_size = 4 # float32 = 4 bytes - - # Extract slot mappings using chunk's k_token_start/end - # These indices point directly into the extended slot mapping array - k_token_start = chunk.k_token_start - 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] - - # Vectorized gather using pre-computed slot mappings - # Gather FP8 data - byte_offsets_fp8 = torch.arange( - head_dim, device=k_cache.device).unsqueeze(0) # [1, head_dim] - gather_indices_fp8 = slot_mapping_fp8_chunk.unsqueeze( - 1) + byte_offsets_fp8 # [num_k_tokens, head_dim] - assert (gather_indices_fp8 - >= k_cache.numel()).sum() == 0, "Out-of-bounds access detected" - gather_indices_fp8 = _unravel_indices(gather_indices_fp8, k_cache.shape) - k_fp8_bytes = k_cache[gather_indices_fp8] - k_fp8 = k_fp8_bytes.view(torch.float8_e4m3fn).view( - num_k_tokens, head_dim) - - # Gather scale data - byte_offsets_scale = torch.arange( - scale_size, device=k_cache.device).unsqueeze(0) # [1, 4] - gather_indices_scale = slot_mapping_scale_chunk.unsqueeze( - 1) + byte_offsets_scale # [num_k_tokens, 4] - assert (gather_indices_scale - >= k_cache.numel()).sum() == 0, "Out-of-bounds access detected" - gather_indices_scale = _unravel_indices(gather_indices_scale, - k_cache.shape) - k_scale_bytes = k_cache[gather_indices_scale] - k_scale = k_scale_bytes.view(torch.float32).view(num_k_tokens, 1) - - return k_fp8, k_scale - - def _call_mqa_logits(self, q_fp8: torch.Tensor, k_fp8: torch.Tensor, - k_scale: torch.Tensor, weights: torch.Tensor, - cu_seqlen_ks: torch.Tensor, cu_seqlen_ke: torch.Tensor, - q_scale: Optional[torch.Tensor]) -> torch.Tensor: - """Dispatch fp8_mqa_logits vs fp8_fp4_mqa_logits based on use_fp4. - - For FP4 the gather output keeps the legacy float8_e4m3fn dtype for - API compatibility; reinterpret the bytes as the int8 / int32 layout - the DeepGEMM kernel expects. The scale tensor is collapsed to 1D for - the kv side and 2D for the q side per the kernel's asserts. - """ - if self.use_fp4: - k_fp4_bytes = k_fp8.view(torch.int8) - k_scale_int32 = k_scale.view(torch.int32).reshape(-1) - # q_scale arrives as (chunk_tokens, n_heads, 1); the FP4 kernel - # asserts q_sf is 2D so collapse the trailing unit axis. - q_scale_2d = q_scale.reshape(-1, self.n_heads) - return fp8_fp4_mqa_logits( - (q_fp8, q_scale_2d), - (k_fp4_bytes, k_scale_int32), - weights, - cu_seqlen_ks, - cu_seqlen_ke, - ) - return fp8_mqa_logits(q_fp8, (k_fp8, k_scale.reshape(-1)), weights, - cu_seqlen_ks, cu_seqlen_ke) - - def _call_paged_mqa_logits(self, q_decode: torch.Tensor, - k_cache: torch.Tensor, - weights_decode: torch.Tensor, - context_lens: torch.Tensor, - block_table: torch.Tensor, - scheduler_metadata_buffer: torch.Tensor, - max_seq_len: int, - q_scale: Optional[torch.Tensor]) -> torch.Tensor: - """Dispatch fp8_paged_mqa_logits vs fp8_fp4_paged_mqa_logits.""" - if self.use_fp4: - return fp8_fp4_paged_mqa_logits( - (q_decode, q_scale), k_cache, weights_decode, context_lens, - block_table, scheduler_metadata_buffer, max_seq_len) - return fp8_paged_mqa_logits(q_decode, k_cache, weights_decode, - context_lens, block_table, - scheduler_metadata_buffer, max_seq_len) - - def sparse_attn_indexer( - self, - metadata: DSAtrtllmAttentionMetadata, - hidden_states: torch.Tensor, - q_fp8: torch.Tensor, - k_fp8: torch.Tensor, - k_scale: torch.Tensor, - weights: torch.Tensor, - use_custom_topk: bool = True, - q_scale: Optional[torch.Tensor] = None, - update_k_cache: bool = True, - ) -> torch.Tensor: - """Run the indexer TopK kernel for both prefill and decode phases. - - q_scale is only consumed by the FP4 dispatch; FP8 path ignores it. - """ - # DSACacheManager / DeepseekV4CacheManager force quant_block_size to - # 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. - assert metadata.kv_cache_manager is None or \ - metadata.kv_cache_manager.quant_block_size in (32, 128), \ - f"Unexpected quant_block_size {metadata.kv_cache_manager.quant_block_size if metadata.kv_cache_manager else 'N/A'}" - # Update the indexer k cache before prefill chunks gather from it. - if update_k_cache: - self._update_k_cache(k_fp8, k_scale, metadata) - - num_contexts = metadata.num_contexts - num_generations = metadata.num_generations - num_ctx_tokens = metadata.num_ctx_tokens - num_tokens = metadata.num_tokens - - has_decode = num_generations > 0 - has_prefill = num_contexts > 0 - num_gen_tokens = num_tokens - num_ctx_tokens - - topk_indices_buffer = metadata.get_empty( - metadata.cuda_graph_buffers, - (hidden_states.shape[0], self.index_topk), - cache_name="indexer_topk_out_buffer", - dtype=torch.int32, - capture_graph=metadata.is_cuda_graph) - if not use_custom_topk: - topk_indices_buffer[:hidden_states.shape[0]] = -1 - - if has_prefill and not metadata.skip_indexer_for_ctx_reqs: - # Use chunked prefill to reduce memory footprint - if metadata.indexer_prefill_chunks is not None: - - sparse_metadata_params = metadata.sparse_metadata_params - q_split_threshold = (sparse_metadata_params.q_split_threshold - if sparse_metadata_params is not None else - 8192) - q_split_eligible = (q_split_threshold >= 0 - and metadata.mapping is not None - and not metadata.mapping.enable_attention_dp - and metadata.mapping.tp_size > 1) - - if q_split_eligible: - tp_rank = metadata.mapping.tp_rank - tp_size = metadata.mapping.tp_size - - 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 - # same in both modes because FP4 packs four UE8M0 exponents - # 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: - # 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: - topk_indices_buffer[ - chunk.token_start:chunk.token_end, :].fill_(-1) - continue - 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, chunk.k_token_start, - num_k_tokens, gather_head_dim) - - chunk_num_token = chunk.token_end - chunk.token_start - apply_q_split = q_split_eligible and chunk_num_token >= q_split_threshold - if apply_q_split: - chunk_q_start = chunk_num_token * tp_rank // tp_size - chunk_q_end = chunk_num_token * (tp_rank + 1) // tp_size - else: - chunk_q_start = 0 - chunk_q_end = chunk_num_token - - global_q_start = chunk.token_start + chunk_q_start - global_q_end = chunk.token_start + chunk_q_end - - # Tile the query dimension so each fp8_mqa_logits call - # allocates at most [q_tile x num_k_tokens] instead of the - # full [local_q x num_k_tokens] (which can reach tens of GB - # on a long context and stall cuMemCreate under - # expandable_segments -> engine hang; see - # _INDEXER_MQA_LOGITS_ELEM_BUDGET). Results are identical: - # each query row's logits/top-k are independent and the KV - # (chunk_k_fp8) is unchanged across tiles, so the per-call - # allocation is the same size and the caching allocator - # reuses one block (peak ~= one tile, no extra sync). - local_q_len = chunk_q_end - chunk_q_start - q_tile = max( - 1, - min( - local_q_len, _INDEXER_MQA_LOGITS_ELEM_BUDGET // - max(1, num_k_tokens))) - for tile_off in range(0, local_q_len, q_tile): - c0 = chunk_q_start + tile_off - c1 = min(c0 + q_tile, chunk_q_end) - g0 = chunk.token_start + c0 - g1 = chunk.token_start + c1 - tile_q_scale = q_scale[g0:g1, - ...] if self.use_fp4 else None - logits = self._call_mqa_logits( - q_fp8[g0:g1, ...], - chunk_k_fp8, - chunk_k_scale, - weights[g0:g1, ...], - chunk.cu_seqlen_ks[c0:c1], - chunk.cu_seqlen_ke[c0:c1], - tile_q_scale, - ) - if use_custom_topk: - torch.ops.trtllm.indexer_topk_prefill( - logits, chunk.cu_seqlen_ks[c0:c1], - chunk.cu_seqlen_ke[c0:c1], - topk_indices_buffer[g0:g1, :], self.index_topk) - else: - topk_indices = logits.topk(min( - self.index_topk, logits.shape[-1]), - dim=-1)[1] - topk_indices -= chunk.cu_seqlen_ks[c0:c1][:, None] - - mask_lo = topk_indices >= 0 - mask_hi = topk_indices - ( - chunk.cu_seqlen_ke[c0:c1] - - chunk.cu_seqlen_ks[c0:c1])[:, None] < 0 - mask = mask_lo & mask_hi - - # local indices per sequence - topk_indices = topk_indices.masked_fill(~mask, -1) - - topk_indices_buffer[ - g0:g1, :topk_indices.shape[-1]] = \ - topk_indices.to(dtype=torch.int32) - - if apply_q_split: - q_sizes = [(r + 1) * chunk_num_token // tp_size - - r * chunk_num_token // tp_size - for r in range(tp_size)] - topk_indices_buffer[ - chunk.token_start:chunk.token_end, :] = allgather( - topk_indices_buffer[ - global_q_start:global_q_end, :], - metadata.mapping, - dim=0, - sizes=q_sizes) - elif metadata.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] - - ctx_q_scale = q_scale[:num_ctx_tokens, - ...] if self.use_fp4 else None - logits = self._call_mqa_logits( - q_fp8[:num_ctx_tokens, ...], - k_fp8[:num_ctx_kv_tokens, ...], - k_scale[:num_ctx_kv_tokens, ...], - weights[:num_ctx_tokens, ...], - cu_seqlen_ks, - cu_seqlen_ke, - ctx_q_scale, - ) - if use_custom_topk: - torch.ops.trtllm.indexer_topk_prefill( - logits, cu_seqlen_ks, cu_seqlen_ke, - topk_indices_buffer[:num_ctx_tokens, :], - self.index_topk) - else: - topk_indices = logits.topk(min(self.index_topk, - logits.shape[-1]), - dim=-1)[1] - topk_indices -= cu_seqlen_ks[:, None] - mask_lo = topk_indices >= 0 - mask_hi = topk_indices - (cu_seqlen_ke - - cu_seqlen_ks)[:, None] < 0 - mask = mask_lo & mask_hi - - # local indices per sequence - topk_indices = topk_indices.masked_fill(~mask, -1) - topk_indices_buffer[:num_ctx_tokens, :topk_indices. - shape[-1]] = topk_indices.to( - dtype=torch.int32) - elif has_prefill and metadata.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, :] - - # Prefill→decode GVR handoff: seed each finishing-prefill sequence's - # heuristic_prev_topk slot with its own last-context-token top-K, so - # the FIRST decode step of that sequence gets a warm-started preIdx - # (~60-75% set-overlap with the eventual decode top-K on this - # workload) instead of the all-zero / all-(-1) cold start that the - # default `heuristic_prev_topk.zero_()` initialization leaves behind. - # Without this, GVR P2 secant on decode step 0 runs from a benign - # but uninformative seed (kernel +1 offset on zeros → all indices - # point at compressed-token position 1), wasting iterations. - # Slot convention (mirrors the existing decode write-back at the - # bottom of the decode block): new gens from finishing prefill - # append after currently-active gens, i.e., slots - # [num_generations : num_generations + num_contexts]. - if (self._enable_heuristic_topk and has_prefill - and not metadata.skip_indexer_for_ctx_reqs): - local_layer = metadata.kv_cache_manager.layer_offsets[ - self.layer_idx] - ctx_seq_lens = metadata.seq_lens[:num_contexts] - # Per-sequence last context-token offset (exclusive cumsum minus 1). - last_ctx_idx = (torch.cumsum(ctx_seq_lens, dim=0) - - 1).to(dtype=torch.long) - metadata.heuristic_prev_topk[ - local_layer, num_generations:num_generations + - num_contexts].copy_(topk_indices_buffer[last_ctx_idx, :]) - - if has_decode and not metadata.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] - max_decode_len = gen_seq_lens.max().item() - min_decode_len = gen_seq_lens.min().item() - assert max_decode_len == min_decode_len, "max_decode_len != min_decode_len, we need padding" - - # Reshape q for decode phase: [num_gen_tokens, ...] -> [batch_size, next_n, ...] - q_decode = q_fp8[num_ctx_tokens:num_ctx_tokens + num_gen_tokens, - ...] - batch_size = num_generations - next_n = num_gen_tokens // num_generations - # 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 not metadata.use_expanded_buffers_for_mtp or next_n == 1: - 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[ - num_contexts:num_contexts + num_generations] - # The 2D-context_lens metadata kernel encodes next_n into the - # schedule (via num_next_n_atoms). MTP forwards alternate - # between the full-window call (next_n == 1+max_draft_tokens) - # and per-token draft calls (next_n == 1), so we must select - # the buffer that was populated for this next_n. The DSL path - # 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 - else: - scheduler_metadata_buffer = metadata.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 - - assert num_gen_tokens == batch_size * next_n - weights_decode = weights[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens, ...] - - # Get k cache and call fp8_paged_mqa_logits / fp8_fp4_paged_mqa_logits - # 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() - - if self.use_cute_dsl_paged_mqa_logits: - # 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()` - # with a (num_gen, 1) input shape, which makes DeepGEMM's wrapper - # compute `num_next_n_atoms = 1`. (DeepGEMM uses the same buffer - # for its own next_n=1 kernel; DSL piggy-backs on it for all - # real next_n values.) All next_n positions of a batch share - # 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 - assert dsl_context_lens is not None - # Wave-aware atom-split: the picker in `_pick_dsl_expand` caches - # (factor, atom) on metadata with invariant - # `factor * atom == 1 + max_draft_tokens` (the target/verify-time - # next_n). MTPEagle reuses the same metadata for its multi-step - # draft loop; after i=0 it mutates seq_lens to 1, so i≥1 - # iterations run with next_n=1. The reshape - # `(num_gen, next_n, ...) -> (num_gen*factor, atom, ...)` is only - # valid when the caller actually supplies next_n == factor * atom - # 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) - if self.use_fp4: - # FP4 DSL signature splits DG's (q, sf_q) tuple into two - # separate args and requires q.dtype == uint8 (q_decode - # came in via the FP8 plumbing as int8; reinterpret with - # no copy). sf_q is the q_scale slice reshaped to - # (B, next_n, H) int32 -- mirrors the non-DSL FP4 branch. - decode_q_scale = q_scale[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens, ...] - decode_q_scale = decode_q_scale.view( - q_decode.shape[0], q_decode.shape[1], self.n_heads) - dsl_q = q_decode.view(torch.uint8) - dsl_block_table = block_table - dsl_schedule_meta = metadata.scheduler_metadata_buffer - if dsl_atom_split: - factor = metadata.dsl_expand_factor - eff_next_n = metadata.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) - - logits_decode = torch.ops.trtllm.cute_dsl_fp4_paged_mqa_logits( - dsl_q, decode_q_scale, k_cache, weights_decode, - dsl_context_lens, dsl_block_table, dsl_schedule_meta, - indexer_max_seq_len) - else: - # FP8 DSL kernel natively supports next_n ∈ {1, 2, 3, 4}. - # Atom-split benefits small-batch / low-ntask configs by - # raising SM utilization at the cost of factorx KV HBM - # re-reads. Context lengths are indexer KV lengths, which - # may be compressed for DeepSeek-V4. - dsl_q = q_decode - fp8_ctx_lens = dsl_context_lens - fp8_block_table = block_table - fp8_schedule_meta = metadata.scheduler_metadata_buffer - if dsl_atom_split: - factor = metadata.dsl_expand_factor - atom = metadata.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) - logits_decode = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( - dsl_q, k_cache, weights_decode, fp8_ctx_lens, - fp8_block_table, fp8_schedule_meta, indexer_max_seq_len) - else: - decode_q_scale = q_scale[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens, - ...] if self.use_fp4 else None - if self.use_fp4: - # q_decode shape is either (num_generations, next_n, n_heads, - # head_dim/2) [non-expanded] or (batch*next_n, 1, n_heads, - # head_dim/2) [expanded]. Match q_scale's batch/next_n dims. - decode_q_scale = decode_q_scale.view( - q_decode.shape[0], q_decode.shape[1], self.n_heads) - logits_decode = self._call_paged_mqa_logits( - q_decode, k_cache, weights_decode, context_lens, - block_table, scheduler_metadata_buffer, indexer_max_seq_len, - decode_q_scale) - - if use_custom_topk: - # Kernel expects kv_lens (total cache length), not seq_lens (new tokens) - # This is because rowEnd = seq_len - next_n + offset + 1 - gen_kv_lens_cuda = metadata.kv_lens_cuda_runtime[ - num_contexts:num_contexts + num_generations] - - pre_idx = None - heuristic_scratch = None - if self._enable_heuristic_topk: - local_layer = metadata.kv_cache_manager.layer_offsets[ - self.layer_idx] - # Pass prev_topk directly; the +1 temporal offset is - # handled inside the C++ kernel (preIdxOffset += 1). - pre_idx = metadata.heuristic_prev_topk[ - local_layer, :num_generations] - # heuristic_scratch is only consumed by the C++ - # indexer_topk_decode path; the GVR DSL op does not take it. - # Guard on the metadata flag so this stays consistent with - # the buffer allocation (also gated on the same flag). - if not metadata.use_cute_dsl_topk: - heuristic_scratch = \ - metadata.heuristic_scratch_values[ - :num_gen_tokens] - - if self.use_cute_dsl_topk and self._enable_heuristic_topk: - # GVR DSL: supports all compress_ratio and next_n values. - torch.ops.trtllm.cute_dsl_gvr_topk_decode( - logits_decode, - pre_idx, - gen_kv_lens_cuda, - topk_indices_buffer[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens, :], - self.index_topk, - next_n=next_n, - compress_ratio=self.compress_ratio, - max_seq_len=indexer_max_seq_len, - order_row=metadata.kv_lens_row_reorder, - ) - # CuTE DSL radix top-k allocates O(num_gen_tokens * kv_len) - # global memory. Beyond 256 tokens the extra memory becomes - # significant, so we cap it at 256 and fall back to C++. - elif (self.use_cute_dsl_topk and num_gen_tokens <= 256 - and (self.compress_ratio == 1 or next_n == 1)): - torch.ops.trtllm.cute_dsl_indexer_topk_decode( - logits_decode, context_lens - if self.compress_ratio > 1 else gen_kv_lens_cuda, - topk_indices_buffer[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens, :], self.index_topk, - next_n) - else: - torch.ops.trtllm.indexer_topk_decode( - logits_decode, - gen_kv_lens_cuda, - topk_indices_buffer[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens, :], - next_n, - self.index_topk, - pre_idx=pre_idx, - heuristic_scratch=heuristic_scratch, - compress_ratio=self.compress_ratio, - radix_aux_indices=metadata.radix_aux_indices, - radix_aux_logits=metadata.radix_aux_logits) - else: - # padded - positions = torch.arange( - logits_decode.shape[-1], - device=q_decode.device).unsqueeze(0).expand( - num_gen_tokens, -1) - row_indices = torch.arange(num_gen_tokens, - device=q_decode.device) // next_n - next_n_offset = torch.arange(num_gen_tokens, - device=q_decode.device) % next_n - index_end_pos = (context_lens[row_indices] - next_n + - next_n_offset).unsqueeze(1) - # index_end_pos: [B * N, 1] - mask = positions <= index_end_pos - # mask: [B * N, L] - logits_decode = logits_decode.masked_fill(~mask, float('-inf')) - topk_indices_decode = logits_decode.topk( - min(self.index_topk, logits_decode.shape[-1]), - dim=-1)[1].to(torch.int32) # [B * N, K] - # ensure we don't set indices for the top k - # that is out of range(masked already) - # this will happen if context length is shorter than K - mask_decode = topk_indices_decode <= index_end_pos - - # local indices per sequence - topk_indices_decode = topk_indices_decode.masked_fill( - ~mask_decode, -1) - # Store in buffer - topk_indices_buffer[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens, :topk_indices_decode. - shape[-1]] = topk_indices_decode.to( - dtype=torch.int32) - - if self._enable_heuristic_topk: - local_layer = metadata.kv_cache_manager.layer_offsets[ - self.layer_idx] - decode_topk = topk_indices_buffer[ - num_ctx_tokens:num_ctx_tokens + num_gen_tokens] - last_mtp_topk = decode_topk[next_n - 1::next_n] - metadata.heuristic_prev_topk[ - local_layer, :num_generations].copy_(last_mtp_topk) - - elif has_decode and metadata.skip_indexer_for_gen_reqs: - # Fill topk_indices_buffer with pre-defined dense topk indices - topk_indices_buffer[num_ctx_tokens:num_tokens, :] = \ - metadata.topk_indices_buffer[num_ctx_tokens:num_tokens, :] - return topk_indices_buffer - - def _weight_scale(self, weights: torch.Tensor, - q_scale: torch.Tensor) -> torch.Tensor: - """Apply quantization scale to indexer attention weights.""" - weights = _scale(weights, q_scale, self.weight_scale_factor) - return weights - - def _qk_projection_and_rope(self, qr: torch.Tensor, indexer_k: torch.Tensor, - position_ids: torch.Tensor): - """Project Q/K and apply RoPE""" - q = self.wq_b(qr) - k = self.k_norm(indexer_k) - q = q.view(-1, self.n_heads, self.head_dim) - q_pe, q_nope = q.split([self.rope_dim, self.head_dim - self.rope_dim], - dim=-1) - k_pe, k_nope = k.split([self.rope_dim, self.head_dim - self.rope_dim], - dim=-1) - q_pe, k_pe = self.rotary_emb(position_ids, [q_pe, k_pe.unsqueeze(1)]) - k_pe = k_pe[:, 0, :] - return q_pe, q_nope, k_pe, k_nope - - def _prep_q_or_k(self, qk_pe: torch.Tensor, qk_nope: torch.Tensor): - """Concatenate and quantize for Q or K. - - FP8 mode: fused cat + FP8 quantize via CUDA kernel. - FP4 mode: fused cat + per-block-32 FP4 E2M1 quantize via CUDA kernel. - The returned packed bytes are int8 (two FP4 codes per byte) and the - scale is int32 (four UE8M0 exponents packed little-endian). - """ - if self.use_fp4: - return torch.ops.trtllm.fused_cat_fp4(qk_pe, qk_nope) - fp8_out, scale = torch.ops.trtllm.fused_cat_fp8( - qk_pe, qk_nope, self.scale_fmt == "ue8m0") - return fp8_out, scale - - def pre_indexer_proj( - self, qr: torch.Tensor, hidden_states: torch.Tensor, - position_ids: torch.Tensor - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, - torch.Tensor]: - """Pure token-wise projections (CUDA-graph-capturable). - - Runs cublas_mm, qk_projection_and_rope, FP8/FP4 quantize, and weight - scaling. Does NOT touch the k cache or any batch-specific metadata, - so this can safely run inside a captured CUDA graph partition. - - Returns (q_fp_bytes, k_fp_bytes, k_scale, weights, q_scale). The last - tensor is only consumed by the FP4 kernel dispatch; the FP8 path - ignores it. It is returned unconditionally so the two-op CUDA graph - split in MLA.forward_dsa_proj sees a stable signature. - """ - assert self._fused_wk_wp_weight is not None, \ - "cache_derived_state() must be called before forward()" - # When the boundary fusion pre-quantized the next layer's kv_a_proj - # NVFP4 input, hidden_states is an Fp4QuantizedTensor that also carries - # the BF16 post-RMSNorm value. Use that BF16 view here -- the indexer - # weight is BF16 and the matmul needs a float input. - if isinstance(hidden_states, Fp4QuantizedTensor): - assert hidden_states.unquantized_hidden_states is not None, ( - "pre_indexer_proj received Fp4QuantizedTensor without bf16 view; " - "the producer fusion must request return_norm_out=True") - hidden_states_bf = hidden_states.unquantized_hidden_states - else: - hidden_states_bf = hidden_states - hidden_float = _to_float(hidden_states_bf) - with _tf32_matmul_enabled(): - # F.linear computes input @ weight.T internally; no explicit .t() needed. - # _fused_wk_wp_weight is [head_dim + n_heads, hidden_size] (nn.Linear convention). - # Goes through PyTorch's cuBLAS handle which respects allow_tf32 and - # dispatches CUBLAS_COMPUTE_32F_FAST_TF32, unlike torch.ops.trtllm.cublas_mm - # which uses its own handle and always falls back to CUDA-core SGEMM. - fused_out = F.linear(hidden_float, self._fused_wk_wp_weight) - indexer_k, weights = fused_out.split([self.head_dim, self.n_heads], - dim=-1) - # Cast indexer_k back to model dtype for downstream ops (k_norm, RoPE, FP8 quantize) - indexer_k = indexer_k.to(hidden_states_bf.dtype) - - q_pe, q_nope, k_pe, k_nope = self._qk_projection_and_rope( - qr, indexer_k, position_ids) - q, k = maybe_execute_in_parallel( - lambda: self._prep_q_or_k(q_pe, q_nope), - lambda: self._prep_q_or_k(k_pe, k_nope), - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - ) - q_fp8, q_scale = q - k_fp8, k_scale = k - if self.use_fp4: - # FP4 packs two codes per byte, so the trailing dim is head_dim // 2. - # fused_cat_fp4 flattens the leading dims to M=N*n_heads; restore - # the (N, n_heads, ...) shape so downstream slicing in - # sparse_attn_indexer (which indexes by token) lines up with - # q_fp8. The DeepGEMM FP4 kernel applies the per-block q_scale - # internally, so weights carry only softmax_scale * n_heads^-0.5. - q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim // 2) - q_scale = q_scale.view(-1, self.n_heads, 1) - # DeepGEMM's fp8_fp4_(paged_)mqa_logits asserts - # `weights.scalar_type() == kFloat`. Unlike the FP8 branch (whose - # `_weight_scale` multiplies by the fp32 `q_scale` tensor and so - # implicitly upcasts), this branch's only multiplier is a Python - # float — `bf16 * float` stays bf16 and trips the kernel assert. - # Cast explicitly so this path doesn't silently rely on an - # upstream `_to_float`. - weights = weights.float() * self.weight_scale_factor - else: - q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim) - q_scale = q_scale.view(-1, self.n_heads, 1) - weights = self._weight_scale(weights, q_scale) - - return q_fp8, k_fp8, k_scale, weights, q_scale - - @torch.inference_mode() - def forward(self, qr: torch.Tensor, hidden_states: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata, - position_ids: torch.Tensor): - q_fp8, k_fp8, k_scale, weights, q_scale = self.pre_indexer_proj( - qr, hidden_states, position_ids) - - # Return topk indices buffer for sparse attention [num_tokens, index_topk] - return self.sparse_attn_indexer(metadata, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - q_scale=q_scale) - - -class DSATrtllmAttention(TrtllmAttention): - """TRT-LLM attention layer with DSA sparse indexer for MLA models.""" - - Metadata = DSAtrtllmAttentionMetadata - - def __init__(self, - layer_idx: int, - num_heads: int, - head_dim: int, - num_kv_heads: Optional[int] = None, - quant_config: Optional[QuantConfig] = None, - q_scaling: Optional[float] = None, - pos_embd_params: Optional[PositionalEmbeddingParams] = None, - mla_params: Optional[MLAParams] = None, - skip_create_weights_in_init: bool = False, - attention_chunk_size: Optional[int] = None, - sparse_params: Optional[DSAParams] = None, - dtype: Optional[torch.dtype] = None, - aux_stream: Optional[torch.cuda.Stream] = None, - **kwargs): - """Initialize DSA attention with an Indexer sub-module for sparse TopK selection.""" - sparse_attention_config = kwargs.pop("sparse_attention_config", None) - self.sparse_attention_config = sparse_attention_config - if (sparse_params is None and sparse_attention_config is not None - and hasattr(sparse_attention_config, "to_sparse_params")): - sparse_params = sparse_attention_config.to_sparse_params( - layer_idx=layer_idx) - if sparse_params is None: - raise ValueError( - "sparse_params is required for DSATrtllmAttention and cannot be None" - ) - TrtllmAttention.__init__( - self, - layer_idx, - num_heads, - head_dim, - sparse_params=sparse_params, - num_kv_heads=num_kv_heads, - quant_config=quant_config, - q_scaling=q_scaling, - pos_embd_params=pos_embd_params, - mla_params=mla_params, - skip_create_weights_in_init=skip_create_weights_in_init, - attention_chunk_size=attention_chunk_size, - **kwargs) - - # Cross-layer indexer sharing: only "full" layers own an indexer; - # "shared" layers reuse the previous full layer's top-k (see - # MLA.forward_dsa_*). Resolved per-layer in to_sparse_params; defaults to - # full (dense per-layer indexer). indexer=None also makes the weight - # loader skip the (absent) shared-layer indexer weights. - self.is_full_indexer_layer = getattr(sparse_params, - 'is_full_indexer_layer', True) - if self.is_full_indexer_layer: - self.indexer = Indexer(quant_config, - pos_embd_params, - mla_params, - skip_create_weights_in_init, - sparse_params, - dtype=dtype, - layer_idx=layer_idx, - aux_stream=aux_stream) - else: - self.indexer = None - - def sparse_attn_predict( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - metadata: DSAtrtllmAttentionMetadata, - forward_args: AttentionForwardArgs, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """Transform local TopK indices to global paged KV cache indices.""" - # Transform the local topk indices to global topk indices in paged kv cache - is_generation = forward_args.attention_input_type == AttentionInputType.generation_only - topk_indices_global, _ = transform_local_topk_and_prepare_pool_view( - forward_args.topk_indices, metadata, - self.get_local_layer_idx(metadata), is_generation) - - # TODO: Use sparse_attn_indexer to predict the indices for DSA attention - # return self.indexer(q, k, metadata, hidden_states, qr, position_ids) - return topk_indices_global, None - - def sparse_kv_predict( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - metadata: DSAtrtllmAttentionMetadata, - forward_args: AttentionForwardArgs, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """No-op KV prediction; DSA uses indexer-based selection instead.""" - return None, None - - def mla_rope_append_paged_kv_assign_q( - self, - q: torch.Tensor, - latent_cache: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata, - is_generation: bool = False, - **kwargs, - ) -> None: - """Apply RoPE, append latent cache to paged KV, and assign query for MLA.""" - if is_generation: - cached_token_indptr = metadata.gen_cached_token_indptr - kv_indptr = metadata.gen_kv_indptr - num_seqs = metadata.num_generations - max_seq_len = metadata.max_gen_seq_len - block_offsets = metadata.kv_cache_block_offsets[:, metadata. - num_contexts:] - else: - cached_token_indptr = metadata.ctx_cached_token_indptr - kv_indptr = metadata.ctx_kv_indptr - num_seqs = metadata.num_contexts - max_seq_len = metadata.max_ctx_seq_len - block_offsets = metadata.kv_cache_block_offsets - assert self.is_mla_enable and self.mla_params is not None - assert metadata.kv_cache_manager is not None - - beam_width = 1 - - torch.ops.trtllm.mla_rope_append_paged_kv_assign_q( - q, - latent_cache, - num_seqs, - cached_token_indptr, - kv_indptr, - max_seq_len, - self.rotary_cos_sin, - self.num_heads, - self.mla_params.qk_nope_head_dim, - self.mla_params.qk_rope_head_dim, - self.mla_params.kv_lora_rank, - block_offsets, - metadata.kv_cache_manager.kv_cache_pool_pointers, - metadata.kv_cache_manager.kv_cache_pool_mapping, - None, # kv_scale_orig_quant - self.get_local_layer_idx(metadata), - metadata.kv_cache_manager.tokens_per_block, - metadata.kv_cache_manager.max_seq_len, - beam_width, - self.quant_mode, - ) - - -class DSACacheManager(KVCacheManager): - """KV cache manager for DSA with additional indexer K-cache pools.""" - - def __init__( - self, - 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, - max_num_tokens: int = 8192, - model_config: Optional[ModelConfig] = None, - max_beam_width: int = 1, - sparse_attention_config: Optional["SparseAttentionConfig"] = None, - pretrained_config=None, - **kwargs, - ) -> None: - """Initialize cache manager with indexer K-cache pool per layer.""" - if sparse_attention_config is None: - sparse_attention_config = kwargs.pop("sparse_attn_config", None) - if sparse_attention_config is None and model_config is not None: - sparse_attention_config = model_config.sparse_attention_config - if sparse_attention_config is None: - raise ValueError( - "sparse_attention_config is required for DSA cache") - sparse_params = sparse_attention_config.to_sparse_params( - pretrained_config=pretrained_config) - if not isinstance(sparse_params, DSAParams): - raise ValueError("DSA cache requires DSA sparse parameters") - self.quant_block_size = 128 - self.index_head_dim = sparse_params.index_head_dim - # FP4 mode packs the indexer K cache as head_dim/2 data bytes + 4 - # scale bytes (vs. head_dim + 4 for FP8). The C++ WindowBlockManager - # allocates the pool with this smaller stride when the flag is set. - self.use_fp4 = sparse_params.indexer_k_dtype == "fp4" - - 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=layer_mask, - max_num_tokens=max_num_tokens, - model_config=model_config, - max_beam_width=max_beam_width, - enable_indexer_k_cache=True, - indexer_k_cache_quant_block_size=128, - indexer_k_cache_index_head_dim=self.index_head_dim, - indexer_k_cache_use_fp4=self.use_fp4, - **kwargs, - ) - self.num_blocks = self.blocks_in_primary_pool - - # Indexer K cache pool for DSA attention - # Shape: [num_blocks, self.tokens_per_block * (index_head_dim + scale_size)] - # Non-interleaved layout: [fp8_tok0 | fp8_tok1 | ... | scale_tok0 | scale_tok1 | ...] - # Store FP8-quantized k values from the indexer - self.indexer_k_cache_pool_per_layer = [ - self.get_indexer_k_cache_pool_data(layer_idx) - for layer_idx in range(self.num_local_layers) - ] - - def get_indexer_k_cache_buffers(self, layer_idx: int): - """Get indexer k cache buffer from a specific layer pool.""" - block_size = self.tokens_per_block - data_bytes = self.index_head_dim // 2 if self.use_fp4 else self.index_head_dim - per_token_size = data_bytes + self.index_head_dim // self.quant_block_size * 4 - layer_offset = self.layer_offsets[layer_idx] - return self.indexer_k_cache_pool_per_layer[layer_offset].view( - self.num_blocks, block_size, 1, per_token_size) - - def get_batch_indexer_k_cache_indices( - self, request_ids: List[int]) -> List[List[int]]: - """ - Get the indices for the indexer k cache for a specific batch of requests. - """ - # All of layers share the same cache indices, so we use layer index 0. - return self.get_batch_cache_indices(request_ids, 0) - - def shutdown(self): - """Release indexer K-cache pool references before C++ buffer cleanup.""" - # Clear Python references BEFORE C++ frees the underlying CUDA buffers - self.indexer_k_cache_pool_per_layer = [] - super().shutdown() - - @staticmethod - def get_cache_size_per_token(model_config: ModelConfig, - mapping: Mapping, - num_layers: Optional[int] = None, - **kwargs): - """Estimate total cache bytes per token including indexer K-cache overhead.""" - config = model_config.pretrained_config - sparse_attention_config = model_config.sparse_attention_config - if sparse_attention_config is None: - raise ValueError( - "sparse_attention_config is required for DSA cache") - sparse_params = sparse_attention_config.to_sparse_params( - pretrained_config=model_config.pretrained_config) - if not isinstance(sparse_params, DSAParams): - raise ValueError("DSA cache requires DSA sparse parameters") - index_head_dim = sparse_params.index_head_dim - quant_block_size = 128 - # Under FP4 the indexer stores two E2M1 codes per byte, so the - # per-token data footprint halves (132 B -> 68 B at index_head_dim=128); - # the scale bytes are unchanged (4 per token, one int32 holding four - # UE8M0 exponents at quant_block_size=32 after packing). - use_fp4 = sparse_params.indexer_k_dtype == "fp4" - indexer_data_dim = index_head_dim // 2 if use_fp4 else index_head_dim - - # get kv cache dtype bytes - mem_per_token = 2 - quant_config = model_config.quant_config - if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache( - ): - mem_per_token = 1 - - # get head dim - head_dim = config.kv_lora_rank + config.qk_rope_head_dim - - num_attention_layers = KVCacheManager._resolve_num_attention_layers( - model_config, mapping, num_layers) - # MLA latent K cache: stored at the KV cache dtype (BF16/FP8). - mem_per_token *= num_attention_layers * head_dim - - # Indexer K cache: physically allocated as raw UINT8 in - # WindowBlockManager::allocatePools (poolDtype = kUINT8), so we assume - # 1 byte/element here -- it is NOT scaled by the KV cache dtype (unlike - # the latent above). The data-portion byte count already reflects fp8 vs - # fp4 via indexer_data_dim. - indexer_bytes_per_token = num_attention_layers * ( - indexer_data_dim + index_head_dim // quant_block_size * 4) - mem_per_token += indexer_bytes_per_token - return mem_per_token - - def get_cache_bytes_per_token(self): - """Compute actual cache bytes per token from instance configuration.""" - # MLA latent K cache: stored at the KV cache dtype (self.dtype). The - # indexer K cache is added separately below. - cache_size_per_token = math.ceil( - self.kv_factor * sum(self.num_kv_heads_per_layer) * self.head_dim) - - if self.dtype not in (DataType.FP8, DataType.HALF, DataType.BF16, - DataType.FLOAT, DataType.NVFP4): - raise ValueError(f'Cannot support {self.dtype} KV cache.') - - cache_size_bytes_per_token = get_size_in_bytes(cache_size_per_token, - self.dtype) - if self.dtype == DataType.NVFP4: - cache_size_bytes_per_token += self.calculate_scaling_factor_size_bytes( - cache_size_per_token, - quant_vector_size=16, - scaling_factor_dtype=DataType.FP8) - - # Indexer K cache: physically allocated as raw UINT8 in - # WindowBlockManager::allocatePools (poolDtype = kUINT8), so we assume - # 1 byte/element here -- it is NOT scaled by the KV cache dtype (unlike - # the latent above). Under FP4 the indexer data portion is halved (two - # E2M1 codes per byte); the scale bytes are unchanged. - indexer_data_dim = self.index_head_dim // 2 if self.use_fp4 else self.index_head_dim - indexer_bytes_per_token = sum(self.num_kv_heads_per_layer) * ( - indexer_data_dim + self.index_head_dim // self.quant_block_size * 4) - cache_size_bytes_per_token += indexer_bytes_per_token - - return cache_size_bytes_per_token diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py new file mode 100644 index 000000000000..44fd65d65e50 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""DeepSeek Sparse Attention backend package.""" + +from .backend import DSATrtllmAttention +from .cache_manager import DSACacheManager +from .indexer import ( + _DG_SCHEDULE_BLOCK_KV, + HAS_FAST_HADAMARD, + Indexer, + IndexerParams, + IndexerPrefillChunkMetadata, + RotaryEmbedding, + _compute_slot_mappings, + _effective_compress_ratio_divisor, + _pick_dsl_expand, + _select_indexer_compress_ratio, + compute_cu_seqlen_kv_bounds_with_cache, + rotate_activation, + split_prefill_chunks, + transform_local_topk_and_prepare_pool_view, + warmup_heuristic_topk_decode, +) +from .metadata import DSAtrtllmAttentionMetadata +from .params import DSABackendForwardArgs, DSAMetadataParams, DSAParams + +__all__ = [ + "HAS_FAST_HADAMARD", + "DSABackendForwardArgs", + "DSACacheManager", + "DSAMetadataParams", + "DSAParams", + "DSATrtllmAttention", + "DSAtrtllmAttentionMetadata", + "Indexer", + "IndexerParams", + "IndexerPrefillChunkMetadata", + "RotaryEmbedding", + "_DG_SCHEDULE_BLOCK_KV", + "_compute_slot_mappings", + "_effective_compress_ratio_divisor", + "_pick_dsl_expand", + "_select_indexer_compress_ratio", + "compute_cu_seqlen_kv_bounds_with_cache", + "rotate_activation", + "split_prefill_chunks", + "transform_local_topk_and_prepare_pool_view", + "warmup_heuristic_topk_decode", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/backend.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/backend.py new file mode 100644 index 000000000000..9a38a29969c0 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/backend.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dense Sparse Attention (DSA) backend for TRT-LLM with indexer-based TopK selection.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionForwardArgs, + AttentionInputType, + MLAParams, + PositionalEmbeddingParams, +) +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm.models.modeling_utils import QuantConfig + +from .indexer import Indexer, transform_local_topk_and_prepare_pool_view +from .metadata import DSAtrtllmAttentionMetadata +from .params import DSAParams + +ModelConfig = tensorrt_llm.bindings.ModelConfig + + +class DSATrtllmAttention(TrtllmAttention): + """TRT-LLM attention layer with DSA sparse indexer for MLA models.""" + + Metadata = DSAtrtllmAttentionMetadata + + def __init__( + self, + layer_idx: int, + num_heads: int, + head_dim: int, + num_kv_heads: Optional[int] = None, + quant_config: Optional[QuantConfig] = None, + q_scaling: Optional[float] = None, + pos_embd_params: Optional[PositionalEmbeddingParams] = None, + mla_params: Optional[MLAParams] = None, + skip_create_weights_in_init: bool = False, + attention_chunk_size: Optional[int] = None, + sparse_params: Optional[DSAParams] = None, + dtype: Optional[torch.dtype] = None, + aux_stream: Optional[torch.cuda.Stream] = None, + **kwargs, + ): + """Initialize DSA attention with an Indexer sub-module for sparse TopK selection.""" + sparse_attention_config = kwargs.pop("sparse_attention_config", None) + self.sparse_attention_config = sparse_attention_config + if ( + sparse_params is None + and sparse_attention_config is not None + and hasattr(sparse_attention_config, "to_sparse_params") + ): + sparse_params = sparse_attention_config.to_sparse_params(layer_idx=layer_idx) + if sparse_params is None: + raise ValueError("sparse_params is required for DSATrtllmAttention and cannot be None") + TrtllmAttention.__init__( + self, + layer_idx, + num_heads, + head_dim, + sparse_params=sparse_params, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + q_scaling=q_scaling, + pos_embd_params=pos_embd_params, + mla_params=mla_params, + skip_create_weights_in_init=skip_create_weights_in_init, + attention_chunk_size=attention_chunk_size, + **kwargs, + ) + + # Cross-layer indexer sharing: only "full" layers own an indexer; + # "shared" layers reuse the previous full layer's top-k (see + # MLA.forward_dsa_*). Resolved per-layer in to_sparse_params; defaults to + # full (dense per-layer indexer). indexer=None also makes the weight + # loader skip the (absent) shared-layer indexer weights. + self.is_full_indexer_layer = getattr(sparse_params, "is_full_indexer_layer", True) + if self.is_full_indexer_layer: + self.indexer = Indexer( + quant_config, + pos_embd_params, + mla_params, + skip_create_weights_in_init, + sparse_params, + dtype=dtype, + layer_idx=layer_idx, + aux_stream=aux_stream, + ) + else: + self.indexer = None + + def sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: DSAtrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Run the DSA indexer and convert its local TopK to paged indices.""" + 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 + if self.indexer is None: + topk_indices = shared_topk_indices[phase_start:phase_end] + else: + indexer_intermediates = sparse_backend_args.indexer_intermediates + topk_indices = self.indexer.forward_from_projected( + metadata, + q, + indexer_intermediates, + is_generation=is_generation, + ) + if shared_topk_indices is not None: + shared_topk_indices[ + phase_start : phase_start + topk_indices.shape[0], + : topk_indices.shape[1], + ].copy_(topk_indices) + + topk_indices_global, _ = transform_local_topk_and_prepare_pool_view( + topk_indices, metadata, self.get_local_layer_idx(metadata), is_generation + ) + + return topk_indices_global, None + + def sparse_kv_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: DSAtrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """No-op KV prediction; DSA uses indexer-based selection instead.""" + return None, None + + def mla_rope_append_paged_kv_assign_q( + self, + q: torch.Tensor, + latent_cache: torch.Tensor, + metadata: DSAtrtllmAttentionMetadata, + is_generation: bool = False, + **kwargs, + ) -> None: + """Apply RoPE, append latent cache to paged KV, and assign query for MLA.""" + if is_generation: + cached_token_indptr = metadata.gen_cached_token_indptr + kv_indptr = metadata.gen_kv_indptr + num_seqs = metadata.num_generations + max_seq_len = metadata.max_gen_seq_len + block_offsets = metadata.kv_cache_block_offsets[:, metadata.num_contexts :] + else: + cached_token_indptr = metadata.ctx_cached_token_indptr + kv_indptr = metadata.ctx_kv_indptr + num_seqs = metadata.num_contexts + max_seq_len = metadata.max_ctx_seq_len + block_offsets = metadata.kv_cache_block_offsets + assert self.is_mla_enable and self.mla_params is not None + assert metadata.kv_cache_manager is not None + + beam_width = 1 + + torch.ops.trtllm.mla_rope_append_paged_kv_assign_q( + q, + latent_cache, + num_seqs, + cached_token_indptr, + kv_indptr, + max_seq_len, + self.rotary_cos_sin, + self.num_heads, + self.mla_params.qk_nope_head_dim, + self.mla_params.qk_rope_head_dim, + self.mla_params.kv_lora_rank, + block_offsets, + metadata.kv_cache_manager.kv_cache_pool_pointers, + metadata.kv_cache_manager.kv_cache_pool_mapping, + None, # kv_scale_orig_quant + self.get_local_layer_idx(metadata), + metadata.kv_cache_manager.tokens_per_block, + metadata.kv_cache_manager.max_seq_len, + beam_width, + self.quant_mode, + ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py new file mode 100644 index 000000000000..6539c48f39da --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dense Sparse Attention (DSA) backend for TRT-LLM with indexer-based TopK selection.""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, List, Optional, Union + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm._utils import get_size_in_bytes +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.mapping import Mapping + +from .params import DSAParams + +ModelConfig = tensorrt_llm.bindings.ModelConfig + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, SparseAttentionConfig + + +class DSACacheManager(KVCacheManager): + """KV cache manager for DSA with additional indexer K-cache pools.""" + + def __init__( + self, + 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, + max_num_tokens: int = 8192, + model_config: Optional[ModelConfig] = None, + max_beam_width: int = 1, + sparse_attention_config: Optional["SparseAttentionConfig"] = None, + pretrained_config=None, + **kwargs, + ) -> None: + """Initialize cache manager with indexer K-cache pool per layer.""" + if sparse_attention_config is None: + sparse_attention_config = kwargs.pop("sparse_attn_config", None) + if sparse_attention_config is None and model_config is not None: + sparse_attention_config = model_config.sparse_attention_config + if sparse_attention_config is None: + raise ValueError("sparse_attention_config is required for DSA cache") + sparse_params = sparse_attention_config.to_sparse_params( + pretrained_config=pretrained_config + ) + if not isinstance(sparse_params, DSAParams): + raise ValueError("DSA cache requires DSA sparse parameters") + self.quant_block_size = 128 + self.index_head_dim = sparse_params.index_head_dim + # FP4 mode packs the indexer K cache as head_dim/2 data bytes + 4 + # scale bytes (vs. head_dim + 4 for FP8). The C++ WindowBlockManager + # allocates the pool with this smaller stride when the flag is set. + self.use_fp4 = sparse_params.indexer_k_dtype == "fp4" + + 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=layer_mask, + max_num_tokens=max_num_tokens, + model_config=model_config, + max_beam_width=max_beam_width, + enable_indexer_k_cache=True, + indexer_k_cache_quant_block_size=128, + indexer_k_cache_index_head_dim=self.index_head_dim, + indexer_k_cache_use_fp4=self.use_fp4, + **kwargs, + ) + self.num_blocks = self.blocks_in_primary_pool + + # Indexer K cache pool for DSA attention + # Shape: [num_blocks, self.tokens_per_block * (index_head_dim + scale_size)] + # Non-interleaved layout: [fp8_tok0 | fp8_tok1 | ... | scale_tok0 | scale_tok1 | ...] + # Store FP8-quantized k values from the indexer + self.indexer_k_cache_pool_per_layer = [ + self.get_indexer_k_cache_pool_data(layer_idx) + for layer_idx in range(self.num_local_layers) + ] + + def get_indexer_k_cache_buffers(self, layer_idx: int): + """Get indexer k cache buffer from a specific layer pool.""" + block_size = self.tokens_per_block + data_bytes = self.index_head_dim // 2 if self.use_fp4 else self.index_head_dim + per_token_size = data_bytes + self.index_head_dim // self.quant_block_size * 4 + layer_offset = self.layer_offsets[layer_idx] + return self.indexer_k_cache_pool_per_layer[layer_offset].view( + self.num_blocks, block_size, 1, per_token_size + ) + + def get_batch_indexer_k_cache_indices(self, request_ids: List[int]) -> List[List[int]]: + """ + Get the indices for the indexer k cache for a specific batch of requests. + """ + # All of layers share the same cache indices, so we use layer index 0. + return self.get_batch_cache_indices(request_ids, 0) + + def shutdown(self): + """Release indexer K-cache pool references before C++ buffer cleanup.""" + # Clear Python references BEFORE C++ frees the underlying CUDA buffers + self.indexer_k_cache_pool_per_layer = [] + super().shutdown() + + @staticmethod + def get_cache_size_per_token( + model_config: ModelConfig, mapping: Mapping, num_layers: Optional[int] = None, **kwargs + ): + """Estimate total cache bytes per token including indexer K-cache overhead.""" + config = model_config.pretrained_config + sparse_attention_config = model_config.sparse_attention_config + if sparse_attention_config is None: + raise ValueError("sparse_attention_config is required for DSA cache") + sparse_params = sparse_attention_config.to_sparse_params( + pretrained_config=model_config.pretrained_config + ) + if not isinstance(sparse_params, DSAParams): + raise ValueError("DSA cache requires DSA sparse parameters") + index_head_dim = sparse_params.index_head_dim + quant_block_size = 128 + # Under FP4 the indexer stores two E2M1 codes per byte, so the + # per-token data footprint halves (132 B -> 68 B at index_head_dim=128); + # the scale bytes are unchanged (4 per token, one int32 holding four + # UE8M0 exponents at quant_block_size=32 after packing). + use_fp4 = sparse_params.indexer_k_dtype == "fp4" + indexer_data_dim = index_head_dim // 2 if use_fp4 else index_head_dim + + # get kv cache dtype bytes + mem_per_token = 2 + quant_config = model_config.quant_config + if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache(): + mem_per_token = 1 + + # get head dim + head_dim = config.kv_lora_rank + config.qk_rope_head_dim + + num_attention_layers = KVCacheManager._resolve_num_attention_layers( + model_config, mapping, num_layers + ) + # MLA latent K cache: stored at the KV cache dtype (BF16/FP8). + mem_per_token *= num_attention_layers * head_dim + + # Indexer K cache: physically allocated as raw UINT8 in + # WindowBlockManager::allocatePools (poolDtype = kUINT8), so we assume + # 1 byte/element here -- it is NOT scaled by the KV cache dtype (unlike + # the latent above). The data-portion byte count already reflects fp8 vs + # fp4 via indexer_data_dim. + indexer_bytes_per_token = num_attention_layers * ( + indexer_data_dim + index_head_dim // quant_block_size * 4 + ) + mem_per_token += indexer_bytes_per_token + return mem_per_token + + def get_cache_bytes_per_token(self): + """Compute actual cache bytes per token from instance configuration.""" + # MLA latent K cache: stored at the KV cache dtype (self.dtype). The + # indexer K cache is added separately below. + cache_size_per_token = math.ceil( + self.kv_factor * sum(self.num_kv_heads_per_layer) * self.head_dim + ) + + if self.dtype not in ( + DataType.FP8, + DataType.HALF, + DataType.BF16, + DataType.FLOAT, + DataType.NVFP4, + ): + raise ValueError(f"Cannot support {self.dtype} KV cache.") + + cache_size_bytes_per_token = get_size_in_bytes(cache_size_per_token, self.dtype) + if self.dtype == DataType.NVFP4: + cache_size_bytes_per_token += self.calculate_scaling_factor_size_bytes( + cache_size_per_token, quant_vector_size=16, scaling_factor_dtype=DataType.FP8 + ) + + # Indexer K cache: physically allocated as raw UINT8 in + # WindowBlockManager::allocatePools (poolDtype = kUINT8), so we assume + # 1 byte/element here -- it is NOT scaled by the KV cache dtype (unlike + # the latent above). Under FP4 the indexer data portion is halved (two + # E2M1 codes per byte); the scale bytes are unchanged. + indexer_data_dim = self.index_head_dim // 2 if self.use_fp4 else self.index_head_dim + indexer_bytes_per_token = sum(self.num_kv_heads_per_layer) * ( + indexer_data_dim + self.index_head_dim // self.quant_block_size * 4 + ) + cache_size_bytes_per_token += indexer_bytes_per_token + + return cache_size_bytes_per_token diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py new file mode 100644 index 000000000000..67c912fba56a --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""DSA custom ops used by piecewise CUDA graph capture.""" + +from typing import List, Optional + +import torch + +from tensorrt_llm._torch.utils import Fp4QuantizedTensor + +from .module import _forward_dsa_attn, forward_dsa_proj + + +def _extract_extra_attrs(layer_idx: str): + from tensorrt_llm._torch.modules.attention import extract_extra_attrs + + return extract_extra_attrs(layer_idx, "mla") + + +@torch.library.custom_op("trtllm::mla_dsa_proj", mutates_args=()) +def mla_dsa_proj( + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + layer_idx: str, + hidden_states_fp4: Optional[torch.Tensor] = None, + hidden_states_sf: Optional[torch.Tensor] = None, +) -> List[torch.Tensor]: + """Token-wise projections for DSA MLA (CUDA-graph-capturable). + + Runs kv_a_proj, layernorms, q_b_proj, and conditionally + indexer.pre_indexer_proj (FP8/FP4 quantize, weight scaling). Does NOT + update the indexer k cache — that happens in Op 2 (mla_dsa_attn_inplace) + because the scatter kernel accesses batch-specific metadata. + + Returns [q, compressed_kv, k_pe, latent_cache] when the short-MHA path + handles all tokens, or [q, compressed_kv, k_pe, latent_cache, q_fp8, + k_fp8, k_scale, weights, q_scale] when the indexer runs. Under torch + compile, _should_use_short_mha returns False so the result is always + length 9, keeping control flow straight-line for CUDA graph capture. + The trailing q_scale is only consumed by the FP4 dispatch; the FP8 + path ignores it in _forward_dsa_attn. + """ + metadata, mla_layer = _extract_extra_attrs(layer_idx) + if hidden_states_fp4 is not None or hidden_states_sf is not None: + assert hidden_states_fp4 is not None and hidden_states_sf is not None, ( + "hidden_states_fp4 and hidden_states_sf must be passed together" + ) + hidden_states = Fp4QuantizedTensor( + fp4_tensor=hidden_states_fp4, + scaling_factor=hidden_states_sf, + unquantized_hidden_states=hidden_states, + ) + return forward_dsa_proj(mla_layer, position_ids, hidden_states, metadata) + + +@mla_dsa_proj.register_fake +def _mla_dsa_proj_fake( + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + layer_idx: str, + hidden_states_fp4: Optional[torch.Tensor] = None, + hidden_states_sf: Optional[torch.Tensor] = None, +) -> List[torch.Tensor]: + # Under torch compile _should_use_short_mha is False, so the result is + # always 9 tensors (4 attention inputs + 5 indexer intermediates, with + # q_scale as the 9th carried for the FP4 dispatch). + metadata, mla_layer = _extract_extra_attrs(layer_idx) + num_tokens = hidden_states.shape[0] + indexer = mla_layer.mqa.indexer + q = hidden_states.new_empty([num_tokens, mla_layer.num_heads_tp * mla_layer.qk_head_dim]) + compressed_kv = hidden_states.new_empty([num_tokens, mla_layer.kv_lora_rank]) + k_pe = hidden_states.new_empty([num_tokens, mla_layer.qk_rope_head_dim]) + latent_cache = hidden_states.new_empty( + [num_tokens, mla_layer.kv_lora_rank + mla_layer.qk_rope_head_dim] + ) + if indexer is None: + # DSA "shared" layer: no indexer, mirror forward_dsa_proj's early + # return of only the 4 base tensors (no indexer intermediates). + return [q, compressed_kv, k_pe, latent_cache] + # Indexer intermediates: q_fp8, k_fp8, k_scale, weights, q_scale. + # Under FP4 q_fp8's trailing dim is head_dim // 2 (two E2M1 codes per + # byte) and q_scale carries one int32 per (token, head) packing four + # UE8M0 exponents; under FP8 q_fp8's trailing dim is head_dim and + # q_scale carries one float32 per (token, head). + if indexer.use_fp4: + q_fp8 = hidden_states.new_empty( + [num_tokens, indexer.n_heads, indexer.head_dim // 2], dtype=torch.int8 + ) + k_fp8 = hidden_states.new_empty([num_tokens, indexer.head_dim // 2], dtype=torch.int8) + k_scale = hidden_states.new_empty([num_tokens, 1], dtype=torch.int32) + q_scale = hidden_states.new_empty([num_tokens, indexer.n_heads, 1], dtype=torch.int32) + else: + q_fp8 = hidden_states.new_empty( + [num_tokens, indexer.n_heads, indexer.head_dim], dtype=torch.float8_e4m3fn + ) + k_fp8 = hidden_states.new_empty([num_tokens, indexer.head_dim], dtype=torch.float8_e4m3fn) + k_scale = hidden_states.new_empty([num_tokens, 1], dtype=torch.float32) + q_scale = hidden_states.new_empty([num_tokens, indexer.n_heads, 1], dtype=torch.float32) + weights = hidden_states.new_empty([num_tokens, indexer.n_heads], dtype=torch.float32) + return [q, compressed_kv, k_pe, latent_cache, q_fp8, k_fp8, k_scale, weights, q_scale] + + +@torch.library.custom_op("trtllm::mla_dsa_attn_inplace", mutates_args=("output",)) +def mla_dsa_attn_inplace( + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + latent_cache: torch.Tensor, + indexer_intermediates: List[torch.Tensor], + position_ids: Optional[torch.Tensor], + layer_idx: str, + output: torch.Tensor, +) -> None: + """Batch-structure-dependent attention dispatch for DSA MLA. + + indexer_intermediates is [q_fp8, k_fp8, k_scale, weights, q_scale] when + the indexer ran in Op 1, or [] when short-MHA handled all tokens. The + trailing q_scale is only consumed by the FP4 dispatch; the FP8 path + ignores it. Runs sparse_attn_indexer then dispatches context/generation + attention. This op is excluded from CUDA graph capture. + """ + metadata, mla_layer = _extract_extra_attrs(layer_idx) + _forward_dsa_attn( + mla_layer, + q, + compressed_kv, + k_pe, + latent_cache, + indexer_intermediates, + position_ids, + metadata, + output, + ) + + +@mla_dsa_attn_inplace.register_fake +def _mla_dsa_attn_inplace_fake( + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + latent_cache: torch.Tensor, + indexer_intermediates: List[torch.Tensor], + position_ids: Optional[torch.Tensor], + layer_idx: str, + output: torch.Tensor, +) -> None: + """Model the in-place output mutation during fake-tensor propagation.""" diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py new file mode 100644 index 000000000000..0f03dd5df8e7 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py @@ -0,0 +1,2006 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dense Sparse Attention (DSA) backend for TRT-LLM with indexer-based TopK selection.""" + +from __future__ import annotations + +import os +import threading +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Optional, Set, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.attention_backend.interface import MLAParams, PositionalEmbeddingParams +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.distributed.ops import allgather +from tensorrt_llm._torch.modules.layer_norm import LayerNorm +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding +from tensorrt_llm._torch.utils import Fp4QuantizedTensor, maybe_compile +from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned +from tensorrt_llm.deep_gemm import ( + fp8_fp4_mqa_logits, + fp8_fp4_paged_mqa_logits, + fp8_mqa_logits, + fp8_paged_mqa_logits, + get_paged_mqa_logits_metadata, +) +from tensorrt_llm.logger import logger +from tensorrt_llm.models.modeling_utils import QuantConfig + +from .params import DSAParams + +ModelConfig = tensorrt_llm.bindings.ModelConfig + +# Cap the per-call indexer MQA-logits transient (in elements). fp8_mqa_logits +# allocates its [q x kv] logits output via torch.empty; the KV dimension is the +# full (compressed) context and is unbounded, so for a large query chunk on a +# long-context prefill this single allocation can reach tens of GB. Under +# PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True such an allocation can stall +# indefinitely in cuMemCreate on the longest-context (attention_dp laggard) rank +# -> GPU idle -> peers block at the next MoE all-to-all -> watchdog hang. Tiling +# the query dimension caps the transient to q_tile x kv with identical results +# (each query row's logits/top-k are independent). Override via env if needed. +_INDEXER_MQA_LOGITS_ELEM_BUDGET = int( + os.environ.get("TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET", 1 << 31) +) + +if TYPE_CHECKING: + from .metadata import DSAtrtllmAttentionMetadata + +# Optional import: fast-hadamard-transform causes CI build issues (requires wheel+torch pre-installed) +try: + from fast_hadamard_transform import hadamard_transform + + HAS_FAST_HADAMARD = True +except ImportError: + hadamard_transform = None + HAS_FAST_HADAMARD = False + +# Idempotency guard for warmup_heuristic_topk_decode — keyed by +# (device_index, top_k, hint_size, num_cols). Prevents repeated allocations +# and synchronizations when multiple Indexer modules invoke the warmup with +# the same parameters during model construction. +_HEURISTIC_TOPK_WARMUP_DONE: Set[Tuple[int, int, int, int]] = set() +_HEURISTIC_TOPK_WARMUP_LOCK = threading.Lock() +_DG_SCHEDULE_BLOCK_KV = 64 + + +def warmup_heuristic_topk_decode( + top_k: int = 2048, hint_size: int = 2048, num_cols: int = 4096 +) -> None: + """Pre-initialize cached hardware attributes in the C++ Scheme X dispatcher. + + The dispatcher inside ``invokeIndexerTopKDecode`` lazily queries + ``cudaDeviceGetAttribute`` for ``MultiProcessorCount`` and + ``L2CacheSize`` on its first call. Those host-side queries must not + be issued during ``cudaStreamBeginCapture / EndCapture``: the values + captured there become frozen into the graph and cannot be refreshed + across replays on a different device. + + This warmup issues one small heuristic decode call so the static + caches are populated before any CUDA Graph capture begins. Must be + called from the Indexer setup hook (``layer_idx == 0``) when + ``enable_heuristic_topk`` is true. + + Repeated invocations with the same ``(device, top_k, hint_size, + num_cols)`` key are short-circuited so that constructing many Indexer + modules in the same process does not re-allocate scratch tensors or + issue redundant synchronizations. + """ + key = (torch.cuda.current_device(), top_k, hint_size, num_cols) + with _HEURISTIC_TOPK_WARMUP_LOCK: + if key in _HEURISTIC_TOPK_WARMUP_DONE: + return + _HEURISTIC_TOPK_WARMUP_DONE.add(key) + + device = torch.device("cuda") + logits = torch.zeros((1, num_cols), dtype=torch.float32, device=device) + seq_lens = torch.tensor([num_cols], dtype=torch.int32, device=device) + indices = torch.empty((1, top_k), dtype=torch.int32, device=device) + pre_idx = torch.zeros((1, hint_size), dtype=torch.int32, device=device) + scratch = torch.empty((top_k,), dtype=torch.float32, device=device) + # The default warmup geometry (num_cols=4096) falls below kSeqSmall=12288 + # and routes to the Radix path with blocks_per_row=2 (num_rows=1 sweeps + # bp ∈ [2, maxByCols=2]). The cpp op rejects blocks_per_row > 1 without + # caller-owned radix aux scratch, so supply worst-case (kMaxBlocksPerRowDecode=10) + # buffers here. Cost is negligible (~80 KB) and the warmup is a one-shot. + _radix_max_bp = 10 + radix_aux_indices = torch.empty((1, _radix_max_bp, top_k), dtype=torch.int32, device=device) + radix_aux_logits = torch.empty((1, _radix_max_bp, top_k), dtype=torch.float32, device=device) + torch.ops.trtllm.indexer_topk_decode( + logits, + seq_lens, + indices, + 1, + top_k, + pre_idx=pre_idx, + heuristic_scratch=scratch, + radix_aux_indices=radix_aux_indices, + radix_aux_logits=radix_aux_logits, + ) + torch.cuda.synchronize() + + +def _pick_dsl_expand( + next_n: int, + num_sms: int, + batch_size: int = 0, + max_ctx: int = 0, + kernel_atoms: Tuple[int, ...] = (1, 2, 3), +) -> Tuple[int, int]: + """Pick (expand_factor, effective_next_n) for the DSL paged kernel + using a wave-aware strategy. Used by both FP4 and FP8 DSL paths. + + The DSL kernel natively supports ``effective_next_n ∈ kernel_atoms`` + (FP4: ``(1, 2, 3)``; FP8: ``(1, 2, 3, 4)``). For ``next_n`` not natively + supported or when SM utilization can be improved, reshape + ``[B, next_n, ...]`` -> ``[B * expand_factor, effective_next_n, ...]`` + caller-side. + + Strategy: enumerate ``(expand_factor, effective_next_n)`` pairs with + ``expand_factor * effective_next_n == next_n`` and ``effective_next_n + in kernel_atoms``. Score each by ``(waves, -expand_factor)`` where + ``waves = ceil(B * expand_factor * ceil(max_ctx/256) / num_sms)``. + Pick min waves; on tie, prefer LARGER expand_factor (more SMs busy per + wave; pays HBM cost of expand_factor x KV re-reads). + + When ``batch_size == 0`` or ``max_ctx == 0`` (workload unknown), fall + back to the legacy HBM-minimizing heuristic: largest effective_next_n + that divides next_n cleanly (still constrained to ``kernel_atoms``). + + Examples (wave-aware, num_sms=148 [B200], SPLIT_KV=256 tokens): + FP4, next_n=4, B=1, ctx=4096 -> (4, 1): ntask=64<148, 1 wave, max factor + FP4, next_n=4, B=32, ctx=4096 -> (2, 2): ntask=1024>148, multi-wave, min factor + FP4, next_n=2, B=1, ctx=4096 -> (2, 1): wave-tie, larger factor + FP8, next_n=4, B=1, ctx=4096 -> (4, 1): kernel_atoms incl. 4 doesn't change small-B pick + """ + # Legacy fallback when workload is unknown. + if batch_size <= 0 or max_ctx <= 0: + for eff in sorted(kernel_atoms, reverse=True): + if next_n % eff == 0: + return next_n // eff, eff + return next_n, 1 + + SPLIT_KV_TOKENS = 256 + cands = [] + for eff in kernel_atoms: + if next_n % eff == 0: + factor = next_n // eff + ntask = batch_size * factor * ((max_ctx + SPLIT_KV_TOKENS - 1) // SPLIT_KV_TOKENS) + waves = (ntask + num_sms - 1) // num_sms + cands.append((waves, factor, eff)) + if not cands: + return next_n, 1 + cands.sort(key=lambda x: (x[0], -x[1])) # min waves, max factor + _, factor, eff = cands[0] + return factor, eff + + +def _compute_slot_mappings( + global_positions: torch.Tensor, + block_offsets: torch.Tensor, + req_indices: torch.Tensor, + head_dim: int, + tokens_per_block: int, + quant_block_size: int, + data_bytes_per_token: Optional[int] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute flat byte indices for FP8/FP4 data and scales from global token positions. + + Shared by Indexer.prepare() (CPU) and on_update_kv_lens() (GPU) to avoid + duplicating the slot mapping arithmetic. + + Args: + global_positions: Per-token absolute position in the KV sequence. + block_offsets: [num_seqs, max_blocks_per_seq] block offset table. + req_indices: Per-token request index. + head_dim: Indexer head dimension (logical element count). + tokens_per_block: Tokens stored per cache block. + quant_block_size: Quantization block size. + data_bytes_per_token: Bytes of quantized data per token in the cache + (default: head_dim, i.e. one byte per FP8 value). For MXFP4 this + should be ``head_dim // 2`` since two E2M1 codes pack into one byte. + The scale layout is identical at head_dim=128: 4 bytes per token + (one float32 for FP8, four packed UE8M0 exponents for FP4). + + Returns: + (fp8_indices, scale_indices): Flat byte offsets into the cache pool. + """ + if data_bytes_per_token is None: + data_bytes_per_token = head_dim + scale_size = head_dim // quant_block_size * 4 # float32 = 4 bytes + block_stride = tokens_per_block * (data_bytes_per_token + scale_size) + scale_base_offset = tokens_per_block * data_bytes_per_token + + block_indices_in_seq = global_positions // tokens_per_block + pos_in_blocks = global_positions % tokens_per_block + + max_blocks = block_offsets.shape[1] + if block_indices_in_seq.is_cuda: + # Clamp to prevent OOB from stale token-to-seq mappings during + # CUDA graph capture/replay with MTP + DSA. + block_indices_in_seq = block_indices_in_seq.clamp(0, max_blocks - 1) + else: + 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()}" + ) + + block_ids = block_offsets[req_indices, block_indices_in_seq].to(torch.int64) + + fp8_indices = block_ids * block_stride + pos_in_blocks * data_bytes_per_token + scale_indices = block_ids * block_stride + scale_base_offset + pos_in_blocks * scale_size + return fp8_indices, scale_indices + + +def _unravel_indices( + flat_indices: torch.Tensor, shape: Tuple[int, ...] +) -> Tuple[torch.Tensor, ...]: + """ + Unravel indices into multiple dimensions. + """ + d3 = shape[3] + i3 = flat_indices % d3 + flat_indices = flat_indices // d3 + d2 = shape[2] + i2 = flat_indices % d2 + flat_indices = flat_indices // d2 + d1 = shape[1] + i1 = flat_indices % d1 + flat_indices = flat_indices // d1 + i0 = flat_indices + return i0, i1, i2, i3 + + +def rotate_activation(x: torch.Tensor) -> torch.Tensor: + """Apply Hadamard rotation to activation tensor for DSA sparse attention.""" + assert x.dtype == torch.bfloat16 + + if not HAS_FAST_HADAMARD: + # Fallback: skip transformation (acceptable for test/dev) + logger.warning_once( + "fast-hadamard-transform not available. Sparse MLA will skip " + "hadamard transformation. Install with: " + "pip install git+https://github.com/Dao-AILab/fast-hadamard-transform.git", + key="fast_hadamard_import_missing", + ) + return x + + hidden_size = x.size(-1) + assert (hidden_size & (hidden_size - 1)) == 0, ( + "Hidden size must be a power of 2 for Hadamard transform." + ) + return hadamard_transform(x, scale=hidden_size**-0.5) + + +def transform_local_topk_and_prepare_pool_view( + topk_indices: torch.Tensor, + attn_metadata: "DSAtrtllmAttentionMetadata", + layer_idx: int, + is_generation: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Convert local topk indices to global pool indices and prepare KV pool. + + Uses cached values from attn_metadata._ensure_pool_view_cached() + to avoid redundant Python/CUDA overhead across layers. + """ + assert topk_indices.dtype == torch.int32 + + attn_metadata._ensure_pool_view_cached() + + if is_generation: + block_table = attn_metadata._cached_block_table_gen + req_idx = attn_metadata._cached_req_idx_gen + else: + block_table = attn_metadata._cached_block_table_ctx + req_idx = attn_metadata._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, + topk_indices.shape[1], + attn_metadata._cached_stride_factor, + layer_idx, + ) + + return global_indices, attn_metadata._cached_pool_view + + +def split_prefill_chunks( + seq_lens: torch.Tensor, + max_chunk_size: int, + start_idx: int = 0, +) -> List[List[Tuple[int, int, int, int]]]: + """ + Split prefill requests into chunks based on max_chunk_size. + Supports two-level chunking: + 1. Request-boundary chunking: group multiple small requests into one chunk + 2. Intra-request chunking: split large requests into multiple Q-block chunks + + Args: + seq_lens: Sequence lengths for all requests + max_chunk_size: Maximum number of tokens per chunk + start_idx: Starting index for prefill requests + + Returns: + List of chunk groups, where each group is a list of chunk specs. + Each chunk spec is (req_idx, token_start_in_req, token_end_in_req, req_cum_start) + + - For multi-request chunks: group contains multiple specs (one per request) + - For intra-request chunks: each Q-block is a separate group with single spec + """ + chunk_groups = [] + num_reqs = len(seq_lens) + + current_req = start_idx + # Compute cumulative token positions + query_start_loc_cpu = torch.cat( + [torch.zeros(1, dtype=torch.int32, device="cpu"), seq_lens.cumsum(dim=0).to(torch.int32)] + ) + + while current_req < num_reqs: + seq_len = seq_lens[current_req].item() + req_cum_start = query_start_loc_cpu[current_req].item() + + if seq_len <= max_chunk_size: + # This request fits in one chunk - try to pack with others + current_size = seq_len + chunk_specs = [(current_req, 0, seq_len, req_cum_start)] + next_req = current_req + 1 + + # Try to add more requests to this chunk + while next_req < num_reqs: + next_seq_len = seq_lens[next_req].item() + if next_seq_len > max_chunk_size: + # Next request is large, stop packing + break + if current_size + next_seq_len <= max_chunk_size: + next_cum_start = query_start_loc_cpu[next_req].item() + chunk_specs.append((next_req, 0, next_seq_len, next_cum_start)) + current_size += next_seq_len + next_req += 1 + else: + break + + # Add as one multi-request chunk group + chunk_groups.append(chunk_specs) + current_req = next_req + else: + # Large request - split into Q-blocks + # Each Q-block is a separate chunk group (processed in separate iteration) + num_q_blocks = (seq_len + max_chunk_size - 1) // max_chunk_size + for q_block_idx in range(num_q_blocks): + token_start = q_block_idx * max_chunk_size + token_end = min(token_start + max_chunk_size, seq_len) + q_block_spec = [(current_req, token_start, token_end, req_cum_start)] + chunk_groups.append(q_block_spec) + + current_req += 1 + + return chunk_groups + + +# Shrink the indexer prefill chunk size for very long requests to bound the +# fp8(_fp4)_mqa_logits activation memory (~ chunk_size * K_compressed), keyed on +# the largest compressed KV length in the batch. Entries are +# (k_compressed_lower_bound_exclusive, chunk_size); >512K -> 8K, [256K, 512K] +# -> 16K, otherwise unchanged. +_INDEXER_CHUNK_SIZE_HEURISTIC = ( + (512 * 1024, 8 * 1024), + (256 * 1024 - 1, 16 * 1024), +) + + +def select_indexer_chunk_size(configured_chunk_size: int, max_k_compressed: int) -> int: + """Pick the indexer prefill chunk size from the batch's largest K_compressed. + + Only reduces ``configured_chunk_size`` (never increases it). + """ + for threshold, chunk_size in _INDEXER_CHUNK_SIZE_HEURISTIC: + if max_k_compressed > threshold: + return min(configured_chunk_size, chunk_size) + return configured_chunk_size + + +def _select_indexer_compress_ratio(compress_ratios: List[int]) -> int: + if 4 in compress_ratios: + return 4 + if 1 in compress_ratios: + return 1 + return 0 + + +def _effective_compress_ratio_divisor(compress_ratio: int) -> int: + return compress_ratio if compress_ratio > 1 else 1 + + +def compute_cu_seqlen_kv_bounds_with_cache( + seq_lens: torch.Tensor, + num_contexts: int, + num_ctx_tokens: int, + cached_token_lens: Optional[torch.Tensor] = None, + kv_lens: Optional[torch.Tensor] = None, + compress_ratio: int = 1, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Compute attention window bounds for batched sequences with causal attention, + accounting for cached KV tokens. + + Args: + seq_lens: current token lengths [num_contexts], dtype=torch.int32 + num_contexts: Number of sequences in the batch + num_ctx_tokens: Total number of context tokens across all sequences in current batch + cached_token_lens: Cached KV token lengths [num_contexts], dtype=torch.int32 (optional) + kv_lens: KV token lengths [num_contexts], dtype=torch.int32 (optional) + compress_ratio: Compression ratio for KV tokens + + Returns: + cu_seqlen_ks: Start index in KV for each Q token [num_ctx_tokens] + cu_seqlen_ke: End index (exclusive) in KV for each Q token [num_ctx_tokens] + """ + device = seq_lens.device + # Total KV lengths per request + if kv_lens is None: + kv_lens = ( + seq_lens if cached_token_lens is None else cached_token_lens + seq_lens + ) # [num_contexts] + kv_lens = kv_lens // compress_ratio + + # Cumulative KV offsets: where each request's KV sequence starts in global KV space + cu_kv_offsets = torch.cat( + [ + torch.zeros(1, device=device, dtype=torch.int32), + torch.cumsum(kv_lens, dim=0).to(torch.int32), + ] + ) # [num_contexts + 1] + + # Map each Q token to its request: [0,0,...,0, 1,1,...,1, ..., B-1,B-1,...,B-1] + batch_ids = torch.repeat_interleave( + torch.arange(num_contexts, device=device, dtype=torch.int32), seq_lens + ) # [num_ctx_tokens] + + # Each Q token's KV window starts at its request's KV sequence start + cu_seqlen_ks = cu_kv_offsets[batch_ids] # [num_ctx_tokens] + + # Compute local Q position within each request (0-based, relative to current batch context tokens) + cu_q_offsets = torch.cat( + [ + torch.zeros(1, device=device, dtype=torch.int32), + torch.cumsum(seq_lens, dim=0).to(torch.int32), + ] + ) # [num_contexts + 1] + + global_q_positions = torch.arange(num_ctx_tokens, device=device, dtype=torch.int32) + local_q_positions = global_q_positions - torch.repeat_interleave( + cu_q_offsets[:-1], seq_lens + ) # [num_ctx_tokens] + + if cached_token_lens is not None: + cached_per_token = torch.repeat_interleave(cached_token_lens, seq_lens) # [num_ctx_tokens] + cu_seqlen_ke = ( + cu_seqlen_ks + (cached_per_token + local_q_positions + 1) // compress_ratio + ) # [num_ctx_tokens] + else: + cu_seqlen_ke = cu_seqlen_ks + (local_q_positions + 1) // compress_ratio # [num_ctx_tokens] + + return cu_seqlen_ks, cu_seqlen_ke + + +@dataclass +class IndexerPrefillChunkMetadata: + """Metadata for a single prefill chunk in the indexer""" + + cu_seqlen_ks: torch.Tensor # Attention window start for each token + cu_seqlen_ke: torch.Tensor # Attention window end for each token + token_start: int # Q token start index in batch + token_end: int # Q token end index in batch + k_token_start: int # K token start index in batch + k_token_end: int # K token end index in batch + + +@maybe_compile(dynamic=True) +def _scale(weights: torch.Tensor, q_scale: torch.Tensor, s: float) -> torch.Tensor: + """Scale attention weights by quantization scale and constant factor.""" + return weights * q_scale.squeeze(-1) * s + + +@maybe_compile(dynamic=True) +def _to_float(hidden_states: torch.Tensor) -> torch.Tensor: + """Cast hidden states to float32 for TF32 GEMM computation.""" + return hidden_states.float() + + +@contextmanager +def _tf32_matmul_enabled(): + """Temporarily enable TF32 tensor cores for FP32 matmul in this scope. + + Forces PyTorch/cuBLASLt to use CUBLAS_COMPUTE_32F_FAST_TF32, which + guarantees TF32 tensor cores. Plain CUBLAS_COMPUTE_32F (used by + torch.ops.trtllm.cublas_mm) falls back to SIMT SGEMM on CUDA cores + based on cuBLASLt heuristics for small M. + """ + prev = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = True + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = prev + + +@dataclass +class IndexerParams: + """ + Parameters for indexer. + """ + + num_contexts: int + num_generations: int + num_ctx_tokens: int + head_dim: int + quant_block_size: int + tokens_per_block: int + compress_ratio: int + request_ids: List[int] + num_past_tokens: List[int] + seq_lens: torch.Tensor + # Bytes of quantized data per token in the indexer K cache; defaults to + # head_dim (one FP8 byte per element). For MXFP4 use head_dim // 2. + data_bytes_per_token: Optional[int] = None + + def __post_init__(self): + # Pre-compute frequently used tensors once instead of on every property access + num_past_tokens_tensor = torch.tensor(self.num_past_tokens, dtype=torch.int32) + self._num_past_tokens_tensor = num_past_tokens_tensor + compress_ratio = self.compress_ratio + self._cached_kv_tokens = num_past_tokens_tensor // compress_ratio + self._all_kv_tokens = (self.seq_lens + num_past_tokens_tensor) // compress_ratio + self._new_kv_tokens = self._all_kv_tokens - self._cached_kv_tokens + self._kv_lens = self._all_kv_tokens + if self.data_bytes_per_token is None: + self.data_bytes_per_token = self.head_dim + self._scale_size = self.head_dim // self.quant_block_size * 4 + self._block_stride = self.tokens_per_block * (self.data_bytes_per_token + self._scale_size) + + @property + def batch_size(self): + return len(self.request_ids) + + @property + def kv_lens(self): + return self._kv_lens + + @property + def total_tokens(self): + return self.seq_lens.sum().item() + + @property + def cached_kv_tokens(self): + return self._cached_kv_tokens + + @property + def all_kv_tokens(self): + return self._all_kv_tokens + + @property + def new_kv_tokens(self): + return self._new_kv_tokens + + @property + def scale_size(self): + return self._scale_size + + @property + def block_stride(self): + return self._block_stride + + +class Indexer(nn.Module): + """DSA sparse attention indexer that selects top-K KV cache entries per token.""" + + def __init__( + self, + quant_config: Optional[QuantConfig], + pos_embd_params: Optional[PositionalEmbeddingParams], + mla_params: Optional[MLAParams], + skip_create_weights_in_init: bool, + sparse_params: DSAParams, + dtype: Optional[torch.dtype], + compress_ratio: int = 1, + layer_idx: int = 0, + aux_stream: Optional[torch.cuda.Stream] = None, + ): + """Initialize indexer with projection weights, norms, and TopK configuration.""" + super().__init__() + self.hidden_size = mla_params.hidden_size + self.q_lora_rank = mla_params.q_lora_rank + self.rope_dim = mla_params.qk_rope_head_dim + self.n_heads = sparse_params.index_n_heads # 64 + self.head_dim = sparse_params.index_head_dim # 128 + self.index_topk = sparse_params.index_topk # 2048 + self.layer_idx = layer_idx + self.compress_ratio = compress_ratio + + self.wq_b = Linear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=skip_create_weights_in_init, + use_custom_cublas_mm=True, + ) + self.wk = Linear( + self.hidden_size, + self.head_dim, + bias=False, + dtype=torch.float32, + quant_config=None, + skip_create_weights_in_init=skip_create_weights_in_init, + use_custom_cublas_mm=True, + ) + self.k_norm = LayerNorm(hidden_size=self.head_dim, eps=1e-6) + self.weights_proj = Linear( + self.hidden_size, + self.n_heads, + bias=False, + dtype=torch.float32, + quant_config=None, + skip_create_weights_in_init=skip_create_weights_in_init, + use_custom_cublas_mm=True, + ) + + # Fused wk + weights_proj weight for single F.linear FP32 GEMM under allow_tf32. + # Maps to TF32 tensor cores on Ampere+. + self._fused_wk_wp_weight: Optional[torch.Tensor] = None + + indexer_rope_interleave = sparse_params.indexer_rope_interleave + self.rotary_emb = RotaryEmbedding( + pos_embd_params.rope, + head_dim=self.rope_dim, + is_neox=not indexer_rope_interleave, + ) + + self.softmax_scale = self.head_dim**-0.5 + # TODO: make it configurable from hf config + self.scale_fmt = "ue8m0" + # indexer_k_dtype controls both Q and K precision. DeepGEMM's + # fp8_fp4_mqa_logits / fp8_fp4_paged_mqa_logits kernels only dispatch + # to FP4xFP4 or FP8xFP8 (no mixed-precision variant). The DeepGEMM + # kernel asserts SM100 + head_dim=128 at launch time under FP4. + self.use_fp4 = sparse_params.indexer_k_dtype == "fp4" + self.aux_stream = aux_stream + self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] + self.use_cute_dsl_topk = sparse_params.use_cute_dsl_topk and IS_CUTLASS_DSL_AVAILABLE + self.use_cute_dsl_paged_mqa_logits = ( + sparse_params.use_cute_dsl_paged_mqa_logits and IS_CUTLASS_DSL_AVAILABLE + ) + self.weight_scale_factor = self.softmax_scale * self.n_heads**-0.5 + + self._enable_heuristic_topk = ( + sparse_params.enable_heuristic_topk and get_sm_version() >= 100 + ) + + if (self.use_cute_dsl_topk or self.use_cute_dsl_paged_mqa_logits) and layer_idx == 0: + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + if self.use_cute_dsl_topk and not self._enable_heuristic_topk: + # the dtype of topk input tensor, which is float32 now. + # Note, need to update it if the dtype of topk input tensor is changed. + cute_dsl_custom_ops.warmup_cute_dsl_indexer_topk( + dtype=torch.float32, top_k=self.index_topk + ) + + if self._enable_heuristic_topk and layer_idx == 0: + # Populate static caches (sm_count, L2 cache size) inside the C++ + # Scheme X dispatcher before any CUDA Graph capture so the host + # attribute queries do not end up frozen into a captured graph. + warmup_heuristic_topk_decode(top_k=self.index_topk) + + # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM + # (populated in cache_derived_state; maps to TF32 tensor cores on Ampere+) + self._fused_wk_wp_weight: Optional[torch.Tensor] = None + + def cache_derived_state(self) -> None: + """Fuse wk and weights_proj for F.linear with TF32 tensor cores on Ampere+.""" + # wk: [head_dim, hidden_size] + weights_proj: [n_heads, hidden_size] + # → fused: [head_dim + n_heads, hidden_size] + self._fused_wk_wp_weight = torch.cat( + [self.wk.weight.data, self.weights_proj.weight.data], dim=0 + ) + + def post_load_weights(self) -> None: + self.cache_derived_state() + + @staticmethod + def prepare_one_prefill_chunk( + metadata: DSAtrtllmAttentionMetadata, + chunk_specs: List[Tuple[int, int, int, int]], + ) -> IndexerPrefillChunkMetadata: + """ + Build metadata for one prefill chunk for indexer forward pass. + Handles both multi-request chunks and intra-request Q-block chunks. + + Args: + metadata: Attention metadata + chunk_specs: List of (req_idx, token_start_in_req, token_end_in_req, req_cum_start) + - token_start_in_req, token_end_in_req are indices into current batch context tokens + - For multi-request: multiple specs from different requests (full requests) + - For intra-request: single spec from one request's Q-block + + Note: Cached token counts are derived from metadata.host_ctx_cached_token_indptr + """ + device = metadata.cu_seqlen_ks.device + compress_ratio = _effective_compress_ratio_divisor( + _select_indexer_compress_ratio(metadata.compress_ratios) + ) + if len(chunk_specs) == 1: + # Single request or intra-request Q-block + req_idx, token_start_in_req, token_end_in_req, req_cum_start = chunk_specs[0] + num_q_tokens = token_end_in_req - token_start_in_req + + # Get cached token count for this request from metadata + num_cached = ( + metadata.host_ctx_cached_token_indptr[req_idx + 1] + - metadata.host_ctx_cached_token_indptr[req_idx] + ).item() + + # Total compressed KV tokens for this request + req_kv_len = (num_cached + token_end_in_req) // compress_ratio + + # For intra-request chunks: Q block attends to all previous K in the request + # Q tokens [token_start_in_req:token_end_in_req] within the request's current tokens + # K tokens [0:req_kv_len] in compressed KV space + cu_seqlen_ks = torch.zeros(num_q_tokens, dtype=torch.int32, device="cpu") + cu_seqlen_ke = ( + torch.arange( + token_start_in_req + 1, token_end_in_req + 1, dtype=torch.int32, device="cpu" + ) + + num_cached + ) // compress_ratio + + # Q token range in batch (indices into context tokens in the current batch) + token_start = req_cum_start + token_start_in_req + token_end = req_cum_start + token_end_in_req + + # K token range: index into full KV slot mapping in compressed KV space + # For req_idx=0 the offset is 0; for other indices compute compressed cumulative offset + kv_offset_in_extended = sum( + (metadata.host_ctx_kv_indptr[j + 1] - metadata.host_ctx_kv_indptr[j]).item() + // compress_ratio + for j in range(req_idx) + ) + total_kv_for_req = req_kv_len + k_token_start = kv_offset_in_extended + k_token_end = kv_offset_in_extended + total_kv_for_req + + else: + # Multi-request chunk: batch multiple full requests together + # Extract sequence lengths for these requests + req_seq_lens = [] + req_num_past_tokens = [] + req_kv_lens = [] + first_req_idx = chunk_specs[0][0] + + for spec in chunk_specs: + req_idx, token_start_in_req, token_end_in_req, _ = spec + req_seq_lens.append(token_end_in_req - token_start_in_req) + # Get cached token count from metadata + num_past_tokens = ( + metadata.host_ctx_cached_token_indptr[req_idx + 1] + - metadata.host_ctx_cached_token_indptr[req_idx] + ).item() + req_num_past_tokens.append(num_past_tokens) + req_kv_lens.append((num_past_tokens + req_seq_lens[-1]) // compress_ratio) + + req_seq_lens_tensor = torch.tensor(req_seq_lens, dtype=torch.int32, device="cpu") + req_num_past_tokens_tensor = torch.tensor( + req_num_past_tokens, dtype=torch.int32, device="cpu" + ) + req_kv_lens_tensor = torch.tensor(req_kv_lens, dtype=torch.int32, device="cpu") + num_q_tokens = sum(req_seq_lens) + + # Compute causal attention bounds for batched requests + cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( + req_seq_lens_tensor, + len(chunk_specs), + num_q_tokens, + req_num_past_tokens_tensor, + req_kv_lens_tensor, + compress_ratio, + ) + + # Global Q token ranges (indices into ctx tokens in the current batch) + token_start = chunk_specs[0][3] # req_cum_start of first request + token_end = token_start + num_q_tokens + + # K token range: index into full kv slot mapping (cached + current ctx tokens within the batch) + # Must use compressed offsets + kv_offset_in_extended = sum( + (metadata.host_ctx_kv_indptr[j + 1] - metadata.host_ctx_kv_indptr[j]).item() + // compress_ratio + for j in range(first_req_idx) + ) + total_kv_len = sum(req_kv_lens) + k_token_start = kv_offset_in_extended + k_token_end = kv_offset_in_extended + total_kv_len + + assert cu_seqlen_ks.shape[0] == num_q_tokens == token_end - token_start, ( + "Indexer.prepare_one_prefill_chunk - cu_seqlen_ks length mismatch: " + f"{cu_seqlen_ks.shape[0]} != {num_q_tokens}" + ) + assert cu_seqlen_ke.shape[0] == num_q_tokens == token_end - token_start, ( + "Indexer.prepare_one_prefill_chunk - cu_seqlen_ke length mismatch: " + f"{cu_seqlen_ke.shape[0]} != {num_q_tokens}" + ) + + return IndexerPrefillChunkMetadata( + cu_seqlen_ks=maybe_pin_memory(cu_seqlen_ks).to(device, non_blocking=True), + cu_seqlen_ke=maybe_pin_memory(cu_seqlen_ke).to(device, non_blocking=True), + token_start=token_start, + token_end=token_end, + k_token_start=k_token_start, + k_token_end=k_token_end, + ) + + @staticmethod + def build_indexer_params(metadata: DSAtrtllmAttentionMetadata) -> Optional[IndexerParams]: + 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 + 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) + ) + 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, + compress_ratio=compress_ratio, + request_ids=metadata.request_ids, + num_past_tokens=metadata.kv_cache_params.num_cached_tokens_per_seq, + seq_lens=metadata.seq_lens, + data_bytes_per_token=data_bytes_per_token, + ) + + @staticmethod + def recompute_slot_mappings( + metadata: DSAtrtllmAttentionMetadata, indexer_params: Optional[IndexerParams] = None + ): + """Recompute slot_mapping_fp8/scale from current block offsets. + + This is the subset of prepare_for_update_k_cache() that maps each new + compressed KV token to its flat cache position. It is safe to call in + isolation after a caller has swapped the active KV cache manager and + block-offset buffers, such as during draft KV-cache replay. + """ + if indexer_params is None: + indexer_params = Indexer.build_indexer_params(metadata) + if indexer_params is None: + return + + batch_size = indexer_params.batch_size + tokens_per_block = indexer_params.tokens_per_block + head_dim = indexer_params.head_dim + new_kv_tokens = indexer_params.new_kv_tokens + total_new_kv_tokens = new_kv_tokens.sum().item() + data_bytes_per_token = indexer_params.data_bytes_per_token + + # Compute global positions for all kv tokens in the batch (fully vectorized) + req_indices = torch.repeat_interleave( + torch.arange(batch_size, dtype=torch.int64, device="cpu"), new_kv_tokens + ) + # Vectorized token_offsets: arange(total) - cumulative start per request + cu_new_kv = torch.zeros(batch_size + 1, dtype=torch.int64, device="cpu") + cu_new_kv[1:] = new_kv_tokens.to(torch.int64).cumsum(0) + token_offsets = torch.arange( + total_new_kv_tokens, dtype=torch.int64, device="cpu" + ) - cu_new_kv[:-1].repeat_interleave(new_kv_tokens) + global_positions = indexer_params.cached_kv_tokens[req_indices] + token_offsets + + # 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] + 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, + req_indices, + head_dim, + tokens_per_block, + indexer_params.quant_block_size, + 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 + + metadata.slot_mapping_fp8[:total_new_kv_tokens].copy_( + metadata.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 + ) + + @staticmethod + def prepare_for_update_k_cache( + metadata: DSAtrtllmAttentionMetadata, indexer_params: IndexerParams + ): + """ + Prepare indexer for the update_k_cache stage. + + Compute slot_mapping for all requests (both context and generation) + This maps each token to its flat cache position for vectorized KV cache updates + """ + Indexer.recompute_slot_mappings(metadata, indexer_params) + + @staticmethod + def prepare_for_chunked_prefill( + metadata: DSAtrtllmAttentionMetadata, indexer_params: IndexerParams + ): + """ + Prepare indexer for the chunked prefill. + """ + num_contexts = indexer_params.num_contexts + seq_lens = indexer_params.seq_lens + tokens_per_block = indexer_params.tokens_per_block + head_dim = indexer_params.head_dim + + # When MLA chunked prefill is active, it already handles chunking + # Indexer should just process the current MLA chunk as a single chunk + has_mla_chunked_prefill = ( + metadata.enable_context_mla_with_cached_kv and metadata.runtime_features.chunked_prefill + ) + if has_mla_chunked_prefill: + chunk_specs = [ + (i, 0, seq_lens[i].item(), seq_lens[:i].sum().item() if i > 0 else 0) + for i in range(num_contexts) + ] + metadata.indexer_prefill_chunks = [ + Indexer.prepare_one_prefill_chunk( + metadata, + chunk_specs, + ) + ] + else: + # Use indexer's chunking to avoid quadratic indexer MQA logits + # computation for long sequences. + # This is only used when MLA chunked prefill is not enabled. + # Adapt chunk size to the batch's largest K_compressed (see + # select_indexer_chunk_size). + max_k_compressed = ( + 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 + ) + chunk_groups = split_prefill_chunks( + seq_lens[:num_contexts], + effective_chunk_size, + start_idx=0, + ) + + if len(chunk_groups) > 1 or metadata.enable_context_mla_with_cached_kv: + metadata.indexer_prefill_chunks = [ + Indexer.prepare_one_prefill_chunk( + metadata, + chunk_specs, + ) + for chunk_specs in chunk_groups + ] + else: + metadata.indexer_prefill_chunks = None + + # Chunked prefill and KV-cache reuse require the full KV for indexer + # logits. The indexer's own chunking gathers only the current chunk. + if metadata.enable_context_mla_with_cached_kv: + # Use kv_lens which correctly computes (raw_past + seq_lens) // compress_ratio. + total_kv_per_request = indexer_params.kv_lens[:num_contexts] + total_kv_len = total_kv_per_request.sum().item() + host_slot_mapping_fp8_fullkv = torch.empty( + total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned() + ) + host_slot_mapping_scale_fullkv = torch.empty( + total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned() + ) + + req_indices = torch.repeat_interleave( + torch.arange(num_contexts, dtype=torch.int64, device="cpu"), total_kv_per_request + ) + + cu_kv = torch.zeros(num_contexts + 1, dtype=torch.int64, device="cpu") + cu_kv[1:] = total_kv_per_request.to(torch.int64).cumsum(0) + kv_positions = torch.arange(total_kv_len, dtype=torch.int64, device="cpu") - cu_kv[ + :-1 + ].repeat_interleave(total_kv_per_request) + + fp8_flat_indices, scale_flat_indices = _compute_slot_mappings( + kv_positions, + metadata.host_indexer_k_cache_block_offsets, + req_indices, + head_dim, + tokens_per_block, + indexer_params.quant_block_size, + data_bytes_per_token=head_dim // 2 + if metadata.kv_cache_manager.use_fp4 + else head_dim, + ) + + host_slot_mapping_fp8_fullkv[:total_kv_len] = fp8_flat_indices + host_slot_mapping_scale_fullkv[:total_kv_len] = scale_flat_indices + + assert len(fp8_flat_indices) == total_kv_len, ( + "host_slot_mapping_fp8_fullkv/host_slot_mapping_scale_fullkv " + f"length mismatch: {len(fp8_flat_indices)} != total_kv_len={total_kv_len}" + ) + + # Store extended mappings for indexer full KV gathering + 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 + ) + else: + metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 + metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + + @staticmethod + def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): + """ + Prepare scheduler metadata for the DeepGEMM decode MQA kernel. + """ + num_contexts = metadata.num_contexts + num_generations = metadata.num_generations + if not metadata.use_expanded_buffers_for_mtp: + gen_seq_lens = metadata.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 + next_n_cap = metadata.kv_lens_cuda_2d.shape[1] + metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_( + gen_seq_lens.unsqueeze(-1).expand(-1, next_n_cap) + ) + scheduler_metadata_buffer = get_paged_mqa_logits_metadata( + gen_seq_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, metadata.num_sms + ) + metadata.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, non_blocking=True) + if metadata.max_draft_tokens > 0: + scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( + metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap], + _DG_SCHEDULE_BLOCK_KV, + metadata.num_sms, + ) + metadata.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 = metadata.num_generations * (1 + metadata.max_draft_tokens) + scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( + metadata.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), + _DG_SCHEDULE_BLOCK_KV, + metadata.num_sms, + ) + metadata.scheduler_metadata_buffer_expanded.copy_( + scheduler_metadata_buffer_expanded, non_blocking=True + ) + + @staticmethod + def prepare(metadata: DSAtrtllmAttentionMetadata): + """ + Prepare indexer for the forward pass. + This should be called during metadata.prepare() stage. + + - Computes slot_mapping for KV cache updates + - Prepares schedule_metadata for fp8_paged_mqa_logits + - Stores generation request IDs for decode phase + """ + + # Skip indexer preparation if the kv_cache_manager doesn't have index_head_dim. + # 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. + indexer_params = Indexer.build_indexer_params(metadata) + if indexer_params is None: + return + + num_contexts = metadata.num_contexts + num_generations = metadata.num_generations + 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) + ) + # Store compressed KV token count for context requests + metadata.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) + + # Prepare for prefill phase if there are context requests + if num_contexts > 0: + # Compute attention window bounds for each query token in batched sequences + # cu_seqlen_ks[i]: start index in global KV for query token i + # cu_seqlen_ke[i]: end index (exclusive) in global KV for query token i + host_seq_lens = seq_lens[:num_contexts] + host_num_past_tokens = indexer_params._num_past_tokens_tensor[:num_contexts] + host_kv_lens = indexer_params.kv_lens[:num_contexts] + host_cu_seqlen_ks, host_cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( + host_seq_lens, + num_contexts, + num_ctx_tokens, + host_num_past_tokens, + host_kv_lens, + compress_ratio, + ) + + metadata.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_( + maybe_pin_memory(host_cu_seqlen_ke), non_blocking=True + ) + Indexer.prepare_for_chunked_prefill(metadata, indexer_params) + + # Prepare for decode phase if there are generation requests + if num_generations > 0: + # Prepare schedule metadata for fp8_paged_mqa_logits + # This is a preprocessing step that computes scheduling information for the kernel + Indexer.prepare_scheduler_metadata(metadata) + + def _update_k_cache( + self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: DSAtrtllmAttentionMetadata + ) -> None: + """ + Insert/append k values and scales into the indexer k cache using pre-computed slot mappings. + Uses flat byte indices with vectorized scatter. + + Args: + 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: + return + + k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(self.layer_idx) + + num_tokens = k_fp8.shape[0] + + # The C++ op reinterprets k_fp8 (FP8) and k_scale (float32) as raw + # bytes internally and only reads the first num_tokens entries from + # the slot mapping buffers, avoiding Python-side view/slice overhead. + if k_scale.element_size() == 1: + # The op expects 4 byte elements. + k_scale = k_scale.view(torch.int32) + torch.ops.trtllm.indexer_k_cache_scatter_op( + k_fp8, + k_scale, + k_cache, + metadata.slot_mapping_fp8, + metadata.slot_mapping_scale, + num_tokens, + ) + + def _gather_k_cache_for_chunk( + self, + metadata: DSAtrtllmAttentionMetadata, + chunk: IndexerPrefillChunkMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Gather K values from indexer cache for a specific chunk. + + Uses pre-computed extended slot mappings that cover cached + current batch context tokens. + chunk.k_token_start/k_token_end directly index into the extended slot mapping. + + Args: + metadata: Attention metadata + chunk: Chunk metadata with k_token_start/end as indices into extended slot mapping + + Returns: + 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, ( + "_gather_k_cache_for_chunk requires extended slot mappings (only available with cached tokens)" + ) + + k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(self.layer_idx) + + head_dim = self.head_dim + scale_size = 4 # float32 = 4 bytes + + # Extract slot mappings using chunk's k_token_start/end + # These indices point directly into the extended slot mapping array + k_token_start = chunk.k_token_start + 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] + + # Vectorized gather using pre-computed slot mappings + # Gather FP8 data + byte_offsets_fp8 = torch.arange(head_dim, device=k_cache.device).unsqueeze( + 0 + ) # [1, head_dim] + gather_indices_fp8 = ( + slot_mapping_fp8_chunk.unsqueeze(1) + byte_offsets_fp8 + ) # [num_k_tokens, head_dim] + assert (gather_indices_fp8 >= k_cache.numel()).sum() == 0, "Out-of-bounds access detected" + gather_indices_fp8 = _unravel_indices(gather_indices_fp8, k_cache.shape) + k_fp8_bytes = k_cache[gather_indices_fp8] + k_fp8 = k_fp8_bytes.view(torch.float8_e4m3fn).view(num_k_tokens, head_dim) + + # Gather scale data + byte_offsets_scale = torch.arange(scale_size, device=k_cache.device).unsqueeze(0) # [1, 4] + gather_indices_scale = ( + slot_mapping_scale_chunk.unsqueeze(1) + byte_offsets_scale + ) # [num_k_tokens, 4] + assert (gather_indices_scale >= k_cache.numel()).sum() == 0, "Out-of-bounds access detected" + gather_indices_scale = _unravel_indices(gather_indices_scale, k_cache.shape) + k_scale_bytes = k_cache[gather_indices_scale] + k_scale = k_scale_bytes.view(torch.float32).view(num_k_tokens, 1) + + return k_fp8, k_scale + + def _call_mqa_logits( + self, + q_fp8: torch.Tensor, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + q_scale: Optional[torch.Tensor], + ) -> torch.Tensor: + """Dispatch fp8_mqa_logits vs fp8_fp4_mqa_logits based on use_fp4. + + For FP4 the gather output keeps the legacy float8_e4m3fn dtype for + API compatibility; reinterpret the bytes as the int8 / int32 layout + the DeepGEMM kernel expects. The scale tensor is collapsed to 1D for + the kv side and 2D for the q side per the kernel's asserts. + """ + if self.use_fp4: + k_fp4_bytes = k_fp8.view(torch.int8) + k_scale_int32 = k_scale.view(torch.int32).reshape(-1) + # q_scale arrives as (chunk_tokens, n_heads, 1); the FP4 kernel + # asserts q_sf is 2D so collapse the trailing unit axis. + q_scale_2d = q_scale.reshape(-1, self.n_heads) + return fp8_fp4_mqa_logits( + (q_fp8, q_scale_2d), + (k_fp4_bytes, k_scale_int32), + weights, + cu_seqlen_ks, + cu_seqlen_ke, + ) + return fp8_mqa_logits( + q_fp8, (k_fp8, k_scale.reshape(-1)), weights, cu_seqlen_ks, cu_seqlen_ke + ) + + def _call_paged_mqa_logits( + self, + q_decode: torch.Tensor, + k_cache: torch.Tensor, + weights_decode: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + scheduler_metadata_buffer: torch.Tensor, + max_seq_len: int, + q_scale: Optional[torch.Tensor], + ) -> torch.Tensor: + """Dispatch fp8_paged_mqa_logits vs fp8_fp4_paged_mqa_logits.""" + if self.use_fp4: + return fp8_fp4_paged_mqa_logits( + (q_decode, q_scale), + k_cache, + weights_decode, + context_lens, + block_table, + scheduler_metadata_buffer, + max_seq_len, + ) + return fp8_paged_mqa_logits( + q_decode, + k_cache, + weights_decode, + context_lens, + block_table, + scheduler_metadata_buffer, + max_seq_len, + ) + + def sparse_attn_indexer( + self, + metadata: DSAtrtllmAttentionMetadata, + hidden_states: torch.Tensor, + q_fp8: torch.Tensor, + k_fp8: torch.Tensor, + k_scale: torch.Tensor, + weights: torch.Tensor, + use_custom_topk: bool = True, + q_scale: Optional[torch.Tensor] = None, + is_generation: Optional[bool] = None, + ) -> torch.Tensor: + """Run the indexer TopK kernel for one phase or the full batch. + + When ``is_generation`` is ``None``, the inputs contain the full mixed + batch. Otherwise they contain only the selected context or generation + phase. ``q_scale`` is only consumed by the FP4 dispatch. + """ + # DSACacheManager / DeepseekV4CacheManager force quant_block_size to + # 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. + assert metadata.kv_cache_manager is None or metadata.kv_cache_manager.quant_block_size in ( + 32, + 128, + ), ( + "Unexpected quant_block_size " + f"{metadata.kv_cache_manager.quant_block_size if metadata.kv_cache_manager else 'N/A'}" + ) + num_contexts = metadata.num_contexts + num_generations = metadata.num_generations + num_ctx_tokens = metadata.num_ctx_tokens + num_tokens = metadata.num_tokens + + num_gen_tokens = num_tokens - num_ctx_tokens + if is_generation is None: + has_prefill = num_contexts > 0 + has_decode = num_generations > 0 + token_offset = num_ctx_tokens + cache_name = "indexer_topk_out_buffer" + else: + has_prefill = not is_generation and num_contexts > 0 + has_decode = is_generation and num_generations > 0 + token_offset = 0 + expected_tokens = num_gen_tokens if is_generation else num_ctx_tokens + assert hidden_states.shape[0] == expected_tokens, ( + "Phase-specific DSA prediction received " + f"{hidden_states.shape[0]} tokens, expected {expected_tokens}." + ) + cache_name = ( + "indexer_topk_out_buffer_gen" if is_generation else "indexer_topk_out_buffer_ctx" + ) + + topk_indices_buffer = metadata.get_empty( + metadata.cuda_graph_buffers, + (hidden_states.shape[0], self.index_topk), + cache_name=cache_name, + dtype=torch.int32, + capture_graph=metadata.is_cuda_graph, + ) + if not use_custom_topk: + topk_indices_buffer[: hidden_states.shape[0]] = -1 + + if has_prefill and not metadata.skip_indexer_for_ctx_reqs: + # Use chunked prefill to reduce memory footprint + if metadata.indexer_prefill_chunks is not None: + sparse_metadata_params = metadata.sparse_metadata_params + q_split_threshold = ( + sparse_metadata_params.q_split_threshold + if sparse_metadata_params is not None + else 8192 + ) + q_split_eligible = ( + q_split_threshold >= 0 + and metadata.mapping is not None + and not metadata.mapping.enable_attention_dp + and metadata.mapping.tp_size > 1 + ) + + if q_split_eligible: + tp_rank = metadata.mapping.tp_rank + tp_size = metadata.mapping.tp_size + + 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 + # same in both modes because FP4 packs four UE8M0 exponents + # 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: + # 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: + topk_indices_buffer[chunk.token_start : chunk.token_end, :].fill_(-1) + continue + 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, + chunk.k_token_start, + num_k_tokens, + gather_head_dim, + ) + + chunk_num_token = chunk.token_end - chunk.token_start + apply_q_split = q_split_eligible and chunk_num_token >= q_split_threshold + if apply_q_split: + chunk_q_start = chunk_num_token * tp_rank // tp_size + chunk_q_end = chunk_num_token * (tp_rank + 1) // tp_size + else: + chunk_q_start = 0 + chunk_q_end = chunk_num_token + + global_q_start = chunk.token_start + chunk_q_start + global_q_end = chunk.token_start + chunk_q_end + + # Tile the query dimension so each fp8_mqa_logits call + # allocates at most [q_tile x num_k_tokens] instead of the + # full [local_q x num_k_tokens] (which can reach tens of GB + # on a long context and stall cuMemCreate under + # expandable_segments -> engine hang; see + # _INDEXER_MQA_LOGITS_ELEM_BUDGET). Results are identical: + # each query row's logits/top-k are independent and the KV + # (chunk_k_fp8) is unchanged across tiles, so the per-call + # allocation is the same size and the caching allocator + # reuses one block (peak ~= one tile, no extra sync). + local_q_len = chunk_q_end - chunk_q_start + q_tile = max( + 1, min(local_q_len, _INDEXER_MQA_LOGITS_ELEM_BUDGET // max(1, num_k_tokens)) + ) + for tile_off in range(0, local_q_len, q_tile): + c0 = chunk_q_start + tile_off + c1 = min(c0 + q_tile, chunk_q_end) + g0 = chunk.token_start + c0 + g1 = chunk.token_start + c1 + tile_q_scale = q_scale[g0:g1, ...] if self.use_fp4 else None + logits = self._call_mqa_logits( + q_fp8[g0:g1, ...], + chunk_k_fp8, + chunk_k_scale, + weights[g0:g1, ...], + chunk.cu_seqlen_ks[c0:c1], + chunk.cu_seqlen_ke[c0:c1], + tile_q_scale, + ) + if use_custom_topk: + torch.ops.trtllm.indexer_topk_prefill( + logits, + chunk.cu_seqlen_ks[c0:c1], + chunk.cu_seqlen_ke[c0:c1], + topk_indices_buffer[g0:g1, :], + self.index_topk, + ) + else: + topk_indices = logits.topk( + min(self.index_topk, logits.shape[-1]), dim=-1 + )[1] + topk_indices -= chunk.cu_seqlen_ks[c0:c1][:, None] + + mask_lo = topk_indices >= 0 + mask_hi = ( + topk_indices + - (chunk.cu_seqlen_ke[c0:c1] - chunk.cu_seqlen_ks[c0:c1])[:, None] + < 0 + ) + mask = mask_lo & mask_hi + + # local indices per sequence + topk_indices = topk_indices.masked_fill(~mask, -1) + + topk_indices_buffer[g0:g1, : topk_indices.shape[-1]] = topk_indices.to( + dtype=torch.int32 + ) + + if apply_q_split: + q_sizes = [ + (r + 1) * chunk_num_token // tp_size - r * chunk_num_token // tp_size + for r in range(tp_size) + ] + topk_indices_buffer[chunk.token_start : chunk.token_end, :] = allgather( + topk_indices_buffer[global_q_start:global_q_end, :], + metadata.mapping, + dim=0, + sizes=q_sizes, + ) + elif metadata.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] + + ctx_q_scale = q_scale[:num_ctx_tokens, ...] if self.use_fp4 else None + logits = self._call_mqa_logits( + q_fp8[:num_ctx_tokens, ...], + k_fp8[:num_ctx_kv_tokens, ...], + k_scale[:num_ctx_kv_tokens, ...], + weights[:num_ctx_tokens, ...], + cu_seqlen_ks, + cu_seqlen_ke, + ctx_q_scale, + ) + if use_custom_topk: + torch.ops.trtllm.indexer_topk_prefill( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + topk_indices_buffer[:num_ctx_tokens, :], + self.index_topk, + ) + else: + topk_indices = logits.topk(min(self.index_topk, logits.shape[-1]), dim=-1)[1] + topk_indices -= cu_seqlen_ks[:, None] + mask_lo = topk_indices >= 0 + mask_hi = topk_indices - (cu_seqlen_ke - cu_seqlen_ks)[:, None] < 0 + mask = mask_lo & mask_hi + + # local indices per sequence + topk_indices = topk_indices.masked_fill(~mask, -1) + topk_indices_buffer[:num_ctx_tokens, : topk_indices.shape[-1]] = ( + topk_indices.to(dtype=torch.int32) + ) + elif has_prefill and metadata.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, : + ] + + # Prefill→decode GVR handoff: seed each finishing-prefill sequence's + # heuristic_prev_topk slot with its own last-context-token top-K, so + # the FIRST decode step of that sequence gets a warm-started preIdx + # (~60-75% set-overlap with the eventual decode top-K on this + # workload) instead of the all-zero / all-(-1) cold start that the + # default `heuristic_prev_topk.zero_()` initialization leaves behind. + # Without this, GVR P2 secant on decode step 0 runs from a benign + # but uninformative seed (kernel +1 offset on zeros → all indices + # point at compressed-token position 1), wasting iterations. + # Slot convention (mirrors the existing decode write-back at the + # bottom of the decode block): new gens from finishing prefill + # append after currently-active gens, i.e., slots + # [num_generations : num_generations + num_contexts]. + if self._enable_heuristic_topk and has_prefill and not metadata.skip_indexer_for_ctx_reqs: + local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] + ctx_seq_lens = metadata.seq_lens[:num_contexts] + # Per-sequence last context-token offset (exclusive cumsum minus 1). + last_ctx_idx = (torch.cumsum(ctx_seq_lens, dim=0) - 1).to(dtype=torch.long) + metadata.heuristic_prev_topk[ + local_layer, num_generations : num_generations + num_contexts + ].copy_(topk_indices_buffer[last_ctx_idx, :]) + + if has_decode and not metadata.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] + max_decode_len = gen_seq_lens.max().item() + min_decode_len = gen_seq_lens.min().item() + assert max_decode_len == min_decode_len, ( + "max_decode_len != min_decode_len, we need padding" + ) + + # Reshape q for decode phase: [num_gen_tokens, ...] -> [batch_size, next_n, ...] + q_decode = q_fp8[token_offset : token_offset + num_gen_tokens, ...] + batch_size = num_generations + next_n = num_gen_tokens // num_generations + # 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 not metadata.use_expanded_buffers_for_mtp or next_n == 1: + 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[ + num_contexts : num_contexts + num_generations + ] + # The 2D-context_lens metadata kernel encodes next_n into the + # schedule (via num_next_n_atoms). MTP forwards alternate + # between the full-window call (next_n == 1+max_draft_tokens) + # and per-token draft calls (next_n == 1), so we must select + # the buffer that was populated for this next_n. The DSL path + # 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 + else: + scheduler_metadata_buffer = metadata.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 + + assert num_gen_tokens == batch_size * next_n + weights_decode = weights[token_offset : token_offset + num_gen_tokens, ...] + + # Get k cache and call fp8_paged_mqa_logits / fp8_fp4_paged_mqa_logits + # 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() + + if self.use_cute_dsl_paged_mqa_logits: + # 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()` + # with a (num_gen, 1) input shape, which makes DeepGEMM's wrapper + # compute `num_next_n_atoms = 1`. (DeepGEMM uses the same buffer + # for its own next_n=1 kernel; DSL piggy-backs on it for all + # real next_n values.) All next_n positions of a batch share + # 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 + assert dsl_context_lens is not None + # Wave-aware atom-split: the picker in `_pick_dsl_expand` caches + # (factor, atom) on metadata with invariant + # `factor * atom == 1 + max_draft_tokens` (the target/verify-time + # next_n). MTPEagle reuses the same metadata for its multi-step + # draft loop; after i=0 it mutates seq_lens to 1, so i≥1 + # iterations run with next_n=1. The reshape + # `(num_gen, next_n, ...) -> (num_gen*factor, atom, ...)` is only + # valid when the caller actually supplies next_n == factor * atom + # 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 + ) + if self.use_fp4: + # FP4 DSL signature splits DG's (q, sf_q) tuple into two + # separate args and requires q.dtype == uint8 (q_decode + # came in via the FP8 plumbing as int8; reinterpret with + # no copy). sf_q is the q_scale slice reshaped to + # (B, next_n, H) int32 -- mirrors the non-DSL FP4 branch. + decode_q_scale = q_scale[token_offset : token_offset + num_gen_tokens, ...] + decode_q_scale = decode_q_scale.view( + q_decode.shape[0], q_decode.shape[1], self.n_heads + ) + dsl_q = q_decode.view(torch.uint8) + dsl_block_table = block_table + dsl_schedule_meta = metadata.scheduler_metadata_buffer + if dsl_atom_split: + factor = metadata.dsl_expand_factor + eff_next_n = metadata.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 + + logits_decode = torch.ops.trtllm.cute_dsl_fp4_paged_mqa_logits( + dsl_q, + decode_q_scale, + k_cache, + weights_decode, + dsl_context_lens, + dsl_block_table, + dsl_schedule_meta, + indexer_max_seq_len, + ) + else: + # FP8 DSL kernel natively supports next_n ∈ {1, 2, 3, 4}. + # Atom-split benefits small-batch / low-ntask configs by + # raising SM utilization at the cost of factorx KV HBM + # re-reads. Context lengths are indexer KV lengths, which + # may be compressed for DeepSeek-V4. + dsl_q = q_decode + fp8_ctx_lens = dsl_context_lens + fp8_block_table = block_table + fp8_schedule_meta = metadata.scheduler_metadata_buffer + if dsl_atom_split: + factor = metadata.dsl_expand_factor + atom = metadata.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 + logits_decode = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( + dsl_q, + k_cache, + weights_decode, + fp8_ctx_lens, + fp8_block_table, + fp8_schedule_meta, + indexer_max_seq_len, + ) + else: + decode_q_scale = ( + q_scale[token_offset : token_offset + num_gen_tokens, ...] + if self.use_fp4 + else None + ) + if self.use_fp4: + # q_decode shape is either (num_generations, next_n, n_heads, + # head_dim/2) [non-expanded] or (batch*next_n, 1, n_heads, + # head_dim/2) [expanded]. Match q_scale's batch/next_n dims. + decode_q_scale = decode_q_scale.view( + q_decode.shape[0], q_decode.shape[1], self.n_heads + ) + logits_decode = self._call_paged_mqa_logits( + q_decode, + k_cache, + weights_decode, + context_lens, + block_table, + scheduler_metadata_buffer, + indexer_max_seq_len, + decode_q_scale, + ) + + if use_custom_topk: + # Kernel expects kv_lens (total cache length), not seq_lens (new tokens) + # This is because rowEnd = seq_len - next_n + offset + 1 + gen_kv_lens_cuda = metadata.kv_lens_cuda_runtime[ + num_contexts : num_contexts + num_generations + ] + + pre_idx = None + heuristic_scratch = None + if self._enable_heuristic_topk: + local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] + # Pass prev_topk directly; the +1 temporal offset is + # handled inside the C++ kernel (preIdxOffset += 1). + pre_idx = metadata.heuristic_prev_topk[local_layer, :num_generations] + if not metadata.use_cute_dsl_topk: + heuristic_scratch = metadata.heuristic_scratch_values[:num_gen_tokens] + + if self.use_cute_dsl_topk and self._enable_heuristic_topk: + torch.ops.trtllm.cute_dsl_gvr_topk_decode( + logits_decode, + pre_idx, + gen_kv_lens_cuda, + topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], + self.index_topk, + next_n=next_n, + compress_ratio=self.compress_ratio, + max_seq_len=indexer_max_seq_len, + order_row=metadata.kv_lens_row_reorder, + ) + # The radix DSL path uses O(num_gen_tokens * kv_len) memory. + elif ( + self.use_cute_dsl_topk + and num_gen_tokens <= 256 + and (self.compress_ratio == 1 or next_n == 1) + ): + torch.ops.trtllm.cute_dsl_indexer_topk_decode( + logits_decode, + context_lens if self.compress_ratio > 1 else gen_kv_lens_cuda, + topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], + self.index_topk, + next_n, + ) + else: + torch.ops.trtllm.indexer_topk_decode( + logits_decode, + gen_kv_lens_cuda, + topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], + next_n, + self.index_topk, + pre_idx=pre_idx, + heuristic_scratch=heuristic_scratch, + compress_ratio=self.compress_ratio, + radix_aux_indices=metadata.radix_aux_indices, + radix_aux_logits=metadata.radix_aux_logits, + ) + else: + # padded + positions = ( + torch.arange(logits_decode.shape[-1], device=q_decode.device) + .unsqueeze(0) + .expand(num_gen_tokens, -1) + ) + row_indices = torch.arange(num_gen_tokens, device=q_decode.device) // next_n + next_n_offset = torch.arange(num_gen_tokens, device=q_decode.device) % next_n + index_end_pos = (context_lens[row_indices] - next_n + next_n_offset).unsqueeze(1) + # index_end_pos: [B * N, 1] + mask = positions <= index_end_pos + # mask: [B * N, L] + logits_decode = logits_decode.masked_fill(~mask, float("-inf")) + topk_indices_decode = logits_decode.topk( + min(self.index_topk, logits_decode.shape[-1]), dim=-1 + )[1].to(torch.int32) # [B * N, K] + # ensure we don't set indices for the top k + # that is out of range(masked already) + # this will happen if context length is shorter than K + mask_decode = topk_indices_decode <= index_end_pos + + # local indices per sequence + topk_indices_decode = topk_indices_decode.masked_fill(~mask_decode, -1) + # Store in buffer + topk_indices_buffer[ + token_offset : token_offset + num_gen_tokens, : topk_indices_decode.shape[-1] + ] = topk_indices_decode.to(dtype=torch.int32) + + if self._enable_heuristic_topk: + local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] + decode_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens] + last_mtp_topk = decode_topk[next_n - 1 :: next_n] + metadata.heuristic_prev_topk[local_layer, :num_generations].copy_(last_mtp_topk) + + elif has_decode and metadata.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, :] + ) + return topk_indices_buffer + + def _weight_scale(self, weights: torch.Tensor, q_scale: torch.Tensor) -> torch.Tensor: + """Apply quantization scale to indexer attention weights.""" + weights = _scale(weights, q_scale, self.weight_scale_factor) + return weights + + def _qk_projection_and_rope( + self, qr: torch.Tensor, indexer_k: torch.Tensor, position_ids: torch.Tensor + ): + """Project Q/K and apply RoPE""" + q = self.wq_b(qr) + k = self.k_norm(indexer_k) + q = q.view(-1, self.n_heads, self.head_dim) + q_pe, q_nope = q.split([self.rope_dim, self.head_dim - self.rope_dim], dim=-1) + k_pe, k_nope = k.split([self.rope_dim, self.head_dim - self.rope_dim], dim=-1) + q_pe, k_pe = self.rotary_emb(position_ids, [q_pe, k_pe.unsqueeze(1)]) + k_pe = k_pe[:, 0, :] + return q_pe, q_nope, k_pe, k_nope + + def _prep_q_or_k(self, qk_pe: torch.Tensor, qk_nope: torch.Tensor): + """Concatenate and quantize for Q or K. + + FP8 mode: fused cat + FP8 quantize via CUDA kernel. + FP4 mode: fused cat + per-block-32 FP4 E2M1 quantize via CUDA kernel. + The returned packed bytes are int8 (two FP4 codes per byte) and the + scale is int32 (four UE8M0 exponents packed little-endian). + """ + if self.use_fp4: + return torch.ops.trtllm.fused_cat_fp4(qk_pe, qk_nope) + fp8_out, scale = torch.ops.trtllm.fused_cat_fp8(qk_pe, qk_nope, self.scale_fmt == "ue8m0") + return fp8_out, scale + + def pre_indexer_proj( + self, qr: torch.Tensor, hidden_states: torch.Tensor, position_ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Pure token-wise projections (CUDA-graph-capturable). + + Runs cublas_mm, qk_projection_and_rope, FP8/FP4 quantize, and weight + scaling. Does NOT touch the k cache or any batch-specific metadata, + so this can safely run inside a captured CUDA graph partition. + + Returns (q_fp_bytes, k_fp_bytes, k_scale, weights, q_scale). The last + tensor is only consumed by the FP4 kernel dispatch; the FP8 path + ignores it. It is returned unconditionally so the two-op CUDA graph + split in dsa.module.forward_dsa_proj sees a stable signature. + """ + assert self._fused_wk_wp_weight is not None, ( + "cache_derived_state() must be called before forward()" + ) + # When the boundary fusion pre-quantized the next layer's kv_a_proj + # NVFP4 input, hidden_states is an Fp4QuantizedTensor that also carries + # the BF16 post-RMSNorm value. Use that BF16 view here -- the indexer + # weight is BF16 and the matmul needs a float input. + if isinstance(hidden_states, Fp4QuantizedTensor): + assert hidden_states.unquantized_hidden_states is not None, ( + "pre_indexer_proj received Fp4QuantizedTensor without bf16 view; " + "the producer fusion must request return_norm_out=True" + ) + hidden_states_bf = hidden_states.unquantized_hidden_states + else: + hidden_states_bf = hidden_states + hidden_float = _to_float(hidden_states_bf) + with _tf32_matmul_enabled(): + # F.linear computes input @ weight.T internally; no explicit .t() needed. + # _fused_wk_wp_weight is [head_dim + n_heads, hidden_size] (nn.Linear convention). + # Goes through PyTorch's cuBLAS handle which respects allow_tf32 and + # dispatches CUBLAS_COMPUTE_32F_FAST_TF32, unlike torch.ops.trtllm.cublas_mm + # which uses its own handle and always falls back to CUDA-core SGEMM. + fused_out = F.linear(hidden_float, self._fused_wk_wp_weight) + indexer_k, weights = fused_out.split([self.head_dim, self.n_heads], dim=-1) + # Cast indexer_k back to model dtype for downstream ops (k_norm, RoPE, FP8 quantize) + indexer_k = indexer_k.to(hidden_states_bf.dtype) + + q_pe, q_nope, k_pe, k_nope = self._qk_projection_and_rope(qr, indexer_k, position_ids) + q, k = maybe_execute_in_parallel( + lambda: self._prep_q_or_k(q_pe, q_nope), + lambda: self._prep_q_or_k(k_pe, k_nope), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + q_fp8, q_scale = q + k_fp8, k_scale = k + if self.use_fp4: + # FP4 packs two codes per byte, so the trailing dim is head_dim // 2. + # fused_cat_fp4 flattens the leading dims to M=N*n_heads; restore + # the (N, n_heads, ...) shape so downstream slicing in + # sparse_attn_indexer (which indexes by token) lines up with + # q_fp8. The DeepGEMM FP4 kernel applies the per-block q_scale + # internally, so weights carry only softmax_scale * n_heads^-0.5. + q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim // 2) + q_scale = q_scale.view(-1, self.n_heads, 1) + # DeepGEMM's fp8_fp4_(paged_)mqa_logits asserts + # `weights.scalar_type() == kFloat`. Unlike the FP8 branch (whose + # `_weight_scale` multiplies by the fp32 `q_scale` tensor and so + # implicitly upcasts), this branch's only multiplier is a Python + # float — `bf16 * float` stays bf16 and trips the kernel assert. + # Cast explicitly so this path doesn't silently rely on an + # upstream `_to_float`. + weights = weights.float() * self.weight_scale_factor + else: + q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim) + q_scale = q_scale.view(-1, self.n_heads, 1) + weights = self._weight_scale(weights, q_scale) + + return q_fp8, k_fp8, k_scale, weights, q_scale + + def forward_from_projected( + self, + metadata: DSAtrtllmAttentionMetadata, + hidden_states: torch.Tensor, + indexer_intermediates: List[torch.Tensor], + is_generation: Optional[bool] = None, + ) -> torch.Tensor: + """Run sparse prediction from graph-captured indexer projections. + + Projected inputs cover the full batch. Phase-specific prediction slices + Q-side inputs here while retaining the full-batch K inputs needed by + context prediction. + """ + if is_generation is None: + phase_start = 0 + phase_end = metadata.num_tokens + elif is_generation: + phase_start = metadata.num_ctx_tokens + phase_end = metadata.num_tokens + else: + phase_start = 0 + phase_end = metadata.num_ctx_tokens + + q_fp8, k_fp8, k_scale, weights, q_scale = indexer_intermediates + q_fp8_phase = q_fp8[phase_start:phase_end] + weights_phase = weights[phase_start:phase_end] + q_scale_phase = q_scale[phase_start:phase_end] if q_scale is not None else None + + return self.sparse_attn_indexer( + metadata, + hidden_states, + q_fp8_phase, + k_fp8, + k_scale, + weights_phase, + q_scale=q_scale_phase, + is_generation=is_generation, + ) + + @torch.inference_mode() + def forward( + self, + qr: torch.Tensor, + hidden_states: torch.Tensor, + metadata: DSAtrtllmAttentionMetadata, + position_ids: torch.Tensor, + ): + q_fp8, k_fp8, k_scale, weights, q_scale = self.pre_indexer_proj( + qr, hidden_states, position_ids + ) + indexer_intermediates = [q_fp8, k_fp8, k_scale, weights, q_scale] + self._update_k_cache(k_fp8, k_scale, metadata) + + # Return topk indices buffer for sparse attention [num_tokens, index_topk] + return self.forward_from_projected(metadata, hidden_states, indexer_intermediates) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/kernels.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/kernels.py new file mode 100644 index 000000000000..7aa85f7d524c --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/kernels.py @@ -0,0 +1,294 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +######################################################## +# Index gather kernel +######################################################## + + +@triton.jit +def _convert_req_index_to_global_index_kernel_with_stride_factor( + req_id_ptr, # int32 [num_tokens] + block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] + token_indices_ptr, # int32 [num_kv_heads, num_tokens, NUM_TOPK_TOKENS] + out_ptr, # int32 [num_kv_heads, num_tokens, kv_factor, NUM_TOPK_TOKENS] + # shapes (compile-time where possible) + max_num_blocks_per_req: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_N: tl.constexpr, # tile width along columns + stride_factor: tl.constexpr, # elements per physical page in pool + layer_id: tl.constexpr, # for multi-layer KV cache + num_kv_heads: tl.constexpr, + kv_factor: tl.constexpr, + bt_stride0, # stride for batch dim + bt_stride1, # stride for page dim + ti_stride0, + ti_stride1, + ti_stride2, + out_stride0, + out_stride1, + out_stride2, + out_stride3, +): + """ + Convert request-local token indices to global KV cache pool indices. + + KV cache pool layout: [max_num_pages, num_layers, kv_factor, num_kv_heads, tokens_per_block, head_dim] + block_table: [num_requests, max_pages_per_req] + - stores memPoolBlockIdx (physical page index, first dim of pool) + - same block_table for K and V (K/V share the same physical page) + + stride_factor = num_layers * kv_factor * num_kv_heads * BLOCK_SIZE + (elements per physical page, excluding head_dim) + + Global index: + memPoolBlockIdx * stride_factor + + layer_id * kv_factor * num_kv_heads * BLOCK_SIZE + + kv_factor_idx * num_kv_heads * BLOCK_SIZE + + kv_head_idx * BLOCK_SIZE + + token_in_page + """ + kv_head_idx = tl.program_id(0) + token_id = tl.program_id(1) + tile_id = tl.program_id(2) + + indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) + + req = tl.load(req_id_ptr + token_id) + + # Load local token indices: token_indices[kv_head_idx, token_id, indice_id] + ti_ptr = ( + token_indices_ptr + + kv_head_idx * ti_stride0 + + token_id * ti_stride1 + + indice_id * ti_stride2 + ) + tok = tl.load(ti_ptr) + + is_invalid_tok = tok < 0 + page_idx = tok // BLOCK_SIZE + token_in_page = tok % BLOCK_SIZE + valid_page = page_idx < max_num_blocks_per_req + + # block_table[req, page_idx] → memPoolBlockIdx (physical page index) + bt_ptr = block_table_ptr + req * bt_stride0 + page_idx * bt_stride1 + mem_pool_idx = tl.load(bt_ptr, mask=valid_page, other=0) + + # Base offset within physical page (invariant across kv_factor loop) + base_off = ( + layer_id * kv_factor * num_kv_heads * BLOCK_SIZE + kv_head_idx * BLOCK_SIZE + token_in_page + ) + + for kv_factor_idx in tl.static_range(kv_factor): + # Within-page offset: layer → kv_factor → kv_head → token + inpage_off = base_off + kv_factor_idx * num_kv_heads * BLOCK_SIZE + + global_idx = mem_pool_idx * stride_factor + inpage_off + + out_val = tl.where(is_invalid_tok | (~valid_page), -1, global_idx) + + out_ptr_ij = ( + out_ptr + + kv_head_idx * out_stride0 + + token_id * out_stride1 + + kv_factor_idx * out_stride2 + + indice_id * out_stride3 + ) + tl.store(out_ptr_ij, out_val) + + +def triton_convert_req_index_to_global_index( + req_id: torch.Tensor, # int32 [num_tokens] + block_table: torch.Tensor, # int32 [num_requests, kv_factor, max_num_blocks_per_req] + token_indices: torch.Tensor, # int32 [num_kv_heads, num_tokens, NUM_TOPK_TOKENS] + BLOCK_SIZE: int, + NUM_TOPK_TOKENS: int = 2048, + BLOCK_N: int = 128, # tile width along columns + stride_factor: int = None, # elements per block in pool + layer_id: int = 0, # for multi-layer KV cache + num_kv_heads: int = 1, + kv_factor: int = 1, +): + """ + Convert request-local token indices to global KV cache pool indices. + + Accepts both 2D and 3D token_indices: + - 2D [num_tokens, topk]: MLA path (num_kv_heads=1 implicit) + - 3D [num_kv_heads, num_tokens, topk]: MQA/GQA path + + Output shape: [num_kv_heads, num_tokens, kv_factor, topk] + + Args: + block_table: [num_requests, max_pages] — physical page indices. + stride_factor: Elements per physical page + (num_layers * kv_factor * num_kv_heads * BLOCK_SIZE). + num_kv_heads: Number of KV heads. + kv_factor: 2 for MQA/GQA, 1 for MLA. + """ + if stride_factor is None: + stride_factor = kv_factor * num_kv_heads * BLOCK_SIZE + assert req_id.dtype == torch.int32 + assert block_table.dtype == torch.int32 + assert token_indices.dtype == torch.int32 + assert token_indices.ndim in (2, 3), ( + f"Expected 2D [num_tokens, topk] or 3D [num_kv_heads, num_tokens, topk], got {token_indices.ndim}D" + ) + assert block_table.ndim == 2, f"Expected 2D [batch, max_pages], got {block_table.ndim}D" + assert NUM_TOPK_TOKENS % BLOCK_N == 0, ( + f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})" + ) + + num_tokens = req_id.shape[0] + max_num_blocks_per_req = block_table.shape[1] + tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N + + req_id_c = req_id.contiguous() + block_table_c = block_table.contiguous() + token_indices_c = token_indices.contiguous() + out = torch.empty( + (num_kv_heads, num_tokens, kv_factor, NUM_TOPK_TOKENS), + dtype=torch.int32, + device=token_indices.device, + ) + + bt_stride0, bt_stride1 = block_table_c.stride() + # 2D input: ti_stride0=0 (kv_head_idx is always 0, term vanishes) + if token_indices_c.ndim == 2: + ti_stride0 = 0 + ti_stride1, ti_stride2 = token_indices_c.stride() + else: + ti_stride0, ti_stride1, ti_stride2 = token_indices_c.stride() + out_stride0, out_stride1, out_stride2, out_stride3 = out.stride() + + grid = (num_kv_heads, num_tokens, tiles_per_row) + + _convert_req_index_to_global_index_kernel_with_stride_factor[grid]( + req_id_c, + block_table_c, + token_indices_c, + out, + max_num_blocks_per_req, + BLOCK_SIZE, + BLOCK_N, + stride_factor, + layer_id, + num_kv_heads, + kv_factor, + bt_stride0, + bt_stride1, + ti_stride0, + ti_stride1, + ti_stride2, + out_stride0, + out_stride1, + out_stride2, + out_stride3, + ) + return out + + +@triton.jit +def _triton_gather_k_cache_kernel( + k_cache_ptr, + slot_fp8_ptr, + slot_scale_ptr, + out_fp8_ptr, + out_scale_ptr, + k_token_start, + num_k_tokens, + HEAD_DIM: tl.constexpr, + SCALE_BYTES: tl.constexpr, + BLOCK_TOKENS: tl.constexpr, +): + pid = tl.program_id(0) + token_offsets = (pid * BLOCK_TOKENS + tl.arange(0, BLOCK_TOKENS)).to(tl.int64) + token_mask = token_offsets < num_k_tokens + + fp8_base = tl.load(slot_fp8_ptr + k_token_start + token_offsets, mask=token_mask, other=0) + scale_base = tl.load(slot_scale_ptr + k_token_start + token_offsets, mask=token_mask, other=0) + + byte_offsets = tl.arange(0, HEAD_DIM).to(tl.int64) + src_fp8 = fp8_base[:, None] + byte_offsets[None, :] + dst_fp8 = token_offsets[:, None] * HEAD_DIM + byte_offsets[None, :] + gather_mask = token_mask[:, None] + + fp8_data = tl.load(k_cache_ptr + src_fp8, mask=gather_mask, other=0) + tl.store(out_fp8_ptr + dst_fp8, fp8_data, mask=gather_mask) + + scale_byte_offsets = tl.arange(0, SCALE_BYTES).to(tl.int64) + src_scale = scale_base[:, None] + scale_byte_offsets[None, :] + dst_scale = token_offsets[:, None] * SCALE_BYTES + scale_byte_offsets[None, :] + + scale_data = tl.load(k_cache_ptr + src_scale, mask=gather_mask, other=0) + tl.store(out_scale_ptr + dst_scale, scale_data, mask=gather_mask) + + +def triton_gather_k_cache( + k_cache: torch.Tensor, + slot_mapping_fp8: torch.Tensor, + slot_mapping_scale: torch.Tensor, + k_token_start: int, + k_token_end: int, + head_dim: int, +): + """Gather K FP8 values and scales from the indexer K cache for a chunk. + + Replaces ``_gather_k_cache_for_chunk``, fusing ~8-12 small PyTorch ops + (arange, unsqueeze, broadcast add, _unravel_indices, advanced indexing) + into a single Triton kernel that directly gathers from flat byte offsets. + This is purely data movement — bit-exact with the original. + + Args: + k_cache: Indexer K cache pool data (2D contiguous), uint8. + slot_mapping_fp8: Flat byte indices for FP8 data + ``[total_kv_len]``, int64. + slot_mapping_scale: Flat byte indices for scale data + ``[total_kv_len]``, int64. + k_token_start: Start index into slot mapping arrays. + k_token_end: End index into slot mapping arrays. + head_dim: FP8 head dimension (typically 128). + + Returns: + Tuple of (k_fp8, k_scale): + k_fp8: ``[num_k_tokens, head_dim]``, float8_e4m3fn. + k_scale: ``[num_k_tokens, 1]``, float32. + """ + num_k_tokens = k_token_end - k_token_start + device = k_cache.device + + if num_k_tokens == 0: + return ( + torch.empty(0, head_dim, dtype=torch.float8_e4m3fn, device=device), + torch.empty(0, 1, dtype=torch.float32, device=device), + ) + + SCALE_BYTES = 4 + BLOCK_TOKENS = 32 + + k_cache_flat = k_cache.reshape(-1) + out_fp8 = torch.empty(num_k_tokens, head_dim, dtype=torch.uint8, device=device) + out_scale = torch.empty(num_k_tokens, SCALE_BYTES, dtype=torch.uint8, device=device) + + grid = (triton.cdiv(num_k_tokens, BLOCK_TOKENS),) + _triton_gather_k_cache_kernel[grid]( + k_cache_flat, + slot_mapping_fp8, + slot_mapping_scale, + out_fp8.view(-1), + out_scale.view(-1), + k_token_start, + num_k_tokens, + HEAD_DIM=head_dim, + SCALE_BYTES=SCALE_BYTES, + BLOCK_TOKENS=BLOCK_TOKENS, + ) + + k_fp8 = out_fp8.view(torch.float8_e4m3fn) + k_scale = out_scale.view(torch.float32).view(num_k_tokens, 1) + return k_fp8, k_scale diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py new file mode 100644 index 000000000000..281158317923 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py @@ -0,0 +1,1080 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dense Sparse Attention (DSA) backend for TRT-LLM with indexer-based TopK selection.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, List, Optional + +import torch + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._torch.utils import maybe_compile +from tensorrt_llm._utils import get_sm_version, prefer_pinned +from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata + +from .indexer import ( + _DG_SCHEDULE_BLOCK_KV, + Indexer, + IndexerPrefillChunkMetadata, + _compute_slot_mappings, + _effective_compress_ratio_divisor, + _pick_dsl_expand, + _select_indexer_compress_ratio, +) +from .params import DSAMetadataParams + +ModelConfig = tensorrt_llm.bindings.ModelConfig + +if TYPE_CHECKING: + from tensorrt_llm._torch.speculative.interface import SpecMetadata + from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager + + +@dataclass(init=False) +class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): + """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" + + sparse_metadata_params: Optional[DSAMetadataParams] = None + # 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 + # 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 + # max number of draft tokens + max_draft_tokens: int = 0 + # Indexer head dimension + indexer_head_dim: int = 128 + # Indexer quant block size + indexer_quant_block_size: int = 128 + # Enable indexer skip for short sequences + enable_indexer_skip: bool = False + # 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 + # 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 + + def __init__(self, *args, **kwargs): + """Initialize DSA metadata with SM count and indexer chunk size.""" + 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_stride_factor = 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 + 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 + + def __post_init__(self): + """Allocate indexer K-cache buffers and heuristic TopK metadata.""" + from .cache_manager import DSACacheManager + + super().__post_init__() + if not isinstance(self.kv_cache_manager, DSACacheManager): + has_deepseek_v4_cache_interface = all( + hasattr(self.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)}" + ) + + 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.num_sparse_topk = sparse_metadata_params.max_sparse_topk + self.sparse_mla_topk = self.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 + self.use_cute_dsl_topk = ( + sparse_metadata_params.use_cute_dsl_topk and IS_CUTLASS_DSL_AVAILABLE + ) + self.kv_lens_row_reorder = None + capture_graph = self.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]) + + # Effective tokens-per-block for the indexer k-cache slot mapping. + # 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 + self._indexer_compress_ratio = _select_indexer_compress_ratio(self.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 + + self.create_buffers_for_mla_rope_append(capture_graph=capture_graph) + self.create_buffers_for_indexer(capture_graph=capture_graph) + + def prepare(self): + super().prepare() + self._invalidate_pool_view_cache() + + # Get kv lengths + assert self.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, + dtype=torch.int, + device="cpu", + ) + if self.enable_helix: + # 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] + kv_lens = cached_token_lens.clone() + kv_lens[active_rank] += self.seq_lens_kv[active_rank] + else: + kv_lens = cached_token_lens + self.seq_lens_kv + + # For mla_rope_append_paged_kv_assign_q + self.prepare_for_mla_rope_append(cached_token_lens, kv_lens) + + # Prepare to support skip indexer + self.prepare_for_skip_indexer(kv_lens) + + # For indices conversion + self.prepare_for_indices_conversion() + + # For indexer k cache + self.prepare_for_indexer_k_cache() + + # For spec decode + self.prepare_for_spec_decode(kv_lens) + + # Prepare metadata for indexer + Indexer.prepare(metadata=self) + + def get_indexer_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor: + if self._indexer_compress_ratio <= 1: + return kv_lens + return kv_lens // self._indexer_compress_ratio + + def get_indexer_max_seq_len(self) -> int: + 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) + + def on_update_kv_lens(self): + # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas. + # Especially for the changes in the _preprocess_inputs() of model_engine.py. + # + # NOTE: + # In overlap scheduler + speculative decoding, kv_lens_cuda can be corrected at runtime + # (inside _preprocess_inputs) to account for variable accepted tokens. The indexer + # 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. + + # _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 + # transform_local_topk_and_prepare_pool_view() call. + self._invalidate_pool_view_cache() + if self.kv_cache_manager is not None and self.num_tokens > 0: + seq_lens = self.seq_lens_cuda[: self.num_seqs] + # Runtime cached lengths after overlap/spec-dec correction. + start_positions = self.kv_lens_cuda[: self.num_seqs] - seq_lens + + # Reuse request-per-token mapping prepared in metadata.prepare(). + # This avoids repeat_interleave in graph-capture mode. + req_indices = self.req_idx_per_token[: self.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) + - seq_starts[req_indices] + ) + + global_positions = start_positions[req_indices] + token_offsets + # 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) + data_bytes_per_token = index_head_dim // 2 if use_fp4 else index_head_dim + fp8_indices, scale_indices = _compute_slot_mappings( + global_positions, + self.indexer_k_cache_block_offsets, + req_indices, + index_head_dim, + self._tokens_per_block, + self.kv_cache_manager.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 + + if self.num_generations > 0: + torch.cumsum( + self.kv_lens_cuda[self.num_contexts : self.num_seqs], # num_contexts should be 0 + dim=0, + dtype=torch.int64, + out=self.gen_kv_indptr[1 : self.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] + ), + dim=0, + dtype=torch.int64, + out=self.gen_cached_token_indptr[1 : self.num_generations + 1], + ) + gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.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 + next_n_cap = self.kv_lens_cuda_2d.shape[1] + self.kv_lens_cuda_2d[: self.num_generations, :next_n_cap].copy_( + gen_indexer_kv_lens.unsqueeze(-1).expand(-1, next_n_cap) + ) + scheduler_metadata_buffer = get_paged_mqa_logits_metadata( + gen_indexer_kv_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, self.num_sms + ) + self.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, non_blocking=True) + if self.max_draft_tokens > 0 and not self.use_expanded_buffers_for_mtp: + scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( + self.kv_lens_cuda_2d[: self.num_generations, :next_n_cap], + _DG_SCHEDULE_BLOCK_KV, + self.num_sms, + ) + self.scheduler_metadata_buffer_full_next_n.copy_( + scheduler_metadata_buffer_full_next_n, non_blocking=True + ) + if self.use_expanded_buffers_for_mtp: + num_draft_tokens = 1 + self.max_draft_tokens + num_tokens = self.num_generations * num_draft_tokens + kv_lens_expanded = torch.stack([gen_indexer_kv_lens] * num_draft_tokens, dim=0) + self.kv_lens_expanded_cuda[:num_tokens] = ( + kv_lens_expanded.transpose(0, 1).contiguous().flatten() + ) + scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( + self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), + _DG_SCHEDULE_BLOCK_KV, + self.num_sms, + ) + self.scheduler_metadata_buffer_expanded.copy_( + scheduler_metadata_buffer_expanded, non_blocking=True + ) + if self.expand_for_dsl and self.dsl_expand_factor > 1: + expand_factor = self.dsl_expand_factor + num_tokens = self.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( + self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), + _DG_SCHEDULE_BLOCK_KV, + self.num_sms, + ) + self.scheduler_metadata_buffer_expanded.copy_( + scheduler_metadata_buffer_expanded, non_blocking=True + ) + self._compute_kv_lens_row_reorder() + self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True) + + def _compute_kv_lens_row_reorder(self): + """Prepare longest-job-first row order for GVR top-k.""" + next_n = 1 + self.max_draft_tokens + if ( + self.enable_heuristic_topk + and self.use_cute_dsl_topk + and self.num_generations * next_n >= 2 * self.num_sms + ): + gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.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] + 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,), + 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, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.ctx_kv_indptr = self.get_empty( + self.cuda_graph_buffers, + (self.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, + 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,), + cache_name="gen_cached_token_indptr", + dtype=torch.int64, + capture_graph=capture_graph, + ) + self.host_gen_cached_token_indptr = torch.zeros_like( + self.gen_cached_token_indptr, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.gen_kv_indptr = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_requests + 1,), + cache_name="gen_kv_indptr", + dtype=torch.int64, + capture_graph=capture_graph, + ) + self.host_gen_kv_indptr = torch.zeros_like( + self.gen_kv_indptr, + device="cpu", + pin_memory=prefer_pinned(), + ) + + def _create_radix_aux_buffers(self, capture_graph=False): + # Persistent scratch for Radix-split-work indexer path (blocks_per_row > 1). + # Mirrors the fix the Heuristic path applied: per-call th::empty inside + # indexer_topk_decode produces stale pointers under CUDA Graph replay when + # the caching allocator is perturbed by chunked prefill at high CONC. + # Sized to the worst case kMaxBlocksPerRowDecode=10 from + # cpp/tensorrt_llm/kernels/indexerTopK.cu, times the max number of + # generation rows (num_seqs * (1 + max_draft_tokens)); the cpp op aborts + # if this is smaller than num_rows*blocks_per_row*index_topk. Allocated + # unconditionally: even with enable_heuristic_topk=True the dispatcher can + # fall back to Radix when canUseHeuristic returns False (small numColumns, + # etc.). MUST be re-created whenever max_draft_tokens changes (see + # update_spec_dec_param) or it is left too small once MTP raises the + # generation-row count. + _radix_max_blocks_per_row = 10 + _radix_max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) + self.radix_aux_indices = self.get_empty( + self.cuda_graph_buffers, + (_radix_max_gen_tokens, _radix_max_blocks_per_row, self.num_sparse_topk), + cache_name="radix_aux_indices", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.radix_aux_logits = self.get_empty( + self.cuda_graph_buffers, + (_radix_max_gen_tokens, _radix_max_blocks_per_row, self.num_sparse_topk), + cache_name="radix_aux_logits", + dtype=torch.float32, + capture_graph=capture_graph, + ) + + def create_buffers_for_indexer(self, capture_graph=False): + 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_k_cache_block_offsets = self.get_empty( + self.cuda_graph_buffers, + [self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq], + cache_name="indexer_k_cache_block_offsets", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.host_indexer_k_cache_block_offsets = torch.zeros_like( + self.indexer_k_cache_block_offsets, + device="cpu", + pin_memory=prefer_pinned(), + ) + # Phase-specific local TopK reused by following shared-indexer layers. + if sparse_metadata_params.has_shared_indexer_layers: + self.shared_topk_indices = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens, self.num_sparse_topk), + cache_name="shared_topk_indices", + dtype=torch.int32, + capture_graph=capture_graph, + ) + else: + self.shared_topk_indices = None + + # 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,), + cache_name="slot_mapping_fp8", + dtype=torch.int64, + capture_graph=capture_graph, + ) + self.host_slot_mapping_fp8 = torch.zeros_like( + self.slot_mapping_fp8, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.slot_mapping_scale = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name="slot_mapping_scale", + dtype=torch.int64, + capture_graph=capture_graph, + ) + self.host_slot_mapping_scale = torch.zeros_like( + self.slot_mapping_scale, + device="cpu", + pin_memory=prefer_pinned(), + ) + # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation. + # These buffers will be allocated dynamically in Indexer.prepare() based on actual total_kv_len to save memory. + if self.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,), + cache_name="req_idx_per_token", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.host_req_idx_per_token = torch.empty_like( + 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], + cache_name="block_table", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.scheduler_metadata_buffer = self.get_empty( + self.cuda_graph_buffers, + (self.num_sms + 1, 2), + cache_name="scheduler_metadata_buffer", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # When MTP runs without the expanded-tokens path, the same forward step + # alternates between full-window calls (next_n == 1 + max_draft_tokens) + # and per-token draft calls (next_n == 1). The 2D DeepGEMM metadata + # API encodes next_n into the schedule, so the precomputed schedule + # for one shape cannot be reused for the other. Maintain a second + # buffer holding the schedule for the full next_n window; the draft + # 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.num_sms + 1, 2), + cache_name="scheduler_metadata_buffer_full_next_n", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # Pre-allocated 2D kv_lens buffer for the new DeepGEMM 2D context_lens + # API. Shape: (max_num_sequences, 1 + max_draft_tokens). Each row + # broadcasts the same kv_len across next_n positions; kernel reads a + # 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,), + 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,), + 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), + cache_name="topk_indices_buffer", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.host_topk_indices_buffer = torch.zeros_like( + self.topk_indices_buffer, + device="cpu", + pin_memory=prefer_pinned(), + ) + # Per-layer persistent buffers for heuristic TopK pre_idx. + # Indexed by [local_layer_idx, generation_position, :]. + # The graph captures reads/writes on these stable-address buffers; + # each replay's write becomes the next replay's read (feedback loop). + self.enable_heuristic_topk = ( + sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100 + ) + if self.enable_heuristic_topk: + num_local_layers = self.kv_cache_manager.num_local_layers + self.heuristic_prev_topk = self.get_empty( + self.cuda_graph_buffers, + (num_local_layers, self.max_num_sequences, self.num_sparse_topk), + cache_name="heuristic_prev_topk", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # Zero-initialize so the first decode step's pre_idx (kernel + # adds +1 offset) points to index 1 — a valid but benign candidate. + # Without this, uninitialized memory produces random hint indices. + self.heuristic_prev_topk.zero_() + # The C++ top-k path needs a stable scratch address. + if not self.use_cute_dsl_topk: + max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) + self.heuristic_scratch_values = self.get_empty( + self.cuda_graph_buffers, + (max_gen_tokens, self.num_sparse_topk), + cache_name="heuristic_scratch_values", + dtype=torch.float32, + capture_graph=capture_graph, + ) + # GVR row order also needs a stable address for CUDA graphs. + if self.use_cute_dsl_topk: + self.kv_lens_row_reorder_buffer = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences,), + cache_name="kv_lens_row_reorder_buffer", + dtype=torch.int32, + capture_graph=capture_graph, + ) + + # Persistent scratch for the Radix-split-work indexer path. Re-created + # in update_spec_dec_param when max_draft_tokens changes so it stays + # large enough for the MTP generation-row count. + self._create_radix_aux_buffers(capture_graph=capture_graph) + + # Create expanded buffers for MTP support + self.create_expanded_buffers(capture_graph=capture_graph) + + def _create_kv_lens_2d_buffer(self, capture_graph=False): + """Pre-allocated buffer for the DeepGEMM 2D context_lens API. + + Avoids per-forward .expand().contiguous() allocations that break CUDA + graphs. The buffer is written in-place via .copy_() inside + on_update_kv_lens so its address stays stable across replays. + """ + self.kv_lens_cuda_2d = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences, 1 + self.max_draft_tokens), + cache_name="kv_lens_cuda_2d", + dtype=torch.int32, + capture_graph=capture_graph, + ) + + # 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.max_draft_tokens),), + cache_name="kv_lens_expanded_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.kv_lens_expanded_host = torch.zeros_like( + self.kv_lens_expanded_cuda, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.block_table_expanded = self.get_empty( + self.cuda_graph_buffers, + [ + self.max_num_sequences * (1 + self.max_draft_tokens), + self.kv_cache_manager.max_blocks_per_seq, + ], + cache_name="block_table_expanded", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.host_block_table_expanded = torch.zeros_like( + self.block_table_expanded, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.scheduler_metadata_buffer_expanded = self.get_empty( + self.cuda_graph_buffers, + (self.num_sms + 1, 2), + cache_name="scheduler_metadata_buffer_expanded", + dtype=torch.int32, + capture_graph=capture_graph, + ) + + # 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( + self, + batch_size, + is_spec_decoding_enabled, + is_spec_dec_tree, + is_spec_dec_dynamic_tree, + max_draft_len, + max_total_draft_tokens, + model_is_wrapped: bool = False, + spec_metadata: Optional["SpecMetadata"] = None, + spec_tree_manager: Optional["SpecTreeManager"] = None, + num_contexts: int = 0, + ): + """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, + model_is_wrapped, + spec_metadata, + spec_tree_manager, + num_contexts=num_contexts, + ) + self.max_draft_tokens = max_draft_len + capture_graph = self.is_cuda_graph + if self.kv_lens_cuda_2d.shape[1] != 1 + self.max_draft_tokens: + self._create_kv_lens_2d_buffer(capture_graph=capture_graph) + init_shape = self.kv_lens_expanded_host.shape[0] + if self.max_num_sequences * (1 + self.max_draft_tokens) != init_shape: + self.create_expanded_buffers(capture_graph=capture_graph) + # Resize heuristic scratch buffer for new max_draft_tokens. + if self.enable_heuristic_topk and not self.use_cute_dsl_topk: + max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) + self.heuristic_scratch_values = self.get_empty( + self.cuda_graph_buffers, + (max_gen_tokens, self.num_sparse_topk), + cache_name="heuristic_scratch_values", + dtype=torch.float32, + capture_graph=capture_graph, + ) + # The Radix-split-work scratch (radix_aux_*) is sized the same way + # (num_seqs * (1 + max_draft_tokens) rows) and is allocated + # unconditionally, so it must be resized here too -- otherwise the + # cpp indexer_topk_decode op aborts once MTP raises max_draft_tokens + # ("radix_aux_* must hold at least num_rows*blocks_per_row*index_topk + # elements"). + self._create_radix_aux_buffers(capture_graph=capture_graph) + + def _get_pool_block_indices(self) -> torch.Tensor: + """Extract memory pool block indices from host_kv_cache_block_offsets. + + The C++ setOffsets() encodes offsets as: + encoded = memPoolBlockIndex * numLayers * kvFactor + For SELFKONLY (MLA/DSA), kvFactor=1, so: + memPoolBlockIndex = encoded // num_local_layers + + Returns a (num_seqs, max_blocks_per_seq) int32 CPU tensor with valid + pool indices clamped to [0, blocks_in_primary_pool - 1]. + """ + num_local_layers = self.kv_cache_manager.num_local_layers + max_pool_idx = self.kv_cache_manager.blocks_in_primary_pool - 1 + # DSA uses SELFKONLY mode where only key cache is stored (kv_factor=1). + # host_kv_cache_block_offsets shape: (num_pools, max_batch*beam, 2, max_blocks_per_seq) + # Note: dim=2 is always 2 in the tensor layout (K and V slots), but for + # SELFKONLY only the K slot (index 0) contains valid data. + assert self.kv_cache_manager.kv_factor == 1, ( + f"DSA requires SELFKONLY mode (kv_factor=1), got kv_factor={self.kv_cache_manager.kv_factor}" + ) + # Pool 0, first num_seqs entries, field 0 (key offsets) + encoded = self.kv_cache_manager.host_kv_cache_block_offsets[0, : self.num_seqs, 0, :] + pool_indices = encoded // num_local_layers + # Clamp for safety: handles garbage padding from torch.empty in uninitialized slots + pool_indices = pool_indices.clamp(min=0, max=max_pool_idx).to(torch.int32) + return pool_indices + + def _invalidate_pool_view_cache(self): + """Invalidate the cached pool view and related step-invariant values. + + Must be called at the start of each forward step (in prepare()) so that + _ensure_pool_view_cached() recomputes them for the new batch. + """ + self._pool_cache_valid = False + + def _ensure_pool_view_cached(self): + """Compute and cache values used by + transform_local_topk_and_prepare_pool_view(). + + These values (pool view, stride factor, block table slices, request + index slices) are constant across all layers sharing the same KV pool + and batch dimensions within a forward pass. Caching them avoids + redundant Python/CUDA overhead per layer. + + Safety: _invalidate_pool_view_cache() is called unconditionally at the + 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): + return + + pool = self.kv_cache_manager.get_unique_primary_pool() + kv_cache_manager = self.kv_cache_manager + num_blocks, num_layers, _, _ = pool.shape + self._cached_tokens_per_block = kv_cache_manager.tokens_per_block + head_dim = kv_cache_manager.head_dim + self._cached_pool_view = pool.squeeze(2).view(-1, 1, head_dim) + self._cached_stride_factor = num_layers * self._cached_tokens_per_block + 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_req_idx_gen = ( + self.req_idx_per_token[self.num_ctx_tokens : self.num_tokens] - self.num_contexts + ) + self._cached_kv_mgr_id = id(kv_cache_manager) + self._pool_cache_valid = True + + @maybe_compile(dynamic=True) + def _get_dense_topk_indices(self, seq_lens, kv_lens, num_tokens): + device = kv_lens.device + past_kv_lens = kv_lens - seq_lens + # get position ids + seq_ends = torch.cumsum(seq_lens, dim=0) + seq_starts = seq_ends - seq_lens + per_seq_offsets = past_kv_lens - seq_starts # Shape: [batch_size] + global_indices = torch.arange(num_tokens, device=device) + batch_indices = torch.searchsorted(seq_ends, global_indices, side="right") + 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) + 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) + 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, + ), + 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, + ) + 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 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, + ), + 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, + ) + self.topk_indices_buffer[gen_range, :].copy_( + self.host_topk_indices_buffer[gen_range, :], non_blocking=True + ) + + def prepare_for_spec_decode(self, kv_lens: torch.Tensor): + # fp8_paged_mqa_logits supports seq_len 1/2/4 on sm100 and 1/2 on + # sm90. Flatten Q and expand kv_lens and block_table for other MTP + # configurations. + # TODO: + # - No distinction between sm90 and sm100 is needed once MTP3 is supported on sm90. + # - Remove this once fp8_paged_mqa_logits supports an arbitrary number of MTP draft tokens. + use_dsl = self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits + self.use_expanded_buffers_for_mtp = not use_dsl and ( + (self.max_draft_tokens > 1 and get_sm_version() == 90) + or ( + (self.max_draft_tokens == 2 or self.max_draft_tokens > 3) + and get_sm_version() >= 100 + ) + ) + if self.use_expanded_buffers_for_mtp: + # Expand kv_lens_cuda (only generation) + num_tokens = self.num_generations * (1 + self.max_draft_tokens) + gen_kv_lens = self.get_indexer_kv_lens(kv_lens[self.num_contexts : self.num_seqs]) + gen_kv_lens_expanded = torch.stack([gen_kv_lens] * (1 + self.max_draft_tokens), dim=0) + gen_kv_lens_expanded = gen_kv_lens_expanded.transpose(0, 1).contiguous().flatten() + 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 + ) + + # Expand indexer_k_cache_block_offsets (only generation) + # host_indexer_k_cache_block_offsets already contains correct pool + # indices from _get_pool_block_indices() in prepare_for_indexer_k_cache(). + if self.kv_cache_manager is not None and self.num_generations > 0: + 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 + ] + expanded_blocks = gen_block_tensor.repeat_interleave( + 1 + self.max_draft_tokens, dim=0 + ) + self.host_block_table_expanded[:num_tokens, :max_len].copy_( + expanded_blocks, non_blocking=True + ) + self.block_table_expanded[:num_tokens].copy_( + self.host_block_table_expanded[:num_tokens], non_blocking=True + ) + self.block_table_expanded.clamp_(min=0) + + self.expand_for_dsl = ( + use_dsl and self.kv_cache_manager is not None and self.max_draft_tokens >= 1 + ) + if self.expand_for_dsl and self.num_generations > 0: + next_n = 1 + self.max_draft_tokens + 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]) + max_ctx = int(gen_kv_lens.max().item()) if gen_kv_lens.numel() else 0 + expand_factor, atom = _pick_dsl_expand( + next_n, + batch_size=self.num_generations, + max_ctx=max_ctx, + num_sms=self.num_sms, + kernel_atoms=kernel_atoms, + ) + self.dsl_expand_factor = expand_factor + self.dsl_atom = atom + if expand_factor > 1: + num_tokens = self.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_( + self.kv_lens_expanded_host[:num_tokens], non_blocking=True + ) + 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 + ] + expanded_blocks = gen_block_tensor.repeat_interleave(expand_factor, dim=0) + self.host_block_table_expanded[:num_tokens, :max_len].copy_( + expanded_blocks, non_blocking=True + ) + self.block_table_expanded[:num_tokens].copy_( + self.host_block_table_expanded[:num_tokens], non_blocking=True + ) + self.block_table_expanded.clamp_(min=0) + else: + self.dsl_expand_factor = 1 + self.dsl_atom = 1 + self.max_draft_tokens + + def prepare_for_indexer_k_cache(self): + # Build indexer_k_cache_block_offsets using pool block indices derived + # from host_kv_cache_block_offsets (populated by super().prepare()). + # This correctly resolves block IDs to memory pool indices, which is + # required when host cache offload is enabled (block IDs != pool indices + # for onboarded secondary blocks). + if self.kv_cache_manager is None: + return + pool_indices = self._get_pool_block_indices() + self.host_indexer_k_cache_block_offsets[: self.num_seqs].copy_(pool_indices) + self.indexer_k_cache_block_offsets[: self.num_seqs].copy_( + self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True + ) + # Safety clamp: prevent OOB from CUDA graph padding entries which + # may contain stale negative or out-of-range values after block + # eviction/onboarding with host cache offload. + self.indexer_k_cache_block_offsets.clamp_(min=0) + + # Build block_table for topk_indices conversion (actual block allocation) + cached_token_lens = torch.tensor( + self.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] + kv_lens = cached_token_lens.clone() + kv_lens[active_rank] += self.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 + # 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. + host_block_table = torch.empty( + (pool_indices.shape[0], max_blocks_used), + dtype=pool_indices.dtype, + pin_memory=prefer_pinned(), + ) + 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 + ) + # Copy to GPU + self.block_table[: self.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: + # 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 + ) + else: + self.skip_indexer_for_ctx_reqs = False + + if self.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 + ) + 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.max_ctx_kv_len = kv_lens[: self.num_contexts].max().item() + self.max_ctx_seq_len = self.seq_lens[: self.num_contexts].max().item() + # context cached token indptr + torch.cumsum( + cached_token_lens[: self.num_contexts], + dim=0, + dtype=torch.int64, + out=self.host_ctx_cached_token_indptr[1 : self.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 + ) + # context kv indptr + torch.cumsum( + kv_lens[: self.num_contexts], + dim=0, + dtype=torch.int64, + out=self.host_ctx_kv_indptr[1 : self.num_contexts + 1], + ) + self.ctx_kv_indptr[: self.num_contexts + 1].copy_( + self.host_ctx_kv_indptr[: self.num_contexts + 1], non_blocking=True + ) + else: + self.num_ctx_cached_tokens = 0 + self.max_ctx_kv_len = 0 + self.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() + # generation cached token indptr + torch.cumsum( + cached_token_lens[self.num_contexts : self.num_seqs], + dim=0, + dtype=torch.int64, + out=self.host_gen_cached_token_indptr[1 : self.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 + ) + # generation kv indptr + torch.cumsum( + kv_lens[self.num_contexts : self.num_seqs], + dim=0, + dtype=torch.int64, + out=self.host_gen_kv_indptr[1 : self.num_generations + 1], + ) + self.gen_kv_indptr[: self.num_generations + 1].copy_( + self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True + ) + else: + self.max_gen_seq_len = 0 + + 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, + dim=0, + ) + self.req_idx_per_token[: self.num_tokens].copy_( + self.host_req_idx_per_token[: self.num_tokens], + non_blocking=True, + ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py new file mode 100644 index 000000000000..743ccf97c1f5 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py @@ -0,0 +1,585 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSA integration for the shared MLA module.""" + +import os +from typing import List, Optional + +import torch + +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionForwardArgs, + AttentionInputType, + AttentionMetadata, +) +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.utils import Fp4QuantizedTensor, is_torch_compiling +from tensorrt_llm._utils import get_sm_version, nvtx_range, nvtx_range_debug +from tensorrt_llm.logger import logger + +from .metadata import DSAtrtllmAttentionMetadata +from .params import DSABackendForwardArgs + +try: + from tensorrt_llm.flash_mla import flash_mla_sparse_fwd +except ImportError: + flash_mla_sparse_fwd = None + + +def initialize_sparse_attn( + self, + *, + config, + mapping, + mapping_o, + rms_norm_eps: float, + quant_config, + q_scaling: float, + bias: bool, + dtype: torch.dtype, + reduce_output: bool, + aux_stream: Optional[torch.cuda.Stream], +) -> None: + """Initialize DSA state and remove the unused dense MHA backend.""" + del mapping, mapping_o, rms_norm_eps, quant_config, q_scaling + del bias, dtype, reduce_output, aux_stream + + threshold = os.environ.get("TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD", "0") + try: + self.short_seq_mha_threshold = int(threshold) + except ValueError as err: + raise ValueError( + f"TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD must be an integer, got {threshold!r}" + ) from err + + use_short_seq_mha = self.short_seq_mha_threshold > 0 and self.mqa.support_fused_rope() + if not use_short_seq_mha: + self.mha = None + self.indexer = getattr(self.mqa, "indexer", None) + + pp_mapping = config.mapping + num_hidden_layers = getattr(config.pretrained_config, "num_hidden_layers", None) + if ( + pp_mapping.has_pp() + and self.layer_idx is not None + and num_hidden_layers is not None + and self.indexer is None + ): + pp_layers = pp_mapping.pp_layers(num_hidden_layers) + if pp_layers and self.layer_idx == pp_layers[0]: + raise ValueError( + f"DSA pipeline stage {pp_mapping.pp_rank} starts with " + f"shared-indexer layer {self.layer_idx}, but shared TopK data " + "is local to a pipeline stage. Choose a pp_partition whose " + "first layer on every stage runs a full indexer." + ) + + +def forward_context_sparse_attn( + self, + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + attn_metadata: AttentionMetadata, + output: torch.Tensor, + latent_cache: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + indexer_intermediates: Optional[List[torch.Tensor]] = None, +) -> torch.Tensor: + if should_use_short_mha(self, attn_metadata, position_ids): + return self.forward_context( + q, compressed_kv, k_pe, position_ids, attn_metadata, output, latent_cache + ) + if get_sm_version() >= 100: + return self.forward_absorption_context( + q, + compressed_kv, + k_pe, + attn_metadata, + output, + position_ids=position_ids, + latent_cache=latent_cache, + sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=indexer_intermediates), + ) + return forward_sparse_mla_kvcache_bf16( + self, + q, + latent_cache, + attn_metadata, + output, + indexer_intermediates, + is_generation=False, + ) + + +def forward_generation_sparse_attn( + self, + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + attn_metadata: AttentionMetadata, + output: torch.Tensor, + position_ids: Optional[torch.Tensor] = None, + latent_cache: Optional[torch.Tensor] = None, + indexer_intermediates: Optional[List[torch.Tensor]] = None, +) -> torch.Tensor: + if get_sm_version() >= 100: + return self.forward_absorption_generation( + q, + compressed_kv, + k_pe, + attn_metadata, + output, + position_ids=position_ids, + latent_cache=latent_cache, + sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=indexer_intermediates), + ) + return forward_sparse_mla_kvcache_bf16( + self, + q, + latent_cache, + attn_metadata, + output, + indexer_intermediates, + is_generation=True, + ) + + +def forward_sparse_attn( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + attn_output: list[torch.Tensor], +) -> None: + """ + Forward pass for the MLA module with DSA (always in MQA mode). + Writes result into output tensor in-place. + + Delegates to forward_dsa_proj (token-wise projections) followed by + _forward_dsa_attn (batch-dependent attention dispatch). + + Args: + position_ids (Optional[torch.IntTensor]): The position IDs. + hidden_states (torch.Tensor): The hidden states. + attn_metadata (AttentionMetadata): The attention metadata. + attn_output: Algorithm-defined output tensors. DSA writes to the + standard output at index zero. + """ + output = attn_output[0] + proj_outputs = forward_dsa_proj(self, position_ids, hidden_states, attn_metadata) + q, compressed_kv, k_pe, latent_cache = proj_outputs[:4] + indexer_intermediates = proj_outputs[4:] + _forward_dsa_attn( + self, + q, + compressed_kv, + k_pe, + latent_cache, + indexer_intermediates, + position_ids, + attn_metadata, + output, + ) + + +def forward_sparse_attn_custom_op( + self, + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_output: list[torch.Tensor], + latent_cache_gen: Optional[torch.Tensor], +) -> None: + """Run DSA's two-stage custom ops for piecewise CUDA graph capture.""" + del latent_cache_gen + from . import custom_ops # noqa: F401 + + if isinstance(hidden_states, Fp4QuantizedTensor): + proj_outputs = torch.ops.trtllm.mla_dsa_proj( + hidden_states.unquantized_hidden_states, + position_ids, + self.layer_idx_str, + hidden_states.fp4_tensor, + hidden_states.scaling_factor, + ) + else: + proj_outputs = torch.ops.trtllm.mla_dsa_proj( + hidden_states, position_ids, self.layer_idx_str + ) + q, compressed_kv, k_pe, latent_cache = proj_outputs[:4] + indexer_intermediates = proj_outputs[4:] + torch.ops.trtllm.mla_dsa_attn_inplace( + q, + compressed_kv, + k_pe, + latent_cache, + indexer_intermediates, + position_ids, + self.layer_idx_str, + attn_output[0], + ) + + +def forward_dsa_proj( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, +) -> List[torch.Tensor]: + """Token-wise projections for DSA MLA (CUDA-graph-capturable Op 1). + + Runs kv_a_proj, layernorms, q_b_proj, and conditionally + indexer.pre_indexer_proj(). + + IMPORTANT: This method must NOT slice tensors by num_tokens or + access batch-specific metadata, so that all operations are + unconditionally straight-line for CUDA graph capture. Slicing + to num_tokens happens in _forward_dsa_attn (Op 2, outside graph). + + Returns [q, compressed_kv, k_pe, latent_cache] when short-MHA + handles all tokens (eager only), or + [q, compressed_kv, k_pe, latent_cache, q_fp8, k_fp8, k_scale, + weights, q_scale] when the indexer runs. q_scale is unused on the + FP8 path but always present so CUDA graph capture sees a stable + 9-tensor shape regardless of indexer dtype. + """ + assert self.mqa is not None, "DSA is only supported in MQA mode" + + q, compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1 + ) + + q_pair, compressed_kv = maybe_execute_in_parallel( + lambda: self._q_a_layernorm_maybe_fused(q, return_norm_out=True), + lambda: self.kv_a_layernorm(compressed_kv), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + q, qr = q_pair + latent_cache = torch.concat([compressed_kv, k_pe], dim=-1) + + q = self.q_b_proj(q) + + use_short_mha_for_ctx = should_use_short_mha(self, attn_metadata, position_ids) + + # Skip the indexer when the short MHA path handles all context + # tokens and there are no generation tokens. + if use_short_mha_for_ctx and attn_metadata.num_generations == 0: + return [q, compressed_kv, k_pe, latent_cache] + + # DSA "shared" layer: no indexer projection; the backend predictor + # reuses the preceding full layer's top-k. + if self.mqa.indexer is None: + return [q, compressed_kv, k_pe, latent_cache] + + # pre_indexer_proj is the CUDA-graph-safe portion: pure token-wise + # compute (cublas_mm, rope, FP4/FP8 quantize, weight scaling) with no + # access to batch-specific metadata or the k cache. Returns q_scale + # as a 5th element so the FP4 dispatch can forward it to the kernel; + # the FP8 path ignores it in _forward_dsa_attn. + q_fp8, k_fp8, k_scale, weights, q_scale = self.mqa.indexer.pre_indexer_proj( + qr, hidden_states, position_ids + ) + + return [ + q, + compressed_kv, + k_pe, + latent_cache, + q_fp8, + k_fp8, + k_scale, + weights, + q_scale, + ] + + +def _forward_dsa_attn( + self, + q: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + latent_cache: torch.Tensor, + indexer_intermediates: List[torch.Tensor], + position_ids: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + output: torch.Tensor, +) -> None: + """Batch-structure-dependent attention for DSA MLA (Op 2, not graph-captured). + + indexer_intermediates is [q_fp8, k_fp8, k_scale, weights, q_scale] + when the indexer ran in Op 1, or [] when short-MHA handled all tokens. + + All num_tokens slicing happens here (not in Op 1) because + num_tokens comes from batch-specific metadata and must not be + baked into CUDA graph capture. + """ + num_contexts = attn_metadata.num_contexts + num_generations = attn_metadata.num_generations + num_ctx_tokens = attn_metadata.num_ctx_tokens + num_tokens = attn_metadata.num_tokens + + # Slice Op 1 outputs to actual num_tokens (Op 1 operates on the + # full padded tensor for CUDA graph compatibility). + q = q[:num_tokens, ...] + compressed_kv = compressed_kv[:num_tokens, ...] + k_pe = k_pe[:num_tokens, ...] + latent_cache = latent_cache[:num_tokens, ...] + if position_ids is not None: + position_ids = position_ids[..., :num_tokens] + + use_short_mha_for_ctx = num_contexts > 0 and should_use_short_mha( + self, attn_metadata, position_ids + ) + + if not (use_short_mha_for_ctx and num_generations == 0) and self.mqa.indexer is not None: + q_fp8, k_fp8, k_scale, weights, q_scale = indexer_intermediates + # Slice indexer intermediates to actual num_tokens (they were + # computed on the full padded tensor in Op 1). + q_fp8 = q_fp8[:num_tokens, ...] + k_fp8 = k_fp8[:num_tokens, ...] + k_scale = k_scale[:num_tokens, ...] + weights = weights[:num_tokens, ...] + q_scale = q_scale[:num_tokens, ...] + indexer_intermediates = [q_fp8, k_fp8, k_scale, weights, q_scale] + # Update the full mixed batch once before context/generation prediction + # splits. Context chunk gathers must observe the newly appended keys. + self.mqa.indexer._update_k_cache(k_fp8, k_scale, attn_metadata) + else: + indexer_intermediates = None + + assert output is not None, "output must be provided" + + if num_contexts > 0: + q_ctx = q[:num_ctx_tokens, ...] + compressed_kv_ctx = compressed_kv[:num_ctx_tokens, ...] + k_pe_ctx = k_pe[:num_ctx_tokens, ...] + latent_cache_ctx = latent_cache[:num_ctx_tokens, ...] + ctx_position_ids = position_ids[..., :num_ctx_tokens] if position_ids is not None else None + if self.apply_rotary_emb: + assert ctx_position_ids is not None + k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids) + + forward_context_sparse_attn( + self, + q_ctx, + compressed_kv_ctx, + k_pe_ctx, + attn_metadata, + output[:num_ctx_tokens, :], + latent_cache_ctx, + position_ids=ctx_position_ids, + indexer_intermediates=indexer_intermediates, + ) + + if num_generations > 0: + q_gen = q[num_ctx_tokens:, ...] + compressed_kv_gen = compressed_kv[num_ctx_tokens:, ...] + k_pe_gen = k_pe[num_ctx_tokens:, ...] + latent_cache_gen = latent_cache[num_ctx_tokens:, ...] + gen_position_ids = ( + position_ids[..., num_ctx_tokens:num_tokens] if position_ids is not None else None + ) + if self.apply_rotary_emb: + assert gen_position_ids is not None + k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids) + + forward_generation_sparse_attn( + self, + q_gen, + compressed_kv_gen, + k_pe_gen, + attn_metadata, + output[num_ctx_tokens:num_tokens, :], + latent_cache=latent_cache_gen, + position_ids=gen_position_ids, + indexer_intermediates=indexer_intermediates, + ) + + +def should_use_short_mha( + self, attn_metadata: AttentionMetadata, position_ids: Optional[torch.Tensor] +) -> bool: + """Check if the short-seq MHA optimization should be used for context. + + Uses max_ctx_kv_len (max total KV length per context sequence, + including cached tokens) when available, to correctly account for + chunked context where the full attention span exceeds the threshold + even if the new token count is small. Falls back to num_ctx_tokens + (total new context tokens) when max_ctx_kv_len is not set. + + Disabled under torch compile so that the split DSA custom ops + (mla_dsa_proj / mla_dsa_attn_inplace) have unconditionally + straight-line control flow for CUDA graph capture. + """ + if is_torch_compiling(): + return False + if not ( + self.short_seq_mha_threshold > 0 + and not self.apply_rotary_emb + and self.mapping.cp_size == 1 + and position_ids is not None + ): + return False + effective_len = getattr(attn_metadata, "max_ctx_kv_len", attn_metadata.num_ctx_tokens) + return effective_len <= self.short_seq_mha_threshold + + +@nvtx_range("forward_sparse_mla_kvcache_bf16") +def forward_sparse_mla_kvcache_bf16( + self, + q: torch.Tensor, + latent_cache: torch.Tensor, + attn_metadata: DSAtrtllmAttentionMetadata, + output: torch.Tensor, + indexer_intermediates: Optional[List[torch.Tensor]], + is_generation: bool = False, +) -> torch.Tensor: + """ + Forward sparse MLA (DSA) for BF16 KV cache for both context and generation phases using FlashMLA kernels + + To form the input for FlashMLA kernel and adapt our KV cache manager, we need to: + 1. Append current tokens to paged cache and apply rope to q/k via mla_rope_append_paged_kv_assign_q + 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" + ) + attention_input_type = ( + AttentionInputType.generation_only if is_generation else AttentionInputType.context_only + ) + topk_indices_pool, _ = self.mqa.sparse_attn_predict( + q, + None, + attn_metadata, + AttentionForwardArgs( + attention_input_type=attention_input_type, + sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=indexer_intermediates), + ), + ) + assert topk_indices_pool is not None + attn_metadata._ensure_pool_view_cached() + kv_cache_pool = attn_metadata._cached_pool_view + assert kv_cache_pool is not None + + # Append current tokens to paged cache and apply RoPE to q + # This writes latent_cache to paged KV and modifies q in-place + trtllm_attention = self.mqa + with nvtx_range_debug(f"mla_rope_append_paged_kv_assign_q_is_generation={is_generation}"): + trtllm_attention.mla_rope_append_paged_kv_assign_q( + q, latent_cache, attn_metadata, is_generation=is_generation + ) + + num_tokens = q.shape[0] + q_nope, q_rope = q.view(-1, self.num_heads_tp, self.qk_head_dim).split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + q_nope_out = torch.empty( + [num_tokens, self.num_heads_tp, (self.kv_lora_rank)], + dtype=q.dtype, + device=q.device, + ) + + if self.k_b_proj_trans.dtype == torch.bfloat16: + # [num_heads, num_tokens, self.qk_nope_head_dim] + q_nope_t = q_nope.transpose(0, 1) + # [num_heads, num_tokens, self.kv_lora_rank] + q_nope_out = q_nope_out.transpose(0, 1) + + # [num_heads, num_tokens, self.qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim] + # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank] + # The output of bmm is written directly into fused_q + self._bmm_bf16_out( + q_nope_t, + self.k_b_proj_trans, + self.k_b_proj_trans.transpose(1, 2), + q_nope_out, + ) + elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn: + # [num_heads, num_tokens, self.kv_lora_rank] + q_nope_out = q_nope_out.transpose(0, 1) + + from tensorrt_llm._torch.modules.mla import fp8_block_scaling_bmm_out + + fp8_block_scaling_bmm_out( + q_nope, + self.k_b_proj_trans, + self.k_b_proj_trans_scale, + q_nope_out, + self.k_b_proj_trans_dequant, + self.use_cute_dsl_blockscaling_bmm, + ) + else: + raise NotImplementedError(f"Missing bmm impl for dtype: {self.k_b_proj_trans.dtype}.") + + q_nope_out = q_nope_out.transpose(0, 1) + q_concat = torch.cat([q_nope_out, q_rope], dim=-1) + + sm_version = get_sm_version() + # FlashMLA sparse kernel (bf16) requires num_heads=128 on sm100 or multiple of 64 on sm90 + if sm_version >= 100: + padding = 128 + assert self.num_heads_tp <= padding, ( + f"SM100 FlashMLA sparse kernel requires exactly {padding} heads, " + f"got {self.num_heads_tp}. Padding from values > {padding} is not supported." + ) + else: # SM90 + padding = ((self.num_heads_tp + 63) // 64) * 64 # multiple of 64 + + if self.num_heads_tp != padding: + logger.warning_once( + f"Padding num_heads from {self.num_heads_tp} to {padding} " + f"due to FlashMLA sparse attention kernel requirement", + key="sparse_mla_padding_warning", + ) + + # Create padded tensor with zeros for extra heads + q_padded = q_concat.new_empty((num_tokens, padding, q_concat.shape[2])) + q_padded[:, : self.num_heads_tp, :] = q_concat + q_concat = q_padded + + topk_indices_pool = topk_indices_pool.view(num_tokens, 1, -1) + if flash_mla_sparse_fwd is not None: + attn_out_latent = flash_mla_sparse_fwd( + q_concat, kv_cache_pool, topk_indices_pool, self.softmax_scale + )[0] + else: + raise RuntimeError( + "flash_mla_sparse_fwd not available. Please ensure FlashMLA module is built." + ) + + # [seq, num_heads, kv_lora_rank], account for padding + attn_out_latent = attn_out_latent[:, : self.num_heads_tp, :] + attn_out_latent = attn_out_latent.view([-1, self.num_heads_tp, self.kv_lora_rank]) + if self.num_heads_tp != padding: + attn_out_latent = attn_out_latent.contiguous() + + assert attn_out_latent.shape[0] == q.shape[0] and attn_out_latent.shape[1] == self.num_heads_tp + + attn_output = output.view([num_tokens, self.num_heads_tp, self.v_head_dim]) + + if self.v_b_proj.dtype == torch.bfloat16: + # [num_heads, seq, kv_lora_rank] x [num_heads, kv_lora_rank, v_head_dim] + # -> [num_heads, seq, v_head_dim] + self._bmm_bf16_out( + attn_out_latent.transpose(0, 1), + self.v_b_proj, + self.v_b_proj.transpose(1, 2), + attn_output.transpose(0, 1), + ) + elif self.v_b_proj.dtype == torch.float8_e4m3fn: + from tensorrt_llm._torch.modules.mla import fp8_block_scaling_bmm_out + + fp8_block_scaling_bmm_out( + attn_out_latent, + self.v_b_proj, + self.v_b_proj_scale, + attn_output.transpose(0, 1), + self.v_b_proj_dequant, + self.use_cute_dsl_blockscaling_bmm, + ) + else: + raise NotImplementedError(f"Missing bmm impl for dtype: {self.v_b_proj.dtype}.") + return output diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py new file mode 100644 index 000000000000..ecaffd89eb35 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSA parameter types.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, List, Literal, Optional + +import torch + +import tensorrt_llm +import tensorrt_llm.bindings + +from ..params import SparseBackendForwardArgs, SparseMetadataParams, SparseParams + +ModelConfig = tensorrt_llm.bindings.ModelConfig + +if TYPE_CHECKING: + pass + + +@dataclass(kw_only=True, slots=True) +class DSABackendForwardArgs(SparseBackendForwardArgs): + """DSA inputs passed from the MLA module to its backend.""" + + indexer_intermediates: Optional[List[torch.Tensor]] = None + + +@dataclass(frozen=True) +class DSAMetadataParams(SparseMetadataParams): + """DSA metadata parameters.""" + + indexer_max_chunk_size: int + max_sparse_topk: Optional[int] + index_head_dim: int + enable_indexer_skip: bool + enable_heuristic_topk: bool + use_cute_dsl_topk: bool + use_cute_dsl_paged_mqa_logits: bool + q_split_threshold: int + has_shared_indexer_layers: bool = False + + +@dataclass(frozen=True) +class DSAParams(SparseParams): + """DSA backend parameters.""" + + algorithm: Literal["dsa"] = field(init=False, default="dsa") + index_n_heads: Optional[int] = None + index_head_dim: Optional[int] = None + index_topk: Optional[int] = None + indexer_max_chunk_size: Optional[int] = None + skip_indexer_for_short_seqs: bool = True + use_cute_dsl_topk: bool = False + use_cute_dsl_paged_mqa_logits: bool = False + q_split_threshold: int = 8192 + indexer_rope_interleave: bool = False + enable_heuristic_topk: bool = False + indexer_k_dtype: Literal["fp8", "fp4"] = "fp8" + # Shared layers reuse the preceding full layer's top-k. + is_full_indexer_layer: bool = True + + @property + def indices_block_size(self) -> int: + return 1 diff --git a/tensorrt_llm/_torch/attention_backend/sparse/hooks.py b/tensorrt_llm/_torch/attention_backend/sparse/hooks.py new file mode 100644 index 000000000000..a7cead7380b5 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/hooks.py @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Module and backend hooks for sparse attention algorithms. + +Algorithms implement only the hooks they need in +``sparse//module.py``. This module validates those hooks and owns +module dispatch. Backend prediction hooks use the backend subclass directly. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from functools import lru_cache +from importlib import import_module +from inspect import Parameter, signature +from types import ModuleType +from typing import TYPE_CHECKING, Callable, Optional + +if TYPE_CHECKING: + import torch + + from ..interface import AttentionForwardArgs, AttentionMetadata + from ..trtllm import TrtllmAttention + from .params import SparseRuntimeParams + +SparseAttnHook = Callable[..., object] + +__all__ = [ + "SparseAttnHooks", + "get_sparse_attn_hooks", + "prepare_sparse_runtime_params", +] + +_SPARSE_ATTN_HOOK_MODULE_PATHS = { + "rocket": ".rocket.module", + "dsa": ".dsa.module", + "deepseek_v4": ".deepseek_v4.module", +} +_HOOK_PARAMETER_NAMES = { + "initialize_sparse_attn": ( + "module", + "config", + "mapping", + "mapping_o", + "rms_norm_eps", + "quant_config", + "q_scaling", + "bias", + "dtype", + "reduce_output", + "aux_stream", + ), + "create_sparse_attn_weights": ("module",), + "transform_sparse_attn_weights": ("module",), + "prepare_sparse_attn_outputs": ("module", "hidden_states", "attn_metadata"), + "forward_sparse_attn": ( + "module", + "position_ids", + "hidden_states", + "attn_metadata", + "attn_output", + ), + "forward_sparse_attn_custom_op": ( + "module", + "hidden_states", + "position_ids", + "attn_output", + "latent_cache_gen", + ), + "project_sparse_attn_output": ( + "module", + "attn_output", + "position_ids", + "attn_metadata", + "all_reduce_params", + ), +} +_ALTERNATE_HOOK_PARAMETER_NAMES = { + "forward_sparse_attn": ( + ( + "module", + "q", + "k", + "v", + "attn_metadata", + "attention_mask", + "attention_window_size", + "attention_mask_data", + "mrope_config", + "attention_sinks", + "relative_attention_bias", + "relative_attention_max_distance", + "has_lora", + "kwargs", + ), + ), + "project_sparse_attn_output": ( + ( + "module", + "attn_output", + "attn_metadata", + "all_reduce_params", + "lora_params", + ), + ), +} +_INITIALIZE_KEYWORD_ONLY_PARAMETERS = frozenset(_HOOK_PARAMETER_NAMES["initialize_sparse_attn"][1:]) + + +def _get_hook( + module: ModuleType, + hook_name: str, + *, + algorithm: str, +) -> Optional[SparseAttnHook]: + hook = getattr(module, hook_name, None) + if hook is None: + return None + if not callable(hook): + raise TypeError( + f"Sparse attention hook {algorithm!r}.{hook_name} must be callable, " + f"got {type(hook).__name__}" + ) + + parameters = tuple(signature(hook).parameters.values()) + parameter_names = tuple( + "module" if index == 0 and parameter.name in ("self", "mla") else parameter.name + for index, parameter in enumerate(parameters) + ) + expected_names = ( + _HOOK_PARAMETER_NAMES[hook_name], + *_ALTERNATE_HOOK_PARAMETER_NAMES.get(hook_name, ()), + ) + if parameter_names not in expected_names: + raise TypeError( + f"Sparse attention hook {algorithm!r}.{hook_name} has parameters {parameter_names}; " + f"expected one of {expected_names}" + ) + if hook_name == "initialize_sparse_attn": + invalid_keyword_only = [ + parameter.name + for parameter in parameters + if parameter.name in _INITIALIZE_KEYWORD_ONLY_PARAMETERS + and parameter.kind != Parameter.KEYWORD_ONLY + ] + if invalid_keyword_only: + raise TypeError( + f"Sparse attention hook {algorithm!r}.{hook_name} must declare these parameters " + f"as keyword-only: {', '.join(invalid_keyword_only)}" + ) + if parameter_names and parameter_names[-1] == "kwargs": + if parameters[-1].kind != Parameter.VAR_KEYWORD: + raise TypeError( + f"Sparse attention hook {algorithm!r}.{hook_name} must declare 'kwargs' as **kwargs" + ) + return hook + + +@dataclass(frozen=True) +class SparseAttnHooks: + """Validated module-layer hooks for one sparse attention algorithm.""" + + algorithm: Optional[str] = None + initialize_sparse_attn: Optional[SparseAttnHook] = None + create_sparse_attn_weights: Optional[SparseAttnHook] = None + transform_sparse_attn_weights: Optional[SparseAttnHook] = None + prepare_sparse_attn_outputs: Optional[SparseAttnHook] = None + forward_sparse_attn: Optional[SparseAttnHook] = None + forward_sparse_attn_custom_op: Optional[SparseAttnHook] = None + project_sparse_attn_output: Optional[SparseAttnHook] = None + + def __bool__(self) -> bool: + """Return whether sparse attention is configured for the module.""" + return self.algorithm is not None + + def require(self, hook_name: str) -> SparseAttnHook: + """Return an implemented hook required by the current module path.""" + if hook_name not in _HOOK_PARAMETER_NAMES: + raise ValueError(f"Unknown sparse attention hook {hook_name!r}") + hook = getattr(self, hook_name) + if hook is None: + raise NotImplementedError( + f"Sparse attention algorithm {self.algorithm!r} does not implement " + f"the {hook_name!r} hook required by this module path" + ) + return hook + + @classmethod + def from_module(cls, algorithm: str, module: ModuleType) -> "SparseAttnHooks": + """Validate and adapt an algorithm ``module.py`` to the hook contract.""" + return cls( + algorithm=algorithm, + initialize_sparse_attn=_get_hook(module, "initialize_sparse_attn", algorithm=algorithm), + create_sparse_attn_weights=_get_hook( + module, "create_sparse_attn_weights", algorithm=algorithm + ), + transform_sparse_attn_weights=_get_hook( + module, "transform_sparse_attn_weights", algorithm=algorithm + ), + prepare_sparse_attn_outputs=_get_hook( + module, "prepare_sparse_attn_outputs", algorithm=algorithm + ), + forward_sparse_attn=_get_hook(module, "forward_sparse_attn", algorithm=algorithm), + forward_sparse_attn_custom_op=_get_hook( + module, "forward_sparse_attn_custom_op", algorithm=algorithm + ), + project_sparse_attn_output=_get_hook( + module, "project_sparse_attn_output", algorithm=algorithm + ), + ) + + +_EMPTY_SPARSE_ATTN_HOOKS = SparseAttnHooks() + + +@lru_cache(maxsize=None) +def _get_sparse_attn_hooks_for_algorithm(algorithm: str) -> SparseAttnHooks: + module_name = _SPARSE_ATTN_HOOK_MODULE_PATHS.get(algorithm) + if module_name is None: + return _EMPTY_SPARSE_ATTN_HOOKS + module = import_module(module_name, package=__package__) + return SparseAttnHooks.from_module(algorithm, module) + + +def get_sparse_attn_hooks(module) -> SparseAttnHooks: + """Return hooks selected by ``module.sparse_params`` or an empty hook set.""" + algorithm = getattr(getattr(module, "sparse_params", None), "algorithm", None) + if algorithm is None: + return _EMPTY_SPARSE_ATTN_HOOKS + return _get_sparse_attn_hooks_for_algorithm(algorithm) + + +def prepare_sparse_runtime_params( + backend: "TrtllmAttention", + q: "torch.Tensor", + k: Optional["torch.Tensor"], + metadata: "AttentionMetadata", + forward_args: "AttentionForwardArgs", +) -> "SparseRuntimeParams": + """Run backend prediction hooks and update attention-op parameters.""" + runtime_params = forward_args.sparse_runtime_params + if backend.sparse_params is None: + return runtime_params + + kv_indices, kv_offsets = backend.sparse_kv_predict(q, k, metadata, forward_args) + attn_indices, attn_offsets = backend.sparse_attn_predict(q, k, metadata, forward_args) + block_size = ( + backend.sparse_params.indices_block_size + if attn_indices is not None or attn_offsets is not None + else runtime_params.sparse_attn_indices_block_size + ) + return replace( + runtime_params, + sparse_kv_indices=kv_indices, + sparse_kv_offsets=kv_offsets, + sparse_attn_indices=attn_indices, + sparse_attn_offsets=attn_offsets, + sparse_attn_indices_block_size=block_size, + ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py index 86cd67fa7272..52bae8974fe5 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -8,7 +8,7 @@ * The main sparse GQA runs through the registered MsaSparseGqaFmha. * The indexer calls fmha_sm100 directly to produce the per-query selected block indices, which the model layer threads through - forward_args.topk_indices. + 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 @@ -720,7 +720,7 @@ def run_indexer( """Write the index-K cache and return the selected block indices. The model layer runs this before forward and threads the result through - forward_args.topk_indices. Returns [total_q, num_kv_heads, topk]. + forward_args.sparse_backend_args. Returns [total_q, num_kv_heads, topk]. Decode uses the prebuilt graph-safe proxy plan; prefill plans eagerly. """ @@ -768,10 +768,11 @@ def sparse_attn_predict( metadata, forward_args: "AttentionForwardArgs", ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - # The model layer runs run_indexer and passes the selected block - # indices through forward_args.topk_indices. Publish them as the - # sparse attention indices MsaSparseGqaFmha reads. - return forward_args.topk_indices, None + # The model layer runs run_indexer and passes the selected blocks + # through the sparse backend payload. + sparse_backend_args = forward_args.sparse_backend_args + topk_indices = sparse_backend_args.topk_indices if sparse_backend_args is not None else None + return topk_indices, None def sparse_kv_predict( self, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/params.py b/tensorrt_llm/_torch/attention_backend/sparse/params.py index cafd26e53b19..99a0648d86c3 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/params.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/params.py @@ -12,49 +12,46 @@ # 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. -"""Base sparse parameter interfaces for attention backend boundaries. - -User-facing SparseAttentionConfig live in LLM / VisualGen args and -ModelConfig. They lower through ``to_sparse_params()`` for per-backend runtime -state and ``to_sparse_metadata_params()`` for metadata allocation/update state. - -Concrete sparse algorithm params live with their backend implementations; -shared kernel-facing carriers live here when they are part of the generic -attention-forward contract. -""" +"""Shared sparse attention parameter types.""" from dataclasses import dataclass +from typing import Optional +import torch -class SparseParams: - """Base for per AttentionBackend instance sparse runtime parameters. - User-facing SparseAttentionConfig objects lower into this type through - ``to_sparse_params()`` before an AttentionBackend is constructed. That - lowering can resolve per-model or per-layer settings without passing - configs into backend instances. - """ +class SparseParams: + """Base parameters for a sparse attention backend.""" algorithm: str class SparseMetadataParams: - """Base for sparse settings needed by AttentionMetadata. + """Base parameters for sparse attention metadata.""" + + +@dataclass(kw_only=True, slots=True) +class SparseBackendForwardArgs: + """Sparse inputs passed from an attention module to its backend.""" + + # Shared by algorithms that accept precomputed top-k indices. + topk_indices: Optional[torch.Tensor] = None - Derived from the same user-facing SparseAttentionConfig through - ``to_sparse_metadata_params()``, but kept separate from SparseParams because - metadata owns batch/runtime buffers rather than per-layer - ``AttentionBackend`` behavior. - """ +@dataclass(kw_only=True, slots=True) +class SparseRuntimeParams: + """Sparse runtime parameters passed from a backend to the attention op.""" -@dataclass -class SkipSoftmaxKernelParams: - """Skip-softmax thresholds passed to attention backend kernels.""" + sparse_kv_indices: Optional[torch.Tensor] = None + sparse_kv_offsets: Optional[torch.Tensor] = None + sparse_attn_indices: Optional[torch.Tensor] = None + sparse_attn_offsets: Optional[torch.Tensor] = None + sparse_attn_indices_block_size: int = 0 + # DeepSeek-V4 compressed-cache inputs. + sparse_mla_topk_lens: Optional[torch.Tensor] = None + compressed_kv_cache_pool_ptr: Optional[int] = None - # The kernel divides this by the context length to get the skip threshold; - # zero turns skip-softmax off. + # SkipSoftmax prefill threshold; kernels divide it by context length. threshold_scale_factor_prefill: float = 0.0 - # Only autoregressive (LLM) decoding has a decode phase; diffusion and - # visual generation leave this at zero. + # SkipSoftmax decode threshold; diffusion models leave it at zero. threshold_scale_factor_decode: float = 0.0 diff --git a/tensorrt_llm/_torch/attention_backend/sparse/utils.py b/tensorrt_llm/_torch/attention_backend/sparse/registry.py similarity index 82% rename from tensorrt_llm/_torch/attention_backend/sparse/utils.py rename to tensorrt_llm/_torch/attention_backend/sparse/registry.py index 1762268bf242..cd2cea5d2db2 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/registry.py @@ -1,16 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from typing import TYPE_CHECKING, Type, Union if TYPE_CHECKING: from tensorrt_llm._torch.attention_backend.interface import AttentionBackend - from tensorrt_llm.llmapi.llm_args import \ - SparseAttentionConfig as LlmSparseAttentionConfig - from tensorrt_llm.visual_gen.args import \ - SparseAttentionConfig as VisualGenSparseAttentionConfig + from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig as LlmSparseAttentionConfig + from tensorrt_llm.visual_gen.args import SparseAttentionConfig as VisualGenSparseAttentionConfig from .params import SparseParams - SparseAttentionConfig = Union[LlmSparseAttentionConfig, - VisualGenSparseAttentionConfig] + SparseAttentionConfig = Union[LlmSparseAttentionConfig, VisualGenSparseAttentionConfig] # Imports of the concrete backends / cache managers are kept local to each # function: they pull in ``trtllm`` and ``resource_manager``, which import @@ -18,14 +18,14 @@ # loaded. -def get_sparse_attn_kv_cache_manager( - sparse_attention_config: "SparseAttentionConfig") -> type: +def get_sparse_attn_kv_cache_manager(sparse_attention_config: "SparseAttentionConfig") -> type: from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from .deepseek_v4 import DeepseekV4CacheManager from .dsa import DSACacheManager from .minimax_m3 import MiniMaxM3KVCacheManagerV2 from .rocket import RocketKVCacheManager + if sparse_attention_config.algorithm == "rocket": return RocketKVCacheManager elif sparse_attention_config.algorithm == "dsa": @@ -42,8 +42,7 @@ def get_sparse_attn_kv_cache_manager( ) -def _resolve_minimax_m3_backend_cls( - sparse_params: "SparseParams") -> Type["AttentionBackend"]: +def _resolve_minimax_m3_backend_cls(sparse_params: "SparseParams") -> Type["AttentionBackend"]: """Select the MiniMax-M3 sparse backend from the lowered params. The Triton reference is the default. When implementation is 'msa' the MSA @@ -51,17 +50,21 @@ def _resolve_minimax_m3_backend_cls( unsupported system fails early rather than at kernel launch. """ from .minimax_m3 import MiniMaxM3SparseRuntimeBackend + if getattr(sparse_params, "implementation", "triton") == "msa": from .minimax_m3 import MiniMaxM3MsaSparseAttention from .minimax_m3.msa_availability import ensure_msa_available + ensure_msa_available() return MiniMaxM3MsaSparseAttention return MiniMaxM3SparseRuntimeBackend def get_vanilla_sparse_attn_attention_backend( - sparse_params: "SparseParams") -> Type["AttentionBackend"]: + sparse_params: "SparseParams", +) -> Type["AttentionBackend"]: from .rocket import RocketVanillaAttention + if sparse_params.algorithm == "rocket": return RocketVanillaAttention elif sparse_params.algorithm == "minimax_m3": @@ -73,12 +76,14 @@ def get_vanilla_sparse_attn_attention_backend( def get_trtllm_sparse_attn_attention_backend( - sparse_params: "SparseParams") -> Type["AttentionBackend"]: + sparse_params: "SparseParams", +) -> Type["AttentionBackend"]: from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention from .deepseek_v4 import DeepseekV4TrtllmAttention from .dsa import DSATrtllmAttention from .rocket import RocketTrtllmAttention + if sparse_params.algorithm == "rocket": return RocketTrtllmAttention elif sparse_params.algorithm == "dsa": @@ -101,7 +106,8 @@ def get_trtllm_sparse_attn_attention_backend( def get_flashinfer_sparse_attn_attention_backend( - sparse_params: "SparseParams") -> Type["AttentionBackend"]: + sparse_params: "SparseParams", +) -> Type["AttentionBackend"]: if sparse_params.algorithm == "minimax_m3": return _resolve_minimax_m3_backend_cls(sparse_params) raise ValueError( diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket.py deleted file mode 100644 index 325cc431dc34..000000000000 --- a/tensorrt_llm/_torch/attention_backend/sparse/rocket.py +++ /dev/null @@ -1,1186 +0,0 @@ -import math -from dataclasses import dataclass, field -from typing import (TYPE_CHECKING, Iterable, List, Literal, Optional, Tuple, - Union) - -import torch -from torch import Tensor -from triton import next_power_of_2 - -import tensorrt_llm -import tensorrt_llm.bindings -from tensorrt_llm._torch.attention_backend.interface import ( - AttentionForwardArgs, AttentionMetadata) -from tensorrt_llm._torch.attention_backend.trtllm import ( - TrtllmAttention, TrtllmAttentionMetadata) -from tensorrt_llm._torch.attention_backend.vanilla import ( - VanillaAttention, VanillaAttentionMetadata, repeat_kv) -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.pyexecutor.resource_manager import (BlockManager, - KVCacheManager) -from tensorrt_llm._utils import get_size_in_bytes, prefer_pinned -from tensorrt_llm.bindings import DataType -from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.bindings.internal.batch_manager import \ - CacheType as CacheTypeCpp -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantConfig - -if TYPE_CHECKING: - from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, - SparseAttentionConfig) - -from .kernel import (triton_bmm, triton_flatten_to_batch, triton_index_gather, - triton_rocket_batch_to_flatten, - triton_rocket_paged_kt_cache_bmm, triton_rocket_qk_split, - triton_rocket_reduce_scores, - triton_rocket_update_kt_cache_ctx, - triton_rocket_update_kt_cache_gen, triton_softmax, - triton_topk) -from .params import SparseMetadataParams, SparseParams - -ModelConfig = tensorrt_llm.bindings.ModelConfig - - -@dataclass(frozen=True) -class RocketKVMetadataParams(SparseMetadataParams): - """RocketKV metadata settings derived from user config.""" - - prompt_budget: int - window_size: int - page_size: int - topk: int - - -@dataclass(frozen=True) -class RocketKVParams(SparseParams): - """RocketKV sparse attention backend parameters.""" - - algorithm: Literal["rocket"] = field(init=False, default="rocket") - window_size: int = 32 - kernel_size: int = 63 - topr: Union[int, float] = 128 - topk: int = 64 - prompt_budget: int = 2048 - page_size: int = 4 - kt_cache_dtype: str = "float8_e5m2" - - @property - def indices_block_size(self) -> int: - return self.page_size - - -class RocketTrtllmAttentionMetadata(TrtllmAttentionMetadata): - - 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.window_size = sparse_metadata_params.window_size - self.page_size = sparse_metadata_params.page_size - self.topk = sparse_metadata_params.topk - - 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 - - # 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, - ) - - self.q_cu_seqlens = torch.zeros_like(self.q_cu_seqlens_cuda, - 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) - - # 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) - - # 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, - ) - - # 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) - - # 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) - - # 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, - ) - - # 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.host_kt_cache_block_offsets = torch.zeros_like( - self.kt_cache_block_offsets, - device='cpu', - pin_memory=prefer_pinned(), - ) - - # Number of KT tokens for each sequence - self.num_kt_tokens = torch.empty( - self.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) - - # 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) - - # 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] - - # 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) - - # -------------------------------- Context phase -------------------------------- - self.context_cumsum[1:self.num_contexts + 1] = torch.cumsum( - self.prompt_lens_cpu[:self.num_contexts], dim=0) - self.context_cumsum_cuda[:self.num_contexts + 1].copy_( - self.context_cumsum[:self.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_seq_indices = torch.where(valid_mask)[0] - invalid_seq_indices = torch.where(~valid_mask)[0] - valid_batch_size = len(valid_seq_indices) - self.valid_seq_indices_cuda[:valid_batch_size].copy_(valid_seq_indices, - non_blocking=True) - - # 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 - 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[valid_seq_indices] = self.prompt_budget - sparse_counts_ctx[invalid_seq_indices] = self.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.q_cu_seqlens[:valid_batch_size + 1] = torch.arange( - valid_batch_size + 1, device='cpu', - dtype=torch.int32) * self.window_size - self.q_cu_seqlens_cuda[:valid_batch_size + 1].copy_( - self.q_cu_seqlens[:valid_batch_size + 1], non_blocking=True) - - self.k_cu_seqlens[1:valid_batch_size + 1] = torch.cumsum( - self.k_context_lens[:valid_batch_size], dim=0) - self.k_cu_seqlens_cuda[:valid_batch_size + 1].copy_( - self.k_cu_seqlens[:valid_batch_size + 1], non_blocking=True) - - if valid_batch_size > 0: - # Maximum context length of RocketKV key for valid sequences for padding - self.max_rocket_k_ctx_len = self.k_context_lens[: - valid_batch_size].max( - ).item() - self.total_rocket_k_ctx_tokens = self.k_cu_seqlens[ - valid_batch_size].item() - else: - self.max_rocket_k_ctx_len = 0 - 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() - - # -------------------------------- Generation phase -------------------------------- - self.num_kt_tokens[:self.num_generations] = ( - self.kv_lens[self.num_contexts:self.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_cuda[:self.num_generations + 1].copy_( - self.cum_kt_lens[:self.num_generations + 1], non_blocking=True) - - self.total_kt_tokens = self.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) - - self.total_sparse_gen_indices = self.topk * self.num_generations - - -class RocketTrtllmAttention(TrtllmAttention): - Metadata = RocketTrtllmAttentionMetadata - - # Access type for different dtype sizes - _access_type = { - 1: torch.int8, - 2: torch.int16, - 4: torch.int32, - 8: torch.int64 - } - - def __init__(self, - layer_idx: int, - num_heads: int, - head_dim: int, - num_kv_heads: Optional[int] = None, - quant_config: Optional[QuantConfig] = None, - q_scaling: Optional[float] = None, - sparse_params: Optional[RocketKVParams] = None, - **kwargs): - if sparse_params is None: - raise ValueError( - "sparse_params is required for RocketTrtllmAttention") - super().__init__(layer_idx, - num_heads, - head_dim, - sparse_params=sparse_params, - num_kv_heads=num_kv_heads, - quant_config=quant_config, - q_scaling=q_scaling, - **kwargs) - self.topr = sparse_params.topr - self.topk = sparse_params.topk - self.prompt_budget = sparse_params.prompt_budget - self.window_size = sparse_params.window_size - self.kernel_size = sparse_params.kernel_size - self.page_size = sparse_params.page_size - - def sparse_kv_predict( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - metadata: TrtllmAttentionMetadata, - forward_args: AttentionForwardArgs, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """ - Predict sparse KV indices using optimized SnapKV algorithm. - - Uses a single Triton kernel to compute attention scores between observation window - and prefix tokens, then selects the most important prefix tokens directly. - """ - - num_ctx_tokens = metadata.num_ctx_tokens - if num_ctx_tokens == 0: - return None, None - - 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: - 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, - self.num_heads, - self.num_kv_heads, - self.head_dim, - self.window_size, - metadata.valid_batch_size, - ) - - scores = triton_bmm(q_window, - k_context, - metadata.q_cu_seqlens_cuda, - metadata.k_cu_seqlens_cuda, - metadata.valid_batch_size, - causal=False) - - scores = triton_softmax(scores, metadata.k_cu_seqlens_cuda, - metadata.valid_batch_size) - - # scores: [num_heads, window_size, total_k_tokens] -> [num_kv_heads, total_k_tokens] - scores = scores.view(self.num_kv_heads, - self.num_heads // self.num_kv_heads, - self.window_size, -1).sum(dim=(1, 2)) - - # Reshape scores to handle variable length sequences with padding using Triton - # 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) - - scores = torch.nn.functional.max_pool1d( - scores, - kernel_size=self.kernel_size, - padding=self.kernel_size // 2, - stride=1) - - # Use indexer topk prefill to select topk prefix indices - total_tasks = metadata.valid_batch_size * self.num_kv_heads - - selected_prefix_indices = torch.empty( - (total_tasks, self.prompt_budget - self.window_size), - device=qkv_input.device, - dtype=torch.int32) - - 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( - self.num_kv_heads) - torch.ops.trtllm.indexer_topk_prefill( - scores, row_starts, row_ends, selected_prefix_indices, - self.prompt_budget - self.window_size) - - # Sort selected prefix indices to keep topk indices in ascending order - selected_prefix_indices = torch.sort(selected_prefix_indices, - dim=-1).values - else: - selected_prefix_indices = torch.empty( - (0, self.prompt_budget - self.window_size), - device=qkv_input.device, - dtype=torch.int32) - - sparse_kv_offsets = metadata.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, sparse_kv_offsets, - metadata.num_contexts, metadata.total_sparse_ctx_indices, - self.window_size, self.prompt_budget, self.num_kv_heads) - - # Update KT cache - kt_cache_tensor = metadata.kv_cache_manager.get_kt_buffers( - self.layer_idx) - - 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], - sparse_kv_indices, - sparse_kv_offsets, - self.num_heads, - self.num_kv_heads, - self.head_dim, - self.page_size, - self.prompt_budget, - metadata.kt_tokens_per_block, - metadata.kv_cache_manager.max_kt_blocks_per_seq, - ) - - # Reduce overhead of post processing - if metadata.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 - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if k is None: - qkv_input = q[metadata.num_ctx_tokens:] - q_hidden_size = self.num_heads * self.head_dim - k_hidden_size = self.num_kv_heads * self.head_dim - q = qkv_input[:, :q_hidden_size] - k = qkv_input[:, q_hidden_size:q_hidden_size + k_hidden_size] - else: - q = q[metadata.num_ctx_tokens:] - k = k[metadata.num_ctx_tokens:] - - q = q.view(-1, self.num_kv_heads, self.num_heads // self.num_kv_heads, - self.head_dim) - - return q, k - - @torch.compile(dynamic=True) - def topr_filter(self, q: torch.Tensor) -> torch.Tensor: - i1 = torch.topk(q.abs().sum(dim=2, keepdim=True), self.topr, - dim=-1).indices - q_mask = torch.zeros_like(q) - q_mask.scatter_(-1, i1.expand_as(q[..., :self.topr]), 1) - return q * q_mask - - def sparse_attn_predict( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - metadata: TrtllmAttentionMetadata, - forward_args: AttentionForwardArgs, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - if metadata.num_generations == 0: - return None, None - - q, k = self.preprocess_for_gen(q, k, metadata) - - if self.topr < self.head_dim: - q = self.topr_filter(q) - - kt_cache_tensor = metadata.kv_cache_manager.get_kt_buffers( - self.layer_idx) - - # Update KT cache with new key values - triton_rocket_update_kt_cache_gen( - k, - kt_cache_tensor, - metadata.kt_cache_block_offsets[metadata.num_contexts:], - metadata.kv_lens_cuda_runtime[metadata.num_contexts:], - metadata.page_size, - metadata.kt_tokens_per_block, - metadata.kv_cache_manager.max_kt_blocks_per_seq, - self.num_kv_heads, - self.head_dim, - ) - - # Perform BMM with updated cache - scores = triton_rocket_paged_kt_cache_bmm( - q, - kt_cache_tensor, - metadata.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, - metadata.kv_cache_manager.max_kt_blocks_per_seq, - metadata.total_kt_tokens, - ) - - scores = triton_softmax(scores, metadata.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, - 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] - - selected_indices = triton_topk(scores, metadata.cum_kt_lens_cuda, - sparse_attn_offsets, - metadata.total_sparse_gen_indices, - metadata.topk) - - return selected_indices, sparse_attn_offsets - - -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.host_kt_cache_block_offsets = torch.zeros_like( - self.kt_cache_block_offsets, - device='cpu', - pin_memory=prefer_pinned(), - ) - - 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) - - -class RocketVanillaAttention(VanillaAttention): - """ - RocketKV sparse attention implementation. - - Context phase: only support sparse kv cache - - Generation phase: only support sparse attention computation - """ - - Metadata = RocketVanillaAttentionMetadata - - def __init__(self, - layer_idx: int, - num_heads: int, - head_dim: int, - num_kv_heads: Optional[int] = None, - quant_config: Optional[QuantConfig] = None, - q_scaling: Optional[float] = None, - sparse_params: Optional[RocketKVParams] = None, - **kwargs): - if sparse_params is None: - raise ValueError( - "sparse_params is required for RocketVanillaAttention") - super().__init__(layer_idx, - num_heads, - head_dim, - sparse_params=sparse_params, - num_kv_heads=num_kv_heads, - quant_config=quant_config, - q_scaling=q_scaling, - **kwargs) - self.topr = sparse_params.topr - self.topk = sparse_params.topk - self.prompt_budget = sparse_params.prompt_budget - self.window_size = sparse_params.window_size - self.kernel_size = sparse_params.kernel_size - self.page_size = sparse_params.page_size - assert sparse_params.kt_cache_dtype == 'bfloat16', "Only bfloat16 kt cache is supported for Vanilla RocketKV" - - def _single_request_sparse_kv_predict( - self, q: Optional[Tensor], k: Optional[Tensor], v: Optional[Tensor], - metadata: RocketVanillaAttentionMetadata, past_seen_token: int, - sample_idx: int, **kwargs) -> tuple[Optional[torch.Tensor], int]: - """ - Predict KV indices for writing new key/value pairs. - For RocketKV: - - Context phase: return SnapKV sparse kv indices - - Generation phase: return None - """ - if k is None or v is None: - return None, 0 - - # Generation phase: do not support sparse kv cache - if past_seen_token > 0: - return None, k.size(1) - - # Context phase: predict SnapKV sparse kv indices - sparse_kv_indices = self._get_snapkv_indices(q, k, sample_idx) - - # Gather the key states using the sparse kv indices - if sparse_kv_indices is not None: - k_snap = triton_index_gather(k, sparse_kv_indices) - else: - k_snap = k - - # Update KT cache - kt_cache_tensor = metadata.kv_cache_manager.get_kt_buffers( - self.layer_idx) - target_seq_len = past_seen_token + k_snap.size(1) - kt_cache_position = torch.arange(past_seen_token // self.page_size, - math.ceil(target_seq_len / - self.page_size), - device=q.device) - self._single_request_update_kt_cache( - k_snap, - kt_cache_tensor, - metadata.kt_cache_block_offsets[sample_idx], - target_seq_len, - kt_cache_position, - update=False) - return sparse_kv_indices, k_snap.size(1) - - def _single_request_sparse_attn_predict( - self, q: Tensor, k: Optional[Tensor], v: Optional[Tensor], - kv_cache_tensor: Tensor, metadata: RocketVanillaAttentionMetadata, - past_seen_token: int, sample_idx: int, - **kwargs) -> tuple[Optional[torch.Tensor], int]: - """ - Predict KV cache indices for sparse attention computation. - For RocketKV: - - Context phase: returns None (use full attention) - - Generation phase: return RocketKV sparse indices for sparse attention computation - """ - if k is None or v is None: - return None, 0 - - # Context phase: use full attention - if past_seen_token == 0: - return None, k.size(1) - - # Get RocketKV sparse indices - sparse_indices = self._rocketkv_selection(q, k, metadata, - past_seen_token, sample_idx) - return sparse_indices, sparse_indices.size(1) - - def _get_snapkv_indices(self, q: Tensor, k: Tensor, - sample_idx: int) -> Tensor: - """Get SnapKV sparse kv indices from the input sequence for context phase.""" - bsz = 1 - seq_len = k.size(1) - - # If the sequence length is less than the prompt budget, do not enable sparse kv cache - if seq_len <= self.prompt_budget: - return None - - # Use last window_size tokens as observation - # (1, num_kv_heads, window_size, head_dim) - q_obs = q[:, :, -self.window_size:] - # (1, num_kv_heads, seq_len, head_dim) - k_pre = repeat_kv(k.transpose(1, 2), self.num_key_value_groups) - - dist = (torch.arange(0, self.window_size, device=q.device)[:, None] - - torch.arange(0, seq_len, device=q.device)[None, :] + seq_len - - self.window_size) - attention_mask = (dist >= 0) - - score = torch.matmul(q_obs, k_pre.transpose(-1, -2)) / math.sqrt( - self.head_dim) - - score = torch.masked_fill( - score, - attention_mask.view(1, 1, self.window_size, seq_len) == False, - torch.scalar_tensor(float("-inf"), - device=score.device, - dtype=score.dtype)) - - score = torch.nn.functional.softmax(score, dim=-1) - - score = torch.masked_fill( - score, - attention_mask.view(1, 1, self.window_size, seq_len) == False, - torch.scalar_tensor(0, device=score.device, dtype=score.dtype)) - - score = score[:, :, -self.window_size:, :-self.window_size].sum(dim=-2) - - score = score.view(bsz, self.num_kv_heads, self.num_key_value_groups, - -1).sum(dim=2) - score = torch.nn.functional.max_pool1d(score, - kernel_size=self.kernel_size, - padding=self.kernel_size // 2, - stride=1) - - # Select top important tokens from prefix - prefix_len = seq_len - self.window_size - selected_prefix_indices = score.topk(self.prompt_budget - - self.window_size, - dim=-1).indices.sort().values - - # Combine selected prefix indices with window indices - window_indices = torch.arange( - prefix_len, seq_len, - device=k.device).unsqueeze(0).unsqueeze(0).expand( - bsz, self.num_kv_heads, -1) - selected_indices = torch.cat([selected_prefix_indices, window_indices], - dim=-1) - - return selected_indices.transpose(1, 2) - - def _single_request_update_kt_cache(self, - k, - kt_cache_tensor, - kt_cache_block_offsets, - seq_len, - cache_position, - update=True): - """Update KT cache for RocketKV algorithm.""" - # (1, num_pages_per_block, num_kv_heads, 2*head_dim) - k_out = kt_cache_tensor[kt_cache_block_offsets[0], :, :, :].unsqueeze(0) - - # k: (1, seq_len, num_kv_heads, head_dim) - if k is not None: - padding_len = self.page_size - ( - (k.size(1) - 1) % self.page_size + 1) - padding_tensor = torch.full( - (1, padding_len, self.num_kv_heads, self.head_dim), - float('inf'), - device=k.device, - dtype=k.dtype) - # (1, seq_len+padding_len, num_kv_heads, head_dim) - k_min = torch.cat([k, padding_tensor], dim=1) - k_min = k_min.reshape(1, -1, self.page_size, self.num_kv_heads, - self.head_dim).amin(dim=2) - k_max = torch.cat([k, -padding_tensor], dim=1) - k_max = k_max.reshape(1, -1, self.page_size, self.num_kv_heads, - self.head_dim).amax(dim=2) - if update and (seq_len - 1) % self.page_size > 0: # gen phase - k_min = torch.min(k_min, - k_out[:, cache_position, :, :self.head_dim]) - k_max = torch.max(k_max, k_out[:, cache_position, :, - self.head_dim:]) - k_value = torch.cat([k_min, k_max], dim=-1) - access_type = self._access_type[k_value.dtype.itemsize] - k_out.view(dtype=access_type).index_copy_( - 1, cache_position, k_value.view(dtype=access_type)) - - return k_out[:, :math.ceil(seq_len / self.page_size), :, :] - - def _rocketkv_selection(self, q: Tensor, k: Tensor, - metadata: RocketVanillaAttentionMetadata, - past_seen_token: int, sample_idx: int) -> Tensor: - """Implement RocketKV's two-stage selection process for generation phase.""" - bsz = 1 - q_len = q.size(2) - - # Helper functions - def _gather(t: Tensor, dim: int, i: Tensor) -> Tensor: - dim += (dim < 0) * t.ndim - return t.gather( - dim, i.expand(*t.shape[:dim], i.shape[dim], *t.shape[dim + 1:])) - - @torch.compile(disable=not torch.cuda.is_available()) - def _scaled_softmax(x: Tensor, divscale: Tensor | float, - dim: int) -> Tensor: - return torch.softmax(x / divscale, dim=dim) - - # Get KT cache for key-token matching - kt_cache_tensor = metadata.kv_cache_manager.get_kt_buffers( - self.layer_idx) - target_seq_len = past_seen_token + 1 # +1 for current token - - # Update KT cache - kt_cache_position = torch.arange(past_seen_token // self.page_size, - math.ceil(target_seq_len / - self.page_size), - device=q.device) - kt_states = self._single_request_update_kt_cache( - k, kt_cache_tensor, metadata.kt_cache_block_offsets[sample_idx], - target_seq_len, kt_cache_position) - - # Reshape query for multi-head processing - qi = q.view(bsz, self.num_kv_heads, self.num_heads // self.num_kv_heads, - q_len, self.head_dim) - qi_abs = torch.abs(qi) - - # Top-r selection on query features - i1 = torch.topk(qi_abs.mean(dim=2, keepdim=True), self.topr, - dim=-1).indices - qi_hat = _gather(qi, -1, i1) - - # Generate signed indices for key-token matching - i1_sign = torch.where( - qi_hat.sum(dim=2, keepdim=True) > 0, i1 + self.head_dim, - i1).transpose(-1, -2) - - # Gather key tokens and compute attention scores - kt_hat = _gather( - kt_states.permute(0, 2, 3, 1).unsqueeze(2), -2, i1_sign) - qk_hat = qi_hat @ kt_hat - qk_hat = qk_hat.repeat_interleave(self.page_size, - dim=-1)[:, :, :, :, :target_seq_len] - scale = torch.sqrt(self.head_dim * - torch.abs(qi_hat).sum(dim=-1, keepdim=True) / - qi_abs.sum(dim=-1, keepdim=True)) - - # (1, num_kv_heads, num_heads, target_seq_len) - s_hat = _scaled_softmax(qk_hat, scale, dim=-1) - - topk = min(self.topk, target_seq_len) - i2 = torch.topk(s_hat.mean(dim=2, keepdim=True), topk, dim=-1).indices - - iKV = i2[:, :, 0, 0, :].transpose(1, 2) # (1, topk, num_kv_heads) - - return iKV - - -class RocketKVCacheManager(KVCacheManager): - - def __init__( - self, - 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, - max_num_tokens: int = 8192, - model_config: Optional[ModelConfig] = None, - max_beam_width: int = 1, - sparse_attention_config: Optional["SparseAttentionConfig"] = None, - pretrained_config=None, - **kwargs, - ) -> None: - - if sparse_attention_config is None: - raise ValueError( - "sparse_attention_config is required for RocketKV cache") - sparse_params = sparse_attention_config.to_sparse_params( - pretrained_config=pretrained_config) - if not isinstance(sparse_params, RocketKVParams): - raise ValueError( - "RocketKV cache requires RocketKV sparse parameters") - assert not kv_cache_config.enable_block_reuse, "RocketKV cache requires block reuse to be disabled in KV cache config" - self.kt_tokens_per_block = next_power_of_2( - math.ceil(tokens_per_block / sparse_params.page_size)) - self.kt_cache_dtype = torch.bfloat16 if sparse_params.kt_cache_dtype == 'bfloat16' else torch.float8_e5m2 - - 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=layer_mask, - max_num_tokens=max_num_tokens, - model_config=model_config, - max_beam_width=max_beam_width, - **kwargs, - ) - self.page_size = sparse_params.page_size - self.prompt_budget = sparse_params.prompt_budget - self.max_batch_size = max_batch_size - - # Per layer KT cache pool - # Use the same number of blocks as the paged kv cache. In this way, the scheduler can use the same number of - # blocks to schedule requests. - # Use kt_tokens_per_block to make sure the KT cache is large enough to hold the kt tokens, - # since kt_tokens_per_block * num_blocks * page_size >= tokens_per_block * num_blocks. - self.num_blocks = self.blocks_in_primary_pool - self.kt_cache_pool_per_layer = [ - torch.empty((self.num_blocks, self.kt_tokens_per_block, - num_kv_heads, head_dim * 2), - device="cuda", - dtype=self.kt_cache_dtype) - for _ in range(self.num_local_layers) - ] - self.max_kt_blocks_per_seq = self.num_blocks - - # Block manager to manage the KT cache blocks for each request. Different layers share the - # same block ids. - self.kt_cache_manager = BlockManager(self.num_blocks, - self.kt_tokens_per_block) - - def add_dummy_requests( - self, - request_ids: List[int], - token_nums: Optional[List[int]] = None, - is_gen: bool = False, - prepare_resource: bool = True, - 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, - num_extra_decoding_steps: int = 0, - draft_kv_cache_manager=None, - ): - requests = super().add_dummy_requests( - request_ids=request_ids, - token_nums=token_nums, - is_gen=is_gen, - prepare_resource=prepare_resource, - 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, - num_extra_decoding_steps=num_extra_decoding_steps, - draft_kv_cache_manager=draft_kv_cache_manager, - ) - if prepare_resource: - for req in requests: - request_id = req.py_request_id - kt_token_num = math.ceil(req.max_beam_num_tokens / - self.page_size) - self.kt_cache_manager.add_tokens(request_id, kt_token_num) - return requests - - def get_kt_buffers(self, layer_idx: int): - return self.kt_cache_pool_per_layer[layer_idx] - - def copy_kt_block_offsets(self, request_ids: List[int], - block_offsets: torch.Tensor) -> torch.Tensor: - self.kt_cache_manager.copy_block_offsets(request_ids, block_offsets) - - def prepare_resources(self, scheduled_batch): - super().prepare_resources(scheduled_batch) - for req in scheduled_batch.context_requests: - request_id = req.py_request_id - num_tokens = req.prompt_len - kt_token_num = math.ceil(num_tokens / self.page_size) - self.kt_cache_manager.add_tokens(request_id, kt_token_num) - - for req in scheduled_batch.generation_requests: - request_id = req.py_request_id - num_tokens = req.max_beam_num_tokens + 1 - if num_tokens % self.page_size == 1: - self.kt_cache_manager.add_tokens(request_id, 1) - - def update_resources(self, - scheduled_batch, - attn_metadata: AttentionMetadata = None, - kv_cache_dtype_byte_size: float = None): - for request in scheduled_batch.context_requests: - if request.state != LlmRequestState.GENERATION_COMPLETE: - seq_len = request.get_num_tokens(0) - rewind_len = max(seq_len - 1 - self.prompt_budget, 0) - self.rewind_kv_cache(request, rewind_len) - # get the rewind length for kt cache - num_tokens = request.max_beam_num_tokens - updated_kt_token_num = num_tokens - rewind_len - rewind_len = (math.ceil(num_tokens / self.page_size) - - math.ceil(updated_kt_token_num / self.page_size)) - self.kt_cache_manager.rewind_cache(request, rewind_len) - - def free_resources(self, request): - super().free_resources(request) - self.kt_cache_manager.free_resources(request) - - @staticmethod - def get_cache_size_per_token(model_config: ModelConfig, - mapping: Mapping, - num_layers: Optional[int] = None, - **kwargs): - # main KV cache dtype bytes (one each for K and V) - kv_dtype_bytes = 2 - quant_config = model_config.quant_config - if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache( - ): - kv_dtype_bytes = 1 - - # get num key value heads - config = model_config.pretrained_config - num_key_value_heads = getattr(config, 'num_key_value_heads', - config.num_attention_heads) - if isinstance(num_key_value_heads, Iterable): - num_key_value_heads = sum(num_key_value_heads) / len( - num_key_value_heads) - - # get head dim - tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size - head_dim = getattr(config, "head_dim", None) - if not isinstance(head_dim, int): - head_dim = config.hidden_size // config.num_attention_heads - head_dim = head_dim * num_key_value_heads // tp_size - - num_attention_layers = KVCacheManager._resolve_num_attention_layers( - model_config, mapping, num_layers) - - # K and V: two tensors stored at the main KV cache dtype. - mem_per_token = kv_dtype_bytes * 2 * num_attention_layers * head_dim - - tokens_per_block = kwargs['tokens_per_block'] - sparse_attention_config = model_config.sparse_attention_config - if sparse_attention_config is None: - raise ValueError( - "sparse_attention_config is required for RocketKV cache") - sparse_params = sparse_attention_config.to_sparse_params( - pretrained_config=model_config.pretrained_config) - if not isinstance(sparse_params, RocketKVParams): - raise ValueError( - "RocketKV cache requires RocketKV sparse parameters") - kt_tokens_per_block = next_power_of_2( - math.ceil(tokens_per_block / sparse_params.page_size)) - # KT (key-landmark) cache: a separate pool stored at kt_cache_dtype, - # NOT the main KV dtype. Per real token it holds - # kt_tokens_per_block / tokens_per_block landmark slots, each storing - # the concatenated (min, max) landmark -- a factor of 2 over head_dim. - # Size it at its own dtype's byte width (1 for float8_e5m2, else 2) so - # the estimate is independent of the main KV dtype and matches both - # get_cache_bytes_per_token and the physical pool allocated in - # RocketKVCacheManager.__init__. - kt_dtype_bytes = 1 if sparse_params.kt_cache_dtype == "float8_e5m2" else 2 - mem_per_token += (kt_dtype_bytes * 2 * kt_tokens_per_block / - tokens_per_block * num_attention_layers * head_dim) - return mem_per_token - - def get_cache_bytes_per_token(self): - # Main K/V cache: stored at the main KV cache dtype (self.dtype). - cache_size_per_token = math.ceil( - self.kv_factor * sum(self.num_kv_heads_per_layer) * self.head_dim) - - if self.dtype not in (DataType.FP8, DataType.HALF, DataType.BF16, - DataType.FLOAT, DataType.NVFP4): - raise ValueError(f'Cannot support {self.dtype} KV cache.') - - cache_size_bytes_per_token = get_size_in_bytes(cache_size_per_token, - self.dtype) - if self.dtype == DataType.NVFP4: - cache_size_bytes_per_token += self.calculate_scaling_factor_size_bytes( - cache_size_per_token, - quant_vector_size=16, - scaling_factor_dtype=DataType.FP8) - - # KT (key-landmark) cache: a separate pool physically allocated at - # kt_cache_dtype, NOT the main KV dtype -- see __init__: - # torch.empty((num_blocks, kt_tokens_per_block, num_kv_heads, - # head_dim * 2), dtype=kt_cache_dtype). - # The trailing factor of 2 is the concatenated (min, max) landmark per - # head_dim. Per real token the pool holds - # kt_tokens_per_block / tokens_per_block landmark slots. Size it at its - # own dtype's byte width (1 for float8_e5m2, else 2) so the estimate - # matches the allocation regardless of self.dtype. - kt_dtype_bytes = 1 if self.kt_cache_dtype == torch.float8_e5m2 else 2 - kt_elements_per_token = math.ceil( - 2 * self.kt_tokens_per_block / self.tokens_per_block * - sum(self.num_kv_heads_per_layer) * self.head_dim) - cache_size_bytes_per_token += kt_elements_per_token * kt_dtype_bytes - - return cache_size_bytes_per_token diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket/__init__.py new file mode 100644 index 000000000000..85c3c9e33f3b --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/rocket/__init__.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""RocketKV sparse attention backend package.""" + +from .backend import RocketTrtllmAttention, RocketVanillaAttention +from .cache_manager import RocketKVCacheManager +from .metadata import RocketTrtllmAttentionMetadata, RocketVanillaAttentionMetadata +from .params import RocketKVMetadataParams, RocketKVParams + +__all__ = [ + "RocketKVCacheManager", + "RocketKVMetadataParams", + "RocketKVParams", + "RocketTrtllmAttention", + "RocketTrtllmAttentionMetadata", + "RocketVanillaAttention", + "RocketVanillaAttentionMetadata", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket/backend.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket/backend.py new file mode 100644 index 000000000000..c5423a33eca8 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/rocket/backend.py @@ -0,0 +1,601 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +from torch import Tensor + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.attention_backend.vanilla import VanillaAttention, repeat_kv +from tensorrt_llm.models.modeling_utils import QuantConfig + +if TYPE_CHECKING: + pass + +from .kernels import ( + triton_bmm, + triton_flatten_to_batch, + triton_index_gather, + triton_rocket_batch_to_flatten, + triton_rocket_paged_kt_cache_bmm, + triton_rocket_qk_split, + triton_rocket_reduce_scores, + triton_rocket_update_kt_cache_ctx, + triton_rocket_update_kt_cache_gen, + triton_softmax, + triton_topk, +) +from .metadata import RocketTrtllmAttentionMetadata, RocketVanillaAttentionMetadata +from .params import RocketKVParams + +ModelConfig = tensorrt_llm.bindings.ModelConfig + + +class RocketTrtllmAttention(TrtllmAttention): + Metadata = RocketTrtllmAttentionMetadata + + # Access type for different dtype sizes + _access_type = {1: torch.int8, 2: torch.int16, 4: torch.int32, 8: torch.int64} + + def __init__( + self, + layer_idx: int, + num_heads: int, + head_dim: int, + num_kv_heads: Optional[int] = None, + quant_config: Optional[QuantConfig] = None, + q_scaling: Optional[float] = None, + sparse_params: Optional[RocketKVParams] = None, + **kwargs, + ): + if sparse_params is None: + raise ValueError("sparse_params is required for RocketTrtllmAttention") + super().__init__( + layer_idx, + num_heads, + head_dim, + sparse_params=sparse_params, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + q_scaling=q_scaling, + **kwargs, + ) + self.topr = sparse_params.topr + self.topk = sparse_params.topk + self.prompt_budget = sparse_params.prompt_budget + self.window_size = sparse_params.window_size + self.kernel_size = sparse_params.kernel_size + self.page_size = sparse_params.page_size + + def sparse_kv_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + Predict sparse KV indices using optimized SnapKV algorithm. + + Uses a single Triton kernel to compute attention scores between observation window + and prefix tokens, then selects the most important prefix tokens directly. + """ + + num_ctx_tokens = metadata.num_ctx_tokens + if num_ctx_tokens == 0: + return None, None + + 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: + 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, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.window_size, + metadata.valid_batch_size, + ) + + scores = triton_bmm( + q_window, + k_context, + metadata.q_cu_seqlens_cuda, + metadata.k_cu_seqlens_cuda, + metadata.valid_batch_size, + causal=False, + ) + + scores = triton_softmax(scores, metadata.k_cu_seqlens_cuda, metadata.valid_batch_size) + + # scores: [num_heads, window_size, total_k_tokens] -> [num_kv_heads, total_k_tokens] + scores = scores.view( + self.num_kv_heads, self.num_heads // self.num_kv_heads, self.window_size, -1 + ).sum(dim=(1, 2)) + + # Reshape scores to handle variable length sequences with padding using Triton + # 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, + ) + + scores = torch.nn.functional.max_pool1d( + scores, kernel_size=self.kernel_size, padding=self.kernel_size // 2, stride=1 + ) + + # Use indexer topk prefill to select topk prefix indices + total_tasks = metadata.valid_batch_size * self.num_kv_heads + + selected_prefix_indices = torch.empty( + (total_tasks, self.prompt_budget - self.window_size), + device=qkv_input.device, + dtype=torch.int32, + ) + + 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( + self.num_kv_heads + ) + torch.ops.trtllm.indexer_topk_prefill( + scores, + row_starts, + row_ends, + selected_prefix_indices, + self.prompt_budget - self.window_size, + ) + + # Sort selected prefix indices to keep topk indices in ascending order + selected_prefix_indices = torch.sort(selected_prefix_indices, dim=-1).values + else: + selected_prefix_indices = torch.empty( + (0, self.prompt_budget - self.window_size), + device=qkv_input.device, + dtype=torch.int32, + ) + + sparse_kv_offsets = metadata.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, + sparse_kv_offsets, + metadata.num_contexts, + metadata.total_sparse_ctx_indices, + self.window_size, + self.prompt_budget, + self.num_kv_heads, + ) + + # Update KT cache + kt_cache_tensor = metadata.kv_cache_manager.get_kt_buffers(self.layer_idx) + + 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], + sparse_kv_indices, + sparse_kv_offsets, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.page_size, + self.prompt_budget, + metadata.kt_tokens_per_block, + metadata.kv_cache_manager.max_kt_blocks_per_seq, + ) + + # Reduce overhead of post processing + if metadata.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 + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if k is None: + qkv_input = q[metadata.num_ctx_tokens :] + q_hidden_size = self.num_heads * self.head_dim + k_hidden_size = self.num_kv_heads * self.head_dim + q = qkv_input[:, :q_hidden_size] + k = qkv_input[:, q_hidden_size : q_hidden_size + k_hidden_size] + else: + q = q[metadata.num_ctx_tokens :] + k = k[metadata.num_ctx_tokens :] + + q = q.view(-1, self.num_kv_heads, self.num_heads // self.num_kv_heads, self.head_dim) + + return q, k + + @torch.compile(dynamic=True) + def topr_filter(self, q: torch.Tensor) -> torch.Tensor: + i1 = torch.topk(q.abs().sum(dim=2, keepdim=True), self.topr, dim=-1).indices + q_mask = torch.zeros_like(q) + q_mask.scatter_(-1, i1.expand_as(q[..., : self.topr]), 1) + return q * q_mask + + def sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + if metadata.num_generations == 0: + return None, None + + q, k = self.preprocess_for_gen(q, k, metadata) + + if self.topr < self.head_dim: + q = self.topr_filter(q) + + kt_cache_tensor = metadata.kv_cache_manager.get_kt_buffers(self.layer_idx) + + # Update KT cache with new key values + triton_rocket_update_kt_cache_gen( + k, + kt_cache_tensor, + metadata.kt_cache_block_offsets[metadata.num_contexts :], + metadata.kv_lens_cuda_runtime[metadata.num_contexts :], + metadata.page_size, + metadata.kt_tokens_per_block, + metadata.kv_cache_manager.max_kt_blocks_per_seq, + self.num_kv_heads, + self.head_dim, + ) + + # Perform BMM with updated cache + scores = triton_rocket_paged_kt_cache_bmm( + q, + kt_cache_tensor, + metadata.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, + metadata.kv_cache_manager.max_kt_blocks_per_seq, + metadata.total_kt_tokens, + ) + + scores = triton_softmax(scores, metadata.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, + 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] + + selected_indices = triton_topk( + scores, + metadata.cum_kt_lens_cuda, + sparse_attn_offsets, + metadata.total_sparse_gen_indices, + metadata.topk, + ) + + return selected_indices, sparse_attn_offsets + + +class RocketVanillaAttention(VanillaAttention): + """ + RocketKV sparse attention implementation. + - Context phase: only support sparse kv cache + - Generation phase: only support sparse attention computation + """ + + Metadata = RocketVanillaAttentionMetadata + + def __init__( + self, + layer_idx: int, + num_heads: int, + head_dim: int, + num_kv_heads: Optional[int] = None, + quant_config: Optional[QuantConfig] = None, + q_scaling: Optional[float] = None, + sparse_params: Optional[RocketKVParams] = None, + **kwargs, + ): + if sparse_params is None: + raise ValueError("sparse_params is required for RocketVanillaAttention") + super().__init__( + layer_idx, + num_heads, + head_dim, + sparse_params=sparse_params, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + q_scaling=q_scaling, + **kwargs, + ) + self.topr = sparse_params.topr + self.topk = sparse_params.topk + self.prompt_budget = sparse_params.prompt_budget + self.window_size = sparse_params.window_size + self.kernel_size = sparse_params.kernel_size + self.page_size = sparse_params.page_size + assert sparse_params.kt_cache_dtype == "bfloat16", ( + "Only bfloat16 kt cache is supported for Vanilla RocketKV" + ) + + def _single_request_sparse_kv_predict( + self, + q: Optional[Tensor], + k: Optional[Tensor], + v: Optional[Tensor], + metadata: RocketVanillaAttentionMetadata, + past_seen_token: int, + sample_idx: int, + **kwargs, + ) -> tuple[Optional[torch.Tensor], int]: + """ + Predict KV indices for writing new key/value pairs. + For RocketKV: + - Context phase: return SnapKV sparse kv indices + - Generation phase: return None + """ + if k is None or v is None: + return None, 0 + + # Generation phase: do not support sparse kv cache + if past_seen_token > 0: + return None, k.size(1) + + # Context phase: predict SnapKV sparse kv indices + sparse_kv_indices = self._get_snapkv_indices(q, k, sample_idx) + + # Gather the key states using the sparse kv indices + if sparse_kv_indices is not None: + k_snap = triton_index_gather(k, sparse_kv_indices) + else: + k_snap = k + + # Update KT cache + kt_cache_tensor = metadata.kv_cache_manager.get_kt_buffers(self.layer_idx) + target_seq_len = past_seen_token + k_snap.size(1) + kt_cache_position = torch.arange( + past_seen_token // self.page_size, + math.ceil(target_seq_len / self.page_size), + device=q.device, + ) + self._single_request_update_kt_cache( + k_snap, + kt_cache_tensor, + metadata.kt_cache_block_offsets[sample_idx], + target_seq_len, + kt_cache_position, + update=False, + ) + return sparse_kv_indices, k_snap.size(1) + + def _single_request_sparse_attn_predict( + self, + q: Tensor, + k: Optional[Tensor], + v: Optional[Tensor], + kv_cache_tensor: Tensor, + metadata: RocketVanillaAttentionMetadata, + past_seen_token: int, + sample_idx: int, + **kwargs, + ) -> tuple[Optional[torch.Tensor], int]: + """ + Predict KV cache indices for sparse attention computation. + For RocketKV: + - Context phase: returns None (use full attention) + - Generation phase: return RocketKV sparse indices for sparse attention computation + """ + if k is None or v is None: + return None, 0 + + # Context phase: use full attention + if past_seen_token == 0: + return None, k.size(1) + + # Get RocketKV sparse indices + sparse_indices = self._rocketkv_selection(q, k, metadata, past_seen_token, sample_idx) + return sparse_indices, sparse_indices.size(1) + + def _get_snapkv_indices(self, q: Tensor, k: Tensor, sample_idx: int) -> Tensor: + """Get SnapKV sparse kv indices from the input sequence for context phase.""" + bsz = 1 + seq_len = k.size(1) + + # If the sequence length is less than the prompt budget, do not enable sparse kv cache + if seq_len <= self.prompt_budget: + return None + + # Use last window_size tokens as observation + # (1, num_kv_heads, window_size, head_dim) + q_obs = q[:, :, -self.window_size :] + # (1, num_kv_heads, seq_len, head_dim) + k_pre = repeat_kv(k.transpose(1, 2), self.num_key_value_groups) + + dist = ( + torch.arange(0, self.window_size, device=q.device)[:, None] + - torch.arange(0, seq_len, device=q.device)[None, :] + + seq_len + - self.window_size + ) + attention_mask = dist >= 0 + + score = torch.matmul(q_obs, k_pre.transpose(-1, -2)) / math.sqrt(self.head_dim) + + score = torch.masked_fill( + score, + ~attention_mask.view(1, 1, self.window_size, seq_len), + torch.scalar_tensor(float("-inf"), device=score.device, dtype=score.dtype), + ) + + score = torch.nn.functional.softmax(score, dim=-1) + + score = torch.masked_fill( + score, + ~attention_mask.view(1, 1, self.window_size, seq_len), + torch.scalar_tensor(0, device=score.device, dtype=score.dtype), + ) + + score = score[:, :, -self.window_size :, : -self.window_size].sum(dim=-2) + + score = score.view(bsz, self.num_kv_heads, self.num_key_value_groups, -1).sum(dim=2) + score = torch.nn.functional.max_pool1d( + score, kernel_size=self.kernel_size, padding=self.kernel_size // 2, stride=1 + ) + + # Select top important tokens from prefix + prefix_len = seq_len - self.window_size + selected_prefix_indices = ( + score.topk(self.prompt_budget - self.window_size, dim=-1).indices.sort().values + ) + + # Combine selected prefix indices with window indices + window_indices = ( + torch.arange(prefix_len, seq_len, device=k.device) + .unsqueeze(0) + .unsqueeze(0) + .expand(bsz, self.num_kv_heads, -1) + ) + selected_indices = torch.cat([selected_prefix_indices, window_indices], dim=-1) + + return selected_indices.transpose(1, 2) + + def _single_request_update_kt_cache( + self, k, kt_cache_tensor, kt_cache_block_offsets, seq_len, cache_position, update=True + ): + """Update KT cache for RocketKV algorithm.""" + # (1, num_pages_per_block, num_kv_heads, 2*head_dim) + k_out = kt_cache_tensor[kt_cache_block_offsets[0], :, :, :].unsqueeze(0) + + # k: (1, seq_len, num_kv_heads, head_dim) + if k is not None: + padding_len = self.page_size - ((k.size(1) - 1) % self.page_size + 1) + padding_tensor = torch.full( + (1, padding_len, self.num_kv_heads, self.head_dim), + float("inf"), + device=k.device, + dtype=k.dtype, + ) + # (1, seq_len+padding_len, num_kv_heads, head_dim) + k_min = torch.cat([k, padding_tensor], dim=1) + k_min = k_min.reshape(1, -1, self.page_size, self.num_kv_heads, self.head_dim).amin( + dim=2 + ) + k_max = torch.cat([k, -padding_tensor], dim=1) + k_max = k_max.reshape(1, -1, self.page_size, self.num_kv_heads, self.head_dim).amax( + dim=2 + ) + if update and (seq_len - 1) % self.page_size > 0: # gen phase + k_min = torch.min(k_min, k_out[:, cache_position, :, : self.head_dim]) + k_max = torch.max(k_max, k_out[:, cache_position, :, self.head_dim :]) + k_value = torch.cat([k_min, k_max], dim=-1) + access_type = self._access_type[k_value.dtype.itemsize] + k_out.view(dtype=access_type).index_copy_( + 1, cache_position, k_value.view(dtype=access_type) + ) + + return k_out[:, : math.ceil(seq_len / self.page_size), :, :] + + def _rocketkv_selection( + self, + q: Tensor, + k: Tensor, + metadata: RocketVanillaAttentionMetadata, + past_seen_token: int, + sample_idx: int, + ) -> Tensor: + """Implement RocketKV's two-stage selection process for generation phase.""" + bsz = 1 + q_len = q.size(2) + + # Helper functions + def _gather(t: Tensor, dim: int, i: Tensor) -> Tensor: + dim += (dim < 0) * t.ndim + return t.gather(dim, i.expand(*t.shape[:dim], i.shape[dim], *t.shape[dim + 1 :])) + + @torch.compile(disable=not torch.cuda.is_available()) + def _scaled_softmax(x: Tensor, divscale: Tensor | float, dim: int) -> Tensor: + return torch.softmax(x / divscale, dim=dim) + + # Get KT cache for key-token matching + kt_cache_tensor = metadata.kv_cache_manager.get_kt_buffers(self.layer_idx) + target_seq_len = past_seen_token + 1 # +1 for current token + + # Update KT cache + kt_cache_position = torch.arange( + past_seen_token // self.page_size, + math.ceil(target_seq_len / self.page_size), + device=q.device, + ) + kt_states = self._single_request_update_kt_cache( + k, + kt_cache_tensor, + metadata.kt_cache_block_offsets[sample_idx], + target_seq_len, + kt_cache_position, + ) + + # Reshape query for multi-head processing + qi = q.view( + bsz, self.num_kv_heads, self.num_heads // self.num_kv_heads, q_len, self.head_dim + ) + qi_abs = torch.abs(qi) + + # Top-r selection on query features + i1 = torch.topk(qi_abs.mean(dim=2, keepdim=True), self.topr, dim=-1).indices + qi_hat = _gather(qi, -1, i1) + + # Generate signed indices for key-token matching + i1_sign = torch.where( + qi_hat.sum(dim=2, keepdim=True) > 0, i1 + self.head_dim, i1 + ).transpose(-1, -2) + + # Gather key tokens and compute attention scores + kt_hat = _gather(kt_states.permute(0, 2, 3, 1).unsqueeze(2), -2, i1_sign) + qk_hat = qi_hat @ kt_hat + qk_hat = qk_hat.repeat_interleave(self.page_size, dim=-1)[:, :, :, :, :target_seq_len] + scale = torch.sqrt( + self.head_dim + * torch.abs(qi_hat).sum(dim=-1, keepdim=True) + / qi_abs.sum(dim=-1, keepdim=True) + ) + + # (1, num_kv_heads, num_heads, target_seq_len) + s_hat = _scaled_softmax(qk_hat, scale, dim=-1) + + topk = min(self.topk, target_seq_len) + i2 = torch.topk(s_hat.mean(dim=2, keepdim=True), topk, dim=-1).indices + + iKV = i2[:, :, 0, 0, :].transpose(1, 2) # (1, topk, num_kv_heads) + + return iKV diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py new file mode 100644 index 000000000000..a711bbe8b53e --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py @@ -0,0 +1,293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Iterable, List, Optional, Union + +import torch +from triton import next_power_of_2 + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.resource_manager import BlockManager, KVCacheManager +from tensorrt_llm._utils import get_size_in_bytes +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.mapping import Mapping + +from .params import RocketKVParams + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, SparseAttentionConfig + +ModelConfig = tensorrt_llm.bindings.ModelConfig + + +class RocketKVCacheManager(KVCacheManager): + def __init__( + self, + 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, + max_num_tokens: int = 8192, + model_config: Optional[ModelConfig] = None, + max_beam_width: int = 1, + sparse_attention_config: Optional["SparseAttentionConfig"] = None, + pretrained_config=None, + **kwargs, + ) -> None: + if sparse_attention_config is None: + raise ValueError("sparse_attention_config is required for RocketKV cache") + sparse_params = sparse_attention_config.to_sparse_params( + pretrained_config=pretrained_config + ) + if not isinstance(sparse_params, RocketKVParams): + raise ValueError("RocketKV cache requires RocketKV sparse parameters") + assert not kv_cache_config.enable_block_reuse, ( + "RocketKV cache requires block reuse to be disabled in KV cache config" + ) + self.kt_tokens_per_block = next_power_of_2( + math.ceil(tokens_per_block / sparse_params.page_size) + ) + self.kt_cache_dtype = ( + torch.bfloat16 if sparse_params.kt_cache_dtype == "bfloat16" else torch.float8_e5m2 + ) + + 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=layer_mask, + max_num_tokens=max_num_tokens, + model_config=model_config, + max_beam_width=max_beam_width, + **kwargs, + ) + self.page_size = sparse_params.page_size + self.prompt_budget = sparse_params.prompt_budget + self.max_batch_size = max_batch_size + + # Per layer KT cache pool + # Use the same number of blocks as the paged kv cache. In this way, the scheduler can use the same number of + # blocks to schedule requests. + # Use kt_tokens_per_block to make sure the KT cache is large enough to hold the kt tokens, + # since kt_tokens_per_block * num_blocks * page_size >= tokens_per_block * num_blocks. + self.num_blocks = self.blocks_in_primary_pool + self.kt_cache_pool_per_layer = [ + torch.empty( + (self.num_blocks, self.kt_tokens_per_block, num_kv_heads, head_dim * 2), + device="cuda", + dtype=self.kt_cache_dtype, + ) + for _ in range(self.num_local_layers) + ] + self.max_kt_blocks_per_seq = self.num_blocks + + # Block manager to manage the KT cache blocks for each request. Different layers share the + # same block ids. + self.kt_cache_manager = BlockManager(self.num_blocks, self.kt_tokens_per_block) + + def add_dummy_requests( + self, + request_ids: List[int], + token_nums: Optional[List[int]] = None, + is_gen: bool = False, + prepare_resource: bool = True, + 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, + num_extra_decoding_steps: int = 0, + draft_kv_cache_manager=None, + ): + requests = super().add_dummy_requests( + request_ids=request_ids, + token_nums=token_nums, + is_gen=is_gen, + prepare_resource=prepare_resource, + 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, + num_extra_decoding_steps=num_extra_decoding_steps, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + if prepare_resource: + for req in requests: + request_id = req.py_request_id + kt_token_num = math.ceil(req.max_beam_num_tokens / self.page_size) + self.kt_cache_manager.add_tokens(request_id, kt_token_num) + return requests + + def get_kt_buffers(self, layer_idx: int): + return self.kt_cache_pool_per_layer[layer_idx] + + def copy_kt_block_offsets( + self, request_ids: List[int], block_offsets: torch.Tensor + ) -> torch.Tensor: + self.kt_cache_manager.copy_block_offsets(request_ids, block_offsets) + + def prepare_resources(self, scheduled_batch): + super().prepare_resources(scheduled_batch) + for req in scheduled_batch.context_requests: + request_id = req.py_request_id + num_tokens = req.prompt_len + kt_token_num = math.ceil(num_tokens / self.page_size) + self.kt_cache_manager.add_tokens(request_id, kt_token_num) + + for req in scheduled_batch.generation_requests: + request_id = req.py_request_id + num_tokens = req.max_beam_num_tokens + 1 + if num_tokens % self.page_size == 1: + self.kt_cache_manager.add_tokens(request_id, 1) + + def update_resources( + self, + scheduled_batch, + attn_metadata: AttentionMetadata = None, + kv_cache_dtype_byte_size: float = None, + ): + for request in scheduled_batch.context_requests: + if request.state != LlmRequestState.GENERATION_COMPLETE: + seq_len = request.get_num_tokens(0) + rewind_len = max(seq_len - 1 - self.prompt_budget, 0) + self.rewind_kv_cache(request, rewind_len) + # get the rewind length for kt cache + num_tokens = request.max_beam_num_tokens + updated_kt_token_num = num_tokens - rewind_len + rewind_len = math.ceil(num_tokens / self.page_size) - math.ceil( + updated_kt_token_num / self.page_size + ) + self.kt_cache_manager.rewind_cache(request, rewind_len) + + def free_resources(self, request): + super().free_resources(request) + self.kt_cache_manager.free_resources(request) + + @staticmethod + def get_cache_size_per_token( + model_config: ModelConfig, mapping: Mapping, num_layers: Optional[int] = None, **kwargs + ): + # main KV cache dtype bytes (one each for K and V) + kv_dtype_bytes = 2 + quant_config = model_config.quant_config + if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache(): + kv_dtype_bytes = 1 + + # get num key value heads + config = model_config.pretrained_config + num_key_value_heads = getattr(config, "num_key_value_heads", config.num_attention_heads) + if isinstance(num_key_value_heads, Iterable): + num_key_value_heads = sum(num_key_value_heads) / len(num_key_value_heads) + + # get head dim + tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size + head_dim = getattr(config, "head_dim", None) + if not isinstance(head_dim, int): + head_dim = config.hidden_size // config.num_attention_heads + head_dim = head_dim * num_key_value_heads // tp_size + + num_attention_layers = KVCacheManager._resolve_num_attention_layers( + model_config, mapping, num_layers + ) + + # K and V: two tensors stored at the main KV cache dtype. + mem_per_token = kv_dtype_bytes * 2 * num_attention_layers * head_dim + + tokens_per_block = kwargs["tokens_per_block"] + sparse_attention_config = model_config.sparse_attention_config + if sparse_attention_config is None: + raise ValueError("sparse_attention_config is required for RocketKV cache") + sparse_params = sparse_attention_config.to_sparse_params( + pretrained_config=model_config.pretrained_config + ) + if not isinstance(sparse_params, RocketKVParams): + raise ValueError("RocketKV cache requires RocketKV sparse parameters") + kt_tokens_per_block = next_power_of_2(math.ceil(tokens_per_block / sparse_params.page_size)) + # KT (key-landmark) cache: a separate pool stored at kt_cache_dtype, + # NOT the main KV dtype. Per real token it holds + # kt_tokens_per_block / tokens_per_block landmark slots, each storing + # the concatenated (min, max) landmark -- a factor of 2 over head_dim. + # Size it at its own dtype's byte width (1 for float8_e5m2, else 2) so + # the estimate is independent of the main KV dtype and matches both + # get_cache_bytes_per_token and the physical pool allocated in + # RocketKVCacheManager.__init__. + kt_dtype_bytes = 1 if sparse_params.kt_cache_dtype == "float8_e5m2" else 2 + mem_per_token += ( + kt_dtype_bytes + * 2 + * kt_tokens_per_block + / tokens_per_block + * num_attention_layers + * head_dim + ) + return mem_per_token + + def get_cache_bytes_per_token(self): + # Main K/V cache: stored at the main KV cache dtype (self.dtype). + cache_size_per_token = math.ceil( + self.kv_factor * sum(self.num_kv_heads_per_layer) * self.head_dim + ) + + if self.dtype not in ( + DataType.FP8, + DataType.HALF, + DataType.BF16, + DataType.FLOAT, + DataType.NVFP4, + ): + raise ValueError(f"Cannot support {self.dtype} KV cache.") + + cache_size_bytes_per_token = get_size_in_bytes(cache_size_per_token, self.dtype) + if self.dtype == DataType.NVFP4: + cache_size_bytes_per_token += self.calculate_scaling_factor_size_bytes( + cache_size_per_token, quant_vector_size=16, scaling_factor_dtype=DataType.FP8 + ) + + # KT (key-landmark) cache: a separate pool physically allocated at + # kt_cache_dtype, NOT the main KV dtype -- see __init__: + # torch.empty((num_blocks, kt_tokens_per_block, num_kv_heads, + # head_dim * 2), dtype=kt_cache_dtype). + # The trailing factor of 2 is the concatenated (min, max) landmark per + # head_dim. Per real token the pool holds + # kt_tokens_per_block / tokens_per_block landmark slots. Size it at its + # own dtype's byte width (1 for float8_e5m2, else 2) so the estimate + # matches the allocation regardless of self.dtype. + kt_dtype_bytes = 1 if self.kt_cache_dtype == torch.float8_e5m2 else 2 + kt_elements_per_token = math.ceil( + 2 + * self.kt_tokens_per_block + / self.tokens_per_block + * sum(self.num_kv_heads_per_layer) + * self.head_dim + ) + cache_size_bytes_per_token += kt_elements_per_token * kt_dtype_bytes + + return cache_size_bytes_per_token diff --git a/tensorrt_llm/_torch/attention_backend/sparse/kernel.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py similarity index 51% rename from tensorrt_llm/_torch/attention_backend/sparse/kernel.py rename to tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py index 6be9cf716325..bcd1ea339e59 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/kernel.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py @@ -1,5 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + import math -import os import torch import triton @@ -13,85 +17,24 @@ @triton.jit -def _index_gather_kernel(output_ptr, input_ptr, index_ptr, in_row_stride, - in_token_stride, in_head_stride, idx_row_stride, - idx_token_stride, idx_head_stride, dim_size, - BLOCK_SIZE: tl.constexpr): - # get program id and block offset - row_pid = tl.program_id(0) - token_pid = tl.program_id(1) - head_pid = tl.program_id(2) - token_block_num = tl.num_programs(1) - head_num = tl.num_programs(2) - - # get index - indices_idx = row_pid * idx_row_stride + token_pid * idx_token_stride + head_pid * idx_head_stride - token_idx = tl.load(index_ptr + indices_idx) - - # get input and output base address - input_base = (row_pid * in_row_stride + token_idx * in_token_stride + - head_pid * in_head_stride) - output_base = (row_pid * token_block_num * head_num * dim_size + - token_pid * head_num * dim_size + head_pid * dim_size) - - # process elements in blocks - for dim_offset in tl.range(0, dim_size, BLOCK_SIZE): - # get offsets - offsets = tl.arange(0, BLOCK_SIZE) - dim_indices = dim_offset + offsets - mask = dim_indices < dim_size - - # load input and store output - input_val = tl.load(input_ptr + input_base + dim_indices, - mask=mask, - other=0.0) - tl.store(output_ptr + output_base + dim_indices, input_val, mask=mask) - - -def triton_index_gather(input, indices): - assert input.ndim == 4, "Input must be a 4D tensor, [row, token, head, dim]" - assert indices.ndim == 3, "Indices must be a 3D tensor, [row, token, head]" - - # shape of input and indices - row_size = input.shape[0] - head_num = input.shape[2] - dim_size = input.shape[3] - num_tokens = indices.shape[1] - - # create output tensor - output = torch.empty((row_size, num_tokens, head_num, dim_size), - device='cuda', - dtype=input.dtype) - - # launch kernel - grid = (row_size, num_tokens, head_num) - _index_gather_kernel[grid](output, - input, - indices, - input.stride(0), - input.stride(1), - input.stride(2), - indices.stride(0), - indices.stride(1), - indices.stride(2), - dim_size, - BLOCK_SIZE=1024) - return output - - -######################################################## -# QK split kernel -######################################################## - - -@triton.jit -def rocket_qk_split_kernel(input_ptr, q_output_ptr, k_output_ptr, - k_output_offsets_ptr, context_lens_ptr, - context_cumsum_ptr, valid_seq_indices_ptr, - window_size, num_heads, num_kv_heads, head_dim, - q_total_output_tokens, k_total_output_tokens, - valid_batch_size, BLOCK_SIZE: tl.constexpr, - BLOCK_SIZE_M: tl.constexpr): +def rocket_qk_split_kernel( + input_ptr, + q_output_ptr, + k_output_ptr, + k_output_offsets_ptr, + context_lens_ptr, + context_cumsum_ptr, + valid_seq_indices_ptr, + window_size, + num_heads, + num_kv_heads, + head_dim, + q_total_output_tokens, + k_total_output_tokens, + valid_batch_size, + BLOCK_SIZE: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, +): valid_seq_idx = tl.program_id(0) head_idx = tl.program_id(1) dim_offset = tl.program_id(2) * BLOCK_SIZE @@ -143,16 +86,17 @@ def rocket_qk_split_kernel(input_ptr, q_output_ptr, k_output_ptr, # Calculate input indices with proper dimension offset for Q/K separation input_dim_indices = input_dim_offset + actual_head_idx * head_dim + dim_indices - src_indices = src_token_pos[:, None] * ( - num_heads * head_dim + - 2 * num_kv_heads * head_dim) + input_dim_indices[None, :] + src_indices = ( + src_token_pos[:, None] * (num_heads * head_dim + 2 * num_kv_heads * head_dim) + + input_dim_indices[None, :] + ) # Calculate output indices - dst_indices = (actual_head_idx * total_output_tokens + - dst_token_pos[:, None]) * head_dim + dim_indices[None, :] + dst_indices = ( + actual_head_idx * total_output_tokens + dst_token_pos[:, None] + ) * head_dim + dim_indices[None, :] - full_mask = token_mask[:, None] & dim_mask[ - None, :] # [BLOCK_SIZE_M, BLOCK_SIZE] + full_mask = token_mask[:, None] & dim_mask[None, :] # [BLOCK_SIZE_M, BLOCK_SIZE] data = tl.load(input_ptr + src_indices, mask=full_mask, other=0.0) tl.store(output_ptr + dst_indices, data, mask=full_mask) @@ -199,12 +143,16 @@ def triton_rocket_qk_split( q_total_output_tokens = window_size * valid_batch_size k_total_output_tokens = total_rocket_k_ctx_tokens - q_window = torch.empty((num_heads, q_total_output_tokens, head_dim), - device=input_tensor.device, - dtype=input_tensor.dtype) - k_context = torch.empty((num_kv_heads, k_total_output_tokens, head_dim), - device=input_tensor.device, - dtype=input_tensor.dtype) + q_window = torch.empty( + (num_heads, q_total_output_tokens, head_dim), + device=input_tensor.device, + dtype=input_tensor.dtype, + ) + k_context = torch.empty( + (num_kv_heads, k_total_output_tokens, head_dim), + device=input_tensor.device, + dtype=input_tensor.dtype, + ) BLOCK_SIZE = 128 # Dimension block size BLOCK_SIZE_M = 128 # Token block size for parallel processing @@ -213,1043 +161,1197 @@ def triton_rocket_qk_split( total_heads = num_heads + num_kv_heads grid = (valid_batch_size, total_heads, triton.cdiv(head_dim, BLOCK_SIZE)) - rocket_qk_split_kernel[grid](input_tensor, - q_window, - k_context, - k_output_offsets, - input_lens, - input_lens_cumsum, - valid_seq_indices, - window_size, - num_heads, - num_kv_heads, - head_dim, - q_total_output_tokens, - k_total_output_tokens, - valid_batch_size, - BLOCK_SIZE=BLOCK_SIZE, - BLOCK_SIZE_M=BLOCK_SIZE_M) + rocket_qk_split_kernel[grid]( + input_tensor, + q_window, + k_context, + k_output_offsets, + input_lens, + input_lens_cumsum, + valid_seq_indices, + window_size, + num_heads, + num_kv_heads, + head_dim, + q_total_output_tokens, + k_total_output_tokens, + valid_batch_size, + BLOCK_SIZE=BLOCK_SIZE, + BLOCK_SIZE_M=BLOCK_SIZE_M, + ) return q_window, k_context -######################################################## -# BMM kernel -######################################################## - - @triton.jit -def bmm_kernel( - q_ptr, - k_ptr, - scores_ptr, - q_cu_seqlens_ptr, - k_cu_seqlens_ptr, - total_q_tokens, - total_k_tokens, - head_dim, +def rocket_batch_to_flatten_kernel( + prefix_indices_ptr, + output_indices_ptr, + context_lens_ptr, + valid_seq_indices_ptr, + sparse_offsets_ptr, batch_size, - num_q_heads, - num_k_heads, - q_len_per_seq, - sm_scale, - causal, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_K: tl.constexpr, + valid_batch_size, + num_kv_heads, + prefix_budget, + window_size, + prompt_budget, + BLOCK_SIZE: tl.constexpr, ): batch_idx = tl.program_id(0) head_idx = tl.program_id(1) - if batch_idx >= batch_size or head_idx >= num_q_heads: + if batch_idx >= batch_size or head_idx >= num_kv_heads: return - # Continuous mapping of query heads to key heads - k_head_idx = head_idx // (num_q_heads // num_k_heads) - - q_seq_start = tl.load(q_cu_seqlens_ptr + batch_idx) - q_seq_end = tl.load(q_cu_seqlens_ptr + batch_idx + 1) - k_seq_start = tl.load(k_cu_seqlens_ptr + batch_idx) - k_seq_end = tl.load(k_cu_seqlens_ptr + batch_idx + 1) - - q_seqlen = q_seq_end - q_seq_start - k_seqlen = k_seq_end - k_seq_start - - if q_seqlen <= 0 or k_seqlen <= 0: - return + # Get context length for this batch + context_len = tl.load(context_lens_ptr + batch_idx) - # Process queries in this batch with BLOCK_M parallelization - for q_block_start in tl.range(0, q_seqlen, BLOCK_M): - q_offsets = q_block_start + tl.arange(0, BLOCK_M) - q_mask = q_offsets < q_seqlen - q_global_offsets = q_seq_start + q_offsets + # Check if this batch is valid (appears in valid_seq_indices) + is_valid = 0 + valid_idx_in_selected = -1 - for k_block_start in tl.range(0, k_seqlen, BLOCK_N): - k_offsets = k_block_start + tl.arange(0, BLOCK_N) - k_mask = k_offsets < k_seqlen - k_global_offsets = k_seq_start + k_offsets + # Search for current batch_idx in valid_seq_indices + for valid_idx in tl.range(0, valid_batch_size): + orig_idx = tl.load(valid_seq_indices_ptr + valid_idx) + if orig_idx == batch_idx: + is_valid = 1 + valid_idx_in_selected = valid_idx - # Initialize QK^T accumulator for this (M, N) block - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + # Calculate output offset for this batch + output_offset = tl.load(sparse_offsets_ptr + batch_idx) + total_sparse_tokens = tl.load(sparse_offsets_ptr + batch_size) - # Tiled matrix multiplication following mm_demo.py pattern - for k_dim_start in tl.range(0, head_dim, BLOCK_K): - k_dim_offsets = k_dim_start + tl.arange(0, BLOCK_K) - k_dim_mask = k_dim_offsets < head_dim + if is_valid: + # Valid batch: copy prefix indices and compute window indices + # Process prefix tokens + for token_block_start in tl.range(0, prefix_budget, BLOCK_SIZE): + token_offsets = token_block_start + tl.arange(0, BLOCK_SIZE) + token_mask = token_offsets < prefix_budget - # Load query chunk [BLOCK_M, BLOCK_K] - q_indices = head_idx * total_q_tokens * head_dim + q_global_offsets[:, None] * head_dim + k_dim_offsets[ - None, :] - q_chunk = tl.load(q_ptr + q_indices, - mask=q_mask[:, None] & k_dim_mask[None, :], - other=0.0) + # Load from prefix_indices + flattened_idx = valid_idx_in_selected * num_kv_heads + head_idx + prefix_indices = flattened_idx * prefix_budget + token_offsets + prefix_values = tl.load(prefix_indices_ptr + prefix_indices, mask=token_mask, other=0) - # Load key chunk [BLOCK_N, BLOCK_K] - k_indices = k_head_idx * total_k_tokens * head_dim + k_global_offsets[:, None] * head_dim + k_dim_offsets[ - None, :] - k_chunk = tl.load(k_ptr + k_indices, - mask=k_mask[:, None] & k_dim_mask[None, :], - other=0.0) + # Store to output + output_indices = head_idx * total_sparse_tokens + output_offset + token_offsets + tl.store(output_indices_ptr + output_indices, prefix_values, mask=token_mask) - # Accumulate QK^T using tl.dot for better performance - qk += tl.dot(q_chunk, tl.trans(k_chunk)) + # Process window tokens + for token_block_start in tl.range(0, window_size, BLOCK_SIZE): + token_offsets = token_block_start + tl.arange(0, BLOCK_SIZE) + token_mask = token_offsets < window_size - # Scale the accumulated QK^T - qk = qk * sm_scale + # Compute window indices: [context_len - window_size, context_len - window_size + 1, ...] + window_values = context_len - window_size + token_offsets - # Apply masking - valid_mask = q_mask[:, None] & k_mask[None, :] - if causal: - # Create causal mask based on positions within this batch's sequence - q_pos_in_seq = q_offsets[:, None] # [BLOCK_M, 1] - k_pos_in_seq = k_offsets[None, :] # [1, BLOCK_N] - causal_mask = q_pos_in_seq >= k_pos_in_seq - qk = tl.where(causal_mask & valid_mask, qk, float("-inf")) - else: - qk = tl.where(valid_mask, qk, float("-inf")) + # Store to output at prefix_budget offset + output_indices = ( + head_idx * total_sparse_tokens + output_offset + prefix_budget + token_offsets + ) + tl.store(output_indices_ptr + output_indices, window_values, mask=token_mask) + else: + # Invalid batch: generate [0, 1, ..., context_len-1] + num_tokens = context_len + for token_block_start in tl.range(0, num_tokens, BLOCK_SIZE): + token_offsets = token_block_start + tl.arange(0, BLOCK_SIZE) + token_mask = token_offsets < num_tokens - # Store results - note that we store in the global k position space - # This matches the original bmm_softmax kernel behavior - output_indices = (head_idx * q_len_per_seq * total_k_tokens + - q_offsets[:, None] * total_k_tokens + - k_global_offsets[None, :]) # [BLOCK_M, BLOCK_N] + # Generate sequential indices + sequential_indices = token_offsets - tl.store(scores_ptr + output_indices, qk, mask=valid_mask) + # Store to output + output_indices = head_idx * total_sparse_tokens + output_offset + token_offsets + tl.store(output_indices_ptr + output_indices, sequential_indices, mask=token_mask) -def triton_bmm(q: torch.Tensor, - k: torch.Tensor, - q_cu_seqlens: torch.Tensor, - k_cu_seqlens: torch.Tensor, - batch_size: int, - sm_scale: float = None, - causal: bool = False) -> torch.Tensor: +def triton_rocket_batch_to_flatten( + prefix_indices: torch.Tensor, + input_lens: torch.Tensor, + valid_seq_indices: torch.Tensor, + output_offsets: torch.Tensor, + batch_size: int, + total_output_tokens: int, + window_size: int, + prompt_budget: int, + num_kv_heads: int, +) -> tuple[torch.Tensor, torch.Tensor]: """ - Compute softmax(QK^T) with flattened input and output. - In this kernel we assume that each sequence in the batch has the same Q length. + Flatten indices considering both valid and invalid batches. + For valid sequences, combines prefix_indices with dynamically computed window indices. + For invalid sequences, generates sequential indices. Args: - q: Query tensor [num_q_heads, total_q_tokens, head_dim] - k: Key tensor [num_kv_heads, total_k_tokens, head_dim] - q_cu_seqlens: Query cumulative sequence lengths [batch_size + 1] - k_cu_seqlens: Key cumulative sequence lengths [batch_size + 1] + prefix_indices: Selected prefix indices [valid_batch_size * num_kv_heads, prefix_budget] + input_lens: Lengths for all sequences [batch_size] + valid_seq_indices: Valid sequence indices [valid_batch_size] + output_offsets: Offset for each batch [batch_size + 1] batch_size: Number of batches - sm_scale: Scaling factor (default: 1/sqrt(head_dim)) - causal: Whether to apply causal masking + total_output_tokens: Total number of output tokens + window_size: Size of sliding window at the end + prompt_budget: Total number of tokens for valid sequences (prefix_budget + window_size) + num_kv_heads: Number of KV heads Returns: - scores: Attention scores [num_q_heads, q_len_per_seq, total_k_tokens] + sparse_indices: Flattened sparse indices [num_kv_heads, total_output_tokens] """ - num_q_heads, total_q_tokens, head_dim = q.shape - num_k_heads, total_k_tokens, _ = k.shape - - assert total_q_tokens % batch_size == 0, "total_q_tokens must be divisible by batch_size" - q_len_per_seq = total_q_tokens // batch_size + total_tasks, prefix_budget = prefix_indices.shape + valid_batch_size = total_tasks // num_kv_heads - if total_k_tokens == 0: - return torch.zeros((num_q_heads, q_len_per_seq, total_k_tokens), - dtype=torch.float32, - device=q.device) + # Create output tensor + sparse_indices = torch.empty( + (num_kv_heads, total_output_tokens), + dtype=prefix_indices.dtype, + device=prefix_indices.device, + ) - if sm_scale is None: - sm_scale = 1.0 / math.sqrt(head_dim) + # Launch kernel + BLOCK_SIZE = 512 + grid = (batch_size, num_kv_heads) - bmm_results = torch.empty((num_q_heads, q_len_per_seq, total_k_tokens), - dtype=torch.float32, - device=q.device) + rocket_batch_to_flatten_kernel[grid]( + prefix_indices, + sparse_indices, + input_lens, + valid_seq_indices, + output_offsets, + batch_size, + valid_batch_size, + num_kv_heads, + prefix_budget, + window_size, + prompt_budget, + BLOCK_SIZE=BLOCK_SIZE, + ) - # BMM kernel configuration - BLOCK_M = 32 - BLOCK_N = 256 - BLOCK_K = 64 - - grid_bmm = lambda meta: (batch_size, num_q_heads) - - bmm_kernel[grid_bmm]( - q, - k, - bmm_results, - q_cu_seqlens, - k_cu_seqlens, - total_q_tokens, - total_k_tokens, - head_dim, - batch_size, - num_q_heads, - num_k_heads, - q_len_per_seq, - sm_scale, - causal, - BLOCK_M=BLOCK_M, - BLOCK_N=BLOCK_N, - BLOCK_K=BLOCK_K, - ) - - return bmm_results - - -######################################################## -# Softmax kernel -######################################################## + return sparse_indices @triton.jit -def softmax_kernel( - input_ptr, - output_ptr, - cu_seq_lens_ptr, - batch_size, - num_heads, - q_len_per_seq, - total_k_tokens, - BLOCK_SIZE: tl.constexpr, +def rocket_update_kt_cache_gen_kernel( + k_ptr, + kt_cache_tensor_ptr, + kt_cache_block_offsets_ptr, + kv_lens_ptr, + num_gen_tokens, + num_kv_heads, + head_dim, + kt_page_size, + tokens_per_block, + max_kt_blocks_per_seq, + k_stride_0, + k_stride_1, + DIM_BLOCK_SIZE: tl.constexpr, ): - head_idx = tl.program_id(0) - q_global_idx = tl.program_id(1) + batch_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + dim_block_idx = tl.program_id(2) - if head_idx >= num_heads or q_global_idx >= (batch_size * q_len_per_seq): + if batch_idx >= num_gen_tokens or kv_head_idx >= num_kv_heads: return - # Determine which batch this q belongs to - batch_idx = q_global_idx // q_len_per_seq - q_local_idx = q_global_idx % q_len_per_seq - - if batch_idx >= batch_size: - return + dim_block_start = dim_block_idx * DIM_BLOCK_SIZE + dim_offsets = tl.arange(0, DIM_BLOCK_SIZE) + dim_indices = dim_block_start + dim_offsets + dim_mask = dim_indices < head_dim - # Get k sequence boundaries for this batch - k_seq_start = tl.load(cu_seq_lens_ptr + batch_idx) - k_seq_end = tl.load(cu_seq_lens_ptr + batch_idx + 1) - k_seqlen = k_seq_end - k_seq_start + k_base = batch_idx * k_stride_0 + kv_head_idx * head_dim * k_stride_1 + k_indices = k_base + dim_indices * k_stride_1 + k_values = tl.load(k_ptr + k_indices, mask=dim_mask, other=0.0) - if k_seqlen <= 0: + kv_len = tl.load(kv_lens_ptr + batch_idx) + if kv_len <= 0: return - # Calculate input/output row start - input_row_start = head_idx * q_len_per_seq * total_k_tokens + q_local_idx * total_k_tokens - output_row_start = input_row_start + # Determine which kt_token to update (the last one) + last_token_idx = kv_len - 1 + last_kt_token_idx = last_token_idx // kt_page_size + kt_token_idx_in_page = last_token_idx % kt_page_size - # Find max value only within the valid k range for this batch - row_max = float('-inf') - for k_block_start in tl.range(0, k_seqlen, BLOCK_SIZE): - k_offsets = k_block_start + tl.arange(0, BLOCK_SIZE) - k_mask = k_offsets < k_seqlen - k_global_offsets = k_seq_start + k_offsets + block_offset_in_seq = last_kt_token_idx // tokens_per_block + if block_offset_in_seq >= max_kt_blocks_per_seq: + return - input_indices = input_row_start + k_global_offsets - values = tl.load(input_ptr + input_indices, - mask=k_mask, - other=float('-inf')) - block_max = tl.max(values, axis=0) - row_max = tl.maximum(row_max, block_max) + block_idx = tl.load( + kt_cache_block_offsets_ptr + batch_idx * max_kt_blocks_per_seq + block_offset_in_seq + ) + token_idx_in_block = last_kt_token_idx % tokens_per_block - # Compute sum of exp(x - max) only within valid k range - row_sum = 0.0 - for k_block_start in tl.range(0, k_seqlen, BLOCK_SIZE): - k_offsets = k_block_start + tl.arange(0, BLOCK_SIZE) - k_mask = k_offsets < k_seqlen - k_global_offsets = k_seq_start + k_offsets + cache_base = ( + block_idx * tokens_per_block + token_idx_in_block + ) * num_kv_heads * 2 * head_dim + kv_head_idx * 2 * head_dim - input_indices = input_row_start + k_global_offsets - values = tl.load(input_ptr + input_indices, - mask=k_mask, - other=float('-inf')) - exp_values = tl.exp(values - row_max) - masked_exp = tl.where(k_mask, exp_values, 0.0) - row_sum += tl.sum(masked_exp, axis=0) + cache_min_indices = cache_base + dim_indices + cache_max_indices = cache_base + head_dim + dim_indices - # Apply softmax and store results - for k_block_start in tl.range(0, k_seqlen, BLOCK_SIZE): - k_offsets = k_block_start + tl.arange(0, BLOCK_SIZE) - k_mask = k_offsets < k_seqlen - k_global_offsets = k_seq_start + k_offsets + kt_mask = dim_mask & (kt_token_idx_in_page > 0) - input_indices = input_row_start + k_global_offsets - output_indices = output_row_start + k_global_offsets + k_min_existing = tl.load( + kt_cache_tensor_ptr + cache_min_indices, mask=kt_mask, other=float("inf") + ) + k_max_existing = tl.load( + kt_cache_tensor_ptr + cache_max_indices, mask=kt_mask, other=float("-inf") + ) - values = tl.load(input_ptr + input_indices, - mask=k_mask, - other=float('-inf')) - exp_values = tl.exp(values - row_max) - masked_exp = tl.where(k_mask, exp_values, 0.0) - softmax_values = masked_exp / row_sum + k_min_new = tl.minimum(k_min_existing, k_values) + k_max_new = tl.maximum(k_max_existing, k_values) + k_min_new = k_min_new.to(kt_cache_tensor_ptr.dtype.element_ty) + k_max_new = k_max_new.to(kt_cache_tensor_ptr.dtype.element_ty) - tl.store(output_ptr + output_indices, softmax_values, mask=k_mask) + tl.store(kt_cache_tensor_ptr + cache_min_indices, k_min_new, mask=dim_mask) + tl.store(kt_cache_tensor_ptr + cache_max_indices, k_max_new, mask=dim_mask) -def triton_softmax( - input_tensor: torch.Tensor, - cum_lens: torch.Tensor, - batch_size: int, -) -> torch.Tensor: +def triton_rocket_update_kt_cache_gen( + k: torch.Tensor, + kt_cache_tensor: torch.Tensor, + kt_cache_block_offsets: torch.Tensor, + kv_lens: torch.Tensor, + kt_page_size: int, + tokens_per_block: int, + max_kt_blocks_per_seq: int, + num_kv_heads: int, + head_dim: int, +) -> None: """ - Apply softmax to flattened input tensor. + Update KT cache with new key values for generation phase. Args: - input_tensor: Input tensor [num_heads, len_per_seq, total_k_tokens] or [num_heads, total_k_tokens] - cum_lens: Cumulative lengths [batch_size + 1] - batch_size: Number of batches - - Returns: - output: Softmax results, shape is like input_tensor + k: Key tensor [num_gen_tokens, num_kv_heads * head_dim] + kt_cache_tensor: KT cache tensor + kt_cache_block_offsets: Block offsets [num_gen_tokens, max_kt_blocks_per_seq] + kv_lens: Sequence lengths [num_gen_tokens] + kt_page_size: Page size for KT tokens + tokens_per_block: Tokens per cache block + max_kt_blocks_per_seq: Maximum KT blocks per sequence """ - if input_tensor.ndim == 2: - num_heads, total_k_tokens = input_tensor.shape - len_per_seq = 1 - else: - num_heads, len_per_seq, total_k_tokens = input_tensor.shape - - output = torch.empty_like(input_tensor, - dtype=input_tensor.dtype, - device=input_tensor.device) + num_gen_tokens = k.shape[0] - BLOCK_SIZE = 512 + grid = (num_gen_tokens, num_kv_heads, 1) - grid = (num_heads, batch_size * len_per_seq) + DIM_BLOCK_SIZE = triton.next_power_of_2(head_dim) - softmax_kernel[grid]( - input_tensor, - output, - cum_lens, - batch_size, - num_heads, - len_per_seq, - total_k_tokens, - BLOCK_SIZE=BLOCK_SIZE, + rocket_update_kt_cache_gen_kernel[grid]( + k, + kt_cache_tensor, + kt_cache_block_offsets, + kv_lens, + num_gen_tokens, + num_kv_heads, + head_dim, + kt_page_size, + tokens_per_block, + max_kt_blocks_per_seq, + k.stride(0), + k.stride(1), + DIM_BLOCK_SIZE=DIM_BLOCK_SIZE, ) - return output - - -######################################################## -# Reshape flatten to batched kernel -######################################################## - @triton.jit -def flatten_to_batch_kernel( - input_ptr, - output_ptr, - input_offsets, +def rocket_update_kt_cache_ctx_kernel( + k_ptr, + kt_cache_tensor_ptr, + kt_cache_block_offsets_ptr, + context_cumsum_ptr, + sparse_kv_indices_ptr, + sparse_kv_offsets_ptr, + batch_size, num_heads, - total_tokens, - padding_to_size, - padding_value, - BLOCK_SIZE: tl.constexpr, + num_kv_heads, + head_dim, + kt_page_size: tl.constexpr, + tokens_per_block, + max_kt_blocks_per_seq, + total_sparse_tokens, + BLOCK_SIZE_KT: tl.constexpr, + BLOCK_SIZE_DIM: tl.constexpr, ): + """ + Triton kernel for updating KT cache during context phase with sparse indices. + """ batch_idx = tl.program_id(0) - head_idx = tl.program_id(1) + kv_head_idx = tl.program_id(1) + kt_block_idx = tl.program_id(2) - k_offset = tl.load(input_offsets + batch_idx) - context_len = tl.load(input_offsets + batch_idx + 1) - k_offset + context_start = tl.load(context_cumsum_ptr + batch_idx) - # Process in blocks - for block_start in tl.range(0, padding_to_size, BLOCK_SIZE): - pos_offsets = block_start + tl.arange(0, BLOCK_SIZE) + sparse_start = tl.load(sparse_kv_offsets_ptr + batch_idx) + sparse_end = tl.load(sparse_kv_offsets_ptr + batch_idx + 1) + num_sparse_tokens = sparse_end - sparse_start - pos_mask = pos_offsets < padding_to_size - valid_mask = pos_offsets < context_len - combined_mask = pos_mask & valid_mask + if num_sparse_tokens <= 0: + return - input_indices = head_idx * total_tokens + k_offset + pos_offsets + q_hidden_size = num_heads * head_dim + kv_hidden_size = num_kv_heads * head_dim + k_dim_base = q_hidden_size + kv_head_idx * head_dim - output_indices = batch_idx * num_heads * padding_to_size + head_idx * padding_to_size + pos_offsets + BLOCK_SIZE_KV: tl.constexpr = kt_page_size * BLOCK_SIZE_KT - values = tl.where( - combined_mask, - tl.load(input_ptr + input_indices, - mask=combined_mask, - other=padding_value), padding_value) + total_kt_tokens = (num_sparse_tokens + kt_page_size - 1) // kt_page_size + kt_offsets = kt_block_idx * BLOCK_SIZE_KT + tl.arange(0, BLOCK_SIZE_KT) + kt_mask = kt_offsets < total_kt_tokens - tl.store(output_ptr + output_indices, values, mask=pos_mask) + kv_start = kt_block_idx * BLOCK_SIZE_KT * kt_page_size + kv_offsets = kv_start + tl.arange(0, BLOCK_SIZE_KV) + kv_mask = kv_offsets < num_sparse_tokens + kv_indices = kv_head_idx * total_sparse_tokens + sparse_start + kv_offsets + for dim_block_start in tl.range(0, head_dim, BLOCK_SIZE_DIM): + dim_offsets = tl.arange(0, BLOCK_SIZE_DIM) + dim_indices = dim_block_start + dim_offsets + dim_mask = dim_indices < head_dim -def triton_flatten_to_batch(input_tensor: torch.Tensor, - input_offsets: torch.Tensor, - batch_size: int, - padding_to_size: int, - padding_value=-1e10) -> torch.Tensor: - """ - Reshape input_tensor from [num_heads, total_tokens] to [batch_size, num_heads, padding_to_size] + kv_token_indices = tl.load(sparse_kv_indices_ptr + kv_indices, mask=kv_mask, other=0) + # Calculate indices for loading keys [BLOCK_SIZE_DIM, BLOCK_SIZE_KV] + k_base_indices = (kv_token_indices[None, :] + context_start) * ( + q_hidden_size + 2 * kv_hidden_size + ) + k_dim_base + k_indices = k_base_indices + dim_indices[:, None] - Args: - input_tensor: Input tensor tensor [num_heads, total_tokens] - input_offsets: Offset for each valid sequence [batch_size + 1] - batch_size: Number of batches - padding_to_size: Target padding size - padding_value: Value to fill for padding - - Returns: - batched_tensor: Output tensor [batch_size, num_heads, padding_to_size] - """ - num_heads, total_tokens = input_tensor.shape - - batched_tensor = torch.empty((batch_size, num_heads, padding_to_size), - device=input_tensor.device, - dtype=input_tensor.dtype) - - grid = lambda meta: (batch_size, num_heads) - flatten_to_batch_kernel[grid](input_tensor, - batched_tensor, - input_offsets, - num_heads, - total_tokens, - padding_to_size, - padding_value, - BLOCK_SIZE=1024) - - return batched_tensor - - -######################################################## -# RocketKV sparse indices batch to flatten flattening kernel -######################################################## - - -@triton.jit -def rocket_batch_to_flatten_kernel( - prefix_indices_ptr, - output_indices_ptr, - context_lens_ptr, - valid_seq_indices_ptr, - sparse_offsets_ptr, - batch_size, - valid_batch_size, - num_kv_heads, - prefix_budget, - window_size, - prompt_budget, - BLOCK_SIZE: tl.constexpr, -): - batch_idx = tl.program_id(0) - head_idx = tl.program_id(1) - - if batch_idx >= batch_size or head_idx >= num_kv_heads: - return - - # Get context length for this batch - context_len = tl.load(context_lens_ptr + batch_idx) - - # Check if this batch is valid (appears in valid_seq_indices) - is_valid = 0 - valid_idx_in_selected = -1 + combined_mask = kv_mask[None, :] & dim_mask[:, None] - # Search for current batch_idx in valid_seq_indices - for valid_idx in tl.range(0, valid_batch_size): - orig_idx = tl.load(valid_seq_indices_ptr + valid_idx) - if orig_idx == batch_idx: - is_valid = 1 - valid_idx_in_selected = valid_idx + # Load key values [BLOCK_SIZE_DIM, BLOCK_SIZE_KV] + k_values = tl.load(k_ptr + k_indices, mask=combined_mask, other=0.0) - # Calculate output offset for this batch - output_offset = tl.load(sparse_offsets_ptr + batch_idx) - total_sparse_tokens = tl.load(sparse_offsets_ptr + batch_size) + k_values = tl.reshape(k_values, (BLOCK_SIZE_DIM, BLOCK_SIZE_KT, kt_page_size)) - if is_valid: - # Valid batch: copy prefix indices and compute window indices - # Process prefix tokens - for token_block_start in tl.range(0, prefix_budget, BLOCK_SIZE): - token_offsets = token_block_start + tl.arange(0, BLOCK_SIZE) - token_mask = token_offsets < prefix_budget + k_min = tl.min(k_values, axis=-1).to(kt_cache_tensor_ptr.dtype.element_ty) + k_max = tl.max(k_values, axis=-1).to(kt_cache_tensor_ptr.dtype.element_ty) - # Load from prefix_indices - flattened_idx = valid_idx_in_selected * num_kv_heads + head_idx - prefix_indices = flattened_idx * prefix_budget + token_offsets - prefix_values = tl.load(prefix_indices_ptr + prefix_indices, - mask=token_mask, - other=0) + # Calculate cache locations [BLOCK_SIZE_KT] + block_offsets_in_seq = kt_offsets // tokens_per_block + valid_block_mask = (block_offsets_in_seq < max_kt_blocks_per_seq) & kt_mask - # Store to output - output_indices = head_idx * total_sparse_tokens + output_offset + token_offsets - tl.store(output_indices_ptr + output_indices, - prefix_values, - mask=token_mask) + # Load block indices [BLOCK_SIZE_KT] + block_offset_addrs = batch_idx * max_kt_blocks_per_seq + block_offsets_in_seq + block_indices = tl.load( + kt_cache_block_offsets_ptr + block_offset_addrs, mask=valid_block_mask, other=0 + ) - # Process window tokens - for token_block_start in tl.range(0, window_size, BLOCK_SIZE): - token_offsets = token_block_start + tl.arange(0, BLOCK_SIZE) - token_mask = token_offsets < window_size + tokens_in_block = kt_offsets % tokens_per_block - # Compute window indices: [context_len - window_size, context_len - window_size + 1, ...] - window_values = context_len - window_size + token_offsets + # Calculate cache base addresses [BLOCK_SIZE_KT] + cache_bases = ( + block_indices * tokens_per_block + tokens_in_block + ) * num_kv_heads * 2 * head_dim + kv_head_idx * 2 * head_dim - # Store to output at prefix_budget offset - output_indices = head_idx * total_sparse_tokens + output_offset + prefix_budget + token_offsets - tl.store(output_indices_ptr + output_indices, - window_values, - mask=token_mask) - else: - # Invalid batch: generate [0, 1, ..., context_len-1] - num_tokens = context_len - for token_block_start in tl.range(0, num_tokens, BLOCK_SIZE): - token_offsets = token_block_start + tl.arange(0, BLOCK_SIZE) - token_mask = token_offsets < num_tokens + cache_min_addrs = cache_bases[None, :] + dim_indices[:, None] + cache_max_addrs = cache_min_addrs + head_dim - # Generate sequential indices - sequential_indices = token_offsets + store_mask = valid_block_mask[None, :] & dim_mask[:, None] - # Store to output - output_indices = head_idx * total_sparse_tokens + output_offset + token_offsets - tl.store(output_indices_ptr + output_indices, - sequential_indices, - mask=token_mask) + tl.store(kt_cache_tensor_ptr + cache_min_addrs, k_min, mask=store_mask) + tl.store(kt_cache_tensor_ptr + cache_max_addrs, k_max, mask=store_mask) -def triton_rocket_batch_to_flatten( - prefix_indices: torch.Tensor, input_lens: torch.Tensor, - valid_seq_indices: torch.Tensor, output_offsets: torch.Tensor, - batch_size: int, total_output_tokens: int, window_size: int, - prompt_budget: int, - num_kv_heads: int) -> tuple[torch.Tensor, torch.Tensor]: +def triton_rocket_update_kt_cache_ctx( + qkv_input: torch.Tensor, + kt_cache_tensor: torch.Tensor, + kt_cache_block_offsets: torch.Tensor, + context_cumsum: torch.Tensor, + sparse_kv_indices: torch.Tensor, + sparse_kv_offsets: torch.Tensor, + num_heads: int, + num_kv_heads: int, + head_dim: int, + kt_page_size: int, + prompt_budget: int, + tokens_per_block: int, + max_kt_blocks_per_seq: int, +): """ - Flatten indices considering both valid and invalid batches. - For valid sequences, combines prefix_indices with dynamically computed window indices. - For invalid sequences, generates sequential indices. + Update KT cache during context phase using sparse indices. Args: - prefix_indices: Selected prefix indices [valid_batch_size * num_kv_heads, prefix_budget] - input_lens: Lengths for all sequences [batch_size] - valid_seq_indices: Valid sequence indices [valid_batch_size] - output_offsets: Offset for each batch [batch_size + 1] - batch_size: Number of batches - total_output_tokens: Total number of output tokens - window_size: Size of sliding window at the end - prompt_budget: Total number of tokens for valid sequences (prefix_budget + window_size) + qkv_input: QKV input tensor [total_sparse_tokens, + num_heads*head_dim + num_kv_heads*head_dim + num_kv_heads*head_dim] + kt_cache_tensor: KT cache [num_blocks, tokens_per_block, num_kv_heads, 2*head_dim] + kt_cache_block_offsets: Block offsets [batch_size, max_kt_blocks_per_seq] + sparse_kv_indices: Sparse KV indices [num_kv_heads, total_sparse_tokens] + sparse_kv_offsets: Sparse offsets [batch_size + 1] + num_heads: Number of Q heads num_kv_heads: Number of KV heads - - Returns: - sparse_indices: Flattened sparse indices [num_kv_heads, total_output_tokens] + head_dim: Head dimension + kt_page_size: Page size for KT tokens + prompt_budget: Prompt budget + tokens_per_block: Tokens per cache block + max_kt_blocks_per_seq: Maximum KT blocks per sequence """ - total_tasks, prefix_budget = prefix_indices.shape - valid_batch_size = total_tasks // num_kv_heads - - # Create output tensor - sparse_indices = torch.empty((num_kv_heads, total_output_tokens), - dtype=prefix_indices.dtype, - device=prefix_indices.device) - - # Launch kernel - BLOCK_SIZE = 512 - grid = (batch_size, num_kv_heads) + batch_size = sparse_kv_offsets.size(0) - 1 + total_sparse_tokens = sparse_kv_indices.size(1) - rocket_batch_to_flatten_kernel[grid](prefix_indices, - sparse_indices, - input_lens, - valid_seq_indices, - output_offsets, - batch_size, - valid_batch_size, - num_kv_heads, - prefix_budget, - window_size, - prompt_budget, - BLOCK_SIZE=BLOCK_SIZE) + total_kt_tokens = (prompt_budget + kt_page_size - 1) // kt_page_size - return sparse_indices + BLOCK_SIZE_KT = 8 + BLOCK_SIZE_DIM = triton.next_power_of_2(head_dim) + grid = (batch_size, num_kv_heads, (total_kt_tokens + BLOCK_SIZE_KT - 1) // BLOCK_SIZE_KT) -######################################################## -# KT cache update kernel -######################################################## + rocket_update_kt_cache_ctx_kernel[grid]( + qkv_input, + kt_cache_tensor, + kt_cache_block_offsets, + context_cumsum, + sparse_kv_indices, + sparse_kv_offsets, + batch_size, + num_heads, + num_kv_heads, + head_dim, + kt_page_size, + tokens_per_block, + max_kt_blocks_per_seq, + total_sparse_tokens, + BLOCK_SIZE_KT=BLOCK_SIZE_KT, + BLOCK_SIZE_DIM=BLOCK_SIZE_DIM, + ) @triton.jit -def rocket_update_kt_cache_gen_kernel( - k_ptr, +def rocket_paged_kt_cache_bmm_kernel( + q_ptr, kt_cache_tensor_ptr, kt_cache_block_offsets_ptr, kv_lens_ptr, + output_ptr, + output_offsets_ptr, num_gen_tokens, num_kv_heads, + num_heads_per_kv: tl.constexpr, head_dim, kt_page_size, tokens_per_block, max_kt_blocks_per_seq, - k_stride_0, - k_stride_1, + total_kt_tokens, + sm_scale, + q_stride_0, + q_stride_1, + q_stride_2, + q_stride_3, + Q_BLOCK_SIZE: tl.constexpr, + KT_BLOCK_SIZE: tl.constexpr, DIM_BLOCK_SIZE: tl.constexpr, ): batch_idx = tl.program_id(0) kv_head_idx = tl.program_id(1) - dim_block_idx = tl.program_id(2) if batch_idx >= num_gen_tokens or kv_head_idx >= num_kv_heads: return + kv_len = tl.load(kv_lens_ptr + batch_idx) + num_kt_tokens = (kv_len + kt_page_size - 1) // kt_page_size - dim_block_start = dim_block_idx * DIM_BLOCK_SIZE - dim_offsets = tl.arange(0, DIM_BLOCK_SIZE) - dim_indices = dim_block_start + dim_offsets + q_base = batch_idx * q_stride_0 + kv_head_idx * q_stride_1 + + q_head_offsets = tl.arange(0, Q_BLOCK_SIZE) + q_head_mask = q_head_offsets < num_heads_per_kv + + output_offset = tl.load(output_offsets_ptr + batch_idx) + + dim_indices = tl.arange(0, DIM_BLOCK_SIZE) dim_mask = dim_indices < head_dim - k_base = batch_idx * k_stride_0 + kv_head_idx * head_dim * k_stride_1 - k_indices = k_base + dim_indices * k_stride_1 - k_values = tl.load(k_ptr + k_indices, mask=dim_mask, other=0.0) + q_indices = q_base + q_head_offsets[:, None] * q_stride_2 + dim_indices[None, :] * q_stride_3 + q_values = tl.load(q_ptr + q_indices, mask=q_head_mask[:, None] & dim_mask[None, :]) - kv_len = tl.load(kv_lens_ptr + batch_idx) - if kv_len <= 0: - return + dim_pos_values = tl.sum(q_values, axis=0) > 0 + dim_pos_values = tl.broadcast_to(dim_pos_values[None, :], (KT_BLOCK_SIZE, DIM_BLOCK_SIZE)) - # Determine which kt_token to update (the last one) - last_token_idx = kv_len - 1 - last_kt_token_idx = last_token_idx // kt_page_size - kt_token_idx_in_page = last_token_idx % kt_page_size + q_values = q_values.to(kt_cache_tensor_ptr.dtype.element_ty) - block_offset_in_seq = last_kt_token_idx // tokens_per_block - if block_offset_in_seq >= max_kt_blocks_per_seq: - return + for kt_block_idx_start in tl.range( + 0, + num_kt_tokens, + KT_BLOCK_SIZE, + ): + kt_block_idx_start = tl.multiple_of(kt_block_idx_start, KT_BLOCK_SIZE) - block_idx = tl.load(kt_cache_block_offsets_ptr + - batch_idx * max_kt_blocks_per_seq + block_offset_in_seq) - token_idx_in_block = last_kt_token_idx % tokens_per_block + kt_token_indices = kt_block_idx_start + tl.arange(0, KT_BLOCK_SIZE) + kt_token_mask = kt_token_indices < num_kt_tokens - cache_base = ((block_idx * tokens_per_block + token_idx_in_block) * - num_kv_heads * 2 * head_dim + kv_head_idx * 2 * head_dim) + block_offsets = batch_idx * max_kt_blocks_per_seq + kt_token_indices // tokens_per_block - cache_min_indices = cache_base + dim_indices - cache_max_indices = cache_base + head_dim + dim_indices + block_indices = tl.load( + kt_cache_block_offsets_ptr + block_offsets, mask=kt_token_mask, other=0 + ) + token_indices_in_block = kt_token_indices % tokens_per_block - kt_mask = dim_mask & (kt_token_idx_in_page > 0) + cache_bases = ( + block_indices * tokens_per_block + token_indices_in_block + ) * num_kv_heads * 2 * head_dim + kv_head_idx * 2 * head_dim - k_min_existing = tl.load(kt_cache_tensor_ptr + cache_min_indices, - mask=kt_mask, - other=float('inf')) - k_max_existing = tl.load(kt_cache_tensor_ptr + cache_max_indices, - mask=kt_mask, - other=float('-inf')) + combined_mask = ( + dim_mask[None, :] & kt_token_mask[:, None] + ) # Shape: [KT_BLOCK_SIZE, DIM_BLOCK_SIZE] - k_min_new = tl.minimum(k_min_existing, k_values) - k_max_new = tl.maximum(k_max_existing, k_values) - k_min_new = k_min_new.to(kt_cache_tensor_ptr.dtype.element_ty) - k_max_new = k_max_new.to(kt_cache_tensor_ptr.dtype.element_ty) + kt_cache_indices_min = cache_bases[:, None] + dim_indices[None, :] + kt_cache_indices_max = kt_cache_indices_min + head_dim + kt_cache_values_min = tl.load( + kt_cache_tensor_ptr + kt_cache_indices_min, mask=combined_mask, other=0.0 + ) + kt_cache_values_max = tl.load( + kt_cache_tensor_ptr + kt_cache_indices_max, mask=combined_mask, other=0.0 + ) - tl.store(kt_cache_tensor_ptr + cache_min_indices, k_min_new, mask=dim_mask) - tl.store(kt_cache_tensor_ptr + cache_max_indices, k_max_new, mask=dim_mask) + kt_cache_values = tl.where(dim_pos_values > 0, kt_cache_values_max, kt_cache_values_min) + results = tl.dot(q_values, kt_cache_values.T) # [Q_BLOCK_SIZE, KT_BLOCK_SIZE] -def triton_rocket_update_kt_cache_gen( - k: torch.Tensor, + output_mask = q_head_mask[:, None] & kt_token_mask[None, :] + output_indices = ( + kv_head_idx * num_heads_per_kv * total_kt_tokens + + q_head_offsets[:, None] * total_kt_tokens + + output_offset + + kt_token_indices[None, :] + ) + + tl.store(output_ptr + output_indices, results * sm_scale, mask=output_mask) + + +def triton_rocket_paged_kt_cache_bmm( + q: torch.Tensor, kt_cache_tensor: torch.Tensor, kt_cache_block_offsets: torch.Tensor, kv_lens: torch.Tensor, + output_offsets: torch.Tensor, kt_page_size: int, tokens_per_block: int, max_kt_blocks_per_seq: int, - num_kv_heads: int, - head_dim: int, -) -> None: + total_kt_tokens: int, + sm_scale: float = None, +) -> torch.Tensor: """ - Update KT cache with new key values for generation phase. + Perform BMM with KT cache for generation phase. Args: - k: Key tensor [num_gen_tokens, num_kv_heads * head_dim] + q: Query tensor [num_gen_tokens, num_kv_heads, num_heads_per_kv, head_dim] kt_cache_tensor: KT cache tensor kt_cache_block_offsets: Block offsets [num_gen_tokens, max_kt_blocks_per_seq] kv_lens: Sequence lengths [num_gen_tokens] + output_offsets: Output offsets [num_gen_tokens + 1] kt_page_size: Page size for KT tokens tokens_per_block: Tokens per cache block max_kt_blocks_per_seq: Maximum KT blocks per sequence + total_kt_tokens: Total number of KT tokens (fixed size) + sm_scale: Scale factor for softmax + Returns: + output: BMM results [num_kv_heads, num_heads_per_kv, total_kt_tokens] """ - num_gen_tokens = k.shape[0] + num_gen_tokens, num_kv_heads, num_heads_per_kv, head_dim = q.shape + num_kv_heads * num_heads_per_kv - grid = (num_gen_tokens, num_kv_heads, 1) + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(head_dim) + + # Create output tensor with shape [num_kv_heads, num_heads_per_kv, total_kt_tokens] + output = torch.empty( + (num_kv_heads, num_heads_per_kv, total_kt_tokens), dtype=torch.float32, device=q.device + ) + + def grid(_meta): + return num_gen_tokens, num_kv_heads + Q_BLOCK_SIZE = num_heads_per_kv + KT_BLOCK_SIZE = 64 DIM_BLOCK_SIZE = triton.next_power_of_2(head_dim) - rocket_update_kt_cache_gen_kernel[grid](k, - kt_cache_tensor, - kt_cache_block_offsets, - kv_lens, - num_gen_tokens, - num_kv_heads, - head_dim, - kt_page_size, - tokens_per_block, - max_kt_blocks_per_seq, - k.stride(0), - k.stride(1), - DIM_BLOCK_SIZE=DIM_BLOCK_SIZE) + rocket_paged_kt_cache_bmm_kernel[grid]( + q, + kt_cache_tensor, + kt_cache_block_offsets, + kv_lens, + output, + output_offsets, + num_gen_tokens, + num_kv_heads, + num_heads_per_kv, + head_dim, + kt_page_size, + tokens_per_block, + max_kt_blocks_per_seq, + total_kt_tokens, + sm_scale, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + Q_BLOCK_SIZE=Q_BLOCK_SIZE, + KT_BLOCK_SIZE=KT_BLOCK_SIZE, + DIM_BLOCK_SIZE=DIM_BLOCK_SIZE, + ) + + return output @triton.jit -def rocket_update_kt_cache_ctx_kernel( - k_ptr, - kt_cache_tensor_ptr, - kt_cache_block_offsets_ptr, - context_cumsum_ptr, - sparse_kv_indices_ptr, - sparse_kv_offsets_ptr, +def rocket_reduce_scores_kernel( + input_ptr, + output_ptr, + cum_lens_ptr, batch_size, - num_heads, num_kv_heads, - head_dim, - kt_page_size: tl.constexpr, - tokens_per_block, - max_kt_blocks_per_seq, - total_sparse_tokens, - BLOCK_SIZE_KT: tl.constexpr, - BLOCK_SIZE_DIM: tl.constexpr, + num_heads_per_kv, + total_kt_tokens, + BLOCK_SIZE: tl.constexpr, ): - """ - Triton kernel for updating KT cache during context phase with sparse indices. - """ batch_idx = tl.program_id(0) kv_head_idx = tl.program_id(1) - kt_block_idx = tl.program_id(2) - - context_start = tl.load(context_cumsum_ptr + batch_idx) - sparse_start = tl.load(sparse_kv_offsets_ptr + batch_idx) - sparse_end = tl.load(sparse_kv_offsets_ptr + batch_idx + 1) - num_sparse_tokens = sparse_end - sparse_start - - if num_sparse_tokens <= 0: + if batch_idx >= batch_size or kv_head_idx >= num_kv_heads: return - q_hidden_size = num_heads * head_dim - kv_hidden_size = num_kv_heads * head_dim - k_dim_base = q_hidden_size + kv_head_idx * head_dim - - BLOCK_SIZE_KV: tl.constexpr = kt_page_size * BLOCK_SIZE_KT - - total_kt_tokens = (num_sparse_tokens + kt_page_size - 1) // kt_page_size - kt_offsets = kt_block_idx * BLOCK_SIZE_KT + tl.arange(0, BLOCK_SIZE_KT) - kt_mask = kt_offsets < total_kt_tokens - - kv_start = kt_block_idx * BLOCK_SIZE_KT * kt_page_size - kv_offsets = kv_start + tl.arange(0, BLOCK_SIZE_KV) - kv_mask = kv_offsets < num_sparse_tokens - kv_indices = kv_head_idx * total_sparse_tokens + sparse_start + kv_offsets - - for dim_block_start in tl.range(0, head_dim, BLOCK_SIZE_DIM): - dim_offsets = tl.arange(0, BLOCK_SIZE_DIM) - dim_indices = dim_block_start + dim_offsets - dim_mask = dim_indices < head_dim - - kv_token_indices = tl.load(sparse_kv_indices_ptr + kv_indices, - mask=kv_mask, - other=0) - # Calculate indices for loading keys [BLOCK_SIZE_DIM, BLOCK_SIZE_KV] - k_base_indices = (kv_token_indices[None, :] + context_start) * ( - q_hidden_size + 2 * kv_hidden_size) + k_dim_base - k_indices = k_base_indices + dim_indices[:, None] - - combined_mask = kv_mask[None, :] & dim_mask[:, None] - - # Load key values [BLOCK_SIZE_DIM, BLOCK_SIZE_KV] - k_values = tl.load(k_ptr + k_indices, mask=combined_mask, other=0.0) - - k_values = tl.reshape(k_values, - (BLOCK_SIZE_DIM, BLOCK_SIZE_KT, kt_page_size)) + # Get KT token boundaries for this batch + kt_start = tl.load(cum_lens_ptr + batch_idx) + kt_end = tl.load(cum_lens_ptr + batch_idx + 1) + kt_len = kt_end - kt_start - k_min = tl.min(k_values, - axis=-1).to(kt_cache_tensor_ptr.dtype.element_ty) - k_max = tl.max(k_values, - axis=-1).to(kt_cache_tensor_ptr.dtype.element_ty) + if kt_len <= 0: + return - # Calculate cache locations [BLOCK_SIZE_KT] - block_offsets_in_seq = kt_offsets // tokens_per_block - valid_block_mask = (block_offsets_in_seq - < max_kt_blocks_per_seq) & kt_mask + # Process KT tokens in blocks + for kt_block_start in tl.range(0, kt_len, BLOCK_SIZE): + kt_offsets = kt_block_start + tl.arange(0, BLOCK_SIZE) + kt_mask = kt_offsets < kt_len + kt_global_offsets = kt_start + kt_offsets - # Load block indices [BLOCK_SIZE_KT] - block_offset_addrs = batch_idx * max_kt_blocks_per_seq + block_offsets_in_seq - block_indices = tl.load(kt_cache_block_offsets_ptr + block_offset_addrs, - mask=valid_block_mask, - other=0) + # Accumulate over num_heads_per_kv + sum_vals = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - tokens_in_block = kt_offsets % tokens_per_block + for head_idx in range(num_heads_per_kv): + global_head_idx = kv_head_idx * num_heads_per_kv + head_idx - # Calculate cache base addresses [BLOCK_SIZE_KT] - cache_bases = ((block_indices * tokens_per_block + tokens_in_block) * - num_kv_heads * 2 * head_dim + kv_head_idx * 2 * head_dim) + # Input indices: [global_head_idx, 1, kt_global_offsets] + input_indices = global_head_idx * total_kt_tokens + kt_global_offsets - cache_min_addrs = cache_bases[None, :] + dim_indices[:, None] - cache_max_addrs = cache_min_addrs + head_dim + values = tl.load(input_ptr + input_indices, mask=kt_mask, other=0.0) + sum_vals += values - store_mask = valid_block_mask[None, :] & dim_mask[:, None] + # Compute mean + mean_vals = sum_vals / num_heads_per_kv - tl.store(kt_cache_tensor_ptr + cache_min_addrs, k_min, mask=store_mask) - tl.store(kt_cache_tensor_ptr + cache_max_addrs, k_max, mask=store_mask) + # Store output: [kv_head_idx, kt_global_offsets] + output_indices = kv_head_idx * total_kt_tokens + kt_global_offsets + tl.store(output_ptr + output_indices, mean_vals, mask=kt_mask) -def triton_rocket_update_kt_cache_ctx( - qkv_input: torch.Tensor, - kt_cache_tensor: torch.Tensor, - kt_cache_block_offsets: torch.Tensor, - context_cumsum: torch.Tensor, - sparse_kv_indices: torch.Tensor, - sparse_kv_offsets: torch.Tensor, - num_heads: int, +def triton_rocket_reduce_scores( + scores: torch.Tensor, + cum_lens: torch.Tensor, + batch_size: int, num_kv_heads: int, - head_dim: int, - kt_page_size: int, - prompt_budget: int, - tokens_per_block: int, - max_kt_blocks_per_seq: int, -): + num_heads_per_kv: int, +) -> torch.Tensor: """ - Update KT cache during context phase using sparse indices. + Reduce scores for generation phase with batch-aware processing. Args: - qkv_input: QKV input tensor [total_sparse_tokens, num_heads*head_dim + num_kv_heads*head_dim + num_kv_heads*head_dim] - kt_cache_tensor: KT cache [num_blocks, tokens_per_block, num_kv_heads, 2*head_dim] - kt_cache_block_offsets: Block offsets [batch_size, max_kt_blocks_per_seq] - sparse_kv_indices: Sparse KV indices [num_kv_heads, total_sparse_tokens] - sparse_kv_offsets: Sparse offsets [batch_size + 1] - num_heads: Number of Q heads + scores: Input scores [num_kv_heads * num_heads_per_kv, 1, total_kt_tokens] + cum_lens: Cumulative scores lengths [batch_size + 1] + batch_size: Number of batches num_kv_heads: Number of KV heads - head_dim: Head dimension - kt_page_size: Page size for KT tokens - prompt_budget: Prompt budget - tokens_per_block: Tokens per cache block - max_kt_blocks_per_seq: Maximum KT blocks per sequence - """ - batch_size = sparse_kv_offsets.size(0) - 1 - total_sparse_tokens = sparse_kv_indices.size(1) + num_heads_per_kv: Number of Q heads per KV head - total_kt_tokens = (prompt_budget + kt_page_size - 1) // kt_page_size + Returns: + output: Reduced scores [num_kv_heads, total_kt_tokens] + """ + total_kt_tokens = scores.shape[-1] - BLOCK_SIZE_KT = 8 - BLOCK_SIZE_DIM = triton.next_power_of_2(head_dim) + output = torch.empty((num_kv_heads, total_kt_tokens), dtype=torch.float32, device=scores.device) - grid = (batch_size, num_kv_heads, - (total_kt_tokens + BLOCK_SIZE_KT - 1) // BLOCK_SIZE_KT) + BLOCK_SIZE = 256 + grid = (batch_size, num_kv_heads) - rocket_update_kt_cache_ctx_kernel[grid]( - qkv_input, - kt_cache_tensor, - kt_cache_block_offsets, - context_cumsum, - sparse_kv_indices, - sparse_kv_offsets, + rocket_reduce_scores_kernel[grid]( + scores, + output, + cum_lens, batch_size, - num_heads, num_kv_heads, - head_dim, - kt_page_size, - tokens_per_block, - max_kt_blocks_per_seq, - total_sparse_tokens, - BLOCK_SIZE_KT=BLOCK_SIZE_KT, - BLOCK_SIZE_DIM=BLOCK_SIZE_DIM, + num_heads_per_kv, + total_kt_tokens, + BLOCK_SIZE=BLOCK_SIZE, ) + return output + ######################################################## -# Paged KT cache BMM kernel +# Index gather kernel ######################################################## @triton.jit -def rocket_paged_kt_cache_bmm_kernel( - q_ptr, - kt_cache_tensor_ptr, - kt_cache_block_offsets_ptr, - kv_lens_ptr, +def _index_gather_kernel( output_ptr, - output_offsets_ptr, - num_gen_tokens, - num_kv_heads, - num_heads_per_kv: tl.constexpr, - head_dim, - kt_page_size, - tokens_per_block, - max_kt_blocks_per_seq, - total_kt_tokens, - sm_scale, - q_stride_0, - q_stride_1, - q_stride_2, - q_stride_3, - Q_BLOCK_SIZE: tl.constexpr, - KT_BLOCK_SIZE: tl.constexpr, - DIM_BLOCK_SIZE: tl.constexpr, + input_ptr, + index_ptr, + in_row_stride, + in_token_stride, + in_head_stride, + idx_row_stride, + idx_token_stride, + idx_head_stride, + dim_size, + BLOCK_SIZE: tl.constexpr, ): - batch_idx = tl.program_id(0) - kv_head_idx = tl.program_id(1) + # get program id and block offset + row_pid = tl.program_id(0) + token_pid = tl.program_id(1) + head_pid = tl.program_id(2) + token_block_num = tl.num_programs(1) + head_num = tl.num_programs(2) - if batch_idx >= num_gen_tokens or kv_head_idx >= num_kv_heads: - return - kv_len = tl.load(kv_lens_ptr + batch_idx) - num_kt_tokens = (kv_len + kt_page_size - 1) // kt_page_size + # get index + indices_idx = ( + row_pid * idx_row_stride + token_pid * idx_token_stride + head_pid * idx_head_stride + ) + token_idx = tl.load(index_ptr + indices_idx) - q_base = batch_idx * q_stride_0 + kv_head_idx * q_stride_1 + # get input and output base address + input_base = row_pid * in_row_stride + token_idx * in_token_stride + head_pid * in_head_stride + output_base = ( + row_pid * token_block_num * head_num * dim_size + + token_pid * head_num * dim_size + + head_pid * dim_size + ) - q_head_offsets = tl.arange(0, Q_BLOCK_SIZE) - q_head_mask = q_head_offsets < num_heads_per_kv + # process elements in blocks + for dim_offset in tl.range(0, dim_size, BLOCK_SIZE): + # get offsets + offsets = tl.arange(0, BLOCK_SIZE) + dim_indices = dim_offset + offsets + mask = dim_indices < dim_size - output_offset = tl.load(output_offsets_ptr + batch_idx) + # load input and store output + input_val = tl.load(input_ptr + input_base + dim_indices, mask=mask, other=0.0) + tl.store(output_ptr + output_base + dim_indices, input_val, mask=mask) - dim_indices = tl.arange(0, DIM_BLOCK_SIZE) - dim_mask = dim_indices < head_dim - q_indices = q_base + q_head_offsets[:, None] * q_stride_2 + dim_indices[ - None, :] * q_stride_3 - q_values = tl.load(q_ptr + q_indices, - mask=q_head_mask[:, None] & dim_mask[None, :]) +def triton_index_gather(input, indices): + assert input.ndim == 4, "Input must be a 4D tensor, [row, token, head, dim]" + assert indices.ndim == 3, "Indices must be a 3D tensor, [row, token, head]" - dim_pos_values = tl.sum(q_values, axis=0) > 0 - dim_pos_values = tl.broadcast_to(dim_pos_values[None, :], - (KT_BLOCK_SIZE, DIM_BLOCK_SIZE)) + # shape of input and indices + row_size = input.shape[0] + head_num = input.shape[2] + dim_size = input.shape[3] + num_tokens = indices.shape[1] - q_values = q_values.to(kt_cache_tensor_ptr.dtype.element_ty) + # create output tensor + output = torch.empty( + (row_size, num_tokens, head_num, dim_size), device="cuda", dtype=input.dtype + ) - for kt_block_idx_start in tl.range( - 0, - num_kt_tokens, - KT_BLOCK_SIZE, - ): - kt_block_idx_start = tl.multiple_of(kt_block_idx_start, KT_BLOCK_SIZE) + # launch kernel + grid = (row_size, num_tokens, head_num) + _index_gather_kernel[grid]( + output, + input, + indices, + input.stride(0), + input.stride(1), + input.stride(2), + indices.stride(0), + indices.stride(1), + indices.stride(2), + dim_size, + BLOCK_SIZE=1024, + ) + return output - kt_token_indices = kt_block_idx_start + tl.arange(0, KT_BLOCK_SIZE) - kt_token_mask = kt_token_indices < num_kt_tokens - block_offsets = batch_idx * max_kt_blocks_per_seq + kt_token_indices // tokens_per_block +@triton.jit +def bmm_kernel( + q_ptr, + k_ptr, + scores_ptr, + q_cu_seqlens_ptr, + k_cu_seqlens_ptr, + total_q_tokens, + total_k_tokens, + head_dim, + batch_size, + num_q_heads, + num_k_heads, + q_len_per_seq, + sm_scale, + causal, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) - block_indices = tl.load(kt_cache_block_offsets_ptr + block_offsets, - mask=kt_token_mask, - other=0) - token_indices_in_block = kt_token_indices % tokens_per_block + if batch_idx >= batch_size or head_idx >= num_q_heads: + return - cache_bases = ( - (block_indices * tokens_per_block + token_indices_in_block) * - num_kv_heads * 2 * head_dim + kv_head_idx * 2 * head_dim) + # Continuous mapping of query heads to key heads + k_head_idx = head_idx // (num_q_heads // num_k_heads) - combined_mask = dim_mask[ - None, :] & kt_token_mask[:, - None] # Shape: [KT_BLOCK_SIZE, DIM_BLOCK_SIZE] + q_seq_start = tl.load(q_cu_seqlens_ptr + batch_idx) + q_seq_end = tl.load(q_cu_seqlens_ptr + batch_idx + 1) + k_seq_start = tl.load(k_cu_seqlens_ptr + batch_idx) + k_seq_end = tl.load(k_cu_seqlens_ptr + batch_idx + 1) - kt_cache_indices_min = cache_bases[:, None] + dim_indices[None, :] - kt_cache_indices_max = kt_cache_indices_min + head_dim - kt_cache_values_min = tl.load(kt_cache_tensor_ptr + - kt_cache_indices_min, - mask=combined_mask, - other=0.0) - kt_cache_values_max = tl.load(kt_cache_tensor_ptr + - kt_cache_indices_max, - mask=combined_mask, - other=0.0) + q_seqlen = q_seq_end - q_seq_start + k_seqlen = k_seq_end - k_seq_start - kt_cache_values = tl.where(dim_pos_values > 0, kt_cache_values_max, - kt_cache_values_min) + if q_seqlen <= 0 or k_seqlen <= 0: + return - results = tl.dot(q_values, - kt_cache_values.T) # [Q_BLOCK_SIZE, KT_BLOCK_SIZE] + # Process queries in this batch with BLOCK_M parallelization + for q_block_start in tl.range(0, q_seqlen, BLOCK_M): + q_offsets = q_block_start + tl.arange(0, BLOCK_M) + q_mask = q_offsets < q_seqlen + q_global_offsets = q_seq_start + q_offsets - output_mask = q_head_mask[:, None] & kt_token_mask[None, :] - output_indices = (kv_head_idx * num_heads_per_kv * total_kt_tokens + - q_head_offsets[:, None] * total_kt_tokens + - output_offset + kt_token_indices[None, :]) + for k_block_start in tl.range(0, k_seqlen, BLOCK_N): + k_offsets = k_block_start + tl.arange(0, BLOCK_N) + k_mask = k_offsets < k_seqlen + k_global_offsets = k_seq_start + k_offsets + + # Initialize QK^T accumulator for this (M, N) block + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + + # Tiled matrix multiplication following mm_demo.py pattern + for k_dim_start in tl.range(0, head_dim, BLOCK_K): + k_dim_offsets = k_dim_start + tl.arange(0, BLOCK_K) + k_dim_mask = k_dim_offsets < head_dim - tl.store(output_ptr + output_indices, - results * sm_scale, - mask=output_mask) + # Load query chunk [BLOCK_M, BLOCK_K] + q_indices = ( + head_idx * total_q_tokens * head_dim + + q_global_offsets[:, None] * head_dim + + k_dim_offsets[None, :] + ) + q_chunk = tl.load( + q_ptr + q_indices, mask=q_mask[:, None] & k_dim_mask[None, :], other=0.0 + ) + # Load key chunk [BLOCK_N, BLOCK_K] + k_indices = ( + k_head_idx * total_k_tokens * head_dim + + k_global_offsets[:, None] * head_dim + + k_dim_offsets[None, :] + ) + k_chunk = tl.load( + k_ptr + k_indices, mask=k_mask[:, None] & k_dim_mask[None, :], other=0.0 + ) -def triton_rocket_paged_kt_cache_bmm( + # Accumulate QK^T using tl.dot for better performance + qk += tl.dot(q_chunk, tl.trans(k_chunk)) + + # Scale the accumulated QK^T + qk = qk * sm_scale + + # Apply masking + valid_mask = q_mask[:, None] & k_mask[None, :] + if causal: + # Create causal mask based on positions within this batch's sequence + q_pos_in_seq = q_offsets[:, None] # [BLOCK_M, 1] + k_pos_in_seq = k_offsets[None, :] # [1, BLOCK_N] + causal_mask = q_pos_in_seq >= k_pos_in_seq + qk = tl.where(causal_mask & valid_mask, qk, float("-inf")) + else: + qk = tl.where(valid_mask, qk, float("-inf")) + + # Store results - note that we store in the global k position space + # This matches the original bmm_softmax kernel behavior + output_indices = ( + head_idx * q_len_per_seq * total_k_tokens + + q_offsets[:, None] * total_k_tokens + + k_global_offsets[None, :] + ) # [BLOCK_M, BLOCK_N] + + tl.store(scores_ptr + output_indices, qk, mask=valid_mask) + + +def triton_bmm( q: torch.Tensor, - kt_cache_tensor: torch.Tensor, - kt_cache_block_offsets: torch.Tensor, - kv_lens: torch.Tensor, - output_offsets: torch.Tensor, - kt_page_size: int, - tokens_per_block: int, - max_kt_blocks_per_seq: int, - total_kt_tokens: int, + k: torch.Tensor, + q_cu_seqlens: torch.Tensor, + k_cu_seqlens: torch.Tensor, + batch_size: int, sm_scale: float = None, + causal: bool = False, +) -> torch.Tensor: + """ + Compute softmax(QK^T) with flattened input and output. + In this kernel we assume that each sequence in the batch has the same Q length. + + Args: + q: Query tensor [num_q_heads, total_q_tokens, head_dim] + k: Key tensor [num_kv_heads, total_k_tokens, head_dim] + q_cu_seqlens: Query cumulative sequence lengths [batch_size + 1] + k_cu_seqlens: Key cumulative sequence lengths [batch_size + 1] + batch_size: Number of batches + sm_scale: Scaling factor (default: 1/sqrt(head_dim)) + causal: Whether to apply causal masking + + Returns: + scores: Attention scores [num_q_heads, q_len_per_seq, total_k_tokens] + """ + num_q_heads, total_q_tokens, head_dim = q.shape + num_k_heads, total_k_tokens, _ = k.shape + + assert total_q_tokens % batch_size == 0, "total_q_tokens must be divisible by batch_size" + q_len_per_seq = total_q_tokens // batch_size + + if total_k_tokens == 0: + return torch.zeros( + (num_q_heads, q_len_per_seq, total_k_tokens), dtype=torch.float32, device=q.device + ) + + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(head_dim) + + bmm_results = torch.empty( + (num_q_heads, q_len_per_seq, total_k_tokens), dtype=torch.float32, device=q.device + ) + + # BMM kernel configuration + BLOCK_M = 32 + BLOCK_N = 256 + BLOCK_K = 64 + + def grid_bmm(_meta): + return batch_size, num_q_heads + + bmm_kernel[grid_bmm]( + q, + k, + bmm_results, + q_cu_seqlens, + k_cu_seqlens, + total_q_tokens, + total_k_tokens, + head_dim, + batch_size, + num_q_heads, + num_k_heads, + q_len_per_seq, + sm_scale, + causal, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, + ) + + return bmm_results + + +@triton.jit +def softmax_kernel( + input_ptr, + output_ptr, + cu_seq_lens_ptr, + batch_size, + num_heads, + q_len_per_seq, + total_k_tokens, + BLOCK_SIZE: tl.constexpr, +): + head_idx = tl.program_id(0) + q_global_idx = tl.program_id(1) + + if head_idx >= num_heads or q_global_idx >= (batch_size * q_len_per_seq): + return + + # Determine which batch this q belongs to + batch_idx = q_global_idx // q_len_per_seq + q_local_idx = q_global_idx % q_len_per_seq + + if batch_idx >= batch_size: + return + + # Get k sequence boundaries for this batch + k_seq_start = tl.load(cu_seq_lens_ptr + batch_idx) + k_seq_end = tl.load(cu_seq_lens_ptr + batch_idx + 1) + k_seqlen = k_seq_end - k_seq_start + + if k_seqlen <= 0: + return + + # Calculate input/output row start + input_row_start = head_idx * q_len_per_seq * total_k_tokens + q_local_idx * total_k_tokens + output_row_start = input_row_start + + # Find max value only within the valid k range for this batch + row_max = float("-inf") + for k_block_start in tl.range(0, k_seqlen, BLOCK_SIZE): + k_offsets = k_block_start + tl.arange(0, BLOCK_SIZE) + k_mask = k_offsets < k_seqlen + k_global_offsets = k_seq_start + k_offsets + + input_indices = input_row_start + k_global_offsets + values = tl.load(input_ptr + input_indices, mask=k_mask, other=float("-inf")) + block_max = tl.max(values, axis=0) + row_max = tl.maximum(row_max, block_max) + + # Compute sum of exp(x - max) only within valid k range + row_sum = 0.0 + for k_block_start in tl.range(0, k_seqlen, BLOCK_SIZE): + k_offsets = k_block_start + tl.arange(0, BLOCK_SIZE) + k_mask = k_offsets < k_seqlen + k_global_offsets = k_seq_start + k_offsets + + input_indices = input_row_start + k_global_offsets + values = tl.load(input_ptr + input_indices, mask=k_mask, other=float("-inf")) + exp_values = tl.exp(values - row_max) + masked_exp = tl.where(k_mask, exp_values, 0.0) + row_sum += tl.sum(masked_exp, axis=0) + + # Apply softmax and store results + for k_block_start in tl.range(0, k_seqlen, BLOCK_SIZE): + k_offsets = k_block_start + tl.arange(0, BLOCK_SIZE) + k_mask = k_offsets < k_seqlen + k_global_offsets = k_seq_start + k_offsets + + input_indices = input_row_start + k_global_offsets + output_indices = output_row_start + k_global_offsets + + values = tl.load(input_ptr + input_indices, mask=k_mask, other=float("-inf")) + exp_values = tl.exp(values - row_max) + masked_exp = tl.where(k_mask, exp_values, 0.0) + softmax_values = masked_exp / row_sum + + tl.store(output_ptr + output_indices, softmax_values, mask=k_mask) + + +def triton_softmax( + input_tensor: torch.Tensor, + cum_lens: torch.Tensor, + batch_size: int, +) -> torch.Tensor: + """ + Apply softmax to flattened input tensor. + + Args: + input_tensor: Input tensor [num_heads, len_per_seq, total_k_tokens] or [num_heads, total_k_tokens] + cum_lens: Cumulative lengths [batch_size + 1] + batch_size: Number of batches + + Returns: + output: Softmax results, shape is like input_tensor + """ + if input_tensor.ndim == 2: + num_heads, total_k_tokens = input_tensor.shape + len_per_seq = 1 + else: + num_heads, len_per_seq, total_k_tokens = input_tensor.shape + + output = torch.empty_like(input_tensor, dtype=input_tensor.dtype, device=input_tensor.device) + + BLOCK_SIZE = 512 + + grid = (num_heads, batch_size * len_per_seq) + + softmax_kernel[grid]( + input_tensor, + output, + cum_lens, + batch_size, + num_heads, + len_per_seq, + total_k_tokens, + BLOCK_SIZE=BLOCK_SIZE, + ) + + return output + + +@triton.jit +def flatten_to_batch_kernel( + input_ptr, + output_ptr, + input_offsets, + num_heads, + total_tokens, + padding_to_size, + padding_value, + BLOCK_SIZE: tl.constexpr, +): + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + + k_offset = tl.load(input_offsets + batch_idx) + context_len = tl.load(input_offsets + batch_idx + 1) - k_offset + + # Process in blocks + for block_start in tl.range(0, padding_to_size, BLOCK_SIZE): + pos_offsets = block_start + tl.arange(0, BLOCK_SIZE) + + pos_mask = pos_offsets < padding_to_size + valid_mask = pos_offsets < context_len + combined_mask = pos_mask & valid_mask + + input_indices = head_idx * total_tokens + k_offset + pos_offsets + + output_indices = ( + batch_idx * num_heads * padding_to_size + head_idx * padding_to_size + pos_offsets + ) + + values = tl.where( + combined_mask, + tl.load(input_ptr + input_indices, mask=combined_mask, other=padding_value), + padding_value, + ) + + tl.store(output_ptr + output_indices, values, mask=pos_mask) + + +def triton_flatten_to_batch( + input_tensor: torch.Tensor, + input_offsets: torch.Tensor, + batch_size: int, + padding_to_size: int, + padding_value=-1e10, ) -> torch.Tensor: """ - Perform BMM with KT cache for generation phase. + Reshape input_tensor from [num_heads, total_tokens] to [batch_size, num_heads, padding_to_size] Args: - q: Query tensor [num_gen_tokens, num_kv_heads, num_heads_per_kv, head_dim] - kt_cache_tensor: KT cache tensor - kt_cache_block_offsets: Block offsets [num_gen_tokens, max_kt_blocks_per_seq] - kv_lens: Sequence lengths [num_gen_tokens] - output_offsets: Output offsets [num_gen_tokens + 1] - kt_page_size: Page size for KT tokens - tokens_per_block: Tokens per cache block - max_kt_blocks_per_seq: Maximum KT blocks per sequence - total_kt_tokens: Total number of KT tokens (fixed size) - sm_scale: Scale factor for softmax + input_tensor: Input tensor tensor [num_heads, total_tokens] + input_offsets: Offset for each valid sequence [batch_size + 1] + batch_size: Number of batches + padding_to_size: Target padding size + padding_value: Value to fill for padding + Returns: - output: BMM results [num_kv_heads, num_heads_per_kv, total_kt_tokens] + batched_tensor: Output tensor [batch_size, num_heads, padding_to_size] """ - num_gen_tokens, num_kv_heads, num_heads_per_kv, head_dim = q.shape - num_kv_heads * num_heads_per_kv - - if sm_scale is None: - sm_scale = 1.0 / math.sqrt(head_dim) - - # Create output tensor with shape [num_kv_heads, num_heads_per_kv, total_kt_tokens] - output = torch.empty((num_kv_heads, num_heads_per_kv, total_kt_tokens), - dtype=torch.float32, - device=q.device) - - grid = lambda meta: (num_gen_tokens, num_kv_heads) - - Q_BLOCK_SIZE = num_heads_per_kv - KT_BLOCK_SIZE = 64 - DIM_BLOCK_SIZE = triton.next_power_of_2(head_dim) + num_heads, total_tokens = input_tensor.shape - rocket_paged_kt_cache_bmm_kernel[grid]( - q, - kt_cache_tensor, - kt_cache_block_offsets, - kv_lens, - output, - output_offsets, - num_gen_tokens, - num_kv_heads, - num_heads_per_kv, - head_dim, - kt_page_size, - tokens_per_block, - max_kt_blocks_per_seq, - total_kt_tokens, - sm_scale, - q.stride(0), - q.stride(1), - q.stride(2), - q.stride(3), - Q_BLOCK_SIZE=Q_BLOCK_SIZE, - KT_BLOCK_SIZE=KT_BLOCK_SIZE, - DIM_BLOCK_SIZE=DIM_BLOCK_SIZE, + batched_tensor = torch.empty( + (batch_size, num_heads, padding_to_size), + device=input_tensor.device, + dtype=input_tensor.dtype, ) - return output + def grid(_meta): + return batch_size, num_heads + flatten_to_batch_kernel[grid]( + input_tensor, + batched_tensor, + input_offsets, + num_heads, + total_tokens, + padding_to_size, + padding_value, + BLOCK_SIZE=1024, + ) -######################################################## -# Triton TopK kernel -######################################################## + return batched_tensor -# Adapted from https://github.com/triton-lang/triton/issues/3698#issuecomment-2067681396 @triton.jit def _compare_and_swap(x, ids, flip, i: core.constexpr, n_dims: core.constexpr): n_outer: core.constexpr = x.numel >> n_dims - shape: core.constexpr = [n_outer * 2**i, 2, 2**(n_dims - i - 1)] + shape: core.constexpr = [n_outer * 2**i, 2, 2 ** (n_dims - i - 1)] y = core.reshape(x, shape) # slice left/right with 'stride' 2**(n_dims - i - 1) mask = core.arange(0, 2)[None, :, None] @@ -1285,13 +1387,12 @@ def _compare_and_swap(x, ids, flip, i: core.constexpr, n_dims: core.constexpr): @triton.jit -def _bitonic_merge(x, ids, stage: core.constexpr, order: core.constexpr, - n_dims: core.constexpr): - ''' +def _bitonic_merge(x, ids, stage: core.constexpr, order: core.constexpr, n_dims: core.constexpr): + """ order_type 0 == ascending order_type 1 == descending order_type 2 == alternating - ''' + """ n_outer: core.constexpr = x.numel >> n_dims core.static_assert(stage <= n_dims) # flip denotes whether to re-arrange sub-sequences of elements in ascending or @@ -1300,11 +1401,11 @@ def _bitonic_merge(x, ids, stage: core.constexpr, order: core.constexpr, # if flip = 00110011... then all the elements will be re-arranged alternatingly (with # a stride of 2) at this stage if order == 2: - shape: core.constexpr = [n_outer * 2**(n_dims - 1 - stage), 2, 2**stage] + shape: core.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2**stage] # Create boolean flip pattern instead of integer - flip = core.reshape( - core.broadcast_to(core.arange(0, 2)[None, :, None], shape), - x.shape) != 0 + flip = ( + core.reshape(core.broadcast_to(core.arange(0, 2)[None, :, None], shape), x.shape) != 0 + ) else: # Ensure flip is boolean for XOR operations flip = order != 0 @@ -1315,20 +1416,15 @@ def _bitonic_merge(x, ids, stage: core.constexpr, order: core.constexpr, @triton.jit -def argsort(x, - ids, - dim: core.constexpr = None, - descending: core.constexpr = core.CONSTEXPR_0): +def argsort(x, ids, dim: core.constexpr = None, descending: core.constexpr = core.CONSTEXPR_0): # handle default dimension or check that it is the most minor dim _dim: core.constexpr = len(x.shape) - 1 if dim is None else dim - core.static_assert(_dim == len(x.shape) - 1, - "only minor dimension is currently supported") + core.static_assert(_dim == len(x.shape) - 1, "only minor dimension is currently supported") # iteratively run bitonic merge-sort steps n_dims: core.constexpr = _log2(x.shape[_dim]) for i in core.static_range(1, n_dims + 1): - x, ids = _bitonic_merge(x, ids, i, 2 if i < n_dims else descending, - n_dims) + x, ids = _bitonic_merge(x, ids, i, 2 if i < n_dims else descending, n_dims) return x, ids @@ -1337,8 +1433,7 @@ def convert_to_sortable_uint16(values): values_fp16 = values.to(tl.float16) values_u16 = values_fp16.to(tl.uint16, bitcast=True) # For negative numbers, flip all bits; for positive, flip sign bit - values_u16 = tl.where(values < 0.0, ~values_u16 & 0xffff, - values_u16 | 0x8000) + values_u16 = tl.where(values < 0.0, ~values_u16 & 0xFFFF, values_u16 | 0x8000) return values_u16 @@ -1347,17 +1442,16 @@ def extract_bin_idx(values_u16, step: tl.constexpr): if step == 0: return 255 - (values_u16 >> 8) else: - return 255 - (values_u16 & 0xff) + return 255 - (values_u16 & 0xFF) @triton.jit -def is_partial_match(values_u16, pattern, step: tl.constexpr, - BLOCK_SIZE: tl.constexpr): +def is_partial_match(values_u16, pattern, step: tl.constexpr, BLOCK_SIZE: tl.constexpr): if step == 0: - return tl.full((BLOCK_SIZE, ), True, dtype=tl.int1) + return tl.full((BLOCK_SIZE,), True, dtype=tl.int1) else: values_high = values_u16 >> 8 - pattern_high = tl.full((BLOCK_SIZE, ), pattern >> 8, dtype=tl.uint16) + pattern_high = tl.full((BLOCK_SIZE,), pattern >> 8, dtype=tl.uint16) return values_high == pattern_high @@ -1389,14 +1483,12 @@ def process_histogram_step( pass # Build histogram - histogram = tl.zeros((NUM_BINS, ), dtype=tl.int32) + histogram = tl.zeros((NUM_BINS,), dtype=tl.int32) for i in tl.range(0, input_len, BLOCK_SIZE): block_offsets = i + tl.arange(0, BLOCK_SIZE) block_mask = block_offsets < input_len - values = tl.load(input_ptr + input_base + block_offsets, - mask=block_mask, - other=-1e10) + values = tl.load(input_ptr + input_base + block_offsets, mask=block_mask, other=-1e10) values_u16 = convert_to_sortable_uint16(values) @@ -1406,8 +1498,7 @@ def process_histogram_step( bin_indices = extract_bin_idx(values_u16, step).to(tl.int32) # Apply pattern mask - bin_indices_masked = tl.where(pattern_match & block_mask, bin_indices, - NUM_BINS - 1) + bin_indices_masked = tl.where(pattern_match & block_mask, bin_indices, NUM_BINS - 1) histogram += tl.histogram(bin_indices_masked, NUM_BINS) @@ -1420,8 +1511,7 @@ def process_histogram_step( threshold_indices = tl.where(crosses_threshold, bin_range, NUM_BINS) threshold_bin_idx = tl.min(threshold_indices) - threshold_bin_size = tl.sum( - tl.where(bin_range == threshold_bin_idx, histogram, 0)) + threshold_bin_size = tl.sum(tl.where(bin_range == threshold_bin_idx, histogram, 0)) if step == 0: new_pattern = (255 - threshold_bin_idx) << 8 @@ -1435,9 +1525,7 @@ def process_histogram_step( for i in tl.range(0, input_len, BLOCK_SIZE): block_offsets = i + tl.arange(0, BLOCK_SIZE) block_mask = block_offsets < input_len - values = tl.load(input_ptr + input_base + block_offsets, - mask=block_mask, - other=-1e10) + values = tl.load(input_ptr + input_base + block_offsets, mask=block_mask, other=-1e10) values_u16 = convert_to_sortable_uint16(values) @@ -1446,40 +1534,34 @@ def process_histogram_step( bin_indices = extract_bin_idx(values_u16, step).to(tl.int32) # Elements smaller than threshold bin are guaranteed top-k - guaranteed_mask = (bin_indices - < threshold_bin_idx) & pattern_match & block_mask + guaranteed_mask = (bin_indices < threshold_bin_idx) & pattern_match & block_mask guaranteed_int = guaranteed_mask.to(tl.int32) write_positions = tl.cumsum(guaranteed_int, axis=0) - guaranteed_int num_guaranteed = tl.sum(guaranteed_int) guaranteed_write_offsets = base_offset + output_base + write_positions - tl.store(output_indices_ptr + guaranteed_write_offsets, - block_offsets, - mask=guaranteed_mask) + tl.store(output_indices_ptr + guaranteed_write_offsets, block_offsets, mask=guaranteed_mask) base_offset += num_guaranteed # Write candidate elements to temp storage (only if threshold bin is small) # Then we can directly sort the temp storage if threshold_bin_size <= FINAL_SIZE: - candidate_mask = (bin_indices - == threshold_bin_idx) & pattern_match & block_mask + candidate_mask = (bin_indices == threshold_bin_idx) & pattern_match & block_mask candidate_int = candidate_mask.to(tl.int32) - candidate_positions = tl.cumsum(candidate_int, - axis=0) - candidate_int + candidate_positions = tl.cumsum(candidate_int, axis=0) - candidate_int num_candidates = tl.sum(candidate_int) candidate_write_offsets = final_offset + temp_base + candidate_positions - candidate_write_mask = (final_offset + candidate_positions - < FINAL_SIZE) & candidate_mask + candidate_write_mask = ( + final_offset + candidate_positions < FINAL_SIZE + ) & candidate_mask - tl.store(temp_values_ptr + candidate_write_offsets, - values, - mask=candidate_write_mask) - tl.store(temp_indices_ptr + candidate_write_offsets, - block_offsets, - mask=candidate_write_mask) + tl.store(temp_values_ptr + candidate_write_offsets, values, mask=candidate_write_mask) + tl.store( + temp_indices_ptr + candidate_write_offsets, block_offsets, mask=candidate_write_mask + ) final_offset += num_candidates @@ -1488,32 +1570,28 @@ def process_histogram_step( for i in tl.range(0, input_len, BLOCK_SIZE): block_offsets = i + tl.arange(0, BLOCK_SIZE) block_mask = block_offsets < input_len - values = tl.load(input_ptr + input_base + block_offsets, - mask=block_mask, - other=-1e10) + values = tl.load(input_ptr + input_base + block_offsets, mask=block_mask, other=-1e10) values_u16 = convert_to_sortable_uint16(values) - pattern_match = is_partial_match(values_u16, pattern, step, - BLOCK_SIZE) + pattern_match = is_partial_match(values_u16, pattern, step, BLOCK_SIZE) bin_indices = extract_bin_idx(values_u16, step).to(tl.int32) - candidate_mask = (bin_indices - == threshold_bin_idx) & pattern_match & block_mask + candidate_mask = (bin_indices == threshold_bin_idx) & pattern_match & block_mask candidate_int = candidate_mask.to(tl.int32) - candidate_positions = tl.cumsum(candidate_int, - axis=0) - candidate_int + candidate_positions = tl.cumsum(candidate_int, axis=0) - candidate_int num_candidates = tl.sum(candidate_int) candidate_write_offsets = base_offset + output_base + candidate_positions - candidate_write_mask = (base_offset + candidate_positions - < topk) & candidate_mask + candidate_write_mask = (base_offset + candidate_positions < topk) & candidate_mask - tl.store(output_indices_ptr + candidate_write_offsets, - block_offsets, - mask=candidate_write_mask) + tl.store( + output_indices_ptr + candidate_write_offsets, + block_offsets, + mask=candidate_write_mask, + ) base_offset += num_candidates @@ -1570,19 +1648,25 @@ def topk_kernel( for i in tl.range(0, input_len, BLOCK_SIZE): block_offsets = i + tl.arange(0, BLOCK_SIZE) block_mask = block_offsets < input_len - tl.store(output_indices_ptr + output_base + block_offsets, - block_offsets, - mask=block_mask) + tl.store( + output_indices_ptr + output_base + block_offsets, block_offsets, mask=block_mask + ) return pattern = 0 found_topk_values = 0 - continue_step, pattern, threshold_bin_size, found_topk_values, final_offset = \ + continue_step, pattern, threshold_bin_size, found_topk_values, final_offset = ( process_histogram_step( - input_ptr, input_base, input_len, - output_indices_ptr, temp_values_ptr, temp_indices_ptr, - output_base, temp_base, topk, + input_ptr, + input_base, + input_len, + output_indices_ptr, + temp_values_ptr, + temp_indices_ptr, + output_base, + temp_base, + topk, step=0, pattern=pattern, found_topk_values=found_topk_values, @@ -1590,13 +1674,20 @@ def topk_kernel( NUM_BINS=NUM_BINS, FINAL_SIZE=FINAL_SIZE, ) + ) if continue_step: - continue_step, pattern, threshold_bin_size, found_topk_values, final_offset = \ + continue_step, pattern, threshold_bin_size, found_topk_values, final_offset = ( process_histogram_step( - input_ptr, input_base, input_len, - output_indices_ptr, temp_values_ptr, temp_indices_ptr, - output_base, temp_base, topk, + input_ptr, + input_base, + input_len, + output_indices_ptr, + temp_values_ptr, + temp_indices_ptr, + output_base, + temp_base, + topk, step=1, pattern=pattern, found_topk_values=found_topk_values, @@ -1604,37 +1695,40 @@ def topk_kernel( NUM_BINS=NUM_BINS, FINAL_SIZE=FINAL_SIZE, ) + ) # Sort the final candidates if not continue_step and final_offset > 0: final_offsets = tl.arange(0, FINAL_SIZE) final_mask = final_offsets < final_offset - final_values = tl.load(temp_values_ptr + temp_base + final_offsets, - mask=final_mask, - other=-1e10) - final_indices = tl.load(temp_indices_ptr + temp_base + final_offsets, - mask=final_mask, - other=-1) - - _, final_sorted_indices = argsort(final_values, - final_indices, - dim=0, - descending=True) + final_values = tl.load( + temp_values_ptr + temp_base + final_offsets, mask=final_mask, other=-1e10 + ) + final_indices = tl.load( + temp_indices_ptr + temp_base + final_offsets, mask=final_mask, other=-1 + ) + + _, final_sorted_indices = argsort(final_values, final_indices, dim=0, descending=True) remain_num = topk - found_topk_values write_offsets = tl.arange(0, FINAL_SIZE) write_mask = write_offsets < remain_num - tl.store(output_indices_ptr + output_base + found_topk_values + - write_offsets, - final_sorted_indices, - mask=write_mask) + tl.store( + output_indices_ptr + output_base + found_topk_values + write_offsets, + final_sorted_indices, + mask=write_mask, + ) -def triton_topk(input_tensor: torch.Tensor, input_offsets: torch.Tensor, - output_offsets: torch.Tensor, total_output_tokens: int, - topk: int) -> torch.Tensor: +def triton_topk( + input_tensor: torch.Tensor, + input_offsets: torch.Tensor, + output_offsets: torch.Tensor, + total_output_tokens: int, + topk: int, +) -> torch.Tensor: """ Perform topk operation on input tensor using two-stage histogram algorithm. Args: @@ -1653,9 +1747,9 @@ def triton_topk(input_tensor: torch.Tensor, input_offsets: torch.Tensor, device = input_tensor.device # Create output tensor - output_indices = torch.empty((num_kv_heads, total_output_tokens), - dtype=torch.int32, - device=device) + output_indices = torch.empty( + (num_kv_heads, total_output_tokens), dtype=torch.int32, device=device + ) grid = (batch_size, num_kv_heads) @@ -1663,12 +1757,12 @@ def triton_topk(input_tensor: torch.Tensor, input_offsets: torch.Tensor, NUM_BINS = 256 FINAL_SIZE = 256 - temp_values = torch.empty((num_kv_heads, batch_size, FINAL_SIZE), - dtype=torch.float32, - device=device) - temp_indices = torch.empty((num_kv_heads, batch_size, FINAL_SIZE), - dtype=torch.int32, - device=device) + temp_values = torch.empty( + (num_kv_heads, batch_size, FINAL_SIZE), dtype=torch.float32, device=device + ) + temp_indices = torch.empty( + (num_kv_heads, batch_size, FINAL_SIZE), dtype=torch.int32, device=device + ) topk_kernel[grid]( input_tensor, @@ -1688,668 +1782,3 @@ def triton_topk(input_tensor: torch.Tensor, input_offsets: torch.Tensor, ) return output_indices - - -######################################################## -# Reduce scores generation kernel -######################################################## - - -@triton.jit -def rocket_reduce_scores_kernel( - input_ptr, - output_ptr, - cum_lens_ptr, - batch_size, - num_kv_heads, - num_heads_per_kv, - total_kt_tokens, - BLOCK_SIZE: tl.constexpr, -): - batch_idx = tl.program_id(0) - kv_head_idx = tl.program_id(1) - - if batch_idx >= batch_size or kv_head_idx >= num_kv_heads: - return - - # Get KT token boundaries for this batch - kt_start = tl.load(cum_lens_ptr + batch_idx) - kt_end = tl.load(cum_lens_ptr + batch_idx + 1) - kt_len = kt_end - kt_start - - if kt_len <= 0: - return - - # Process KT tokens in blocks - for kt_block_start in tl.range(0, kt_len, BLOCK_SIZE): - kt_offsets = kt_block_start + tl.arange(0, BLOCK_SIZE) - kt_mask = kt_offsets < kt_len - kt_global_offsets = kt_start + kt_offsets - - # Accumulate over num_heads_per_kv - sum_vals = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - - for head_idx in range(num_heads_per_kv): - global_head_idx = kv_head_idx * num_heads_per_kv + head_idx - - # Input indices: [global_head_idx, 1, kt_global_offsets] - input_indices = global_head_idx * total_kt_tokens + kt_global_offsets - - values = tl.load(input_ptr + input_indices, mask=kt_mask, other=0.0) - sum_vals += values - - # Compute mean - mean_vals = sum_vals / num_heads_per_kv - - # Store output: [kv_head_idx, kt_global_offsets] - output_indices = kv_head_idx * total_kt_tokens + kt_global_offsets - tl.store(output_ptr + output_indices, mean_vals, mask=kt_mask) - - -def triton_rocket_reduce_scores( - scores: torch.Tensor, - cum_lens: torch.Tensor, - batch_size: int, - num_kv_heads: int, - num_heads_per_kv: int, -) -> torch.Tensor: - """ - Reduce scores for generation phase with batch-aware processing. - - Args: - scores: Input scores [num_kv_heads * num_heads_per_kv, 1, total_kt_tokens] - cum_lens: Cumulative scores lengths [batch_size + 1] - batch_size: Number of batches - num_kv_heads: Number of KV heads - num_heads_per_kv: Number of Q heads per KV head - - Returns: - output: Reduced scores [num_kv_heads, total_kt_tokens] - """ - total_kt_tokens = scores.shape[-1] - - output = torch.empty((num_kv_heads, total_kt_tokens), - dtype=torch.float32, - device=scores.device) - - BLOCK_SIZE = 256 - grid = (batch_size, num_kv_heads) - - rocket_reduce_scores_kernel[grid]( - scores, - output, - cum_lens, - batch_size, - num_kv_heads, - num_heads_per_kv, - total_kt_tokens, - BLOCK_SIZE=BLOCK_SIZE, - ) - - return output - - -######################################################## -# Convert request-local indices to global KV cache pool indices -######################################################## - - -@triton.jit -def _convert_req_index_to_global_index_kernel_with_stride_factor( - req_id_ptr, # int32 [num_tokens] - block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] - token_indices_ptr, # int32 [num_kv_heads, num_tokens, NUM_TOPK_TOKENS] - out_ptr, # int32 [num_kv_heads, num_tokens, kv_factor, NUM_TOPK_TOKENS] - # shapes (compile-time where possible) - max_num_blocks_per_req: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - BLOCK_N: tl.constexpr, # tile width along columns - stride_factor: tl.constexpr, # elements per physical page in pool - layer_id: tl.constexpr, # for multi-layer KV cache - num_kv_heads: tl.constexpr, - kv_factor: tl.constexpr, - bt_stride0, # stride for batch dim - bt_stride1, # stride for page dim - ti_stride0, - ti_stride1, - ti_stride2, - out_stride0, - out_stride1, - out_stride2, - out_stride3, -): - """ - Convert request-local token indices to global KV cache pool indices. - - KV cache pool layout: [max_num_pages, num_layers, kv_factor, num_kv_heads, tokens_per_block, head_dim] - block_table: [num_requests, max_pages_per_req] - - stores memPoolBlockIdx (physical page index, first dim of pool) - - same block_table for K and V (K/V share the same physical page) - - stride_factor = num_layers * kv_factor * num_kv_heads * BLOCK_SIZE - (elements per physical page, excluding head_dim) - - Global index: - memPoolBlockIdx * stride_factor - + layer_id * kv_factor * num_kv_heads * BLOCK_SIZE - + kv_factor_idx * num_kv_heads * BLOCK_SIZE - + kv_head_idx * BLOCK_SIZE - + token_in_page - """ - kv_head_idx = tl.program_id(0) - token_id = tl.program_id(1) - tile_id = tl.program_id(2) - - indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) - - req = tl.load(req_id_ptr + token_id) - - # Load local token indices: token_indices[kv_head_idx, token_id, indice_id] - ti_ptr = (token_indices_ptr + kv_head_idx * ti_stride0 + - token_id * ti_stride1 + indice_id * ti_stride2) - tok = tl.load(ti_ptr) - - is_invalid_tok = tok < 0 - page_idx = tok // BLOCK_SIZE - token_in_page = tok % BLOCK_SIZE - valid_page = page_idx < max_num_blocks_per_req - - # block_table[req, page_idx] → memPoolBlockIdx (physical page index) - bt_ptr = (block_table_ptr + req * bt_stride0 + page_idx * bt_stride1) - mem_pool_idx = tl.load(bt_ptr, mask=valid_page, other=0) - - # Base offset within physical page (invariant across kv_factor loop) - base_off = (layer_id * kv_factor * num_kv_heads * BLOCK_SIZE + - kv_head_idx * BLOCK_SIZE + token_in_page) - - for kv_factor_idx in tl.static_range(kv_factor): - # Within-page offset: layer → kv_factor → kv_head → token - inpage_off = base_off + kv_factor_idx * num_kv_heads * BLOCK_SIZE - - global_idx = mem_pool_idx * stride_factor + inpage_off - - out_val = tl.where(is_invalid_tok | (~valid_page), -1, global_idx) - - out_ptr_ij = (out_ptr + kv_head_idx * out_stride0 + - token_id * out_stride1 + kv_factor_idx * out_stride2 + - indice_id * out_stride3) - tl.store(out_ptr_ij, out_val) - - -def triton_convert_req_index_to_global_index( - req_id: torch.Tensor, # int32 [num_tokens] - block_table: torch. - Tensor, # int32 [num_requests, kv_factor, max_num_blocks_per_req] - token_indices: torch. - Tensor, # int32 [num_kv_heads, num_tokens, NUM_TOPK_TOKENS] - BLOCK_SIZE: int, - NUM_TOPK_TOKENS: int = 2048, - BLOCK_N: int = 128, # tile width along columns - stride_factor: int = None, # elements per block in pool - layer_id: int = 0, # for multi-layer KV cache - num_kv_heads: int = 1, - kv_factor: int = 1, -): - """ - Convert request-local token indices to global KV cache pool indices. - - Accepts both 2D and 3D token_indices: - - 2D [num_tokens, topk]: MLA path (num_kv_heads=1 implicit) - - 3D [num_kv_heads, num_tokens, topk]: MQA/GQA path - - Output shape: [num_kv_heads, num_tokens, kv_factor, topk] - - Args: - block_table: [num_requests, max_pages] — physical page indices. - stride_factor: Elements per physical page - (num_layers * kv_factor * num_kv_heads * BLOCK_SIZE). - num_kv_heads: Number of KV heads. - kv_factor: 2 for MQA/GQA, 1 for MLA. - """ - if stride_factor is None: - stride_factor = kv_factor * num_kv_heads * BLOCK_SIZE - assert req_id.dtype == torch.int32 - assert block_table.dtype == torch.int32 - assert token_indices.dtype == torch.int32 - assert token_indices.ndim in (2, 3), \ - f"Expected 2D [num_tokens, topk] or 3D [num_kv_heads, num_tokens, topk], got {token_indices.ndim}D" - assert block_table.ndim == 2, f"Expected 2D [batch, max_pages], got {block_table.ndim}D" - assert NUM_TOPK_TOKENS % BLOCK_N == 0, \ - f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})" - - num_tokens = req_id.shape[0] - max_num_blocks_per_req = block_table.shape[1] - tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N - - req_id_c = req_id.contiguous() - block_table_c = block_table.contiguous() - token_indices_c = token_indices.contiguous() - out = torch.empty((num_kv_heads, num_tokens, kv_factor, NUM_TOPK_TOKENS), - dtype=torch.int32, - device=token_indices.device) - - bt_stride0, bt_stride1 = block_table_c.stride() - # 2D input: ti_stride0=0 (kv_head_idx is always 0, term vanishes) - if token_indices_c.ndim == 2: - ti_stride0 = 0 - ti_stride1, ti_stride2 = token_indices_c.stride() - else: - ti_stride0, ti_stride1, ti_stride2 = token_indices_c.stride() - out_stride0, out_stride1, out_stride2, out_stride3 = out.stride() - - grid = (num_kv_heads, num_tokens, tiles_per_row) - - _convert_req_index_to_global_index_kernel_with_stride_factor[grid]( - req_id_c, - block_table_c, - token_indices_c, - out, - max_num_blocks_per_req, - BLOCK_SIZE, - BLOCK_N, - stride_factor, - layer_id, - num_kv_heads, - kv_factor, - bt_stride0, - bt_stride1, - ti_stride0, - ti_stride1, - ti_stride2, - out_stride0, - out_stride1, - out_stride2, - out_stride3, - ) - return out - - -######################################################## -# Fused K cache gather kernel -######################################################## - - -@triton.jit -def _triton_gather_k_cache_kernel( - k_cache_ptr, - slot_fp8_ptr, - slot_scale_ptr, - out_fp8_ptr, - out_scale_ptr, - k_token_start, - num_k_tokens, - HEAD_DIM: tl.constexpr, - SCALE_BYTES: tl.constexpr, - BLOCK_TOKENS: tl.constexpr, -): - pid = tl.program_id(0) - token_offsets = (pid * BLOCK_TOKENS + tl.arange(0, BLOCK_TOKENS)).to( - tl.int64) - token_mask = token_offsets < num_k_tokens - - fp8_base = tl.load(slot_fp8_ptr + k_token_start + token_offsets, - mask=token_mask, - other=0) - scale_base = tl.load(slot_scale_ptr + k_token_start + token_offsets, - mask=token_mask, - other=0) - - byte_offsets = tl.arange(0, HEAD_DIM).to(tl.int64) - src_fp8 = fp8_base[:, None] + byte_offsets[None, :] - dst_fp8 = token_offsets[:, None] * HEAD_DIM + byte_offsets[None, :] - gather_mask = token_mask[:, None] - - fp8_data = tl.load(k_cache_ptr + src_fp8, mask=gather_mask, other=0) - tl.store(out_fp8_ptr + dst_fp8, fp8_data, mask=gather_mask) - - scale_byte_offsets = tl.arange(0, SCALE_BYTES).to(tl.int64) - src_scale = scale_base[:, None] + scale_byte_offsets[None, :] - dst_scale = token_offsets[:, - None] * SCALE_BYTES + scale_byte_offsets[None, :] - - scale_data = tl.load(k_cache_ptr + src_scale, mask=gather_mask, other=0) - tl.store(out_scale_ptr + dst_scale, scale_data, mask=gather_mask) - - -def triton_gather_k_cache( - k_cache: torch.Tensor, - slot_mapping_fp8: torch.Tensor, - slot_mapping_scale: torch.Tensor, - k_token_start: int, - k_token_end: int, - head_dim: int, -): - """Gather K FP8 values and scales from the indexer K cache for a chunk. - - Replaces ``_gather_k_cache_for_chunk``, fusing ~8-12 small PyTorch ops - (arange, unsqueeze, broadcast add, _unravel_indices, advanced indexing) - into a single Triton kernel that directly gathers from flat byte offsets. - This is purely data movement — bit-exact with the original. - - Args: - k_cache: Indexer K cache pool data (2D contiguous), uint8. - slot_mapping_fp8: Flat byte indices for FP8 data - ``[total_kv_len]``, int64. - slot_mapping_scale: Flat byte indices for scale data - ``[total_kv_len]``, int64. - k_token_start: Start index into slot mapping arrays. - k_token_end: End index into slot mapping arrays. - head_dim: FP8 head dimension (typically 128). - - Returns: - Tuple of (k_fp8, k_scale): - k_fp8: ``[num_k_tokens, head_dim]``, float8_e4m3fn. - k_scale: ``[num_k_tokens, 1]``, float32. - """ - num_k_tokens = k_token_end - k_token_start - device = k_cache.device - - if num_k_tokens == 0: - return ( - torch.empty(0, head_dim, dtype=torch.float8_e4m3fn, device=device), - torch.empty(0, 1, dtype=torch.float32, device=device), - ) - - SCALE_BYTES = 4 - BLOCK_TOKENS = 32 - - k_cache_flat = k_cache.reshape(-1) - out_fp8 = torch.empty(num_k_tokens, - head_dim, - dtype=torch.uint8, - device=device) - out_scale = torch.empty(num_k_tokens, - SCALE_BYTES, - dtype=torch.uint8, - device=device) - - grid = (triton.cdiv(num_k_tokens, BLOCK_TOKENS), ) - _triton_gather_k_cache_kernel[grid]( - k_cache_flat, - slot_mapping_fp8, - slot_mapping_scale, - out_fp8.view(-1), - out_scale.view(-1), - k_token_start, - num_k_tokens, - HEAD_DIM=head_dim, - SCALE_BYTES=SCALE_BYTES, - BLOCK_TOKENS=BLOCK_TOKENS, - ) - - k_fp8 = out_fp8.view(torch.float8_e4m3fn) - k_scale = out_scale.view(torch.float32).view(num_k_tokens, 1) - return k_fp8, k_scale - - -######################################################## -# DeepSeek-V4 local to global index transformation kernel -######################################################## - - -@triton.jit -def _deepseek_v4_local_to_global_kernel( - req_id_ptr, - block_table_swa_ptr, - block_table_compressed_ptr, - swa_local_indices_ptr, - compressed_local_indices_ptr, - out_ptr, - swa_buffer_offset_in_tokens, - compressed_buffer_offset_in_tokens, - tokens_per_block_swa: tl.constexpr, - tokens_per_block_compressed: tl.constexpr, - max_blocks_swa, - max_blocks_compressed, - num_swa_indices: tl.constexpr, - num_compressed_indices: tl.constexpr, - total_output_indices, - has_compressed: tl.constexpr, - bt_swa_stride0, - bt_swa_stride1, - bt_compressed_stride0, - bt_compressed_stride1, - swa_indices_stride0, - swa_indices_stride1, - compressed_indices_stride0, - compressed_indices_stride1, - out_stride0, - out_stride1, - LAUNCH_WITH_PDL: tl.constexpr, -): - """ - Triton kernel for converting local indices to global KV cache pool indices. - - Dual-pool output layout: - - SWA region [0, num_swa_indices): indices relative to swa_pool_base_ptr - - Compress region [num_swa_indices, num_swa_indices + num_compressed_indices): - indices relative to compress_pool_base_ptr - - Invalid positions padded with -1 at their fixed positions. - - This enables the FMHA kernel to determine which TMA descriptor to use based - solely on tile index (tile 0 = SWA via tmaKSecondary_, rest = compress via tmaK_). - """ - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - token_id = tl.program_id(0) - - # Load request ID for this token - req = tl.load(req_id_ptr + token_id) - - # Load all SWA local indices for this token - swa_ids = tl.arange(0, num_swa_indices) - swa_ptr = swa_local_indices_ptr + token_id * swa_indices_stride0 + swa_ids * swa_indices_stride1 - swa_local_idx = tl.load(swa_ptr) - - # Compute global indices for all SWA positions - swa_valid_mask = swa_local_idx >= 0 - swa_block_ordinal = swa_local_idx // tokens_per_block_swa - swa_token_in_block = swa_local_idx % tokens_per_block_swa - swa_valid_block = swa_block_ordinal < max_blocks_swa - swa_full_mask = swa_valid_mask & swa_valid_block - - swa_bt_ptr = block_table_swa_ptr + req * bt_swa_stride0 + swa_block_ordinal * bt_swa_stride1 - swa_page_index = tl.load(swa_bt_ptr, mask=swa_full_mask, other=0) - - swa_global_index = swa_buffer_offset_in_tokens + swa_page_index * tokens_per_block_swa + swa_token_in_block - swa_global_index = tl.where(swa_full_mask, swa_global_index, -1) - - # Store SWA results at fixed positions [0, num_swa_indices) - swa_out_ptr = out_ptr + token_id * out_stride0 + swa_ids * out_stride1 - tl.store(swa_out_ptr, swa_global_index) - - if has_compressed: - # Load all compressed local indices for this token - compressed_ids = tl.arange(0, num_compressed_indices) - compressed_ptr = (compressed_local_indices_ptr + - token_id * compressed_indices_stride0 + - compressed_ids * compressed_indices_stride1) - compressed_local_idx = tl.load(compressed_ptr) - - # Compute global indices for all compressed positions - compressed_valid_mask = compressed_local_idx >= 0 - compressed_block_ordinal = compressed_local_idx // tokens_per_block_compressed - compressed_token_in_block = compressed_local_idx % tokens_per_block_compressed - compressed_valid_block = compressed_block_ordinal < max_blocks_compressed - compressed_full_mask = compressed_valid_mask & compressed_valid_block - - compressed_bt_ptr = (block_table_compressed_ptr + - req * bt_compressed_stride0 + - compressed_block_ordinal * bt_compressed_stride1) - compressed_page_index = tl.load(compressed_bt_ptr, - mask=compressed_full_mask, - other=0) - - compressed_global_index = ( - compressed_buffer_offset_in_tokens + - compressed_page_index * tokens_per_block_compressed + - compressed_token_in_block) - compressed_global_index = tl.where(compressed_full_mask, - compressed_global_index, -1) - - # Store compressed results at fixed positions [num_swa_indices, total) - compressed_write_pos = num_swa_indices + compressed_ids - compressed_out_ptr = out_ptr + token_id * out_stride0 + compressed_write_pos * out_stride1 - tl.store(compressed_out_ptr, compressed_global_index) - - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -def deepseek_v4_local_to_global_indices( - req_id: torch.Tensor, # int32 [num_tokens] - block_table_swa: torch.Tensor, # int32 [num_requests, max_blocks_swa] - swa_local_indices: torch.Tensor, # int32 [num_tokens, num_swa_indices] - swa_pool_base_ptr: int, # int64: base address of SWA pool - swa_buffer_ptr: int, # int64: base address of SWA buffer - tokens_per_block: int, # tokens per block for SWA - token_stride: int, # bytes per token (use SWA token stride) - # Optional compressed arguments (for compress_ratio > 1) - block_table_compressed: torch.Tensor | None = None, - compressed_local_indices: torch.Tensor | None = None, - compress_pool_base_ptr: int = 0, # int64: base address of compress pool - compressed_buffer_ptr: int = 0, - compress_ratio: int = 1, - num_compressed_indices: int = 0, # max number of compressed indices -) -> torch.Tensor: - """ - Convert local token indices to global KV cache pool indices. - - Dual-pool output layout: - - SWA region [0, window_size): indices relative to swa_pool_base_ptr - - Compress region [window_size, window_size + num_compressed_indices): - indices relative to compress_pool_base_ptr - - Invalid positions padded with -1 at their fixed positions. - - For compress_ratio=1: Only SWA region, all indices relative to swa_pool_base_ptr. - For compress_ratio>1: SWA region + compress region with separate base pointers. - - Args: - req_id: Request ID per token [num_tokens], int32 - block_table_swa: SWA block table [num_requests, max_blocks_swa], int32 - swa_local_indices: Local indices for SWA cache [num_tokens, num_swa_indices], int32 - Use -1 for invalid/padding indices. - swa_pool_base_ptr: Base address of SWA pool - swa_buffer_ptr: Base address of SWA buffer (base_pool_ptr + buffer_offset_in_slot) - tokens_per_block: Number of tokens per block for SWA cache - token_stride: Bytes per token (use SWA token stride) - block_table_compressed: Compressed block table [num_requests, max_blocks_compressed], int32 (optional) - compressed_local_indices: Local indices for compressed cache [num_tokens, num_compressed_indices], int32 (optional) - Use -1 for invalid/padding indices. - compress_pool_base_ptr: Base address of compress pool - compressed_buffer_ptr: Base address of compressed buffer (optional) - compress_ratio: Compression ratio (1: no compression, >1: with compression) - num_compressed_indices: Max number of compressed indices for CUDA graph compatibility - Output width = num_swa_indices + num_compressed_indices. - - Returns: - global_indices: int32 [num_tokens, num_swa_indices + num_compressed_indices] - """ - assert req_id.dtype == torch.int32, f"req_id must be int32, got {req_id.dtype}" - assert block_table_swa.dtype == torch.int32, f"block_table_swa must be int32, got {block_table_swa.dtype}" - assert swa_local_indices.dtype == torch.int32, f"swa_local_indices must be int32, got {swa_local_indices.dtype}" - - num_tokens = req_id.shape[0] - num_swa_indices = swa_local_indices.shape[1] - - assert swa_local_indices.shape[0] == num_tokens - - has_compressed = compress_ratio > 1 - - # Compute SWA buffer offset relative to swa_pool_base_ptr in tokens - swa_buffer_offset_in_tokens = (swa_buffer_ptr - - swa_pool_base_ptr) // token_stride - - if has_compressed: - assert block_table_compressed is not None, "block_table_compressed required when compress_ratio > 1" - assert compressed_local_indices is not None, "compressed_local_indices required when compress_ratio > 1" - assert ( - block_table_compressed.dtype == torch.int32 - ), f"block_table_compressed must be int32, got {block_table_compressed.dtype}" - assert ( - compressed_local_indices.dtype == torch.int32 - ), f"compressed_local_indices must be int32, got {compressed_local_indices.dtype}" - assert compressed_local_indices.shape[0] == num_tokens - - tokens_per_block_compressed = tokens_per_block // compress_ratio - # Compute compressed buffer offset relative to compress_pool_base_ptr in tokens - assert ( - compressed_buffer_ptr - compress_pool_base_ptr - ) % token_stride == 0, "compressed_buffer_ptr must be aligned to token_stride" - compressed_buffer_offset_in_tokens = ( - compressed_buffer_ptr - compress_pool_base_ptr) // token_stride - _, max_blocks_compressed = block_table_compressed.shape - block_table_compressed_c = block_table_compressed.contiguous() - compressed_local_indices_c = compressed_local_indices.contiguous() - else: - # Dummy values - tokens_per_block_compressed = tokens_per_block - compressed_buffer_offset_in_tokens = 0 - max_blocks_compressed = 1 - block_table_compressed_c = torch.zeros((1, 1), - dtype=torch.int32, - device=req_id.device) - compressed_local_indices_c = torch.zeros((1, 1), - dtype=torch.int32, - device=req_id.device) - - total_output_indices = num_swa_indices + num_compressed_indices - _, max_blocks_swa = block_table_swa.shape - - # Ensure contiguous tensors - req_id_c = req_id.contiguous() - block_table_swa_c = block_table_swa.contiguous() - swa_local_indices_c = swa_local_indices.contiguous() - - # Create output tensor - out = torch.empty((num_tokens, total_output_indices), - dtype=torch.int32, - device=req_id.device) - - # Grid: one program per token - grid = (num_tokens, ) - - # Get strides - bt_swa_stride0, bt_swa_stride1 = block_table_swa_c.stride() - bt_compressed_stride0, bt_compressed_stride1 = block_table_compressed_c.stride( - ) - swa_indices_stride0, swa_indices_stride1 = swa_local_indices_c.stride() - compressed_indices_stride0, compressed_indices_stride1 = compressed_local_indices_c.stride( - ) - out_stride0, out_stride1 = out.stride() - launch_with_pdl = os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" - - # Launch kernel - _deepseek_v4_local_to_global_kernel[grid]( - req_id_c, - block_table_swa_c, - block_table_compressed_c, - swa_local_indices_c, - compressed_local_indices_c, - out, - swa_buffer_offset_in_tokens, - compressed_buffer_offset_in_tokens, - tokens_per_block, - tokens_per_block_compressed, - max_blocks_swa, - max_blocks_compressed, - num_swa_indices, - num_compressed_indices, - total_output_indices, - has_compressed, - bt_swa_stride0, - bt_swa_stride1, - bt_compressed_stride0, - bt_compressed_stride1, - swa_indices_stride0, - swa_indices_stride1, - compressed_indices_stride0, - compressed_indices_stride1, - out_stride0, - out_stride1, - LAUNCH_WITH_PDL=launch_with_pdl, - launch_pdl=launch_with_pdl, - ) - - return out diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.py new file mode 100644 index 000000000000..c5967b713b28 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.py @@ -0,0 +1,337 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Optional + +import torch +from triton import next_power_of_2 + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.attention_backend.vanilla import VanillaAttentionMetadata +from tensorrt_llm._utils import prefer_pinned + +from .params import RocketKVMetadataParams + +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): + 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 + + 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 + + # 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, + ) + + self.q_cu_seqlens = torch.zeros_like( + self.q_cu_seqlens_cuda, 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 + ) + + # 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 + ) + + # 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, + ) + + # 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 + ) + + # 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 + ) + + # 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, + ) + + # 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.host_kt_cache_block_offsets = torch.zeros_like( + self.kt_cache_block_offsets, + device="cpu", + pin_memory=prefer_pinned(), + ) + + # Number of KT tokens for each sequence + self.num_kt_tokens = torch.empty( + self.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) + + # 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 + ) + + # 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] + + # 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 + ) + + # -------------------------------- Context phase -------------------------------- + self.context_cumsum[1 : self.num_contexts + 1] = torch.cumsum( + self.prompt_lens_cpu[: self.num_contexts], dim=0 + ) + self.context_cumsum_cuda[: self.num_contexts + 1].copy_( + self.context_cumsum[: self.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_seq_indices = torch.where(valid_mask)[0] + invalid_seq_indices = torch.where(~valid_mask)[0] + valid_batch_size = len(valid_seq_indices) + self.valid_seq_indices_cuda[:valid_batch_size].copy_(valid_seq_indices, non_blocking=True) + + # 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 + ) + 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[valid_seq_indices] = self.prompt_budget + sparse_counts_ctx[invalid_seq_indices] = self.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.q_cu_seqlens[: valid_batch_size + 1] = ( + torch.arange(valid_batch_size + 1, device="cpu", dtype=torch.int32) * self.window_size + ) + self.q_cu_seqlens_cuda[: valid_batch_size + 1].copy_( + self.q_cu_seqlens[: valid_batch_size + 1], non_blocking=True + ) + + self.k_cu_seqlens[1 : valid_batch_size + 1] = torch.cumsum( + self.k_context_lens[:valid_batch_size], dim=0 + ) + self.k_cu_seqlens_cuda[: valid_batch_size + 1].copy_( + self.k_cu_seqlens[: valid_batch_size + 1], non_blocking=True + ) + + if valid_batch_size > 0: + # Maximum context length of RocketKV key for valid sequences for padding + self.max_rocket_k_ctx_len = self.k_context_lens[:valid_batch_size].max().item() + self.total_rocket_k_ctx_tokens = self.k_cu_seqlens[valid_batch_size].item() + else: + self.max_rocket_k_ctx_len = 0 + 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() + + # -------------------------------- Generation phase -------------------------------- + self.num_kt_tokens[: self.num_generations] = ( + self.kv_lens[self.num_contexts : self.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_cuda[: self.num_generations + 1].copy_( + self.cum_kt_lens[: self.num_generations + 1], non_blocking=True + ) + + self.total_kt_tokens = self.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 + ) + + 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.host_kt_cache_block_offsets = torch.zeros_like( + self.kt_cache_block_offsets, + device="cpu", + pin_memory=prefer_pinned(), + ) + + 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 + ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket/module.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket/module.py new file mode 100644 index 000000000000..5301199a2473 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/rocket/module.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RocketKV module-layer integration.""" + +from typing import Optional + +import torch + +from tensorrt_llm.logger import logger + + +def initialize_sparse_attn( + self, + *, + config, + mapping, + mapping_o, + rms_norm_eps: float, + quant_config, + q_scaling: float, + bias: bool, + dtype: torch.dtype, + reduce_output: bool, + aux_stream: Optional[torch.cuda.Stream], +) -> None: + """Configure the Attention module for RocketKV.""" + del config, mapping, mapping_o, rms_norm_eps, quant_config, q_scaling + del bias, dtype, reduce_output, aux_stream + + logger.warning_once( + "disable rope_fusion for RocketKV.", + key="disable_rope_fusion_for_rocketkv", + ) + self.rope_fusion = False diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket/params.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket/params.py new file mode 100644 index 000000000000..940a840a8686 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/rocket/params.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""RocketKV parameter types.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Literal, Union + +import tensorrt_llm +import tensorrt_llm.bindings + +if TYPE_CHECKING: + pass + +from ..params import SparseMetadataParams, SparseParams + +ModelConfig = tensorrt_llm.bindings.ModelConfig + + +@dataclass(frozen=True) +class RocketKVMetadataParams(SparseMetadataParams): + """RocketKV metadata parameters.""" + + prompt_budget: int + window_size: int + page_size: int + topk: int + + +@dataclass(frozen=True) +class RocketKVParams(SparseParams): + """RocketKV backend parameters.""" + + algorithm: Literal["rocket"] = field(init=False, default="rocket") + window_size: int = 32 + kernel_size: int = 63 + topr: Union[int, float] = 128 + topk: int = 64 + prompt_budget: int = 2048 + page_size: int = 4 + kt_cache_dtype: str = "float8_e5m2" + + @property + def indices_block_size(self) -> int: + return self.page_size diff --git a/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/__init__.py new file mode 100644 index 000000000000..13afc5a9b238 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/__init__.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Skip-softmax configuration, scheduling, and kernel parameters.""" + +from .params import ( + SkipSoftmaxFormula, + SkipSoftmaxParams, + SkipSoftmaxScheduler, + skip_softmax_config_from_ckpt_sparse_attention_config, + skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config, + skip_softmax_formula_from_ckpt_sparse_attention_config, + skip_softmax_ignore_from_ckpt_sparse_attention_config, + skip_softmax_target_sparsity_from_ckpt_sparse_attention_config, + skip_softmax_threshold_scale_factor_config_from_ckpt_sparse_attention_config, +) + +__all__ = [ + "SkipSoftmaxFormula", + "SkipSoftmaxParams", + "SkipSoftmaxScheduler", + "skip_softmax_config_from_ckpt_sparse_attention_config", + "skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config", + "skip_softmax_formula_from_ckpt_sparse_attention_config", + "skip_softmax_ignore_from_ckpt_sparse_attention_config", + "skip_softmax_target_sparsity_from_ckpt_sparse_attention_config", + "skip_softmax_threshold_scale_factor_config_from_ckpt_sparse_attention_config", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax.py b/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/params.py similarity index 88% rename from tensorrt_llm/_torch/attention_backend/sparse/skip_softmax.py rename to tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/params.py index 1290632b6a54..feefcc2e21b1 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/params.py @@ -12,16 +12,9 @@ # 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. -"""Skip-softmax calibration helpers shared by the LLM and VisualGen pipelines. +"""Skip-softmax parameter and calibration types.""" -The kernel consumes a scalar ``threshold_scale_factor`` (combined with -the sequence length at runtime) to decide which KV blocks to skip. -This module owns the calibration side that produces that scalar from a -semantic ``target_sparsity`` via a formula shipped with the checkpoint. -The helpers are shared by the LLM and VisualGen pipelines. -""" - -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any, Dict, Literal, Optional, Union import numexpr @@ -31,7 +24,7 @@ from tensorrt_llm.llmapi.utils import StrictBaseModel -from .params import SkipSoftmaxKernelParams, SparseParams +from ..params import SparseParams, SparseRuntimeParams _RESERVED_FORMULA_KEYS = frozenset({"formula", "target_sparsity"}) _SKIP_SOFTMAX_ALGORITHMS = frozenset({"skip_softmax", "softmax_skip"}) @@ -63,8 +56,7 @@ def skip_softmax_config_from_ckpt_sparse_attention_config( config_groups = ckpt_sparse_attention_config.get("config_groups") if isinstance(config_groups, dict): - # ModelOpt may emit many sparse-attention groups. The shared - # skip-softmax helpers operate on the single skip-softmax group. + # ModelOpt may emit multiple sparse-attention groups. groups = [ group for group in config_groups.values() @@ -121,18 +113,9 @@ def skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config( class SkipSoftmaxFormula(StrictBaseModel): - """Numexpr formula + coefficients that map target_sparsity to threshold_scale_factor. - - ``formula`` references ``target_sparsity`` and one or more named - coefficients (e.g. ``"a * exp(b * target_sparsity)"``); - ``coefficients`` supplies their values. - """ - - # Frozen: validated once at construction and never mutated, so the - # construction-time formula/coefficient checks always hold. - # NOTE: ``model_config`` here is the Pydantic model setting (this - # class's own validation behavior) — unrelated to the LLM runtime's - # ``ModelConfig`` / model config plumbing. + """Formula mapping target sparsity to a threshold scale factor.""" + + # Keep validated formulas immutable. model_config = ConfigDict(frozen=True) formula: str = PydanticField( @@ -236,8 +219,7 @@ def _extend(patterns: Any) -> None: ignore.append(pattern) if isinstance(ckpt_sparse_attention_config, dict): - # The input may already be the skip-softmax group rather than the - # outer sparse_attention_config. + # The input may already be the skip-softmax group. _extend(ckpt_sparse_attention_config.get("ignore")) config = skip_softmax_config_from_ckpt_sparse_attention_config(ckpt_sparse_attention_config) if isinstance(config, dict) and config is not ckpt_sparse_attention_config: @@ -341,8 +323,7 @@ def _compute(phase: str, sparsity: Optional[float]) -> Optional[float]: if threshold_config is None: phase_formula = None else: - # Prefer phase-specific coefficients for LLM, then shared - # coefficients, then inline coefficients in the threshold dict. + # Prefer phase-specific, shared, then inline coefficients. coefficient_config = threshold_config.get(phase) if coefficient_config is None: coefficient_config = threshold_config.get("coefficients") @@ -385,11 +366,7 @@ def get_graph_phase_for_timestep( *, disabled_until_timestep: Optional[float], ) -> Optional[int]: - """Return the graph-key phase for the timestep disablement boundary. - - VisualGen denoising timesteps descend from 1 to 0, so skip-softmax is - disabled while ``timestep >= disabled_until_timestep``. - """ + """Return 1 after descending timesteps cross the cutoff, otherwise 0.""" if disabled_until_timestep is None: return None timestep_value = cls._as_float(timestep) @@ -397,8 +374,15 @@ def get_graph_phase_for_timestep( return None return int(timestep_value < disabled_until_timestep) - def get_kernel_params(self, *, timestep: Any = None): - """Return kernel params, applying timestep-based disablement when provided.""" + def get_runtime_params( + self, + *, + runtime_params: Optional[SparseRuntimeParams] = None, + timestep: Any = None, + ) -> SparseRuntimeParams: + """Return runtime parameters with skip-softmax thresholds.""" + if runtime_params is None: + runtime_params = SparseRuntimeParams() if ( self.get_graph_phase_for_timestep( timestep, @@ -406,8 +390,13 @@ def get_kernel_params(self, *, timestep: Any = None): ) == 0 ): - return SkipSoftmaxKernelParams() - return SkipSoftmaxKernelParams( + return replace( + runtime_params, + threshold_scale_factor_prefill=0.0, + threshold_scale_factor_decode=0.0, + ) + return replace( + runtime_params, threshold_scale_factor_prefill=self.threshold_scale_factor_prefill, threshold_scale_factor_decode=self.threshold_scale_factor_decode, ) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index a3f4e89d8d85..b6a365c055cf 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -17,7 +17,7 @@ import math import os import weakref -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -41,6 +41,7 @@ KVCacheParams, MLAParams, PositionalEmbeddingParams, PredefinedAttentionMask, RopeParams, merge_attention_forward_args) +from .sparse.hooks import prepare_sparse_runtime_params from .sparse.params import SparseParams from .sparse.skip_softmax import SkipSoftmaxParams @@ -1602,23 +1603,8 @@ def forward( seq_start=num_ctx, ) - # RocketKV and DSA predict which blocks to keep, so build their sparse - # index tensors here. Skip-softmax needs no prediction. - sparse_params = self.sparse_params - if (sparse_params is not None - and not isinstance(sparse_params, SkipSoftmaxParams)): - kv_idx, kv_off = self.sparse_kv_predict(q, k, metadata, - forward_args) - at_idx, at_off = self.sparse_attn_predict(q, k, metadata, - forward_args) - forward_args.sparse_prediction = replace( - forward_args.sparse_prediction, - sparse_kv_indices=kv_idx, - sparse_kv_offsets=kv_off, - sparse_attn_indices=at_idx, - sparse_attn_offsets=at_off, - sparse_attn_indices_block_size=sparse_params.indices_block_size, - ) + forward_args.sparse_runtime_params = prepare_sparse_runtime_params( + self, q, k, metadata, forward_args) # Compute FlashMLA tile-scheduler metadata once per forward pass. # The flag is reset in prepare_flash_mla() and update_for_spec_dec() to trigger @@ -1675,7 +1661,8 @@ def forward( assert k.shape[0] == num_tokens assert v.shape[0] == num_tokens else: - is_sparse_attn = forward_args.sparse_prediction.sparse_attn_indices is not None and forward_args.sparse_prediction.sparse_attn_indices.numel( + sparse_attn_indices = forward_args.sparse_runtime_params.sparse_attn_indices + is_sparse_attn = sparse_attn_indices is not None and sparse_attn_indices.numel( ) > 0 if attention_input_type == AttentionInputType.context_only and is_sparse_attn: assert forward_args.is_fused_qkv @@ -1751,9 +1738,11 @@ def forward( sparse_params = self.sparse_params if isinstance(sparse_params, SkipSoftmaxParams): - forward_args.skip_softmax_kernel_params = ( - sparse_params.scheduler.get_kernel_params( - timestep=forward_args.timestep)) + forward_args.sparse_runtime_params = ( + sparse_params.scheduler.get_runtime_params( + runtime_params=forward_args.sparse_runtime_params, + timestep=forward_args.timestep, + )) # max_context_q_len_override is only set when encoder CUDA graphs are enabled. if metadata.max_context_q_len_override is not None: @@ -1974,10 +1963,8 @@ def sparse_kv_predict( metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """ - Predict sparse kv indices. It's implemented in the derived class. - """ - raise NotImplementedError + """Predict sparse KV indices when required by an algorithm.""" + return None, None def sparse_attn_predict( self, @@ -1986,10 +1973,8 @@ def sparse_attn_predict( metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """ - Predict sparse attn indices. It's implemented in the derived class. - """ - raise NotImplementedError + """Predict sparse attention indices when required by an algorithm.""" + return None, None def mla_rope_generation( self, diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index eef2fc520906..799de82668ea 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/attention_backend/vanilla.py @@ -17,7 +17,6 @@ from .interface import (AttentionBackend, AttentionForwardArgs, AttentionInputType, AttentionMask, AttentionMetadata, PredefinedAttentionMask, merge_attention_forward_args) -from .sparse.kernel import triton_index_gather from .sparse.params import SparseParams @@ -142,6 +141,8 @@ def _single_request_update_kv_cache(self, sparse_kv_indices=None): # select tokens using the sparse kv indices if sparse_kv_indices is not None: + from .sparse.rocket.kernels import triton_index_gather + k_selected = triton_index_gather(k, sparse_kv_indices) v_selected = triton_index_gather(v, sparse_kv_indices) else: @@ -252,6 +253,8 @@ def _single_request_attn_forward(self, """ # select the key and value states using the sparse indices if sparse_indices is not None: + from .sparse.rocket.kernels import triton_index_gather + key_states = triton_index_gather(key_states, sparse_indices) value_states = triton_index_gather(value_states, sparse_indices) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index c95aac6c8147..de2ac10e1bfe 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + # -------------------------------------------------- # Portions of this code were derived from DeepSeek‑V3: # https://github.com/deepseek-ai/DeepSeek-V3 @@ -831,8 +834,6 @@ def __init__( mapping_with_cp=mapping_with_cp, reduce_output=reduce_output) - self.indexer = self.mqa.indexer - self.kv_a_proj_with_mqa = DeepseekV3Linear( config.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim + self.q_lora_rank, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 5724c1052d8c..b44d1fc4f49b 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -52,7 +52,7 @@ from tensorrt_llm.quantization.mode import QuantAlgo from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams -from ..attention_backend.sparse.deepseek_v4.deepseek_v4 import DeepseekV4TrtllmAttentionMetadata +from ..attention_backend.sparse.deepseek_v4 import DeepseekV4TrtllmAttentionMetadata from ..distributed import ( AllReduce, AllReduceFusionOp, @@ -1355,9 +1355,6 @@ def __init__( reduce_output=reduce_output, ) - self.indexer = getattr(self.mqa, "indexer", None) - self.compressor = getattr(self.mqa, "compressor", None) - class DeepseekV4Gate(nn.Module): def __init__( diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index f10696c97bb9..8cc26ba95035 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -46,6 +46,7 @@ _gather_paged_batched, _write_main_kv_slots_to_pool, ) +from ..attention_backend.sparse.params import SparseBackendForwardArgs from ..distributed import AllReduce, AllReduceParams, MiniMaxAllReduceRMS from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer @@ -1132,7 +1133,10 @@ def _attention_core( assert idx_q is not None and idx_k is not None # Publish the selected blocks so the FMHA runs the sparse path. kv_block_indexes = self.attn.run_indexer(idx_q, idx_k, attn_metadata) - forward_args = AttentionForwardArgs(output=output, topk_indices=kv_block_indexes) + forward_args = AttentionForwardArgs( + output=output, + sparse_backend_args=SparseBackendForwardArgs(topk_indices=kv_block_indexes), + ) else: assert idx_q is None and idx_k is None # No top-k selection means the FMHA attends the full page table. diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index 94e8f089e045..48366251339d 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -108,17 +108,61 @@ KV-cache contract. At a high level, it owns: (`is_lite == True`). In lite mode there is no separate Q low-rank compression stage. `is_lite` changes the projection structure, not just a small code path. -Dense MLA and current sparse MLA variants still use the same module/backend/ -metadata/KV-cache split described above. Sparse-specific routing may currently -pass through `MLA`, but that should be treated as an implementation detail -rather than a stable design boundary. +Dense and sparse MLA variants use the same `MLA` module. `MLA.forward_impl()` +selects the dense implementation or the sparse facade in +`attention_backend/sparse/hooks.py`. The hook facade then dispatches by sparse +algorithm to its `module.py`; MLA-specific dispatch is intentionally not part +of the generic `AttentionBackend` interface. + +The sparse hook facade defines the uniform contract for module initialization, +weight lifecycle, forward, optional output preparation, custom-op bridging, +and output projection. Every hook is optional at registration time: an +algorithm implements only the module paths it needs, while a caller that +requires a hook fails explicitly if it is absent. Shared `MLA` code creates the +sparse-aware MQA and the ordinary dense MHA exactly once, then invokes the +sparse initialization hook. The hook removes dense modules that the algorithm +does not use and initializes only algorithm-specific module state. Dense +create/transform behavior remains inline in `MLA` and is used whenever the +algorithm does not override it. + +`Attention` and `MLA` resolve and cache `self.sparse_attn_hooks` during module +initialization. Later lifecycle methods use that cached contract instead of +redispatching from `sparse_params`. Dense modules and algorithms without +module-layer overrides receive an empty hook set, so optional call sites only +need to test the relevant callable. A module path that requires an override +uses `require()` and fails explicitly when that hook is absent. + +`Attention` invokes the same lifecycle contract at initialization, +`forward_impl()`, and output projection. Its initialization hook runs before +rotary embedding and backend construction so an algorithm can configure +module-level choices that affect both. Forward and projection signatures are +module-specific and validated as either the `Attention` or `MLA` contract. + +Ordinary sparse variants use `attention_output_hidden_size` and the shared +output allocation. DeepSeek-V4's fused epilogue is the exception: it requires +two typed output buffers before the attention op runs, so it implements the +optional output-preparation hook. `_create_outputs()` always returns a tensor +list whose first entry is the standard attention output. The same list flows +through forward, the registered custom op, and output projection; algorithms +own the meaning of any additional entries. Context- and generation-phase +helpers stay within each algorithm module and are not part of the generic hook +facade. + +Sparse prediction inputs stay out of shared MLA APIs. Algorithm modules wrap +their module-to-backend inputs in a `SparseBackendForwardArgs` subclass and +pass it through the registered `AttentionForwardArgs.sparse_backend_args` +field. For example, DSA owns `DSABackendForwardArgs`, whose indexer +intermediates are consumed by `DSATrtllmAttention.sparse_attn_predict`. +Shared sparse carriers, including `SparseBackendForwardArgs.topk_indices` and +the backend-to-AttentionOp `SparseRuntimeParams`, live in +`attention_backend/sparse/params.py`. For MLA-related tasks, first check whether the work fits the current projection structure, can stay on an existing backend and metadata family, and can preserve the current latent-cache / paged-KV contract. If it can, the task usually stays within the existing MLA stack. If it depends on sparse -helper-level control flow, read `mla.py` and the relevant sparse -backend code directly. +helper-level control flow, read `mla.py`, `attention_backend/sparse/hooks.py`, +and the relevant algorithm's `module.py` directly. ## 2. Backend Layer Reference @@ -148,7 +192,7 @@ managers stay model-scope and consume the user-facing config directly. Sparse metadata consumes `SparseMetadataParams`, derived independently from the same user-facing config. -Sparse registrations are defined in `attention_backend/sparse/utils.py`. Check +Sparse registrations are defined in `attention_backend/sparse/registry.py`. Check that file for the current supported combinations, as they may change over time. ### 2.3 Backend contract @@ -345,7 +389,8 @@ Working rules: - Stay on `Attention` or `MLA` plus an existing backend family when possible. - Extend the `TRTLLM` backend path before adding a new backend. -- Extend module-level hooks before adding a new backend. +- Extend the sparse hook facade and the algorithm's `module.py` for + sparse MLA module-side behavior. - Follow an existing sparse family pattern before adding a new sparse abstraction. - Treat cache-manager mismatch as a real blocker. @@ -362,7 +407,9 @@ Working rules: | `tensorrt_llm/_torch/attention_backend/fmha/` | Internal TRTLLM FMHA libraries | | `tensorrt_llm/_torch/attention_backend/vanilla.py` | Torch fallback backend and metadata | | `tensorrt_llm/_torch/attention_backend/flashinfer.py` | FlashInfer backend and metadata | -| `tensorrt_llm/_torch/attention_backend/sparse/` | DSA, Rocket sparse backends, metadata, cache managers | +| `tensorrt_llm/_torch/attention_backend/sparse/hooks.py` | Sparse module hooks and backend prediction orchestration | +| `tensorrt_llm/_torch/attention_backend/sparse//module.py` | Algorithm-specific module-hook implementations | +| `tensorrt_llm/_torch/attention_backend/sparse/` | Sparse prediction backends, metadata, cache managers, and kernels | ## 6. Testing Notes diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index e902a892af3e..16cb0548399d 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -17,6 +17,7 @@ from ..attention_backend.interface import (AttentionMask, CustomAttentionMask, PositionalEmbeddingParams, PredefinedAttentionMask) +from ..attention_backend.sparse.hooks import get_sparse_attn_hooks from ..attention_backend.utils import create_attention, get_attention_backend from ..distributed import (AllReduceParams, HelixAllToAllNative, alltoall_helix, cp_allgather, reducescatter) @@ -602,6 +603,8 @@ def __init__( sparse_params = (sparse_attn_cfg.to_sparse_params( pretrained_config=config.pretrained_config, layer_idx=self.layer_idx) if sparse_attn_cfg is not None else None) + self.sparse_params = sparse_params + self.sparse_attn_hooks = get_sparse_attn_hooks(self) attn_cls = get_attention_backend(self.attn_backend, sparse_params=sparse_params) @@ -631,10 +634,21 @@ def __init__( logger.info_once(f"Using sparse attention: {algo} {cfg_dump}", key="sparse_attention_config") - if config.sparse_attention_config.algorithm == "rocket": - logger.warning_once("disable rope_fusion for RocketKV.", - key="disable_rope_fusion_for_rocketkv") - self.rope_fusion = False + if initialize_sparse_attn := self.sparse_attn_hooks.initialize_sparse_attn: + initialize_sparse_attn( + self, + config=config, + mapping=mapping, + mapping_o=mapping_o, + rms_norm_eps=getattr(config.pretrained_config, "rms_norm_eps", + 1e-6), + quant_config=self.quant_config, + q_scaling=self.q_scaling, + bias=bias, + dtype=dtype, + reduce_output=reduce_output, + aux_stream=None, + ) if self.rope_fusion and not attn_cls.support_fused_rope(): logger.warning_once( @@ -915,7 +929,26 @@ def forward_impl( relative_attention_bias: Optional[torch.Tensor] = None, relative_attention_max_distance: int = 0, has_lora: bool = False, + **kwargs, ): + if forward_sparse_attn := self.sparse_attn_hooks.forward_sparse_attn: + return forward_sparse_attn( + self, + q, + k, + v, + attn_metadata, + attention_mask, + attention_window_size, + attention_mask_data, + mrope_config, + attention_sinks, + relative_attention_bias, + relative_attention_max_distance, + has_lora, + **kwargs, + ) + mrope_rotary_cos_sin = None mrope_position_deltas = None if mrope_config is not None: @@ -1058,11 +1091,21 @@ def forward( relative_attention_bias=relative_attention_bias, relative_attention_max_distance=relative_attention_max_distance, has_lora=bool(lora_params), + **kwargs, ) if self.attn_output_gate: attn_output = self.apply_output_gate(attn_output, gate) + if project_sparse_attn_output := self.sparse_attn_hooks.project_sparse_attn_output: + return project_sparse_attn_output( + self, + attn_output, + attn_metadata, + all_reduce_params, + lora_params, + ) + attn_output = _helix_cp_output_projection(self.o_proj, attn_output, attn_metadata, all_reduce_params, diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 63f9a1bd65d3..d8f1cee8f794 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -15,15 +15,14 @@ import functools import math -import os import weakref -from typing import List, Optional, cast +from typing import Optional, cast import torch from torch import nn import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils -from tensorrt_llm._utils import get_sm_version, is_sm_100f, nvtx_range, nvtx_range_debug +from tensorrt_llm._utils import get_sm_version, is_sm_100f from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.utils.fp4_utils import NVFP4_SF_VEC_SIZE @@ -41,17 +40,13 @@ PositionalEmbeddingParams, PredefinedAttentionMask, ) -from ..attention_backend.sparse.dsa import ( - DSAtrtllmAttentionMetadata, - transform_local_topk_and_prepare_pool_view, -) +from ..attention_backend.sparse.hooks import get_sparse_attn_hooks from ..attention_backend.utils import create_attention from ..distributed import AllReduceParams from ..model_config import ModelConfig from ..utils import ( Fp4QuantizedTensor, compute_swizzled_sf_shape, - is_torch_compiling, maybe_compiled_cat, maybe_compiled_copy_, ) @@ -63,20 +58,10 @@ extract_extra_attrs, ) from .linear import Linear, TensorParallelMode, is_static_nvfp4_input_eligible -from .multi_stream_utils import do_multi_stream, maybe_execute_in_parallel +from .multi_stream_utils import maybe_execute_in_parallel from .rms_norm import RMSNorm from .rotary_embedding import RotaryEmbedding -# Import FlashMLA sparse attention kernel -try: - from tensorrt_llm.flash_mla import flash_mla_sparse_fwd -except ImportError: - flash_mla_sparse_fwd = None - - -def _is_env_truthy(name: str) -> bool: - return os.environ.get(name, "").strip().lower() in ("1", "true", "on") - def _slice_hidden_states_to_num_tokens(hidden_states, num_tokens: int): """Drop CUDA-graph padding by slicing hidden_states to num_tokens rows. @@ -97,20 +82,12 @@ def _slice_hidden_states_to_num_tokens(hidden_states, num_tokens: int): fp4 = hidden_states.fp4_tensor if fp4.shape[0] == num_tokens: - return hidden_states # no padding to strip - # The row-slice math below relies on the swizzled 128-row-tile layout; a - # non-swizzled (linear) scale buffer would be sliced incorrectly and then - # silently relabeled as swizzled. Reject it rather than corrupt the SF. + return hidden_states if not hidden_states.is_sf_swizzled: raise ValueError( "_slice_hidden_states_to_num_tokens only supports swizzled FP4 scaling factors" ) sf = hidden_states.scaling_factor - # fp4 packs 2 elements/byte, so the logical hidden dim is fp4_cols*2 and - # the scale-factor column count is that / NVFP4_SF_VEC_SIZE. The swizzled - # SF is laid out in independent 128-row tiles, so the leading num_tokens - # rows' scale factors occupy exactly the first padded_rows*padded_cols - # bytes. sf_cols = fp4.shape[-1] * 2 // NVFP4_SF_VEC_SIZE padded_rows, padded_cols = compute_swizzled_sf_shape(num_tokens, sf_cols) sf_len = padded_rows * padded_cols @@ -133,21 +110,13 @@ def _extract_mla_extra_attrs(layer_idx: str): return metadata, mla_layer -def create_mla_outputs_impl(hidden_states: torch.Tensor, layer_idx: str) -> List[torch.Tensor]: +def create_mla_outputs_impl(hidden_states: torch.Tensor, layer_idx: str) -> list[torch.Tensor]: metadata, mla_layer = _extract_mla_extra_attrs(layer_idx) - enable_dsv4_epilogue_fusion = mla_layer._should_use_dsv4_epilogue_fusion( - metadata.num_contexts, metadata.num_generations - ) - output_input = hidden_states[:0] if enable_dsv4_epilogue_fusion else hidden_states - attn_output = mla_layer.create_output(output_input, metadata.num_contexts) - outputs = [attn_output] - if enable_dsv4_epilogue_fusion: - outputs.extend(mla_layer._create_dsv4_epilogue_buffers(hidden_states, metadata.num_tokens)) - return outputs + return mla_layer._create_outputs(hidden_states, metadata) @torch.library.custom_op("trtllm::create_mla_outputs", mutates_args=()) -def create_mla_outputs(hidden_states: torch.Tensor, layer_idx: str) -> List[torch.Tensor]: +def create_mla_outputs(hidden_states: torch.Tensor, layer_idx: str) -> list[torch.Tensor]: return create_mla_outputs_impl(hidden_states, layer_idx) @@ -158,7 +127,7 @@ def _create_mla_outputs_fake(hidden_states, layer_idx): @torch.library.custom_op( "trtllm::mla_custom_op_inplace", - mutates_args=("output", "dsv4_output", "dsv4_output_sf"), + mutates_args=("output", "sparse_output", "sparse_output_sf"), ) def mla_custom_op_inplace( hidden_states: torch.Tensor, @@ -166,17 +135,11 @@ def mla_custom_op_inplace( layer_idx: str, output: torch.Tensor, latent_cache_gen: Optional[torch.Tensor], - dsv4_output: Optional[torch.Tensor], - dsv4_output_sf: Optional[torch.Tensor], - enable_dsv4_epilogue_fusion: bool, + sparse_output: Optional[torch.Tensor], + sparse_output_sf: Optional[torch.Tensor], hidden_states_fp4: Optional[torch.Tensor] = None, hidden_states_sf: Optional[torch.Tensor] = None, ) -> None: - # When hidden_states_fp4/_sf are provided, the previous layer's boundary - # fusion pre-quantized this layer's kv_a_proj NVFP4 input. Custom ops can't - # take dataclasses, so the Fp4QuantizedTensor is passed as explicit tensors - # and reconstructed here (mirrors mla_dsa_proj). unquantized_hidden_states carries - # the un-quantized view for the o_proj output sizing in create_output. metadata, mla_layer = _extract_mla_extra_attrs(layer_idx) if hidden_states_fp4 is not None or hidden_states_sf is not None: assert hidden_states_fp4 is not None and hidden_states_sf is not None, ( @@ -187,152 +150,19 @@ def mla_custom_op_inplace( scaling_factor=hidden_states_sf, unquantized_hidden_states=hidden_states, ) - if mla_layer.is_deepseek_v4: - # DeepSeek-V4 uses MQA mode and has no residual-less RMSNorm+quant - # fusion entry point, so it cannot be reached with a pre-quantized - # Fp4QuantizedTensor input; the call site passes plain hidden_states. - if enable_dsv4_epilogue_fusion: - if dsv4_output is None or dsv4_output_sf is None: - raise RuntimeError( - "DSv4 fused epilogue requires caller-provided output and output_sf buffers." - ) - dsv4_epilogue_output = (dsv4_output, dsv4_output_sf) - else: - if dsv4_output is not None or dsv4_output_sf is not None: - raise RuntimeError( - "DSv4 fused epilogue buffers require epilogue fusion to be enabled." - ) - dsv4_epilogue_output = None - mla_layer.forward_impl_with_deepseek_v4( - position_ids, - hidden_states, - metadata, - output=output, - dsv4_epilogue_output=dsv4_epilogue_output, - ) - else: - if enable_dsv4_epilogue_fusion: - raise RuntimeError("DSv4 fused epilogue cannot be enabled for non-DeepSeek-V4 MLA.") - mla_layer.forward_impl( - position_ids, hidden_states, metadata, output=output, latent_cache_gen=latent_cache_gen - ) - - -@torch.library.custom_op("trtllm::mla_dsa_proj", mutates_args=()) -def mla_dsa_proj( - hidden_states: torch.Tensor, - position_ids: Optional[torch.Tensor], - layer_idx: str, - hidden_states_fp4: Optional[torch.Tensor] = None, - hidden_states_sf: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - """Token-wise projections for DSA MLA (CUDA-graph-capturable). - - Runs kv_a_proj, layernorms, q_b_proj, and conditionally - indexer.pre_indexer_proj (FP8/FP4 quantize, weight scaling). Does NOT - update the indexer k cache — that happens in Op 2 (mla_dsa_attn_inplace) - because the scatter kernel accesses batch-specific metadata. - - When ``hidden_states_fp4`` / ``hidden_states_sf`` are provided, the - previous layer's boundary fusion already produced a fused NVFP4 view of - the post-RMSNorm hidden states (via fused_add_rmsnorm_fp4_quantize). The - op rebuilds an ``Fp4QuantizedTensor`` carrying the BF16 view in - ``hidden_states`` (for DSA's ``pre_indexer_proj``) plus the - pre-quantized FP4 form (consumed by ``kv_a_proj_with_mqa``). - Both must be passed together; passing either alone is rejected. - - Returns [q, compressed_kv, k_pe, latent_cache] when the short-MHA path - handles all tokens, or [q, compressed_kv, k_pe, latent_cache, q_fp8, - k_fp8, k_scale, weights, q_scale] when the indexer runs. Under torch - compile, _should_use_short_mha returns False so the result is always - length 9, keeping control flow straight-line for CUDA graph capture. - The trailing q_scale is only consumed by the FP4 dispatch; the FP8 - path ignores it in forward_dsa_attn. - """ - metadata, mla_layer = _extract_mla_extra_attrs(layer_idx) - if hidden_states_fp4 is not None or hidden_states_sf is not None: - assert hidden_states_fp4 is not None and hidden_states_sf is not None, ( - "hidden_states_fp4 and hidden_states_sf must be passed together" - ) - hs = Fp4QuantizedTensor( - fp4_tensor=hidden_states_fp4, - scaling_factor=hidden_states_sf, - unquantized_hidden_states=hidden_states, - ) - else: - hs = hidden_states - return mla_layer.forward_dsa_proj(position_ids, hs, metadata) - - -@mla_dsa_proj.register_fake -def _mla_dsa_proj_fake( - hidden_states: torch.Tensor, - position_ids: Optional[torch.Tensor], - layer_idx: str, - hidden_states_fp4: Optional[torch.Tensor] = None, - hidden_states_sf: Optional[torch.Tensor] = None, -) -> List[torch.Tensor]: - # Under torch compile _should_use_short_mha is False, so the result is - # always 9 tensors (4 attention inputs + 5 indexer intermediates, with - # q_scale as the 9th carried for the FP4 dispatch). - metadata, mla_layer = _extract_mla_extra_attrs(layer_idx) - num_tokens = hidden_states.shape[0] - indexer = mla_layer.mqa.indexer - q = hidden_states.new_empty([num_tokens, mla_layer.num_heads_tp * mla_layer.qk_head_dim]) - compressed_kv = hidden_states.new_empty([num_tokens, mla_layer.kv_lora_rank]) - k_pe = hidden_states.new_empty([num_tokens, mla_layer.qk_rope_head_dim]) - latent_cache = hidden_states.new_empty( - [num_tokens, mla_layer.kv_lora_rank + mla_layer.qk_rope_head_dim] - ) - if indexer is None: - # DSA "shared" layer: no indexer, mirror forward_dsa_proj's early - # return of only the 4 base tensors (no indexer intermediates). - return [q, compressed_kv, k_pe, latent_cache] - # Indexer intermediates: q_fp8, k_fp8, k_scale, weights, q_scale. - # Under FP4 q_fp8's trailing dim is head_dim // 2 (two E2M1 codes per - # byte) and q_scale carries one int32 per (token, head) packing four - # UE8M0 exponents; under FP8 q_fp8's trailing dim is head_dim and - # q_scale carries one float32 per (token, head). - if indexer.use_fp4: - q_fp8 = hidden_states.new_empty( - [num_tokens, indexer.n_heads, indexer.head_dim // 2], dtype=torch.int8 - ) - k_fp8 = hidden_states.new_empty([num_tokens, indexer.head_dim // 2], dtype=torch.int8) - k_scale = hidden_states.new_empty([num_tokens, 1], dtype=torch.int32) - q_scale = hidden_states.new_empty([num_tokens, indexer.n_heads, 1], dtype=torch.int32) - else: - q_fp8 = hidden_states.new_empty( - [num_tokens, indexer.n_heads, indexer.head_dim], dtype=torch.float8_e4m3fn - ) - k_fp8 = hidden_states.new_empty([num_tokens, indexer.head_dim], dtype=torch.float8_e4m3fn) - k_scale = hidden_states.new_empty([num_tokens, 1], dtype=torch.float32) - q_scale = hidden_states.new_empty([num_tokens, indexer.n_heads, 1], dtype=torch.float32) - weights = hidden_states.new_empty([num_tokens, indexer.n_heads], dtype=torch.float32) - return [q, compressed_kv, k_pe, latent_cache, q_fp8, k_fp8, k_scale, weights, q_scale] - - -@torch.library.custom_op("trtllm::mla_dsa_attn_inplace", mutates_args=("output",)) -def mla_dsa_attn_inplace( - q: torch.Tensor, - compressed_kv: torch.Tensor, - k_pe: torch.Tensor, - latent_cache: torch.Tensor, - indexer_intermediates: List[torch.Tensor], - position_ids: Optional[torch.Tensor], - layer_idx: str, - output: torch.Tensor, -) -> None: - """Batch-structure-dependent attention dispatch for DSA MLA. - - indexer_intermediates is [q_fp8, k_fp8, k_scale, weights, q_scale] when - the indexer ran in Op 1, or [] when short-MHA handled all tokens. The - trailing q_scale is only consumed by the FP4 dispatch; the FP8 path - ignores it. Runs sparse_attn_indexer then dispatches context/generation - attention. This op is excluded from CUDA graph capture. - """ - metadata, mla_layer = _extract_mla_extra_attrs(layer_idx) - mla_layer.forward_dsa_attn( - q, compressed_kv, k_pe, latent_cache, indexer_intermediates, position_ids, metadata, output + attn_output = [output] + if sparse_output is not None: + attn_output.append(sparse_output) + if sparse_output_sf is not None: + if sparse_output is None: + raise RuntimeError("sparse_output_sf requires sparse_output") + attn_output.append(sparse_output_sf) + mla_layer.forward_impl( + position_ids, + hidden_states, + metadata, + attn_output=attn_output, + latent_cache_gen=latent_cache_gen, ) @@ -373,40 +203,6 @@ def fp8_block_scaling_bmm_out( raise NotImplementedError(f"SM{sm_version} is not supported") -_q_b_proj_cute_dsl_import_ok: Optional[bool] = None - - -def _q_b_proj_cute_dsl_bf16(q: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - """BF16 dense GEMM via CuTe DSL. - - Computes ``q @ weight.T`` for [M, K] @ [N, K]^T -> [M, N]. - - Delegates to ``torch.ops.trtllm.cute_dsl_bf16_gemm_blackwell`` (which - runs its own autotune over (use_2cta, mma_tiler, cluster_shape)). Falls - back to ``torch.nn.functional.linear`` if CuTe DSL is unavailable. - """ - global _q_b_proj_cute_dsl_import_ok - if _q_b_proj_cute_dsl_import_ok is None: - try: - from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE - - _q_b_proj_cute_dsl_import_ok = IS_CUTLASS_DSL_AVAILABLE - except ImportError: - _q_b_proj_cute_dsl_import_ok = False - if not _q_b_proj_cute_dsl_import_ok or not is_sm_100f(): - return torch.nn.functional.linear(q, weight) - - assert q.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( - "q_b_proj cute_dsl path requires bfloat16 inputs" - ) - q = q.contiguous() - weight = weight.contiguous() - m, n = q.shape[0], weight.shape[0] - out = q.new_empty((m, n), dtype=torch.bfloat16) - torch.ops.trtllm.cute_dsl_bf16_gemm_blackwell(q, weight, out) - return out - - class MLA(nn.Module): def __init__( self, @@ -517,20 +313,12 @@ def __init__( if sparse_attn_cfg is not None else None ) + self.sparse_params = sparse_params + self.sparse_attn_hooks = get_sparse_attn_hooks(self) - sparse_algorithm = getattr(sparse_params, "algorithm", None) - self.is_dsa = sparse_algorithm == "dsa" - self.is_deepseek_v4 = sparse_algorithm == "deepseek_v4" - self._disable_dsv4_epilogue_fusion = self.is_deepseek_v4 and _is_env_truthy( - "TRTLLM_DSV4_DISABLE_FMHA_EPILOGUE_FUSION" - ) - - # Fold of the residual-less q_a_layernorm -> q_b_proj NVFP4 input - # quantize into a single fused (rmsnorm + fp4_quantize) kernel. Self- - # gating on q_b_proj being static-NVFP4 (see _resolve_qa_fused_scale); - # resolved lazily on first forward because q_b_proj.input_scale is only - # finalized after weight loading. - # None = not yet resolved; False = disabled; tensor = q_b_proj input_scale. + # Fold the residual-less q_a_layernorm -> q_b_proj NVFP4 input + # quantization into one fused RMSNorm + FP4-quantize kernel. Resolve + # lazily because q_b_proj.input_scale is finalized after weight loading. self._qa_fused_scale = None # tensor parallel @@ -571,16 +359,6 @@ def __init__( self.num_heads_tp = self.num_heads // tp_size self.num_heads_tp_cp = self.num_heads_tp // cp_size self.num_key_value_heads_tp = (self.num_key_value_heads + tp_size - 1) // tp_size - if self.is_deepseek_v4: - if self.num_groups % tp_size != 0: - raise ValueError( - f"DeepSeek-V4 num_groups ({self.num_groups}) must be divisible by tp_size ({tp_size})." - ) - if self.num_heads % self.num_groups != 0: - raise ValueError( - f"DeepSeek-V4 num_heads ({self.num_heads}) must be divisible by num_groups ({self.num_groups})." - ) - self.n_local_groups = self.num_groups // tp_size rms_norm_eps = getattr(config.pretrained_config, "rms_norm_eps", 1e-6) quant_config = config.get_quant_config() @@ -623,11 +401,6 @@ def __init__( use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, ) - if self.is_deepseek_v4: - # V4 unweighted per-head RMS on Q post-wq_b; q.view(-1, head_dim) at call site. - self.q_b_layernorm = RMSNorm( - hidden_size=self.qk_head_dim, eps=rms_norm_eps, dtype=dtype, has_weights=False - ) else: self.kv_a_proj_with_mqa = Linear( hidden_size, @@ -658,39 +431,6 @@ def __init__( ) self.q_b_proj = self.q_proj - kv_a_layernorm_hidden_size = ( - self.kv_lora_rank + self.qk_rope_head_dim if self.is_deepseek_v4 else kv_lora_rank - ) - self.kv_a_layernorm = RMSNorm( - hidden_size=kv_a_layernorm_hidden_size, dtype=dtype, eps=rms_norm_eps - ) - - if not self.is_deepseek_v4: - self.kv_b_proj = Linear( - self.kv_lora_rank, - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), - bias=bias, - dtype=dtype, - mapping=mapping, - tensor_parallel_mode=TensorParallelMode.COLUMN, - quant_config=quant_config, - skip_create_weights_in_init=config.skip_create_weights_in_init, - allreduce_strategy=config.allreduce_strategy, - force_dynamic_quantization=config.force_dynamic_quantization, - use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, - use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, - ) - # This parameter will view into self.kv_b_proj.weight after loading weights. - # For dummy weight initialization, this parameter is initialized with empty tensor. - # Used in forward_absorption only - self.v_b_proj = nn.Parameter( - torch.empty( - (self.num_heads_tp_cp, self.v_head_dim, self.kv_lora_rank), - dtype=dtype, - ), - requires_grad=False, - ) - mapping_o = Mapping( world_size=pp_size * dp_size * tp_size * cp_size, tp_size=tp_size * cp_size, @@ -701,49 +441,6 @@ def __init__( enable_attention_dp=self.mapping.enable_attention_dp, ) self.mapping_o = mapping_o - if self.is_deepseek_v4: - self.o_a_proj = nn.Parameter( - torch.empty( - ( - self.n_local_groups, - self.o_lora_rank, - self.num_heads * self.qk_head_dim // self.num_groups, - ), - dtype=dtype, - ), - requires_grad=False, - ) - self.o_b_proj = Linear( - self.num_groups * self.o_lora_rank, - self.hidden_size, - bias=False, - dtype=dtype, - mapping=mapping_o, - tensor_parallel_mode=TensorParallelMode.ROW, - quant_config=quant_config, - skip_create_weights_in_init=config.skip_create_weights_in_init, - reduce_output=reduce_output, - allreduce_strategy=config.allreduce_strategy, - force_dynamic_quantization=config.force_dynamic_quantization, - use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, - use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, - ) - else: - self.o_proj = Linear( - self.num_key_value_heads * self.v_head_dim, - self.hidden_size, - bias=self.dense_bias, - dtype=dtype, - mapping=mapping_o, - tensor_parallel_mode=TensorParallelMode.ROW, - quant_config=quant_config, - skip_create_weights_in_init=config.skip_create_weights_in_init, - reduce_output=reduce_output, - allreduce_strategy=config.allreduce_strategy, - force_dynamic_quantization=config.force_dynamic_quantization, - use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, - use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, - ) def yarn_get_mscale(scale=1, mscale=1): if scale <= 1: @@ -755,22 +452,46 @@ def yarn_get_mscale(scale=1, mscale=1): mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) q_scaling = 1.0 / (mscale * mscale) - self.has_dsv4_indexer = ( - self.is_deepseek_v4 - and layer_idx is not None - and sparse_params is not None - and sparse_params.compress_ratios[layer_idx] == 4 + self.aux_stream = aux_stream + self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] + self.kv_a_layernorm = RMSNorm(hidden_size=self.kv_lora_rank, dtype=dtype, eps=rms_norm_eps) + self.kv_b_proj = Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=bias, + dtype=dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + quant_config=quant_config, + skip_create_weights_in_init=config.skip_create_weights_in_init, + allreduce_strategy=config.allreduce_strategy, + force_dynamic_quantization=config.force_dynamic_quantization, + use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, + use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, + ) + self.v_b_proj = nn.Parameter( + torch.empty( + (self.num_heads_tp_cp, self.v_head_dim, self.kv_lora_rank), + dtype=dtype, + ), + requires_grad=False, ) - self.indexer_stream = None - self.indexer_aux_stream = None - self.compressor_stream = None - if self.has_dsv4_indexer and aux_stream is not None: - self.indexer_stream = torch.cuda.Stream(device=aux_stream.device) - self.indexer_aux_stream = torch.cuda.Stream(device=aux_stream.device) - self.compressor_stream = torch.cuda.Stream(device=aux_stream.device) - mqa_aux_stream = ( - self.indexer_aux_stream if self.indexer_aux_stream is not None else aux_stream + self.o_proj = Linear( + self.num_key_value_heads * self.v_head_dim, + self.hidden_size, + bias=self.dense_bias, + dtype=dtype, + mapping=mapping_o, + tensor_parallel_mode=TensorParallelMode.ROW, + quant_config=quant_config, + skip_create_weights_in_init=config.skip_create_weights_in_init, + reduce_output=reduce_output, + allreduce_strategy=config.allreduce_strategy, + force_dynamic_quantization=config.force_dynamic_quantization, + use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, + use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, ) + self.attention_output_hidden_size = self.o_proj.in_features self.mqa = create_attention( config.attn_backend, @@ -778,7 +499,7 @@ def yarn_get_mscale(scale=1, mscale=1): self.num_heads_tp, head_dim=self.kv_lora_rank + self.qk_rope_head_dim, num_kv_heads=1, - pos_embd_params=pos_embd_params, + pos_embd_params=self.pos_embd_params, quant_config=quant_config, q_scaling=q_scaling, is_mla_enable=True, @@ -786,26 +507,58 @@ def yarn_get_mscale(scale=1, mscale=1): kv_lora_rank=self.kv_lora_rank, qk_nope_head_dim=self.qk_nope_head_dim, qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim if self.is_deepseek_v4 else self.kv_lora_rank, hidden_size=self.hidden_size, predicted_tokens_per_seq=self.predicted_tokens_per_seq, skip_create_weights_in_init=config.skip_create_weights_in_init, - sparse_params=sparse_params, + v_head_dim=self.kv_lora_rank, + sparse_params=self.sparse_params, dtype=dtype, - aux_stream=mqa_aux_stream, - rope_append=not self.is_deepseek_v4, + aux_stream=aux_stream, + rope_append=True, ) - self.compressor = getattr(self.mqa, "compressor", None) - self.indexer = getattr(self.mqa, "indexer", None) - self.softmax_scale = 1.0 / (math.sqrt(self.qk_head_dim) * q_scaling) + # MHA is the dense expanded-KV path. Algorithms that do not use it + # remove it in their initialization hook. + mha_sparse_params = self.sparse_params if not self.sparse_attn_hooks else None + self.mha = create_attention( + config.attn_backend, + self.layer_idx, + self.num_heads_tp, + head_dim=self.qk_head_dim, + num_kv_heads=self.num_key_value_heads_tp, + pos_embd_params=self.pos_embd_params, + quant_config=quant_config, + q_scaling=q_scaling, + is_mla_enable=True, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + qk_nope_head_dim=self.qk_nope_head_dim, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + predicted_tokens_per_seq=self.predicted_tokens_per_seq, + skip_create_weights_in_init=config.skip_create_weights_in_init, + sparse_params=mha_sparse_params, + ) - self.aux_stream = aux_stream - self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] - self.dsv4_overlap_start_event = torch.cuda.Event() - self.dsv4_compressor_start_event = torch.cuda.Event() - self.dsv4_compressor_event = torch.cuda.Event() - self.dsv4_indexer_event = torch.cuda.Event() + if initialize_sparse_attn := self.sparse_attn_hooks.initialize_sparse_attn: + initialize_sparse_attn( + self, + config=config, + mapping=mapping, + mapping_o=mapping_o, + rms_norm_eps=rms_norm_eps, + quant_config=quant_config, + q_scaling=q_scaling, + bias=bias, + dtype=dtype, + reduce_output=reduce_output, + aux_stream=aux_stream, + ) + + if self.mqa is None: + raise RuntimeError("MLA requires a non-null MQA attention backend") + + self.softmax_scale = 1.0 / (math.sqrt(self.qk_head_dim) * q_scaling) self.rope_fusion = self.mqa.support_fused_rope() self.rotary_emb = None @@ -817,61 +570,13 @@ def yarn_get_mscale(scale=1, mscale=1): is_neox=pos_embd_params.is_neox, ) - if self.is_deepseek_v4: - self.inverse_rotary_emb = RotaryEmbedding( - pos_embd_params.rope, - head_dim=self.qk_rope_head_dim, - is_neox=pos_embd_params.is_neox, - inverse=True, - ) - - # Short-sequence MHA optimization for DSA models: - # For short prefill sequences, use MHA (kv_b_proj expansion + standard - # attention) instead of the absorption path, which has overhead from - # extra BMMs and larger head_dim (kv_lora_rank + qk_rope_head_dim). - # Only active when rope_fusion is True (DSA with TrtllmAttention). - _threshold_str = os.environ.get("TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD", "0") - try: - self.short_seq_mha_threshold = int(_threshold_str) - except ValueError as err: - raise ValueError( - f"TRTLLM_MLA_SHORT_SEQ_MHA_THRESHOLD must be an integer, got '{_threshold_str}'" - ) from err - - # MHA attention backend: used by non-DSA (standard MLA) and optionally - # by DSA for the short-seq path (dense attention, no sparse config). - _short_seq_mha = ( - self.is_dsa and self.short_seq_mha_threshold > 0 and not self.apply_rotary_emb - ) - if (not self.is_dsa or _short_seq_mha) and not self.is_deepseek_v4: - mha_sparse_params = None if _short_seq_mha else sparse_params - self.mha = create_attention( - config.attn_backend, - self.layer_idx, - self.num_heads_tp, - head_dim=self.qk_head_dim, - num_kv_heads=self.num_key_value_heads_tp, - pos_embd_params=pos_embd_params, - quant_config=quant_config, - q_scaling=q_scaling, - is_mla_enable=True, - q_lora_rank=self.q_lora_rank, - kv_lora_rank=self.kv_lora_rank, - qk_nope_head_dim=self.qk_nope_head_dim, - qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim, - predicted_tokens_per_seq=self.predicted_tokens_per_seq, - skip_create_weights_in_init=config.skip_create_weights_in_init, - sparse_params=mha_sparse_params, - ) - else: - self.mha = None - self.llama_4_scaling = False if hasattr(config.pretrained_config, "llama_4_scaling"): self.llama_4_scaling = True self.floor_scale = getattr( - config.pretrained_config.llama_4_scaling, "original_max_position_embeddings", 8192 + config.pretrained_config.llama_4_scaling, + "original_max_position_embeddings", + 8192, ) self.attn_scale = getattr(config.pretrained_config.llama_4_scaling, "beta", 0.1) @@ -881,41 +586,35 @@ def yarn_get_mscale(scale=1, mscale=1): def create_weights(self): # self.mha/mqa has no weights but has states that are related to # quant_config, which could be modified after __init__. - # self.mha is non-None for non-DSA models (standard MHA) and for DSA - # models when the short-seq MHA optimization is active. if self.mha is not None: self.mha.update_quant_config(self.quant_config) self.mqa.update_quant_config(self.quant_config) # Although we use FP8 MLA for context/generation phase, the output is still in BF16 self.out_scale = None + if create_sparse_attn_weights := self.sparse_attn_hooks.create_sparse_attn_weights: + create_sparse_attn_weights(self) + self._weights_transformed = False + return # k_b_proj_trans's dtype must be consistent with self.kv_b_proj, # which can be modified after __init__ - if self.is_deepseek_v4: - has_fp8_block_scales = ( - self.o_b_proj.quant_config - and self.o_b_proj.quant_config.quant_mode.has_fp8_block_scales() - ) - else: - has_fp8_block_scales = ( - self.kv_b_proj.quant_config - and self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales() - ) - - mla_weight_dtype = torch.float8_e4m3fn if has_fp8_block_scales else self.dtype - self.k_b_proj_trans = nn.Parameter( - torch.empty( - (self.num_heads_tp, self.kv_lora_rank, self.qk_nope_head_dim), - dtype=mla_weight_dtype, - ), - requires_grad=False, - ) + has_fp8_block_scales = bool( + self.kv_b_proj.quant_config + and self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales() + ) + mla_weight_dtype = torch.float8_e4m3fn if has_fp8_block_scales else self.dtype + self.k_b_proj_trans = nn.Parameter( + torch.empty( + (self.num_heads_tp, self.kv_lora_rank, self.qk_nope_head_dim), + dtype=mla_weight_dtype, + ), + requires_grad=False, + ) self.k_b_proj_trans_dequant = None self.v_b_proj_dequant = None - self.o_a_proj_dequant = None - if has_fp8_block_scales and not self.is_deepseek_v4: + if has_fp8_block_scales: self.k_b_proj_trans_scale = nn.Parameter( torch.empty( ( @@ -956,55 +655,9 @@ def create_weights(self): ), requires_grad=False, ) - elif has_fp8_block_scales: - self.o_a_proj_scale = nn.Parameter( - torch.empty( - ( - self.n_local_groups, - self.o_lora_rank // 128, - self.num_heads * self.qk_head_dim // self.num_groups // 128, - ), - dtype=torch.float32, - ), - requires_grad=False, - ) - if is_sm_100f(): - # DSv4 always keeps o_a_proj in its native FP8 e4m3 form so - # cute_dsl_fp8_bmm_blackwell + fused_inv_rope_fp8_quant can - # consume it directly. Decoupled from - # use_cute_dsl_blockscaling_bmm: only DSv4 has o_a_proj, and - # the fused inv-RoPE -> FP8 quant -> cute-dsl BMM chain is the - # only viable path for it on SM100; gating on the global - # bmm-config flag was conflating two independent kernel - # choices (K/V absorption BMM vs. DSv4 o_a_proj BMM). - if self.is_deepseek_v4: - self.o_a_proj = nn.Parameter( - torch.empty( - ( - self.n_local_groups, - self.o_lora_rank, - self.num_heads * self.qk_head_dim // self.num_groups, - ), - dtype=torch.float8_e4m3fn, - ), - requires_grad=False, - ) - else: - self.o_a_proj_dequant = nn.Parameter( - torch.empty( - ( - self.n_local_groups, - self.o_lora_rank, - self.num_heads * self.qk_head_dim // self.num_groups, - ), - dtype=self.dtype, - ), - requires_grad=False, - ) else: self.k_b_proj_trans_scale = None self.v_b_proj_scale = None - self.o_a_proj_scale = None self._weights_transformed = False def apply_rope( @@ -1021,68 +674,6 @@ def apply_rope( q[..., self.qk_nope_head_dim :] = q_pe.view(-1, self.num_heads_tp, self.qk_rope_head_dim) return k_pe - def _deepseek_v4_q_b_layernorm(self, q: torch.Tensor) -> torch.Tensor: - assert q.dim() == 2 and q.shape[1] == self.num_heads_tp * self.qk_head_dim - return torch.ops.trtllm.deepseek_v4_q_norm( - q, self.num_heads_tp, self.qk_head_dim, float(self.q_b_layernorm.variance_epsilon) - ) - - def _is_fused_q_fp8_quant_enabled(self, num_generations: int = 0) -> bool: - # Context-only batches: the fused path leaves a placeholder bf16 q_buf - # that forward_generation_sparse_mla would read uninitialized, so - # mixed/gen batches must take the legacy unfused path. - # `TRTLLM_DISABLE_FUSED_Q_FP8_QUANT=1` opts back into the legacy - # two-kernel Q-quant path as a kill switch. - if os.environ.get("TRTLLM_DISABLE_FUSED_Q_FP8_QUANT", "0") == "1": - return False - if not self.is_deepseek_v4: - return False - if self.qk_head_dim != 512 or self.kv_lora_rank != 448: - return False - if num_generations > 0: - return False - return bool(getattr(self.mqa, "has_fp8_kv_cache", False)) - - def _deepseek_v4_q_b_layernorm_fused_fp8(self, q_proj: torch.Tensor): - # Returns (placeholder_q, quant_q_buffer, q_pe, quant_scale_qkv). - # `placeholder_q` keeps the [num_tokens, num_heads*head_dim] bf16 layout - # the downstream `forward_absorption_context` needs for its `q.shape[0]` - # check and `q.view().split()` call. Its contents are never read on the - # fused FP8 path: the nope segment lives in `quant_q_buffer`, the rope - # segment is passed in `q_pe`, and the split's `q_nope`/`q_pe` outputs - # are either overridden by the caller or discarded by the DSv4 branch. - # Reusing `q_proj` (q_b_proj output) avoids a ~num_tokens x hidden bf16 - # allocation per forward. - assert q_proj.dim() == 2 - assert q_proj.shape[1] == self.num_heads_tp * self.qk_head_dim - if getattr(self, "_quant_scale_qkv", None) is None: - self._quant_scale_qkv = torch.tensor([1.0], dtype=torch.float32, device=q_proj.device) - # q_pe is 3D so thop.attention's sparse-MLA context branch passes its - # q_pe->dim() == 3 check; the kernel op consumes the flat 2D view. - num_tokens = q_proj.shape[0] - rope_dim = self.qk_head_dim - self.kv_lora_rank - quant_q_buffer = q_proj.new_empty( - (num_tokens, self.num_heads_tp * self.qk_head_dim), dtype=torch.float8_e4m3fn - ) - q_pe = q_proj.new_empty((num_tokens, self.num_heads_tp, rope_dim)) - torch.ops.trtllm.deepseek_v4_q_norm_fused_fp8( - q_proj, - quant_q_buffer, - q_pe.view(num_tokens, self.num_heads_tp * rope_dim), - self.num_heads_tp, - self.qk_head_dim, - self.kv_lora_rank, - float(self.q_b_layernorm.variance_epsilon), - self._quant_scale_qkv, - ) - # Both buffers must be live for the fused path; the downstream - # absorption-context op switches on `quant_scale_qkv is not None` - # to enable the C++ fusion (see trtllm.py `thop.attention` call). - assert self._quant_scale_qkv is not None, ( - "fused FP8-Q quant requires _quant_scale_qkv to be set" - ) - return q_proj, quant_q_buffer, q_pe, self._quant_scale_qkv - def _attn_forward_gen( self, attn_backend: AttentionBackend, @@ -1139,10 +730,6 @@ def _attn_forward_gen( return attn_output def create_output(self, hidden_states: torch.Tensor, num_contexts: int): - # Upstream POST_MoE/MLP fusion (or attention-DP no-fusion fold) may pass - # an Fp4QuantizedTensor here; unpack to the BF16 view for sizing. The - # producing fold must have requested return_norm_out so the BF16 view - # is present (folds feeding an MLA always do). if isinstance(hidden_states, Fp4QuantizedTensor): assert hidden_states.unquantized_hidden_states is not None, ( "MLA.create_output received an Fp4QuantizedTensor without a " @@ -1151,11 +738,17 @@ def create_output(self, hidden_states: torch.Tensor, num_contexts: int): ) hidden_states = hidden_states.unquantized_hidden_states num_tokens = hidden_states.shape[0] - if self.is_deepseek_v4: - hidden_size = self.num_heads_tp_cp * self.v_head_dim - else: - hidden_size = self.o_proj.in_features - return hidden_states.new_empty([num_tokens, hidden_size], dtype=hidden_states.dtype) + return hidden_states.new_empty( + [num_tokens, self.attention_output_hidden_size], dtype=hidden_states.dtype + ) + + def _create_outputs( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> list[torch.Tensor]: + """Create the standard output and any algorithm-specific output buffers.""" + if prepare_outputs := self.sparse_attn_hooks.prepare_sparse_attn_outputs: + return prepare_outputs(self, hidden_states, attn_metadata) + return [self.create_output(hidden_states, attn_metadata.num_contexts)] def _attention_scaling(self, q, position_ids): def _get_attn_scale(position_ids: torch.Tensor) -> torch.Tensor: @@ -1168,186 +761,9 @@ def _get_attn_scale(position_ids: torch.Tensor) -> torch.Tensor: q = (q * attn_scale).to(q.dtype) return q - def _should_use_dsv4_epilogue_fusion(self, num_contexts: int, num_generations: int) -> bool: - if self._disable_dsv4_epilogue_fusion: - return False - if not self.is_deepseek_v4: - return False - if num_contexts == 0 and num_generations == 0: - return False - if num_contexts > 0 and num_generations > 0: - # Context and generation use separate FMHA calls, but the fused - # buffers do not carry token offsets for a mixed batch. - return False - if self.mapping.has_cp_helix(): - return False - if not is_sm_100f(): - return False - if not getattr(self.mapping, "enable_attention_dp", False): - return False - if self.num_heads != 128 or self.num_heads_tp != 128: - return False - if getattr(self.mqa, "sparse_params", None) is None: - return False - if not getattr(self.mqa, "has_fp8_kv_cache", False): - return False - if self.o_a_proj.dtype != torch.float8_e4m3fn: - return False - if self.kv_lora_rank != 448 or self.qk_rope_head_dim != 64: - return False - if self.qk_head_dim != 512 or self.v_head_dim != 512: - return False - if self.n_local_groups <= 0 or self.num_heads_tp % self.n_local_groups != 0: - return False - return not self.inverse_rotary_emb.is_neox - - def _create_dsv4_epilogue_buffers( - self, q: torch.Tensor, num_tokens: int - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.n_local_groups <= 0 or self.num_heads_tp % self.n_local_groups != 0: - raise ValueError( - "DSv4 fused epilogue requires num_heads_tp to be divisible by n_local_groups." - ) - heads_per_group = self.num_heads_tp // self.n_local_groups - scale_buf_m = (num_tokens + 3) // 4 * 4 - fp8_o = q.new_empty( - (self.n_local_groups, num_tokens, heads_per_group * self.v_head_dim), - dtype=torch.float8_e4m3fn, - ) - output_sf = q.new_empty( - ( - self.n_local_groups, - heads_per_group * (self.v_head_dim // 128), - scale_buf_m, - ), - dtype=torch.float32, - ) - return fp8_o, output_sf - - def _validate_dsv4_epilogue_buffers( - self, - num_tokens: int, - dsv4_epilogue_output: tuple[torch.Tensor, torch.Tensor], - ) -> tuple[torch.Tensor, torch.Tensor]: - fp8_o, output_sf = dsv4_epilogue_output - scale_buf_m = (num_tokens + 3) // 4 * 4 - if fp8_o.shape[1] != num_tokens or output_sf.shape[2] != scale_buf_m: - raise RuntimeError("Invalid DSv4 fused epilogue buffers for current token count.") - return fp8_o, output_sf - - def _deepseek_v4_o_proj( - self, - attn_out_latent: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - position_ids: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if isinstance(attn_out_latent, tuple): - attn_fp8, attn_scale = attn_out_latent - num_tokens = attn_fp8.shape[1] - o_lora = torch.empty( - [num_tokens, self.n_local_groups, self.o_lora_rank], - device=attn_fp8.device, - dtype=self.dtype, - ) - torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( - attn_fp8, - self.o_a_proj, - attn_scale, - self.o_a_proj_scale, - o_lora.transpose(0, 1), - ) - return self.o_b_proj(o_lora.flatten(1)) - - assert position_ids is not None - num_tokens = attn_out_latent.shape[0] - attn_out_latent = attn_out_latent.view(num_tokens, self.num_heads_tp, -1) - - # When o_a_proj is FP8 on SM100 (which is always the case for DSv4 - # under FP8 block-scales after init), fuse the inverse-RoPE into the - # FP8-quant epilogue (vLLM-ported Triton kernel) and call - # cute_dsl_fp8_bmm_blackwell directly. Saves one BF16 read+write of - # the latent vs the mla_rope_inplace + - # fp8_batched_quantize_1x128_permute102 pair. Decoupled from - # use_cute_dsl_blockscaling_bmm (which gates the separate K/V - # absorption BMM kernel choice). - fused_inv_rope_fp8 = self.o_a_proj.dtype == torch.float8_e4m3fn and is_sm_100f() - if fused_inv_rope_fp8: - heads_per_group = self.num_heads_tp // self.n_local_groups - attn_fp8, attn_scale = torch.ops.trtllm.fused_inv_rope_fp8_quant_vllm_port( - attn_out_latent, - position_ids.view(-1), - self.inverse_rotary_emb.rotary_cos_sin, - self.n_local_groups, - heads_per_group, - self.qk_nope_head_dim, - self.qk_rope_head_dim, - 128, - self.inverse_rotary_emb.is_neox, - ) - o_lora = torch.empty( - [num_tokens, self.n_local_groups, self.o_lora_rank], - device=attn_out_latent.device, - dtype=self.dtype, - ) - torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell( - attn_fp8, self.o_a_proj, attn_scale, self.o_a_proj_scale, o_lora.transpose(0, 1) - ) - o_lora = o_lora.flatten(1) - return self.o_b_proj(o_lora) - - # Fused in-place inverse RoPE on the rope portion of each head - torch.ops.trtllm.mla_rope_inplace( - attn_out_latent, - position_ids.view(-1), - self.inverse_rotary_emb.rotary_cos_sin, - self.num_heads_tp, - self.qk_nope_head_dim, - self.qk_rope_head_dim, - True, - self.inverse_rotary_emb.is_neox, - ) - - # Output projections - o_lora = torch.empty( - [num_tokens, self.n_local_groups, self.o_lora_rank], - device=attn_out_latent.device, - dtype=attn_out_latent.dtype, - ) - if self.o_a_proj.dtype == torch.bfloat16: - # dim = head_dim * num_head // num_group - # [num_groups, num_tokens, dim] x [num_groups, dim, o_lora_rank] - # -> [num_groups, num_tokens, o_lora_rank] - torch.ops.trtllm.bmm_out( - attn_out_latent.view(num_tokens, self.n_local_groups, -1).transpose(0, 1), - self.o_a_proj.transpose(1, 2), - o_lora.transpose(0, 1), - ) - elif self.o_a_proj.dtype == torch.float8_e4m3fn: - fp8_block_scaling_bmm_out( - attn_out_latent.view(num_tokens, self.n_local_groups, -1), - self.o_a_proj, - self.o_a_proj_scale, - o_lora.transpose(0, 1), - self.o_a_proj_dequant, - self.use_cute_dsl_blockscaling_bmm, - ) - else: - raise NotImplementedError(f"Missing bmm impl for dtype: {self.o_a_proj.dtype}.") - o_lora = o_lora.flatten(1) - output = self.o_b_proj(o_lora) - return output - def _resolve_qa_fused_scale(self): - """Lazily decide whether the residual-less q_a_layernorm -> q_b_proj - NVFP4 fusion can run, caching q_b_proj's static input_scale. - - Returns the input_scale tensor when eligible, else None. Eligible iff: - (1) model is non-lite (has a distinct q_a_layernorm / q_b_proj), - (2) q_b_proj is static-NVFP4 (calibrated input_scale, no - pre_quant_scale / AWQ, not forced-dynamic), - (3) q_a_layernorm is a plain (non-gemma, non-quantizing) RMSNorm. - """ + """Resolve whether q_a_layernorm can fuse q_b_proj NVFP4 input quantization.""" if self._qa_fused_scale is not None: - # Already resolved: tensor (enabled) or False (disabled). return self._qa_fused_scale if self._qa_fused_scale is not False else None eligible = False if ( @@ -1357,42 +773,20 @@ def _resolve_qa_fused_scale(self): ): qb = self.q_b_proj qa = self.q_a_layernorm - # q_b_proj must be static-NVFP4 (shared predicate) and q_a_layernorm - # a plain (non-gemma, non-quantizing) RMSNorm. if is_static_nvfp4_input_eligible(qb) and not qa.use_gemma and not qa.is_nvfp4: eligible = True if eligible: - # Mark q_a_layernorm as an NVFP4 norm and attach q_b_proj's static - # input_scale, so RMSNorm.forward folds the residual-less - # q_a_layernorm + q_b_proj NVFP4 input-quant into one kernel (the - # size gate routes this N=q_lora_rank<2048 edge to reduce_fusion). self.q_a_layernorm.is_nvfp4 = True self.q_a_layernorm.nvfp4_scale = qb.input_scale self._qa_fused_scale = qb.input_scale if eligible else False return self._qa_fused_scale if eligible else None def _q_a_layernorm_maybe_fused(self, q: torch.Tensor, return_norm_out: bool = False): - """Apply q_a_layernorm, optionally folding the q_b_proj NVFP4 input - quantize into the same kernel. - - When fusion is disabled, returns the BF16/FP16 normed tensor (and, if - return_norm_out, the same tensor again so callers have a stable arity). - - When fusion is enabled, returns an Fp4QuantizedTensor (consumed by - q_b_proj without re-quantizing). If return_norm_out is set, also returns - the BF16 post-RMSNorm value, needed by the DSA indexer's - pre_indexer_proj which requires a float q.""" + """Apply q_a_layernorm and fuse static NVFP4 input quantization when eligible.""" scale = self._resolve_qa_fused_scale() if scale is None: normed = self.q_a_layernorm(q) return (normed, normed) if return_norm_out else normed - # q is typically the leading q_lora_rank columns of the - # kv_a_proj_with_mqa output split, i.e. a column slice whose last dim is - # unit-stride but whose row pitch is the full projection width. The - # fused RMSNorm+NVFP4-quant kernel reads that strided layout directly, - # so no .contiguous() copy is needed. _resolve_qa_fused_scale has - # attached .nvfp4_scale to q_a_layernorm, so RMSNorm.forward folds the - # residual-less RMSNorm + NVFP4 quant automatically. return self.q_a_layernorm(q, return_norm_out=return_norm_out) def forward_impl( @@ -1400,534 +794,82 @@ def forward_impl( position_ids: Optional[torch.Tensor], hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, - output: torch.Tensor, + attn_output: list[torch.Tensor], latent_cache_gen: Optional[torch.Tensor] = None, ) -> None: - """ - Forward pass for the MLA module. Writes result into output tensor in-place. - - Args: - position_ids (Optional[torch.IntTensor]): The position IDs. - hidden_states (torch.Tensor): The hidden states. - attn_metadata (AttentionMetadata): The attention metadata. - output (torch.Tensor): The output tensor to write results into. - latent_cache_gen (Optional[torch.Tensor]): The latent cache used in generation. - """ - # split q, k, v into context and gen batches - num_contexts = attn_metadata.num_contexts - num_generations = attn_metadata.num_generations - num_ctx_tokens = attn_metadata.num_ctx_tokens - num_tokens = attn_metadata.num_tokens - - hidden_states = _slice_hidden_states_to_num_tokens(hidden_states, num_tokens) - if position_ids is not None: - position_ids = position_ids[..., :num_tokens] - - if self.is_lite: - compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( - [self.kv_lora_rank, self.qk_rope_head_dim], -1 - ) - compressed_kv = self.kv_a_layernorm(compressed_kv) - q = hidden_states - else: - q, compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( - [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1 - ) - - q, compressed_kv = maybe_execute_in_parallel( - lambda: self._q_a_layernorm_maybe_fused(q), - lambda: self.kv_a_layernorm(compressed_kv), - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - ) - - q, latent_cache = maybe_execute_in_parallel( - lambda: self.q_b_proj(q), - lambda: torch.concat([compressed_kv, k_pe], dim=-1), - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - ) - - assert q.shape[0] == num_tokens, ( - f"Expect q.shape[0] to be {num_tokens}, but got {q.shape[0]}" - ) - - assert output is not None, "output must be provided" - - if num_contexts > 0: - q_ctx = q[:num_ctx_tokens, ...] - compressed_kv_ctx = compressed_kv[:num_ctx_tokens, ...] - k_pe_ctx = k_pe[:num_ctx_tokens, ...] - latent_cache_ctx = latent_cache[:num_ctx_tokens, ...] - if self.apply_rotary_emb: - assert position_ids is not None - # position_ids spans [ctx..., gen...] in mixed batches; slice to - # match q_ctx/k_pe_ctx so external RoPE uses ctx positions. - ctx_position_ids = position_ids[..., :num_ctx_tokens] - k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids) - # External RoPE is only used by backends that do not handle - # fused RoPE internally, so keep latent_cache in sync. - latent_cache_ctx = torch.cat([compressed_kv_ctx, k_pe_ctx], dim=-1) - - if self.llama_4_scaling: - q_ctx = self._attention_scaling(q_ctx, position_ids[..., :num_ctx_tokens]) - - self.forward_context( - q_ctx, - compressed_kv_ctx, - k_pe_ctx, + """Run the dense or sparse implementation of the shared MLA module.""" + if forward_sparse_attn := self.sparse_attn_hooks.forward_sparse_attn: + forward_sparse_attn( + self, position_ids, + hidden_states, attn_metadata, - output[:num_ctx_tokens, :], - latent_cache_ctx, - ) - - if num_generations > 0: - q_gen = q[num_ctx_tokens:, ...] - compressed_kv_gen = compressed_kv[num_ctx_tokens:, ...] - k_pe_gen = k_pe[num_ctx_tokens:, ...] - if latent_cache_gen is None: - latent_cache_gen = latent_cache[num_ctx_tokens:, ...] - - if self.llama_4_scaling: - q_gen = self._attention_scaling(q_gen, position_ids[..., num_ctx_tokens:]) - - self.forward_absorption_generation( - q_gen, - compressed_kv_gen, - k_pe_gen, - attn_metadata, - output[num_ctx_tokens:num_tokens, :], - position_ids=position_ids, - latent_cache=latent_cache_gen, + attn_output, ) + return - def forward_impl_with_dsa( - self, - position_ids: Optional[torch.Tensor], - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - output: torch.Tensor, - ) -> None: - """ - Forward pass for the MLA module with DSA (always in MQA mode). - Writes result into output tensor in-place. - - Delegates to forward_dsa_proj (token-wise projections) followed by - forward_dsa_attn (batch-dependent attention dispatch). - - Args: - position_ids (Optional[torch.IntTensor]): The position IDs. - hidden_states (torch.Tensor): The hidden states. - attn_metadata (AttentionMetadata): The attention metadata. - output (torch.Tensor): The output tensor to write results into. - """ - proj_outputs = self.forward_dsa_proj(position_ids, hidden_states, attn_metadata) - q, compressed_kv, k_pe, latent_cache = proj_outputs[:4] - indexer_intermediates = proj_outputs[4:] - self.forward_dsa_attn( - q, - compressed_kv, - k_pe, - latent_cache, - indexer_intermediates, + self._forward_impl( position_ids, - attn_metadata, - output, - ) - - def forward_dsa_proj( - self, - position_ids: Optional[torch.Tensor], - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> List[torch.Tensor]: - """Token-wise projections for DSA MLA (CUDA-graph-capturable Op 1). - - Runs kv_a_proj, layernorms, q_b_proj, and conditionally - indexer.pre_indexer_proj(). - - IMPORTANT: This method must NOT slice tensors by num_tokens or - access batch-specific metadata, so that all operations are - unconditionally straight-line for CUDA graph capture. Slicing - to num_tokens happens in forward_dsa_attn (Op 2, outside graph). - - Returns [q, compressed_kv, k_pe, latent_cache] when short-MHA - handles all tokens (eager only), or - [q, compressed_kv, k_pe, latent_cache, q_fp8, k_fp8, k_scale, - weights, q_scale] when the indexer runs. q_scale is unused on the - FP8 path but always present so CUDA graph capture sees a stable - 9-tensor shape regardless of indexer dtype. - """ - assert self.mqa is not None, "DSA is only supported in MQA mode" - - q, compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( - [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1 - ) - - # When the q_a_layernorm -> q_b_proj NVFP4 fusion is active, q becomes an - # Fp4QuantizedTensor consumed by q_b_proj; qr (fed to the indexer) needs - # the BF16 post-RMSNorm value, so request return_norm_out here. - q_pair, compressed_kv = maybe_execute_in_parallel( - lambda: self._q_a_layernorm_maybe_fused(q, return_norm_out=True), - lambda: self.kv_a_layernorm(compressed_kv), - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - ) - q, qr = q_pair - latent_cache = torch.concat([compressed_kv, k_pe], dim=-1) - - q = self.q_b_proj(q) - - use_short_mha_for_ctx = self._should_use_short_mha(attn_metadata, position_ids) - - # Skip the indexer when the short MHA path handles all context - # tokens and there are no generation tokens. - if use_short_mha_for_ctx and attn_metadata.num_generations == 0: - return [q, compressed_kv, k_pe, latent_cache] - - # DSA "shared" layer: no indexer; reuses the previous full layer's - # top-k (in forward_dsa_attn), so skip the projection. - if self.mqa.indexer is None: - return [q, compressed_kv, k_pe, latent_cache] - - # pre_indexer_proj is the CUDA-graph-safe portion: pure token-wise - # compute (cublas_mm, rope, FP4/FP8 quantize, weight scaling) with no - # access to batch-specific metadata or the k cache. Returns q_scale - # as a 5th element so the FP4 dispatch can forward it to the kernel; - # the FP8 path ignores it in forward_dsa_attn. - q_fp8, k_fp8, k_scale, weights, q_scale = self.mqa.indexer.pre_indexer_proj( - qr, hidden_states, position_ids - ) - - return [q, compressed_kv, k_pe, latent_cache, q_fp8, k_fp8, k_scale, weights, q_scale] - - def forward_dsa_attn( - self, - q: torch.Tensor, - compressed_kv: torch.Tensor, - k_pe: torch.Tensor, - latent_cache: torch.Tensor, - indexer_intermediates: List[torch.Tensor], - position_ids: Optional[torch.Tensor], - attn_metadata: AttentionMetadata, - output: torch.Tensor, - ) -> None: - """Batch-structure-dependent attention for DSA MLA (Op 2, not graph-captured). - - indexer_intermediates is [q_fp8, k_fp8, k_scale, weights, q_scale] - when the indexer ran in Op 1, or [] when short-MHA handled all tokens. - - All num_tokens slicing happens here (not in Op 1) because - num_tokens comes from batch-specific metadata and must not be - baked into CUDA graph capture. - """ - num_contexts = attn_metadata.num_contexts - num_generations = attn_metadata.num_generations - num_ctx_tokens = attn_metadata.num_ctx_tokens - num_tokens = attn_metadata.num_tokens - - # Slice Op 1 outputs to actual num_tokens (Op 1 operates on the - # full padded tensor for CUDA graph compatibility). - q = q[:num_tokens, ...] - compressed_kv = compressed_kv[:num_tokens, ...] - k_pe = k_pe[:num_tokens, ...] - latent_cache = latent_cache[:num_tokens, ...] - if position_ids is not None: - position_ids = position_ids[..., :num_tokens] - - use_short_mha_for_ctx = num_contexts > 0 and self._should_use_short_mha( - attn_metadata, position_ids - ) - - if use_short_mha_for_ctx and num_generations == 0: - topk_indices = None - elif self.mqa.indexer is None: - # DSA "shared" layer: reuse the previous full layer's top-k. These - # are local token positions, so they are layer-agnostic; each layer - # applies its own paged-KV transform downstream. - topk_indices = getattr(attn_metadata, "shared_topk_indices", None) - assert topk_indices is not None, ( - "DSA shared layer has no top-k from a preceding full indexer " - "layer; check the index_topk_pattern/freq schedule." - ) - else: - q_fp8, k_fp8, k_scale, weights, q_scale = indexer_intermediates - # Slice indexer intermediates to actual num_tokens (they were - # computed on the full padded tensor in Op 1). - q_fp8 = q_fp8[:num_tokens, ...] - k_fp8 = k_fp8[:num_tokens, ...] - k_scale = k_scale[:num_tokens, ...] - weights = weights[:num_tokens, ...] - q_scale = q_scale[:num_tokens, ...] - topk_indices = self.mqa.indexer.sparse_attn_indexer( - attn_metadata, - q, # only used for shape/device in buffer allocation - q_fp8, - k_fp8, - k_scale, - weights, - q_scale=q_scale, - ) - # Stash for subsequent DSA "shared" layers (full -> shared reuse); - # unused by a dense per-layer indexer. - attn_metadata.shared_topk_indices = topk_indices - - assert output is not None, "output must be provided" - - if num_contexts > 0: - q_ctx = q[:num_ctx_tokens, ...] - compressed_kv_ctx = compressed_kv[:num_ctx_tokens, ...] - k_pe_ctx = k_pe[:num_ctx_tokens, ...] - latent_cache_ctx = latent_cache[:num_ctx_tokens, ...] - ctx_position_ids = ( - position_ids[..., :num_ctx_tokens] if position_ids is not None else None - ) - if self.apply_rotary_emb: - assert ctx_position_ids is not None - k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids) - - self.forward_context_sparse_mla( - q_ctx, - compressed_kv_ctx, - k_pe_ctx, - attn_metadata, - output[:num_ctx_tokens, :], - latent_cache_ctx, - topk_indices=topk_indices[:num_ctx_tokens, :] if topk_indices is not None else None, - position_ids=ctx_position_ids, - ) - - if num_generations > 0: - q_gen = q[num_ctx_tokens:, ...] - compressed_kv_gen = compressed_kv[num_ctx_tokens:, ...] - k_pe_gen = k_pe[num_ctx_tokens:, ...] - latent_cache_gen = latent_cache[num_ctx_tokens:, ...] - gen_position_ids = ( - position_ids[..., num_ctx_tokens:num_tokens] if position_ids is not None else None - ) - if self.apply_rotary_emb: - assert gen_position_ids is not None - k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids) - - self.forward_generation_sparse_mla( - q_gen, - compressed_kv_gen, - k_pe_gen, - attn_metadata, - output[num_ctx_tokens:num_tokens, :], - latent_cache=latent_cache_gen, - topk_indices=topk_indices[num_ctx_tokens:num_tokens, :], - position_ids=gen_position_ids, - ) + hidden_states, + attn_metadata, + attn_output[0], + latent_cache_gen=latent_cache_gen, + ) - def forward_impl_with_deepseek_v4( + def _forward_impl( self, position_ids: Optional[torch.Tensor], hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, output: torch.Tensor, - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + latent_cache_gen: Optional[torch.Tensor] = None, ) -> None: """ - Forward pass for the MLA module with DeepSeek-V4 (always in MQA mode). + Forward pass for the MLA module. Writes result into output tensor in-place. Args: position_ids (Optional[torch.IntTensor]): The position IDs. hidden_states (torch.Tensor): The hidden states. attn_metadata (AttentionMetadata): The attention metadata. - output (torch.Tensor): Pre-allocated output tensor, written in-place - when epilogue fusion is disabled. - dsv4_epilogue_output: Caller-provided ``(fp8_o, output_sf)`` - buffers, written in-place when epilogue fusion is enabled. + output (torch.Tensor): The output tensor to write results into. + latent_cache_gen (Optional[torch.Tensor]): The latent cache used in generation. """ - assert self.mha is None and self.mqa is not None, ( - "DeepSeek-V4 is only supported in MQA mode" - ) # split q, k, v into context and gen batches num_contexts = attn_metadata.num_contexts num_generations = attn_metadata.num_generations num_ctx_tokens = attn_metadata.num_ctx_tokens num_tokens = attn_metadata.num_tokens - enable_dsv4_epilogue_fusion = dsv4_epilogue_output is not None - if enable_dsv4_epilogue_fusion and ((num_contexts > 0) == (num_generations > 0)): - raise RuntimeError( - "DSv4 epilogue fusion requires a context-only or generation-only batch." - ) - hidden_states = hidden_states[:num_tokens, ...] + hidden_states = _slice_hidden_states_to_num_tokens(hidden_states, num_tokens) if position_ids is not None: position_ids = position_ids[..., :num_tokens] - # TRTLLM_MLA_EXTRA_OVERLAP=1 reorders the V4 attention prologue so the - # outer compressor and the ratio-4 indexer can execute concurrently - # with q_b_proj + q_b_layernorm. The indexer is launched on a - # dedicated stream and still uses a different aux stream for its - # internal q-proj/weights-proj split. - _v4_extra_overlap = ( - os.environ.get("TRTLLM_MLA_EXTRA_OVERLAP", "1") == "1" - and self.compressor is not None - and self.aux_stream is not None - ) - _use_indexer_overlap = ( - _v4_extra_overlap - and do_multi_stream() - and self.indexer is not None - and self.indexer_stream is not None - ) + if self.is_lite: + compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( + [self.kv_lora_rank, self.qk_rope_head_dim], -1 + ) + compressed_kv = self.kv_a_layernorm(compressed_kv) + q = hidden_states + else: + q, compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1 + ) - # Pre-launch the outer compressor on compressor_stream BEFORE - # kv_a_proj_with_mqa. The compressor only reads hidden_states + - # attn_metadata, so it has no data dependency on the kv_a_proj GEMM or - # the downstream q_a/kv_a LN split. A dedicated stream (not aux_stream) - # keeps kv_a_layernorm free to run on aux_stream in parallel. - # _q_branch will be queued onto this same stream further down so it - # runs strictly serial after the compressor; dsv4_compressor_event is - # recorded only at the end of _q_branch, gating the caller's downstream - # waits on both compressor + _q_branch completion. - if _use_indexer_overlap: - self.dsv4_compressor_start_event.record() - with torch.cuda.stream(self.compressor_stream): - self.dsv4_compressor_start_event.wait() - self.compressor(hidden_states, attn_metadata) - - # Pre-launch the qr-independent half of the indexer prepare phase - # (weights_proj + internal compressor + k_cache_update) on the - # indexer's aux stream (self.indexer_aux_stream, wired into the - # indexer module as its aux_stream). Only reads hidden_states + - # attn_metadata, so it can overlap with the kv_a_proj -> LN -> split - # chain on the caller stream and the outer compressor on - # compressor_stream. The returned tuple is fed back into - # self.indexer() via pre_aux so the later _indexer_branch skips its - # own aux-stream launch. - _indexer_pre_aux = None - if _use_indexer_overlap: - _indexer_pre_aux = self.indexer.precompute_aux(hidden_states, attn_metadata) - - q, kv = self.kv_a_proj_with_mqa(hidden_states).split( - [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], -1 - ) + q, compressed_kv = maybe_execute_in_parallel( + lambda: self._q_a_layernorm_maybe_fused(q), + lambda: self.kv_a_layernorm(compressed_kv), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) - q, kv = maybe_execute_in_parallel( - lambda: self.q_a_layernorm(q), - lambda: self.kv_a_layernorm(kv), + q, latent_cache = maybe_execute_in_parallel( + lambda: self.q_b_proj(q), + lambda: torch.concat([compressed_kv, k_pe], dim=-1), self.ln_events[0], self.ln_events[1], self.aux_stream, ) - compressed_kv, k_pe = kv.split([self.kv_lora_rank, self.qk_rope_head_dim], -1) - qr = q - latent_cache = torch.concat([compressed_kv, k_pe], dim=-1) - - # CuTe DSL path for q_b_proj (hardware-default cluster count). - # Restricted to DSv4 CSA layers with compress_ratio=4 so the kernel - # swap only kicks in where the prologue overlap is exercised; other - # layers keep the cuBLAS path. Set TRTLLM_MLA_Q_B_PROJ_USE_CUTE_DSL=0 - # to disable. Bias and quantization are not handled. - _use_q_b_cute = ( - self.has_dsv4_indexer - and os.environ.get("TRTLLM_MLA_Q_B_PROJ_USE_CUTE_DSL", "1") == "1" - and self.q_b_proj.bias is None - and self.q_b_proj.weight.dtype == torch.bfloat16 - ) - - def _q_branch(): - # CuTe DSL bf16 path is bench-only and intentionally bypasses the - # FP8-fused-quant branch (weights are bf16, so the fused FP8 path - # would never apply anyway, but assert to make the contract - # explicit and catch any future config drift). - if _use_q_b_cute: - assert not self._is_fused_q_fp8_quant_enabled(num_generations=num_generations), ( - "CuTe DSL q_b_proj path is incompatible with the fused FP8 q-quant branch" - ) - q_proj = _q_b_proj_cute_dsl_bf16(q, self.q_b_proj.weight) - # Cross-iter cleanup: forward_absorption_* downstream gates - # the fused-FP8 attention path on these attrs being non-None. - # The FP8 path cannot trigger when weights are bf16, but clear - # them so stale buffers cannot silently re-enable fusion. - self._fused_quant_q_buffer = None - self._fused_q_pe = None - return self._deepseek_v4_q_b_layernorm(q_proj) - q_proj = self.q_b_proj(q) - if self._is_fused_q_fp8_quant_enabled(num_generations=num_generations): - placeholder_q, quant_q_buffer, q_pe, quant_scale_qkv = ( - self._deepseek_v4_q_b_layernorm_fused_fp8(q_proj) - ) - self._fused_quant_q_buffer = quant_q_buffer - self._fused_q_pe = q_pe - self._quant_scale_qkv = quant_scale_qkv - return placeholder_q - self._fused_quant_q_buffer = None - self._fused_q_pe = None - return self._deepseek_v4_q_b_layernorm(q_proj) - - def _compressor_branch(): - self.compressor(hidden_states, attn_metadata) - return None - - def _indexer_branch(): - return self.indexer( - qr, - hidden_states, - attn_metadata, - position_ids, - pre_aux=_indexer_pre_aux, - ) - - topk_indices = None - indexer_ran = False - if _v4_extra_overlap: - if _use_indexer_overlap: - # Compressor + indexer-aux are already in flight from the - # pre-launch block above. The outer compressor's tail is - # deferred until after _q_branch so the single wait below - # gates both compressor and _q_branch completion. - self.dsv4_overlap_start_event.record() - - with torch.cuda.stream(self.indexer_stream): - self.dsv4_overlap_start_event.wait() - topk_indices = _indexer_branch() - indexer_ran = True - self.dsv4_indexer_event.record() - - # _q_branch reads qr (post-q_a_layernorm), so it must wait for - # dsv4_overlap_start_event. Queuing it on compressor_stream - # serializes compressor -> q_b_proj -> q_b_layernorm while - # freeing the caller stream during the prologue window. - with torch.cuda.stream(self.compressor_stream): - self.dsv4_overlap_start_event.wait() - q = _q_branch() - self.dsv4_compressor_event.record() - - self.dsv4_compressor_event.wait() - self.dsv4_indexer_event.wait() - - # q/topk_indices were produced on other streams; record on the - # consuming stream so the caching allocator cannot recycle them mid-use. - cur_stream = torch.cuda.current_stream() - if q is not None: - q.record_stream(cur_stream) - if topk_indices is not None: - topk_indices.record_stream(cur_stream) - else: - q, _ = maybe_execute_in_parallel( - _q_branch, - _compressor_branch, - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - ) - else: - q = _q_branch() - if self.compressor is not None: - self.compressor(hidden_states, attn_metadata) - - if self.indexer is not None: - if not indexer_ran: - topk_indices = _indexer_branch() assert q.shape[0] == num_tokens, ( f"Expect q.shape[0] to be {num_tokens}, but got {q.shape[0]}" @@ -1937,58 +879,50 @@ def _indexer_branch(): if num_contexts > 0: q_ctx = q[:num_ctx_tokens, ...] - topk_indices_ctx = ( - topk_indices[:num_ctx_tokens, :] if topk_indices is not None else None - ) compressed_kv_ctx = compressed_kv[:num_ctx_tokens, ...] k_pe_ctx = k_pe[:num_ctx_tokens, ...] latent_cache_ctx = latent_cache[:num_ctx_tokens, ...] - ctx_position_ids = ( - position_ids[..., :num_ctx_tokens] if position_ids is not None else None - ) if self.apply_rotary_emb: - assert ctx_position_ids is not None + assert position_ids is not None + # position_ids spans [ctx..., gen...] in mixed batches; slice to + # match q_ctx/k_pe_ctx so external RoPE uses ctx positions. + ctx_position_ids = position_ids[..., :num_ctx_tokens] k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids) + # External RoPE is only used by backends that do not handle + # fused RoPE internally, so keep latent_cache in sync. + latent_cache_ctx = torch.cat([compressed_kv_ctx, k_pe_ctx], dim=-1) - self.forward_context_sparse_mla( + if self.llama_4_scaling: + q_ctx = self._attention_scaling(q_ctx, position_ids[..., :num_ctx_tokens]) + + self.forward_context( q_ctx, compressed_kv_ctx, k_pe_ctx, + position_ids, attn_metadata, output[:num_ctx_tokens, :], - position_ids=ctx_position_ids, - latent_cache=latent_cache_ctx, - topk_indices=topk_indices_ctx, - enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, - dsv4_epilogue_output=dsv4_epilogue_output, + latent_cache_ctx, ) if num_generations > 0: q_gen = q[num_ctx_tokens:, ...] - topk_indices_gen = ( - topk_indices[num_ctx_tokens:num_tokens, :] if topk_indices is not None else None - ) compressed_kv_gen = compressed_kv[num_ctx_tokens:, ...] k_pe_gen = k_pe[num_ctx_tokens:, ...] - latent_cache_gen = latent_cache[num_ctx_tokens:, ...] - gen_position_ids = ( - position_ids[..., num_ctx_tokens:num_tokens] if position_ids is not None else None - ) - if self.apply_rotary_emb: - assert gen_position_ids is not None - k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids) + if latent_cache_gen is None: + latent_cache_gen = latent_cache[num_ctx_tokens:, ...] + + if self.llama_4_scaling: + q_gen = self._attention_scaling(q_gen, position_ids[..., num_ctx_tokens:]) - self.forward_generation_sparse_mla( + self.forward_absorption_generation( q_gen, compressed_kv_gen, k_pe_gen, attn_metadata, output[num_ctx_tokens:num_tokens, :], - position_ids=gen_position_ids, + position_ids=position_ids, latent_cache=latent_cache_gen, - topk_indices=topk_indices_gen, - enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, - dsv4_epilogue_output=dsv4_epilogue_output, ) def forward_context_default( @@ -2001,13 +935,13 @@ def forward_context_default( output: torch.Tensor, latent_cache: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Dense MHA context path: expand KV via kv_b_proj and run attention. - - Used by non-DSA models and as the short-seq MHA fallback for DSA models. - """ + """Dense MHA context path: expand KV via kv_b_proj and run attention.""" kv = self.kv_b_proj(compressed_kv) k_nope, v = kv.split( - [self.num_heads_tp * self.qk_nope_head_dim, self.num_heads_tp * self.v_head_dim], + [ + self.num_heads_tp * self.qk_nope_head_dim, + self.num_heads_tp * self.v_head_dim, + ], -1, ) @@ -2038,130 +972,6 @@ def forward_context_default( return attn_output - def _should_use_short_mha( - self, attn_metadata: AttentionMetadata, position_ids: Optional[torch.Tensor] - ) -> bool: - """Check if the short-seq MHA optimization should be used for context. - - Uses max_ctx_kv_len (max total KV length per context sequence, - including cached tokens) when available, to correctly account for - chunked context where the full attention span exceeds the threshold - even if the new token count is small. Falls back to num_ctx_tokens - (total new context tokens) when max_ctx_kv_len is not set. - - Disabled under torch compile so that the split DSA custom ops - (mla_dsa_proj / mla_dsa_attn_inplace) have unconditionally - straight-line control flow for CUDA graph capture. - """ - if is_torch_compiling(): - return False - if not ( - self.short_seq_mha_threshold > 0 - and not self.apply_rotary_emb - and self.mapping.cp_size == 1 - and position_ids is not None - ): - return False - effective_len = getattr(attn_metadata, "max_ctx_kv_len", attn_metadata.num_ctx_tokens) - return effective_len <= self.short_seq_mha_threshold - - def forward_context_sparse_mla( - self, - q: torch.Tensor, - compressed_kv: torch.Tensor, - k_pe: torch.Tensor, - attn_metadata: AttentionMetadata, - output: torch.Tensor, - latent_cache: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - enable_dsv4_epilogue_fusion: bool = False, - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """Run context-phase attention for DSA models. - - Dispatches to the short-seq MHA path (forward_context) when the max - per-sequence KV length (including cached tokens) is within the - threshold, or falls through to the absorption/sparse MLA path - otherwise. forward_context() further dispatches to the appropriate - handler (forward_context_default, forward_context_with_cached_kv, or - forward_context_with_chunked_prefill) based on cached-KV state. - - Args: - q: Query tensor, shape [num_ctx_tokens, num_heads * qk_head_dim]. - compressed_kv: Latent KV, shape [num_ctx_tokens, kv_lora_rank]. - k_pe: RoPE key portion, shape [num_ctx_tokens, qk_rope_head_dim]. - attn_metadata: Attention metadata for the current batch. - output: Pre-allocated output tensor, written in-place. - latent_cache: Concatenated [compressed_kv, k_pe] for KV cache. - topk_indices: Sparse routing indices from the indexer (None when - the short-seq MHA path is used). - position_ids: Token position IDs (required for short-seq MHA). - """ - # Short-sequence MHA: bypass absorption path for short prefills, - # using kv_b_proj expansion + standard attention instead. - # See __init__ comment for rationale. topk_indices is not used - # because dense attention is faster than sparse routing at this scale. - # forward_context() handles cached tokens by dispatching to - # forward_context_with_cached_kv or forward_context_with_chunked_prefill. - if not enable_dsv4_epilogue_fusion and self._should_use_short_mha( - attn_metadata, position_ids - ): - return self.forward_context( - q, compressed_kv, k_pe, position_ids, attn_metadata, output, latent_cache - ) - - if get_sm_version() >= 100: - return self.forward_absorption_context( - q, - compressed_kv, - k_pe, - attn_metadata, - output, - position_ids=position_ids, - latent_cache=latent_cache, - topk_indices=topk_indices, - enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, - dsv4_epilogue_output=dsv4_epilogue_output, - ) - else: - assert not self.is_deepseek_v4, "DeepSeek-V4 is not supported on pre-blackwell GPUs." - return self.forward_sparse_mla_kvcache_bf16( - q, latent_cache, attn_metadata, output, topk_indices, is_generation=False - ) - - def forward_generation_sparse_mla( - self, - q: torch.Tensor, - compressed_kv: torch.Tensor, - k_pe: torch.Tensor, - attn_metadata: AttentionMetadata, - output: torch.Tensor, - position_ids: Optional[torch.Tensor] = None, - latent_cache: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, - enable_dsv4_epilogue_fusion: bool = False, - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - if get_sm_version() >= 100: - return self.forward_absorption_generation( - q, - compressed_kv, - k_pe, - attn_metadata, - output, - position_ids=position_ids, - latent_cache=latent_cache, - topk_indices=topk_indices, - enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, - dsv4_epilogue_output=dsv4_epilogue_output, - ) - else: - assert not self.is_deepseek_v4, "DeepSeek-V4 is not supported on pre-blackwell GPUs." - return self.forward_sparse_mla_kvcache_bf16( - q, latent_cache, attn_metadata, output, topk_indices, is_generation=True - ) - def forward_context_with_cached_kv( self, q: torch.Tensor, @@ -2194,7 +1004,10 @@ def forward_context_with_cached_kv( # compute full_k_nope and full_v from full_compressed_kv full_kv = self.kv_b_proj(full_compressed_kv) full_k_nope, full_v = full_kv.split( - [self.num_heads_tp * self.qk_nope_head_dim, self.num_heads_tp * self.v_head_dim], + [ + self.num_heads_tp * self.qk_nope_head_dim, + self.num_heads_tp * self.v_head_dim, + ], -1, ) @@ -2288,7 +1101,10 @@ def forward_context_with_chunked_prefill( # [tokens, 2, h, kv_dim], without rope_dim chunked_kv = self.kv_b_proj(chunked_compressed_kv) chunked_k_nope, chunked_v = chunked_kv.split( - [self.num_heads_tp * self.qk_nope_head_dim, self.num_heads_tp * self.v_head_dim], + [ + self.num_heads_tp * self.qk_nope_head_dim, + self.num_heads_tp * self.v_head_dim, + ], -1, ) @@ -2347,7 +1163,10 @@ def forward_context_with_chunked_prefill( # final round of attention k_nope, v = kv.split( - [self.num_heads_tp * self.qk_nope_head_dim, self.num_heads_tp * self.v_head_dim], + [ + self.num_heads_tp * self.qk_nope_head_dim, + self.num_heads_tp * self.v_head_dim, + ], -1, ) @@ -2398,7 +1217,13 @@ def forward_context_with_chunked_prefill( @staticmethod @functools.cache def cached_warmup_forward_context_with_cached_kv( - num_heads_tp, qk_nope_head_dim, qk_rope_head_dim, kv_lora_rank, v_head_dim, dtype, device + num_heads_tp, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, ): """Warmup torch.compile for cat operations with different tensor layouts. @@ -2422,7 +1247,11 @@ def warmup(num_tokens): num_tokens, 1, qk_rope_head_dim, dtype=dtype, device=device ).expand(-1, num_heads_tp, -1) k_pe = torch.empty( - num_tokens, 1, kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device + num_tokens, + 1, + kv_lora_rank + qk_rope_head_dim, + dtype=dtype, + device=device, )[:, :, -qk_rope_head_dim:].expand(-1, num_heads_tp, -1) torch._dynamo.maybe_mark_dynamic(chunked_k_nope, 0) torch._dynamo.maybe_mark_dynamic(chunked_k_pe, 0) @@ -2510,10 +1339,8 @@ def forward_absorption_generation( output: torch.Tensor, position_ids: Optional[torch.Tensor] = None, latent_cache: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, - enable_dsv4_epilogue_fusion: bool = False, - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + **kwargs, + ) -> torch.Tensor: num_tokens = q.shape[0] q_nope, q_pe = q.view([-1, self.num_heads_tp, self.qk_head_dim]).split( [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 @@ -2544,23 +1371,13 @@ def forward_absorption_generation( device=q.device, ) - if self.is_deepseek_v4: - fused_q = q - self.mqa.mla_rope_generation( - fused_q, - q_pe, - latent_cache, - attn_metadata, - cu_q_seqlens, - cu_kv_seqlens, - fmha_scheduler_counter, - mla_bmm1_scale, - mla_bmm2_scale, - quant_q_buffer, - ) - else: + if hasattr(self, "k_b_proj_trans"): fused_q = torch.empty( - [num_tokens, self.num_heads_tp, (self.kv_lora_rank + self.qk_rope_head_dim)], + [ + num_tokens, + self.num_heads_tp, + (self.kv_lora_rank + self.qk_rope_head_dim), + ], dtype=q.dtype, device=q.device, ) @@ -2640,21 +1457,13 @@ def _mla_gen_rope(): ) fused_q = fused_q.view( - [num_tokens, self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim)] + [ + num_tokens, + self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim), + ] ) - - # Use generation_only for generation phase and context_only for context phase in DSA attention - attention_input_type = AttentionInputType.generation_only - dsv4_output = output if self.is_deepseek_v4 else None - dsv4_output_sf = None - dsv4_cos_sin_cache = None - if enable_dsv4_epilogue_fusion: - assert self.is_deepseek_v4 - assert dsv4_epilogue_output is not None - dsv4_output, dsv4_output_sf = self._validate_dsv4_epilogue_buffers( - num_tokens, dsv4_epilogue_output - ) - dsv4_cos_sin_cache = self.inverse_rotary_emb.rotary_cos_sin + else: + raise RuntimeError("MLA absorption requires k_b_proj_trans") attn_out_latent = self._attn_forward_gen( self.mqa, @@ -2663,39 +1472,20 @@ def _mla_gen_rope(): None, position_ids, attn_metadata, - attention_input_type=attention_input_type, + attention_input_type=AttentionInputType.generation_only, out_scale=self.out_scale, - output=dsv4_output, - output_sf=dsv4_output_sf, latent_cache=latent_cache, # kvcache and k_pe q_pe=q_pe, # used by `invokeMLARopeGeneration` - topk_indices=topk_indices, # used by DSA attention cu_q_seqlens=cu_q_seqlens, # used by `mlaGeneration` cu_kv_seqlens=cu_kv_seqlens, # used by `mlaGeneration` fmha_scheduler_counter=fmha_scheduler_counter, # used by `mlaGeneration` mla_bmm1_scale=mla_bmm1_scale, # used by `mlaGeneration` mla_bmm2_scale=mla_bmm2_scale, # used by `mlaGeneration` quant_q_buffer=quant_q_buffer, # used by `mlaGeneration` - dsv4_inv_rope_cos_sin_cache=dsv4_cos_sin_cache, - enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, + **kwargs, ) fused_q = None - if enable_dsv4_epilogue_fusion: - return attn_out_latent - - if self.is_deepseek_v4: - if self.mapping.has_cp_helix(): - raise RuntimeError( - "DeepSeek-V4 + CP Helix is not supported: " - "_helix_post_process returns a different tensor, " - "bypassing the pre-allocated output buffer." - ) - assert attn_out_latent.data_ptr() == output.data_ptr(), ( - "Attention backend did not write into the provided output buffer." - ) - return output - # note: if we do not have CP, then num_heads_tp_cp == num_heads_tp assert ( attn_out_latent.shape[0] == q.shape[0] @@ -2739,23 +1529,23 @@ def forward_absorption_context( output: torch.Tensor, position_ids: Optional[torch.Tensor] = None, latent_cache: Optional[torch.Tensor] = None, - topk_indices: Optional[torch.Tensor] = None, - enable_dsv4_epilogue_fusion: bool = False, - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + **kwargs, + ) -> torch.Tensor: num_tokens = q.shape[0] q_nope, q_pe = q.view([-1, self.num_heads_tp, self.qk_head_dim]).split( [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 ) - if self.is_deepseek_v4: - fused_q = q - else: + if hasattr(self, "k_b_proj_trans"): # fused_q contains 1) the result of the following bmm with shape [num_tokens, num_heads, kv_lora_rank] # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp fused_q = torch.empty( - [num_tokens, self.num_heads_tp, (self.kv_lora_rank + self.qk_rope_head_dim)], + [ + num_tokens, + self.num_heads_tp, + (self.kv_lora_rank + self.qk_rope_head_dim), + ], dtype=q.dtype, device=q.device, ) @@ -2770,7 +1560,10 @@ def forward_absorption_context( # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank] # The output of bmm is written directly into fused_q self._bmm_bf16_out( - q_nope_t, self.k_b_proj_trans, self.k_b_proj_trans.transpose(1, 2), q_nope_out + q_nope_t, + self.k_b_proj_trans, + self.k_b_proj_trans.transpose(1, 2), + q_nope_out, ) elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn: # [num_heads, num_tokens, self.kv_lora_rank] @@ -2792,46 +1585,13 @@ def forward_absorption_context( if self.apply_rotary_emb: fused_q[..., self.kv_lora_rank :] = q_pe fused_q = fused_q.view( - [num_tokens, self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim)] - ) - - # Use generation_only for generation phase and context_only for context phase in DSA attention - attention_input_type = AttentionInputType.context_only - - # Fused FP8-Q path: forward the pre-quantized buffers stashed in - # `_q_branch`; the C++ op enables fusion when both are non-None. - quant_q_buffer = getattr(self, "_fused_quant_q_buffer", None) - fused_q_pe = getattr(self, "_fused_q_pe", None) - quant_scale_qkv = getattr(self, "_quant_scale_qkv", None) - use_fused_q_fp8 = ( - self.is_deepseek_v4 - and quant_q_buffer is not None - and fused_q_pe is not None - and quant_scale_qkv is not None - ) - - if use_fused_q_fp8: - # Defensive prefix slicing: context-only batches today, mixed-batch later. - q_pe = fused_q_pe[:num_tokens] - quant_q_buffer = quant_q_buffer[:num_tokens].view( - num_tokens, - self.num_heads_tp, - self.kv_lora_rank + self.qk_rope_head_dim, + [ + num_tokens, + self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim), + ] ) else: - quant_q_buffer = None - quant_scale_qkv = None - - dsv4_output = output if self.is_deepseek_v4 else None - dsv4_output_sf = None - dsv4_cos_sin_cache = None - if enable_dsv4_epilogue_fusion: - assert self.is_deepseek_v4 - assert dsv4_epilogue_output is not None - dsv4_output, dsv4_output_sf = self._validate_dsv4_epilogue_buffers( - num_tokens, dsv4_epilogue_output - ) - dsv4_cos_sin_cache = self.inverse_rotary_emb.rotary_cos_sin + raise RuntimeError("MLA absorption requires k_b_proj_trans") attn_out_latent = self._attn_forward_gen( self.mqa, @@ -2840,36 +1600,13 @@ def forward_absorption_context( None, position_ids, attn_metadata, - attention_input_type=attention_input_type, + attention_input_type=AttentionInputType.context_only, out_scale=self.out_scale, - output=dsv4_output, - output_sf=dsv4_output_sf, latent_cache=latent_cache, # kvcache and k_pe q_pe=q_pe, # used by applyMLARopeAndAssignQKVKernelOptContext - quant_q_buffer=quant_q_buffer, # fused-FP8 path only - quant_scale_qkv=quant_scale_qkv, # fused-FP8 path only - topk_indices=topk_indices, # used by DSA attention - dsv4_inv_rope_cos_sin_cache=dsv4_cos_sin_cache, - enable_dsv4_epilogue_fusion=enable_dsv4_epilogue_fusion, + **kwargs, ) fused_q = None - self._fused_quant_q_buffer = None - self._fused_q_pe = None - - if enable_dsv4_epilogue_fusion: - return attn_out_latent - - if self.is_deepseek_v4: - if self.mapping.has_cp_helix(): - raise RuntimeError( - "DeepSeek-V4 + CP Helix is not supported: " - "_helix_post_process returns a different tensor, " - "bypassing the pre-allocated output buffer." - ) - assert attn_out_latent.data_ptr() == output.data_ptr(), ( - "Attention backend did not write into the provided output buffer." - ) - return output # note: if we do not have CP, then num_heads_tp_cp == num_heads_tp assert ( @@ -2905,154 +1642,90 @@ def forward_absorption_context( return output - @nvtx_range("forward_sparse_mla_kvcache_bf16") - def forward_sparse_mla_kvcache_bf16( + def _forward_custom_op( self, - q: torch.Tensor, - latent_cache: torch.Tensor, - attn_metadata: DSAtrtllmAttentionMetadata, - output: torch.Tensor, - topk_indices: torch.Tensor, - is_generation: bool = False, - ) -> torch.Tensor: - """ - Forward sparse MLA (DSA) for BF16 KV cache for both context and generation phases using FlashMLA kernels - - To form the input for FlashMLA kernel and adapt our KV cache manager, we need to: - 1. Append current tokens to paged cache and apply rope to q/k via mla_rope_append_paged_kv_assign_q - 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" - ) - # Append current tokens to paged cache and apply RoPE to q - # This writes latent_cache to paged KV and modifies q in-place - trtllm_attention = self.mqa - with nvtx_range_debug(f"mla_rope_append_paged_kv_assign_q_is_generation={is_generation}"): - trtllm_attention.mla_rope_append_paged_kv_assign_q( - q, latent_cache, attn_metadata, is_generation=is_generation + hidden_states: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_output: list[torch.Tensor], + latent_cache_gen: Optional[torch.Tensor], + ) -> None: + """Run the dense or sparse registered custom-op implementation.""" + if custom_op := self.sparse_attn_hooks.forward_sparse_attn_custom_op: + custom_op( + self, + hidden_states, + position_ids, + attn_output, + latent_cache_gen, ) + return - num_tokens = q.shape[0] - q_nope, q_rope = q.view(-1, self.num_heads_tp, self.qk_head_dim).split( - [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 - ) - q_nope_out = torch.empty( - [num_tokens, self.num_heads_tp, (self.kv_lora_rank)], - dtype=q.dtype, - device=q.device, - ) - - if self.k_b_proj_trans.dtype == torch.bfloat16: - # [num_heads, num_tokens, self.qk_nope_head_dim] - q_nope_t = q_nope.transpose(0, 1) - # [num_heads, num_tokens, self.kv_lora_rank] - q_nope_out = q_nope_out.transpose(0, 1) - - # [num_heads, num_tokens, self.qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim] - # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank] - # The output of bmm is written directly into fused_q - self._bmm_bf16_out( - q_nope_t, self.k_b_proj_trans, self.k_b_proj_trans.transpose(1, 2), q_nope_out - ) - elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn: - # [num_heads, num_tokens, self.kv_lora_rank] - q_nope_out = q_nope_out.transpose(0, 1) + output = attn_output[0] + sparse_output = None + sparse_output_sf = None + if len(attn_output) > 3: + raise RuntimeError("MLA output hooks may return at most two sparse output buffers.") + if len(attn_output) > 1: + sparse_output = attn_output[1] + if len(attn_output) > 2: + sparse_output_sf = attn_output[2] - fp8_block_scaling_bmm_out( - q_nope, - self.k_b_proj_trans, - self.k_b_proj_trans_scale, - q_nope_out, - self.k_b_proj_trans_dequant, - self.use_cute_dsl_blockscaling_bmm, + if isinstance(hidden_states, Fp4QuantizedTensor): + torch.ops.trtllm.mla_custom_op_inplace( + hidden_states.unquantized_hidden_states, + position_ids, + self.layer_idx_str, + output, + latent_cache_gen, + sparse_output, + sparse_output_sf, + hidden_states.fp4_tensor, + hidden_states.scaling_factor, ) else: - raise NotImplementedError(f"Missing bmm impl for dtype: {self.k_b_proj_trans.dtype}.") - - q_nope_out = q_nope_out.transpose(0, 1) - q_concat = torch.cat([q_nope_out, q_rope], dim=-1) - - sm_version = get_sm_version() - # FlashMLA sparse kernel (bf16) requires num_heads=128 on sm100 or a - # multiple of 64 on sm90. - if sm_version >= 100: - padding = 128 - assert self.num_heads_tp <= padding, ( - f"SM100 FlashMLA sparse kernel requires exactly {padding} heads, " - f"got {self.num_heads_tp}. Padding from values > {padding} is not supported." + torch.ops.trtllm.mla_custom_op_inplace( + hidden_states, + position_ids, + self.layer_idx_str, + output, + latent_cache_gen, + sparse_output, + sparse_output_sf, ) - else: # SM90 - padding = ((self.num_heads_tp + 63) // 64) * 64 # multiple of 64 - if self.num_heads_tp != padding: - logger.warning_once( - f"Padding num_heads from {self.num_heads_tp} to {padding} " - f"due to FlashMLA sparse attention kernel requirement", - key="sparse_mla_padding_warning", - ) + def _project_output( + self, + attn_output: list[torch.Tensor], + position_ids: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams], + ) -> torch.Tensor: + """Apply the MLA output projection.""" + if project_output := self.sparse_attn_hooks.project_sparse_attn_output: + return project_output(self, attn_output, position_ids, attn_metadata, all_reduce_params) - # Create padded tensor with zeros for extra heads - q_padded = q_concat.new_empty((num_tokens, padding, q_concat.shape[2])) - q_padded[:, : self.num_heads_tp, :] = q_concat - q_concat = q_padded - - # Convert indices and return all-layer KV pool - # Note: underlying pool is layer-interleaved: [num_blocks, num_layers, - # kv_factor, tokens_per_block, num_kv_heads, head_dim]. To avoid a - # reshape(copy) per-layer KV cache, return the all-layer KV pool with - # topk indices adjusted by stride_factor=num_layers*tokens_per_block. - topk_indices_pool, kv_cache_pool = transform_local_topk_and_prepare_pool_view( - topk_indices, - attn_metadata, - layer_idx=self.layer_idx, - is_generation=is_generation, + return self._project_output_impl( + attn_output[0], position_ids, attn_metadata, all_reduce_params ) - topk_indices_pool = topk_indices_pool.view(num_tokens, 1, -1) - if flash_mla_sparse_fwd is not None: - attn_out_latent = flash_mla_sparse_fwd( - q_concat, kv_cache_pool, topk_indices_pool, self.softmax_scale - )[0] - else: - raise RuntimeError( - "flash_mla_sparse_fwd not available. Please ensure FlashMLA module is built." - ) - - # [seq, num_heads, kv_lora_rank], account for padding - attn_out_latent = attn_out_latent[:, : self.num_heads_tp, :] - attn_out_latent = attn_out_latent.view([-1, self.num_heads_tp, self.kv_lora_rank]) - if self.num_heads_tp != padding: - attn_out_latent = attn_out_latent.contiguous() - assert ( - attn_out_latent.shape[0] == q.shape[0] and attn_out_latent.shape[1] == self.num_heads_tp + def _project_output_impl( + self, + attn_output: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams], + ) -> torch.Tensor: + """Apply the default MLA output projection implementation.""" + return _helix_cp_output_projection( + self.o_proj, + attn_output, + attn_metadata, + all_reduce_params, + self.mapping, + self.mapping_o, + self.layer_idx, ) - attn_output = output.view([num_tokens, self.num_heads_tp, self.v_head_dim]) - - if self.v_b_proj.dtype == torch.bfloat16: - # [num_heads, seq, kv_lora_rank] x [num_heads, kv_lora_rank, v_head_dim] - # -> [num_heads, seq, v_head_dim] - self._bmm_bf16_out( - attn_out_latent.transpose(0, 1), - self.v_b_proj, - self.v_b_proj.transpose(1, 2), - attn_output.transpose(0, 1), - ) - elif self.v_b_proj.dtype == torch.float8_e4m3fn: - fp8_block_scaling_bmm_out( - attn_out_latent, - self.v_b_proj, - self.v_b_proj_scale, - attn_output.transpose(0, 1), - self.v_b_proj_dequant, - self.use_cute_dsl_blockscaling_bmm, - ) - else: - raise NotImplementedError(f"Missing bmm impl for dtype: {self.v_b_proj.dtype}.") - return output - def forward( self, position_ids: Optional[torch.Tensor], @@ -3065,143 +1738,34 @@ def forward( hidden_states, attn_metadata, self.mapping, self.layer_idx ) - dsv4_epilogue_output: Optional[tuple[torch.Tensor, torch.Tensor]] = None if self.register_to_config: - if self.is_deepseek_v4: - outputs = torch.ops.trtllm.create_mla_outputs(hidden_states, self.layer_idx_str) - attn_output = outputs[0] - dsv4_output = None - dsv4_output_sf = None - if len(outputs) == 3: - dsv4_output, dsv4_output_sf = outputs[1], outputs[2] - dsv4_epilogue_output = (dsv4_output, dsv4_output_sf) - elif len(outputs) != 1: - raise RuntimeError( - "create_mla_outputs must return either legacy output or " - "legacy output plus DSv4 fused epilogue buffers." - ) - - torch.ops.trtllm.mla_custom_op_inplace( - hidden_states, - position_ids, - self.layer_idx_str, - attn_output, - latent_cache_gen, - dsv4_output, - dsv4_output_sf, - dsv4_epilogue_output is not None, - ) - else: - attn_output = self.create_output(hidden_states, attn_metadata.num_contexts) - if self.is_dsa: - # When the previous layer's boundary fusion pre-quantized - # this layer's kv_a_proj NVFP4 input, hidden_states is an - # Fp4QuantizedTensor with both BF16 + FP4 views. Custom ops - # can't take dataclasses, so pass tensors explicitly. - if isinstance(hidden_states, Fp4QuantizedTensor): - proj_outputs = torch.ops.trtllm.mla_dsa_proj( - hidden_states.unquantized_hidden_states, - position_ids, - self.layer_idx_str, - hidden_states.fp4_tensor, - hidden_states.scaling_factor, - ) - else: - proj_outputs = torch.ops.trtllm.mla_dsa_proj( - hidden_states, position_ids, self.layer_idx_str - ) - q, compressed_kv, k_pe, latent_cache = proj_outputs[:4] - indexer_intermediates = proj_outputs[4:] - torch.ops.trtllm.mla_dsa_attn_inplace( - q, - compressed_kv, - k_pe, - latent_cache, - indexer_intermediates, - position_ids, - self.layer_idx_str, - attn_output, - ) - else: - # Vanilla MLA (e.g. Kimi-K2.5): when the previous layer's - # boundary fusion pre-quantized this layer's input, - # hidden_states is an Fp4QuantizedTensor. Custom ops can't - # take dataclasses, so pass the BF16 + FP4 + SF views as - # explicit tensors. - if isinstance(hidden_states, Fp4QuantizedTensor): - torch.ops.trtllm.mla_custom_op_inplace( - hidden_states.unquantized_hidden_states, - position_ids, - self.layer_idx_str, - attn_output, - latent_cache_gen, - None, - None, - False, - hidden_states.fp4_tensor, - hidden_states.scaling_factor, - ) - else: - torch.ops.trtllm.mla_custom_op_inplace( - hidden_states, - position_ids, - self.layer_idx_str, - attn_output, - latent_cache_gen, - None, - None, - False, - ) - else: - enable_dsv4_epilogue_fusion = ( - self.is_deepseek_v4 - and self._should_use_dsv4_epilogue_fusion( - attn_metadata.num_contexts, attn_metadata.num_generations + output_hidden_states = hidden_states + if isinstance(hidden_states, Fp4QuantizedTensor): + assert hidden_states.unquantized_hidden_states is not None, ( + "MLA.forward received an Fp4QuantizedTensor without a " + "unquantized_hidden_states view" ) + output_hidden_states = hidden_states.unquantized_hidden_states + attn_output = torch.ops.trtllm.create_mla_outputs( + output_hidden_states, self.layer_idx_str ) - if enable_dsv4_epilogue_fusion: - dsv4_epilogue_output = self._create_dsv4_epilogue_buffers( - hidden_states, attn_metadata.num_tokens - ) - output_input = hidden_states[:0] if enable_dsv4_epilogue_fusion else hidden_states - attn_output = self.create_output(output_input, attn_metadata.num_contexts) - if self.is_dsa: - self.forward_impl_with_dsa( - position_ids, hidden_states, attn_metadata, output=attn_output - ) - elif self.is_deepseek_v4: - self.forward_impl_with_deepseek_v4( - position_ids, - hidden_states, - attn_metadata, - output=attn_output, - dsv4_epilogue_output=dsv4_epilogue_output, - ) - else: - self.forward_impl( - position_ids, - hidden_states, - attn_metadata, - output=attn_output, - latent_cache_gen=latent_cache_gen, - ) - - if self.is_deepseek_v4: - if dsv4_epilogue_output is not None: - attn_output = self._deepseek_v4_o_proj(dsv4_epilogue_output) - else: - attn_output = self._deepseek_v4_o_proj(attn_output, position_ids) - else: - attn_output = _helix_cp_output_projection( - self.o_proj, + self._forward_custom_op( + hidden_states, + position_ids, attn_output, + latent_cache_gen, + ) + else: + attn_output = self._create_outputs(hidden_states, attn_metadata) + self.forward_impl( + position_ids, + hidden_states, attn_metadata, - all_reduce_params, - self.mapping, - self.mapping_o, - self.layer_idx, + attn_output=attn_output, + latent_cache_gen=latent_cache_gen, ) - return attn_output + + return self._project_output(attn_output, position_ids, attn_metadata, all_reduce_params) def resmooth_parameters(self, module_weight, module_weight_scale, recipe=(1, 128, 128)): weight, weight_scale = fp8_utils.resmooth_to_fp8_e8m0(module_weight, module_weight_scale) @@ -3223,16 +1787,18 @@ def resmooth_parameters(self, module_weight, module_weight_scale, recipe=(1, 128 def transform_weights(self) -> None: if self._weights_transformed: return - # In DeepSeek-V4 mode, kv_b_proj doesn't exist - if getattr(self, "is_deepseek_v4", False): - has_fp8_block_scales = False - else: - has_fp8_block_scales = ( - self.kv_b_proj.quant_config - and self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales() - ) - is_sm120 = get_sm_version() == 120 - if is_sm120 and has_fp8_block_scales: + if transform_sparse_attn_weights := getattr( + getattr(self, "sparse_attn_hooks", None), "transform_sparse_attn_weights", None + ): + transform_sparse_attn_weights(self) + self._weights_transformed = True + return + + has_fp8_block_scales = bool( + self.kv_b_proj.quant_config + and self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales() + ) + if get_sm_version() == 120 and has_fp8_block_scales: self.k_b_proj_trans, self.k_b_proj_trans_scale = self.resmooth_parameters( self.k_b_proj_trans, self.k_b_proj_trans_scale, recipe=(1, 128, 128) ) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d09e0b576b5b..70106037ec64 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1064,6 +1064,11 @@ def _value(name: str, default=None): return getattr(pretrained_config, name, default) return default + num_layers = getattr(pretrained_config, "num_hidden_layers", None) + has_shared_indexer_layers = (num_layers is not None and any( + not self._is_full_indexer_layer(pretrained_config, layer_idx) + for layer_idx in range(num_layers))) + return DSAMetadataParams( indexer_max_chunk_size=self.indexer_max_chunk_size or 32768, max_sparse_topk=_value("index_topk"), @@ -1073,6 +1078,7 @@ def _value(name: str, default=None): use_cute_dsl_topk=self.use_cute_dsl_topk, use_cute_dsl_paged_mqa_logits=(self.use_cute_dsl_paged_mqa_logits), q_split_threshold=self.q_split_threshold, + has_shared_indexer_layers=has_shared_indexer_layers, ) 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 3ba703006935..b8649e8621f8 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 @@ -30,16 +30,18 @@ PositionEmbeddingType, RotaryScalingType, ) -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import DeepseekV4CacheManager +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import ( + DeepseekV4CacheManager, + DeepseekV4Indexer, + DeepseekV4TrtllmAttentionMetadata, +) from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor import ( Compressor, KVCacheDtype, ) -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.params import ( DEEPSEEK_V4_SLIDING_ATTENTION, DeepseekV4AttentionType, - DeepseekV4Indexer, - DeepseekV4TrtllmAttentionMetadata, ) from tensorrt_llm._torch.modules.rotary_embedding import RopeParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index bcba3baf1b24..4331fed0d4e3 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -21,7 +21,7 @@ from utils.util import skip_pre_blackwell from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import DeepseekV4CacheManager -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.params import ( DEEPSEEK_V4_SLIDING_ATTENTION, DeepseekV4AttentionType, compress_ratio_has_attention, 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 d7802af211d2..285418898479 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 @@ -27,10 +27,8 @@ import pytest import torch -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( - DeepseekV4AttentionType, - get_token_bytes, -) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.cache_manager import get_token_bytes +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.params import DeepseekV4AttentionType from tensorrt_llm.llmapi.llm_args import ( DeepSeekSparseAttentionConfig, DeepSeekV4SparseAttentionConfig, diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py index b60aff4c96a1..916e6cdca480 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py @@ -28,8 +28,10 @@ DeepseekV4AttentionType, DeepseekV4CacheManager, ) -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import get_token_bytes -from tensorrt_llm._torch.attention_backend.sparse.kernel import deepseek_v4_local_to_global_indices +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.cache_manager import get_token_bytes +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.kernels import ( + deepseek_v4_local_to_global_indices, +) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType, SamplingConfig diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py index 7325b3cecf25..196c06531831 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py @@ -25,6 +25,9 @@ from utils.util import skip_pre_blackwell from tensorrt_llm._torch.attention_backend.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module import ( + project_sparse_attn_output, +) from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_deepseekv3 import weight_dequant from tensorrt_llm._torch.modules.mla import MLA @@ -209,6 +212,10 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): num_groups=num_groups, o_lora_rank=o_lora_rank, ).to(device) + assert mla.mha is None + assert not hasattr(mla, "kv_b_proj") + assert not hasattr(mla, "v_b_proj") + assert not hasattr(mla, "o_proj") # Initialize weights nn_init_std = 0.02 @@ -271,7 +278,7 @@ def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): # Call the deepseek_v4 output projection (mla_rope_inplace modifies attn_out_latent # in-place, so clone before passing to preserve original for reference) - output = mla._deepseek_v4_o_proj(attn_out_latent.clone(), position_ids) + output = project_sparse_attn_output(mla, [attn_out_latent.clone()], position_ids) # Calculate reference output if dtype_str == "bf16": 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 e7df3560db53..77e1e9dcd5ff 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 @@ -37,11 +37,10 @@ DeepseekV4AttentionType, DeepseekV4CacheManager, DeepseekV4TrtllmAttention, -) -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( DeepseekV4TrtllmAttentionMetadata, - get_token_bytes, ) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.cache_manager import get_token_bytes +from tensorrt_llm._torch.attention_backend.sparse.params import SparseBackendForwardArgs from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -781,7 +780,6 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): attention_input_type=AttentionInputType.context_only, latent_cache=latent_cache, q_pe=q_pe, - topk_indices=None, softmax_stats_tensor=softmax_stats_tensor, ) @@ -1087,7 +1085,7 @@ def yarn_get_mscale(scale=1, mscale=1): attention_input_type=AttentionInputType.context_only, latent_cache=latent_cache, q_pe=q_pe, - topk_indices=topk_indices, + sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk_indices), ) # Reference computation @@ -1237,7 +1235,7 @@ def yarn_get_mscale(scale=1, mscale=1): cu_q_seqlens=cu_q_seqlens, cu_kv_seqlens=cu_kv_seqlens, fmha_scheduler_counter=fmha_scheduler_counter, - topk_indices=topk_indices, + sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk_indices), ) # Reference: apply RoPE separately, then compute attention @@ -1432,7 +1430,7 @@ def test_deepseek_v4_sparse_mla_mixed_batch(context_lengths: List[int]): attention_input_type=AttentionInputType.context_only, latent_cache=prefill_latent, q_pe=prefill_q_pe, - topk_indices=prefill_topk, + sparse_backend_args=SparseBackendForwardArgs(topk_indices=prefill_topk), ) gen_requests = requests[1:] @@ -1494,7 +1492,7 @@ def test_deepseek_v4_sparse_mla_mixed_batch(context_lengths: List[int]): ) mixed_metadata.prepare() - # 5. Per-layer forward + verify (mirrors forward_impl_with_deepseek_v4). + # 5. Per-layer forward + verify (mirrors forward_sparse_mla). for li in TEST_LAYERS: ratio = scenario.compress_ratios[li] print(f"\n--- Mixed: layer {li}, compress_ratio={ratio} ---") @@ -1557,7 +1555,7 @@ def test_deepseek_v4_sparse_mla_mixed_batch(context_lengths: List[int]): output=output[:total_ctx_tokens], latent_cache=ctx_latent, q_pe=ctx_q_pe, - topk_indices=ctx_topk, + sparse_backend_args=SparseBackendForwardArgs(topk_indices=ctx_topk), ) # Generation forward → output[total_ctx_tokens:] @@ -1589,7 +1587,7 @@ def test_deepseek_v4_sparse_mla_mixed_batch(context_lengths: List[int]): cu_q_seqlens=cu_q, cu_kv_seqlens=cu_kv, fmha_scheduler_counter=counter, - topk_indices=gen_topk, + sparse_backend_args=SparseBackendForwardArgs(topk_indices=gen_topk), ) # Context reference 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 cf4e9af27de4..22addacdec10 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -34,9 +34,16 @@ from utils.util import check_accuracy, skip_pre_blackwell, skip_pre_hopper from tensorrt_llm import deep_gemm -from tensorrt_llm._torch.attention_backend.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionForwardArgs, + AttentionInputType, + PositionalEmbeddingParams, + RopeParams, +) from tensorrt_llm._torch.attention_backend.sparse.dsa import ( + DSABackendForwardArgs, DSACacheManager, + DSATrtllmAttention, DSAtrtllmAttentionMetadata, Indexer, _effective_compress_ratio_divisor, @@ -89,7 +96,7 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): metadata.create_buffers_for_indexer = Mock() with patch( - "tensorrt_llm._torch.attention_backend.sparse.dsa.TrtllmAttentionMetadata.__post_init__" + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.TrtllmAttentionMetadata.__post_init__" ): DSAtrtllmAttentionMetadata.__post_init__(metadata) @@ -99,6 +106,107 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(): assert metadata._tokens_per_block == 64 +def test_shared_topk_lifecycle(): + sparse_config = DeepSeekSparseAttentionConfig( + index_n_heads=1, + index_head_dim=8, + index_topk=3, + skip_indexer_for_short_seqs=False, + ) + pretrained_config = SimpleNamespace( + num_hidden_layers=2, + index_topk_pattern=["F", "S"], + ) + sparse_metadata_params = sparse_config.to_sparse_metadata_params( + pretrained_config=pretrained_config + ) + 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.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 + metadata.enable_indexer_skip = False + metadata.get_empty = Mock( + side_effect=lambda _, shape, **kwargs: torch.empty(tuple(shape), dtype=kwargs["dtype"]) + ) + metadata._create_kv_lens_2d_buffer = Mock() + metadata._create_radix_aux_buffers = Mock() + metadata.create_expanded_buffers = Mock() + + with patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.prefer_pinned", + return_value=False, + ): + metadata.create_buffers_for_indexer() + + buffer = metadata.shared_topk_indices + assert buffer.shape == (4, 3) + + context_topk = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.int32) + generation_topk = torch.tensor([[2, 1, 0], [3, 2, 1]], dtype=torch.int32) + full_backend = SimpleNamespace( + indexer=SimpleNamespace( + forward_from_projected=Mock(side_effect=[context_topk, generation_topk]) + ), + get_local_layer_idx=Mock(return_value=0), + ) + shared_backend = SimpleNamespace( + indexer=None, + get_local_layer_idx=Mock(return_value=1), + ) + metadata._num_ctx_tokens = context_topk.shape[0] + metadata._num_tokens = context_topk.shape[0] + generation_topk.shape[0] + backend_args = DSABackendForwardArgs(indexer_intermediates=[]) + + def _predict(backend, input_type): + return DSATrtllmAttention.sparse_attn_predict( + backend, + torch.empty((2, 1)), + None, + metadata, + AttentionForwardArgs( + attention_input_type=input_type, + sparse_backend_args=backend_args, + ), + )[0] + + with patch( + "tensorrt_llm._torch.attention_backend.sparse.dsa.backend." + "transform_local_topk_and_prepare_pool_view", + side_effect=lambda topk, *_: (topk.clone(), None), + ): + _predict(full_backend, AttentionInputType.context_only) + _predict(full_backend, AttentionInputType.generation_only) + torch.testing.assert_close( + metadata.shared_topk_indices, + torch.cat([context_topk, generation_topk]), + ) + torch.testing.assert_close( + _predict(shared_backend, AttentionInputType.context_only), + context_topk, + ) + torch.testing.assert_close( + _predict(shared_backend, AttentionInputType.generation_only), + generation_topk, + ) + + 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) + metadata.prepare_dense_topk_indices = Mock() + metadata.on_update_kv_lens() + + assert metadata.shared_topk_indices is buffer + + def test_indexer_post_load_weights_caches_fused_weight(): indexer = Indexer.__new__(Indexer) torch.nn.Module.__init__(indexer) @@ -3314,7 +3422,6 @@ def _run_indexer(): k_fp8=dummy, k_scale=dummy, weights=dummy, - update_k_cache=False, ) out1 = _run_indexer() 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 efbe2e2e56d2..fec8746fcc48 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 @@ -21,6 +21,7 @@ from dataclasses import dataclass, field from types import SimpleNamespace from typing import List, Optional +from unittest.mock import Mock import pytest import torch @@ -33,7 +34,7 @@ PositionalEmbeddingParams, RopeParams, ) -from tensorrt_llm._torch.attention_backend.sparse.dsa import DSACacheManager +from tensorrt_llm._torch.attention_backend.sparse.dsa import DSABackendForwardArgs, DSACacheManager from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig @@ -867,6 +868,9 @@ def yarn_get_mscale(scale=1, mscale=1): topk_indices = _build_sparse_topk_indices_context( context_sequence_lengths, SPARSE_TOPK, device ) + ctx_layers[layer_idx].indexer.forward_from_projected = Mock( + return_value=topk_indices + ) result = ctx_layers[layer_idx].forward( fused_q.clone(), None, @@ -875,7 +879,7 @@ def yarn_get_mscale(scale=1, mscale=1): attention_input_type=AttentionInputType.context_only, latent_cache=latent_cache, q_pe=q_pe, - topk_indices=topk_indices, + sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=[]), ) k_pe_ref = _rotate_k_pe_for_ctx(k_pe, rope_cos_sin, context_sequence_lengths) latent_cache_ref = torch.cat([compressed_kv, k_pe_ref], dim=-1) @@ -958,6 +962,9 @@ def yarn_get_mscale(scale=1, mscale=1): topk_indices = _build_sparse_topk_indices_generation( cached_lens, generation_seq_len_q, SPARSE_TOPK, device ) + gen_layers[layer_idx].indexer.forward_from_projected = Mock( + return_value=topk_indices + ) result = gen_layers[layer_idx].forward( fused_q, @@ -973,7 +980,7 @@ def yarn_get_mscale(scale=1, mscale=1): mla_bmm1_scale=mla_bmm1_scale, mla_bmm2_scale=mla_bmm2_scale, quant_q_buffer=quant_q_buffer, - topk_indices=topk_indices, + sparse_backend_args=DSABackendForwardArgs(indexer_intermediates=[]), ) ref_result, latent_cache_ref = calculate_ref_result_gen( fused_q, diff --git a/tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py b/tests/unittest/_torch/attention/sparse/dsa/test_kernels.py similarity index 95% rename from tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py rename to tests/unittest/_torch/attention/sparse/dsa/test_kernels.py index ed0ff4a10357..98f3d6734acb 100644 --- a/tests/unittest/_torch/attention/sparse/test_triton_gather_k_cache.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_kernels.py @@ -1,7 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import pytest import torch -from tensorrt_llm._torch.attention_backend.sparse.kernel import triton_gather_k_cache +from tensorrt_llm._torch.attention_backend.sparse.dsa.kernels import triton_gather_k_cache def reference_gather_k_cache( 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 36a788d3fafa..40cf00bf0f74 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 @@ -36,6 +36,10 @@ RopeParams, ) from tensorrt_llm._torch.attention_backend.sparse.dsa import DSACacheManager +from tensorrt_llm._torch.attention_backend.sparse.dsa.module import ( + _forward_dsa_attn, + should_use_short_mha, +) from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig @@ -384,19 +388,20 @@ def _make_metadata( def _run_forward( - mla, q, compressed_kv, k_pe, latent_cache, position_ids, metadata, topk_indices=None + mla, q, compressed_kv, k_pe, latent_cache, position_ids, metadata, indexer_inputs=None ): - """Run forward_context_sparse_mla on cloned inputs and return the output tensor.""" + """Run the production DSA attention dispatcher on cloned inputs.""" output = torch.empty(q.shape[0], NUM_HEADS * V_HEAD_DIM, dtype=q.dtype, device=q.device) - mla.forward_context_sparse_mla( + _forward_dsa_attn( + mla, q=q.clone(), compressed_kv=compressed_kv.clone(), k_pe=k_pe.clone(), - attn_metadata=metadata, - output=output, latent_cache=latent_cache.clone(), - topk_indices=topk_indices, + indexer_intermediates=indexer_inputs or [], position_ids=position_ids.clone(), + attn_metadata=metadata, + output=output, ) return output @@ -450,6 +455,7 @@ def test_threshold_zero_disables_mha(): @pytest.mark.skipif(get_sm_version() < 90, reason="MLA requires SM90+") +@torch.inference_mode() def test_standard_path_when_exceeds_threshold(): """When total tokens > threshold, the standard absorption path is used.""" device = torch.device("cuda") @@ -470,21 +476,15 @@ def test_standard_path_when_exceeds_threshold(): qr = torch.randn(total_tokens, Q_LORA_RANK, dtype=torch.bfloat16, device=device) metadata = _make_metadata(attn_cls, seq_lens, kv_mgr, mapping, sparse_config) - topk_indices = mla.mqa.indexer( - qr, - hidden_states, - metadata, - position_ids, - ) + inputs = list(mla.mqa.indexer.pre_indexer_proj(qr, hidden_states, position_ids)) - output = _run_forward( - mla, q, compressed_kv, k_pe, latent_cache, position_ids, metadata, topk_indices - ) + output = _run_forward(mla, q, compressed_kv, k_pe, latent_cache, position_ids, metadata, inputs) assert torch.isfinite(output).all() kv_mgr.shutdown() @pytest.mark.skipif(get_sm_version() < 90, reason="MLA requires SM90+") +@torch.inference_mode() def test_agrees_with_absorption_path(): """Short MHA and absorption paths produce numerically close results.""" device = torch.device("cuda") @@ -525,16 +525,15 @@ def _run(mla_module): mla_module.short_seq_mha_threshold > 0 and total_tokens <= mla_module.short_seq_mha_threshold ) - topk = None + inputs = None if not use_short: - topk = mla_module.mqa.indexer( - qr.clone(), - hidden_states.clone(), - meta, - position_ids.clone(), + inputs = list( + mla_module.mqa.indexer.pre_indexer_proj( + qr.clone(), hidden_states.clone(), position_ids.clone() + ) ) out = _run_forward( - mla_module, q, compressed_kv, k_pe, latent_cache, position_ids, meta, topk + mla_module, q, compressed_kv, k_pe, latent_cache, position_ids, meta, inputs ) kv_mgr.shutdown() return out @@ -665,10 +664,10 @@ def test_chunked_context_rejects_when_kv_exceeds_threshold(): pos_c2 = torch.arange(cached_per_seq[0], total_per_seq[0], device=device, dtype=torch.int32) # max_ctx_kv_len (96) > threshold (80) -> short MHA should NOT be used. - assert not mla._should_use_short_mha(meta_c2, pos_c2) + assert not should_use_short_mha(mla, meta_c2, pos_c2) # With threshold large enough for the full KV -> short MHA IS used. mla.short_seq_mha_threshold = total_per_seq[0] + 100 - assert mla._should_use_short_mha(meta_c2, pos_c2) + assert should_use_short_mha(mla, meta_c2, pos_c2) kv_mgr.shutdown() diff --git a/tests/unittest/_torch/attention/sparse/kernel/test_triton.py b/tests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py similarity index 98% rename from tests/unittest/_torch/attention/sparse/kernel/test_triton.py rename to tests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py index 089513644ee8..abaffde6b686 100644 --- a/tests/unittest/_torch/attention/sparse/kernel/test_triton.py +++ b/tests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py @@ -1,9 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import math import pytest import torch -from tensorrt_llm._torch.attention_backend.sparse.kernel import ( +from tensorrt_llm._torch.attention_backend.sparse.rocket.kernels import ( triton_bmm, triton_rocket_paged_kt_cache_bmm, triton_topk, diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index 58d651631608..330edb63a57a 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -11,7 +11,7 @@ import torch from tensorrt_llm._torch.attention_backend.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention -from tensorrt_llm._torch.attention_backend.sparse.utils import _resolve_minimax_m3_backend_cls +from tensorrt_llm._torch.attention_backend.sparse.registry import _resolve_minimax_m3_backend_cls from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py index 2b436084ebc6..8d390399ed73 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py @@ -1,10 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Unit tests for sparse attention with TrtllmAttention backend. """ import math from dataclasses import dataclass +from types import ModuleType from typing import List, Optional, Tuple +from unittest.mock import Mock import pytest import torch @@ -12,12 +29,18 @@ import tensorrt_llm from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs -from tensorrt_llm._torch.attention_backend.sparse.kernel import ( +from tensorrt_llm._torch.attention_backend.sparse.dsa.kernels import ( triton_convert_req_index_to_global_index, ) -from tensorrt_llm._torch.attention_backend.sparse.params import SparseParams +from tensorrt_llm._torch.attention_backend.sparse.hooks import ( + SparseAttnHooks, + get_sparse_attn_hooks, + prepare_sparse_runtime_params, +) +from tensorrt_llm._torch.attention_backend.sparse.params import SparseParams, SparseRuntimeParams from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.modules.mla import MLA from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str from tensorrt_llm.bindings.executor import KvCacheConfig @@ -151,6 +174,76 @@ def sparse_attn_predict(self, q, k, metadata, forward_args: AttentionForwardArgs return self._sparse_attn_indices, self._sparse_attn_offsets +def test_sparse_runtime_params() -> None: + attention = TestSparseAttention.__new__(TestSparseAttention) + attention.sparse_params = MockSparseParams() + attention._sparse_kv_indices = torch.tensor([1], dtype=torch.int32) + attention._sparse_kv_offsets = torch.tensor([0, 1], dtype=torch.int32) + attention._sparse_attn_indices = torch.tensor([2], dtype=torch.int32) + attention._sparse_attn_offsets = None + forward_args = AttentionForwardArgs( + sparse_runtime_params=SparseRuntimeParams(sparse_mla_topk_lens=torch.tensor([3])) + ) + + runtime_params = prepare_sparse_runtime_params( + attention, torch.empty(0), None, None, forward_args + ) + + assert runtime_params.sparse_kv_indices is attention._sparse_kv_indices + assert runtime_params.sparse_kv_offsets is attention._sparse_kv_offsets + assert runtime_params.sparse_attn_indices is attention._sparse_attn_indices + assert runtime_params.sparse_attn_offsets is None + assert runtime_params.sparse_attn_indices_block_size == 1 + assert ( + runtime_params.sparse_mla_topk_lens + is forward_args.sparse_runtime_params.sparse_mla_topk_lens + ) + + +def test_invalid_sparse_attn_hook() -> None: + hook_module = ModuleType("invalid_sparse_attn_hook") + hook_module.forward_sparse_attn = lambda module: None + with pytest.raises(TypeError, match="expected"): + SparseAttnHooks.from_module("test", hook_module) + + +def test_mla_backend_only_forward() -> None: + backend_only_module = ModuleType("backend_only_sparse_attention") + backend_only_module.sparse_params = MockSparseParams() + backend_only_module.sparse_params.algorithm = "skip_softmax" + hooks = get_sparse_attn_hooks(backend_only_module) + assert not hooks + + mla = MLA.__new__(MLA) + torch.nn.Module.__init__(mla) + mla.sparse_attn_hooks = hooks + default_forward = Mock() + mla._forward_impl = default_forward + hidden_states = torch.empty(0) + attn_output = [torch.empty(0)] + + MLA.forward_impl(mla, None, hidden_states, None, attn_output) + + default_forward.assert_called_once_with( + None, + hidden_states, + None, + attn_output[0], + latent_cache_gen=None, + ) + + +def test_sparse_runtime_params_without_prediction() -> None: + attention = TrtllmAttention.__new__(TrtllmAttention) + attention.sparse_params = MockSparseParams() + + runtime_params = prepare_sparse_runtime_params( + attention, torch.empty(0), None, None, AttentionForwardArgs() + ) + + assert runtime_params == SparseRuntimeParams() + + def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """Repeat kv heads to match query heads.""" batch, num_key_value_heads, slen, head_dim = hidden_states.shape 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 72454df6fe86..120771e3a884 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py @@ -28,10 +28,16 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.attention_backend.interface import ( AttentionBackend, PositionalEmbeddingParams, RopeParams) -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import ( - DeepseekV4CacheManager, make_deepseek_v4_sparse_metadata_params) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import \ + DeepseekV4CacheManager +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module import \ + forward_context_sparse_attn as forward_context_deepseek_v4_mla +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module import \ + forward_generation_sparse_attn as forward_generation_deepseek_v4_mla from tensorrt_llm._torch.attention_backend.sparse.dsa import (HAS_FAST_HADAMARD, DSACacheManager) +from tensorrt_llm._torch.attention_backend.sparse.dsa.module import \ + _forward_dsa_attn from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig @@ -1297,9 +1303,8 @@ def populate_gen_dsa_kv_cache(mla: MLA, AttentionCls: AttentionBackend, offset:offset + cached_len, pretrained_config.kv_lora_rank:].clone() offset += cached_len - print( - f" - Allocated+populated cache for gen request {req_idx}: {cached_len} cached + 1 new = {cached_len+1} tokens" - ) + print(f" - Allocated+populated cache for gen request {req_idx}: " + f"{cached_len} cached + 1 new = {cached_len + 1} tokens") kv_cache_for_ref = {} kv_cache_for_ref["compressed_kv"] = all_cached_compressed_kv @@ -1494,6 +1499,7 @@ def populate_gen_deepseek_v4_kv_cache( return kv_cache_for_ref +@torch.inference_mode() def mla_forward_impl_with_dsa_wo_linear(mla, attn_metadata, q, qr, compressed_kv, k_pe, latent_cache, hidden_states, position_ids, dtype, @@ -1509,44 +1515,25 @@ def mla_forward_impl_with_dsa_wo_linear(mla, attn_metadata, q, qr, dtype=dtype, device=device) - topk_indices_local = mla.mqa.indexer( - qr, - hidden_states, - attn_metadata, + indexer_intermediates = list( + mla.mqa.indexer.pre_indexer_proj(qr, hidden_states, position_ids)) + _forward_dsa_attn( + mla, + q, + compressed_kv, + k_pe, + latent_cache, + indexer_intermediates, position_ids, + attn_metadata, + output, ) - - # Validate indexer output against expected causal indices (since seq_len < topk=2048) if num_contexts > 0: - # Transform context indices from local to global - ctx_topk_local = topk_indices_local[:num_ctx_tokens] - - mla.forward_context_sparse_mla( - q=q[:num_ctx_tokens], - compressed_kv=compressed_kv[:num_ctx_tokens], - k_pe=k_pe[:num_ctx_tokens], - attn_metadata=attn_metadata, - output=output[:num_ctx_tokens], - latent_cache=latent_cache[:num_ctx_tokens], - topk_indices=ctx_topk_local, # Use global indices - ) print( f" ✓ Context forward: {num_ctx_tokens} tokens from {num_contexts} requests" ) if num_generations > 0: - # Transform generation indices from local to global - gen_topk_local = topk_indices_local[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens] - mla.forward_generation_sparse_mla( - q=q[num_ctx_tokens:], - compressed_kv=compressed_kv[num_ctx_tokens:], - k_pe=k_pe[num_ctx_tokens:], - attn_metadata=attn_metadata, - output=output[num_ctx_tokens:], - latent_cache=latent_cache[num_ctx_tokens:], - topk_indices=gen_topk_local, # Use global indices - ) print( f" ✓ Generation forward: {num_gen_tokens} tokens from {num_generations} requests" ) @@ -1615,7 +1602,8 @@ def mla_forward_impl_with_deepseek_v4_wo_linear(mla, if num_contexts > 0: ctx_topk_local = topk_indices_for_forward[:num_ctx_tokens] - mla.forward_context_sparse_mla( + forward_context_deepseek_v4_mla( + mla, q=q[:num_ctx_tokens], compressed_kv=compressed_kv[:num_ctx_tokens], k_pe=k_pe[:num_ctx_tokens], @@ -1633,7 +1621,8 @@ def mla_forward_impl_with_deepseek_v4_wo_linear(mla, if num_generations > 0: gen_topk_local = topk_indices_for_forward[num_ctx_tokens:num_tokens] - mla.forward_generation_sparse_mla( + forward_generation_deepseek_v4_mla( + mla, q=q[num_ctx_tokens:], compressed_kv=compressed_kv[num_ctx_tokens:], k_pe=k_pe[num_ctx_tokens:], @@ -1661,8 +1650,8 @@ def test_forward_sparse_mla_unified(batch_name, kv_cache_dtype: str, sparse_attn_algo: str): """Test sparse MLA attention for pure prefill, pure decode, and mixed batches.""" print( - f"\n{'='*80}\nTesting: {batch_name}, sparse_attn_algo: {sparse_attn_algo}, kv_cache_dtype: {kv_cache_dtype}\n{'='*80}" - ) + f"\n{'=' * 80}\nTesting: {batch_name}, sparse_attn_algo: {sparse_attn_algo}, " + f"kv_cache_dtype: {kv_cache_dtype}\n{'=' * 80}") if sparse_attn_algo == "deepseek_v4" and get_sm_version() < 100: pytest.skip( "DeepSeek-V4 sparse MLA unittest is not supported on pre-Blackwell architectures" @@ -1914,7 +1903,7 @@ def yarn_get_mscale(scale=1, mscale=1): if num_layers > 1: print(f" Testing layer {layer_idx} of {num_layers} (multi-layer pool)") else: - print(f" Testing single layer (baseline)") + print(" Testing single layer (baseline)") # For deepseek_v4: derive freqs_cis from the production RotaryEmbedding to # guarantee matching frequencies (the reference precompute_freqs_cis may @@ -2007,7 +1996,7 @@ def yarn_get_mscale(scale=1, mscale=1): # Allocate and pre-populate KV cache in batch order [context...][generation...] - print(f" Allocating and pre-populating cache...") + print(" Allocating and pre-populating cache...") # Allocate context requests first for req_idx in ctx_indices: @@ -2028,9 +2017,6 @@ def yarn_get_mscale(scale=1, mscale=1): # Allocate and pre-populate generation requests (batched) if gen_indices: - gen_cached_lens = [ - cached_lens[i] for i in gen_indices if cached_lens[i] > 0 - ] gen_with_cache = [i for i in gen_indices if cached_lens[i] > 0] # Allocate all generation caches for req_idx in gen_with_cache: @@ -2059,7 +2045,7 @@ def yarn_get_mscale(scale=1, mscale=1): raise ValueError( f"Invalid sparse attention algorithm: {sparse_attn_algo}") - print(f" ✓ KV cache allocated and pre-populated") + print(" ✓ KV cache allocated and pre-populated") # Generate inputs directly in batch order [context...][generation...] batch_order = ctx_indices + gen_indices @@ -2099,11 +2085,7 @@ def yarn_get_mscale(scale=1, mscale=1): num_heads * v_head_dim, dtype=dtype, device=device) - if sparse_config.algorithm == "deepseek_v4": - sparse_metadata_params = make_deepseek_v4_sparse_metadata_params( - sparse_config) - else: - sparse_metadata_params = sparse_config.to_sparse_metadata_params() + sparse_metadata_params = sparse_config.to_sparse_metadata_params() attn_metadata = AttentionCls.Metadata( seq_lens=torch.tensor(batch_query_lens, dtype=torch.int), @@ -2349,20 +2331,21 @@ def yarn_get_mscale(scale=1, mscale=1): def test_forward_sparse_mla_unified_fused_q_fp8(monkeypatch, batch_name): """Regression test for the sparse-MLA context-branch fused FP8 Q-quant wiring. The wo-linear helper bypasses `_q_branch`, so we monkey-patch - `forward_context_sparse_mla` to manually populate the fused buffers - (mimicking `_deepseek_v4_q_b_layernorm_fused_fp8`) and replace q with + `forward_context_sparse_attn` to manually populate the fused buffers + (mimicking `_q_b_layernorm_fused_fp8`) and replace q with NaN. Without the C++ wiring fix the legacy standalone quantize runs over the NaN placeholder and the output diverges from the reference. """ monkeypatch.setenv("TRTLLM_DISABLE_FUSED_Q_FP8_QUANT", "0") - from tensorrt_llm._torch.modules.mla import MLA - original_fwd = MLA.forward_context_sparse_mla + from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import \ + module as deepseek_v4_module + + original_fwd = deepseek_v4_module.forward_context_sparse_attn def patched_fwd(self, q, compressed_kv, k_pe, attn_metadata, output, **kwargs): - if (self.is_deepseek_v4 and self.qk_head_dim == 512 - and self.kv_lora_rank == 448 + if (self.qk_head_dim == 512 and self.kv_lora_rank == 448 and bool(getattr(self.mqa, "has_fp8_kv_cache", False))): num_tokens = q.shape[0] num_heads = self.num_heads_tp @@ -2385,7 +2368,8 @@ def patched_fwd(self, q, compressed_kv, k_pe, attn_metadata, output, return original_fwd(self, q, compressed_kv, k_pe, attn_metadata, output, **kwargs) - monkeypatch.setattr(MLA, 'forward_context_sparse_mla', patched_fwd) + monkeypatch.setattr(deepseek_v4_module, 'forward_context_sparse_attn', + patched_fwd) test_forward_sparse_mla_unified(batch_name=batch_name, kv_cache_dtype="fp8", sparse_attn_algo="deepseek_v4") diff --git a/tests/unittest/_torch/attention/sparse/test_triton_topk.py b/tests/unittest/_torch/attention/sparse/test_triton_topk.py deleted file mode 100644 index 5375619bb411..000000000000 --- a/tests/unittest/_torch/attention/sparse/test_triton_topk.py +++ /dev/null @@ -1,228 +0,0 @@ -import pytest -import torch - -from tensorrt_llm._torch.attention_backend.sparse.kernel import triton_topk - - -def pytorch_reference_topk( - input_tensor: torch.Tensor, - kv_offsets: torch.Tensor, - kv_lens: torch.Tensor, - topk: int, -) -> torch.Tensor: - """ - Args: - input_tensor: Input scores [num_kv_heads, sum(kv_lens)] - kv_offsets: KV offsets [batch_size + 1] - kv_lens: KV lengths [batch_size] - topk: TopK parameter - - Returns: - output_indices: Padded indices [num_kv_heads, batch_size, topk] - """ - num_kv_heads = input_tensor.shape[0] - batch_size = len(kv_lens) - device = input_tensor.device - - # Compute max sequence length for padding - max_seq_len = kv_lens.max().item() - - # Ensure padding size >= topk - pad_size = max(max_seq_len, topk) - - # Create padded tensor [num_kv_heads, batch_size, pad_size] - padded_tensor = torch.full( - (num_kv_heads, batch_size, pad_size), float("-inf"), dtype=input_tensor.dtype, device=device - ) - - # Fill in actual values - for batch_idx in range(batch_size): - start = kv_offsets[batch_idx].item() - end = kv_offsets[batch_idx + 1].item() - seq_len = kv_lens[batch_idx].item() - - for head_idx in range(num_kv_heads): - padded_tensor[head_idx, batch_idx, :seq_len] = input_tensor[head_idx, start:end] - - # Perform batch topk: [num_kv_heads, batch_size, pad_size] -> [num_kv_heads, batch_size, topk] - topk_values, topk_indices = torch.topk( - padded_tensor, - k=topk, - dim=-1, - largest=True, - ) - - # Mask out invalid indices based on each batch's seq_len - seq_lens_expanded = kv_lens.to(device).unsqueeze(0).unsqueeze(-1) # [1, batch_size, 1] - # topk_indices: [num_kv_heads, batch_size, topk] - mask = topk_indices >= seq_lens_expanded - topk_indices.masked_fill_(mask, -1) - - return topk_indices - - -def triton_topk_wrapper( - input_tensor: torch.Tensor, - kv_offsets: torch.Tensor, - kv_lens: torch.Tensor, - topk: int, -) -> torch.Tensor: - """ - Args: - input_tensor: Input scores [num_kv_heads, sum(kv_lens)] - kv_offsets: KV offsets [batch_size + 1] - kv_lens: KV lengths [batch_size] - topk: TopK parameter - - Returns: - output_indices: Padded indices [num_kv_heads, batch_size, topk] - """ - num_kv_heads = input_tensor.shape[0] - batch_size = len(kv_lens) - device = input_tensor.device - - sparse_lens = torch.tensor( - [min(topk, seq_len.item()) for seq_len in kv_lens], dtype=torch.int32, device=device - ) - sparse_offsets = torch.cat( - [torch.zeros(1, dtype=torch.int32, device=device), torch.cumsum(sparse_lens, dim=0)] - ).to(device) - total_sparse_indices = sparse_offsets[-1].item() - - output_indices_flat = triton_topk( - input_tensor, kv_offsets, sparse_offsets, total_sparse_indices, topk - ) - - # Convert flat format to padded format [num_kv_heads, batch_size, topk] - output_indices_padded = torch.full( - (num_kv_heads, batch_size, topk), -1, dtype=torch.int32, device=device - ) - - for batch_idx in range(batch_size): - start = sparse_offsets[batch_idx].item() - end = sparse_offsets[batch_idx + 1].item() - actual_len = end - start - - for head_idx in range(num_kv_heads): - output_indices_padded[head_idx, batch_idx, :actual_len] = output_indices_flat[ - head_idx, start:end - ] - - return output_indices_padded - - -def compute_overlap_ratio( - triton_indices: torch.Tensor, - reference_indices: torch.Tensor, - kv_lens: torch.Tensor, -) -> float: - """ - Args: - triton_indices: Triton topk results [num_kv_heads, batch_size, topk] - reference_indices: Reference topk results [num_kv_heads, batch_size, topk] - kv_lens: KV lengths [batch_size] - - Returns: - Average overlap ratio across all batches and heads - """ - num_kv_heads = triton_indices.shape[0] - batch_size = triton_indices.shape[1] - - overlap_ratios = [] - - # Compare batch by batch - for batch_idx in range(batch_size): - for head_idx in range(num_kv_heads): - # Extract indices for this batch and head - triton_batch = triton_indices[head_idx, batch_idx, :].cpu().tolist() - reference_batch = reference_indices[head_idx, batch_idx, :].cpu().tolist() - - # Filter out -1 (invalid/padding indices) - triton_set = set([x for x in triton_batch if x >= 0]) - reference_set = set([x for x in reference_batch if x >= 0]) - - if len(reference_set) > 0: - overlap = len(triton_set & reference_set) - overlap_ratio = overlap / len(reference_set) - overlap_ratios.append(overlap_ratio) - - if len(overlap_ratios) == 0: - return 1.0 - - return sum(overlap_ratios) / len(overlap_ratios) - - -@pytest.mark.parametrize( - "batch_size,seq_lens,num_kv_heads,topk", - [ - # Single batch, seq_len > topk - (1, [3000], 1, 2048), - # Single batch, seq_len < topk - (1, [1000], 8, 2048), - # Multiple batches, mixed seq_len (some < topk, some > topk) - (6, [50, 150, 80, 300, 100, 256], 8, 128), - (6, [1000, 2500, 1500, 3000, 1800, 4000], 8, 2048), - ], -) -def test_topk_kernel(batch_size, seq_lens, num_kv_heads, topk): - device = torch.device("cuda") - dtype = torch.float32 - - kv_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) - - kv_offsets = torch.cat( - [torch.zeros(1, dtype=torch.int32, device=device), torch.cumsum(kv_lens, dim=0)] - ).to(device) - - total_tokens = kv_offsets[-1].item() - - input_tensor = torch.randn((num_kv_heads, total_tokens), dtype=dtype, device=device) - - triton_output = triton_topk_wrapper( - input_tensor=input_tensor, - kv_offsets=kv_offsets, - kv_lens=kv_lens, - topk=topk, - ) - - reference_output = pytorch_reference_topk( - input_tensor=input_tensor, - kv_offsets=kv_offsets, - kv_lens=kv_lens, - topk=topk, - ) - - overlap_ratio = compute_overlap_ratio( - triton_output, - reference_output, - kv_lens, - ) - - min_threshold = 0.99 - - print(f"overlap_ratio: {overlap_ratio}") - - assert overlap_ratio >= min_threshold, ( - f"Overlap ratio {overlap_ratio:.4f} is too low (< {min_threshold})" - ) - - -if __name__ == "__main__": - print("\n" + "=" * 80) - print("Testing Triton TopK Kernel") - print("=" * 80) - - # Single batch tests - print("\n--- Single batch, seq_len > topk ---") - test_topk_kernel(1, [3000], 8, 2048) - - print("\n--- Single batch, seq_len < topk ---") - test_topk_kernel(1, [1000], 8, 2048) - - print("\n--- Multiple batches, mixed seq_len ---") - test_topk_kernel(6, [50, 150, 80, 300, 100, 256], 8, 128) - test_topk_kernel(6, [1000, 2500, 1500, 3000, 1800, 4000], 8, 2048) - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) diff --git a/tests/unittest/_torch/attention/test_attention_op_sync.py b/tests/unittest/_torch/attention/test_attention_op_sync.py index cfd56978b848..7782d47af5b2 100644 --- a/tests/unittest/_torch/attention/test_attention_op_sync.py +++ b/tests/unittest/_torch/attention/test_attention_op_sync.py @@ -24,7 +24,7 @@ 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 - nested dataclass sub-bags like ``SparsePrediction``) is consumed at + nested dataclass sub-bags like ``SparseRuntimeParams``) is consumed at the call site — directly, transitively via a @property of the containing class, or listed in ``_THOP_EXCLUDED_FIELDS``. 4. Every kwarg passed as a literal constant matches an entry in @@ -73,14 +73,14 @@ "skip_softmax_threshold_scale_factor_decode": ( "forward_args", ( - "skip_softmax_kernel_params", + "sparse_runtime_params", "threshold_scale_factor_decode", ), ), "skip_softmax_threshold_scale_factor_prefill": ( "forward_args", ( - "skip_softmax_kernel_params", + "sparse_runtime_params", "threshold_scale_factor_prefill", ), ), @@ -571,7 +571,7 @@ def _verify_consumed(cls, chains: set[tuple[str, ...]], excluded=frozenset()): def test_every_forward_args_field_is_consumed(): """Recursively check that every dataclass field reachable from ``AttentionForwardArgs`` (including nested sub-bags such as - ``SparsePrediction``) is consumed at the call site, transitively + ``SparseRuntimeParams``) is consumed at the call site, transitively via @property where applicable, or listed in ``_THOP_EXCLUDED_FIELDS``. """ _verify_consumed( diff --git a/tests/unittest/_torch/compilation/test_remove_copy_pass.py b/tests/unittest/_torch/compilation/test_remove_copy_pass.py index ab5d7afad17a..f13a1043bbf4 100644 --- a/tests/unittest/_torch/compilation/test_remove_copy_pass.py +++ b/tests/unittest/_torch/compilation/test_remove_copy_pass.py @@ -70,11 +70,10 @@ def test_remove_copy_for_mutates_args_restores_optional_none() -> None: "position_ids": None, "layer_idx": "0", "latent_cache_gen": None, - "enable_dsv4_epilogue_fusion": False, "_all_bases": (output,), "_output_base_index": 0, - "_dsv4_output_base_index": None, - "_dsv4_output_sf_base_index": None, + "_sparse_output_base_index": None, + "_sparse_output_sf_base_index": None, }, ) mutated_output = graph.call_function(getitem, args=(functionalized, 1)) @@ -86,8 +85,8 @@ def test_remove_copy_for_mutates_args_restores_optional_none() -> None: inplace_nodes = [node for node in graph.nodes if node.target == inplace_func] assert len(inplace_nodes) == 1 assert inplace_nodes[0].kwargs["output"] is output - assert inplace_nodes[0].kwargs["dsv4_output"] is None - assert inplace_nodes[0].kwargs["dsv4_output_sf"] is None + assert inplace_nodes[0].kwargs["sparse_output"] is None + assert inplace_nodes[0].kwargs["sparse_output_sf"] is None assert clone.args[0] is output graph.lint() @@ -107,11 +106,10 @@ def test_remove_copy_for_mutates_args_rejects_getitem_for_optional_none( "position_ids": None, "layer_idx": "0", "latent_cache_gen": None, - "enable_dsv4_epilogue_fusion": False, "_all_bases": (output,), "_output_base_index": 0, - "_dsv4_output_base_index": None, - "_dsv4_output_sf_base_index": None, + "_sparse_output_base_index": None, + "_sparse_output_sf_base_index": None, }, ) optional_output = graph.call_function(getitem, args=(functionalized, 2)) @@ -121,13 +119,13 @@ def test_remove_copy_for_mutates_args_rejects_getitem_for_optional_none( monkeypatch.setattr( remove_copy_pass, "inplace_info", - lambda: {inplace_func: {1: "output", 2: "dsv4_output"}}, + lambda: {inplace_func: {1: "output", 2: "sparse_output"}}, ) with pytest.raises( AssertionError, match=( - "getitem user for optional output 'dsv4_output' has no " + "getitem user for optional output 'sparse_output' has no " "base tensor -- graph is malformed" ), ): diff --git a/tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py b/tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py index e29d45b94a20..0309b9dcf6cd 100644 --- a/tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py +++ b/tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py @@ -205,37 +205,37 @@ def test_deepseek_v4_q_norm_fused_fp8_zero_rows(): @pytest.mark.parametrize("num_tokens", [1, 7, 129]) -def test_deepseek_v4_q_b_layernorm_fused_fp8_returns_3d_q_pe(num_tokens): +def test_dsv4_fused_q_norm_3d_q_pe(num_tokens): """Lock down the q_pe dim==3 contract expected by thop.attention's sparse-MLA context branch (TORCH_CHECK(q_pe->dim() == 3)).""" - import types - from types import SimpleNamespace - - from tensorrt_llm._torch.modules.mla import MLA - num_heads = 16 qk_head_dim = 512 kv_lora_rank = 448 rope_dim = qk_head_dim - kv_lora_rank - stub = SimpleNamespace( - num_heads_tp=num_heads, - qk_head_dim=qk_head_dim, - kv_lora_rank=kv_lora_rank, - q_b_layernorm=SimpleNamespace(variance_epsilon=1e-6), - ) - fused = types.MethodType(MLA._deepseek_v4_q_b_layernorm_fused_fp8, stub) q_proj = torch.randn( num_tokens, num_heads * qk_head_dim, dtype=torch.bfloat16, device="cuda" ).contiguous() - - placeholder_q, quant_q_buffer, q_pe, scale = fused(q_proj) + quant_q_buffer = q_proj.new_empty( + (num_tokens, num_heads * qk_head_dim), dtype=torch.float8_e4m3fn + ) + q_pe = q_proj.new_empty((num_tokens, num_heads, rope_dim)) + scale = torch.tensor([1.0], dtype=torch.float32, device="cuda") + torch.ops.trtllm.deepseek_v4_q_norm_fused_fp8( + q_proj, + quant_q_buffer, + q_pe.view(num_tokens, num_heads * rope_dim), + num_heads, + qk_head_dim, + kv_lora_rank, + 1e-6, + scale, + ) assert q_pe.shape == (num_tokens, num_heads, rope_dim) assert q_pe.stride(2) == 1 assert q_pe.is_contiguous() assert quant_q_buffer.shape == (num_tokens, num_heads * qk_head_dim) assert quant_q_buffer.dtype == torch.float8_e4m3fn - assert placeholder_q.data_ptr() == q_proj.data_ptr() assert scale.shape == (1,) and scale.dtype == torch.float32 assert float(scale.item()) == 1.0 diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index b46ed3e477b9..5ca4314a4bd4 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -16,15 +16,13 @@ # from utils.util import default_dtype import tensorrt_llm from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.cache_manager import ( +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import ( DeepseekV4CacheManager, -) -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor import Compressor -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( DeepseekV4Indexer, DeepseekV4TrtllmAttention, DeepseekV4TrtllmAttentionMetadata, ) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor import Compressor from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention from tensorrt_llm._torch.configs.deepseekv4 import DeepseekV4Config from tensorrt_llm._torch.metadata import KVCacheParams @@ -371,18 +369,39 @@ def test_deepseek_v4_q_b_layernorm_differs_from_joint_flat_rms(): def test_deepseek_v4_mla_q_b_layernorm_init_and_forward_shape(): + from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.module import ( + forward_sparse_attn, + initialize_sparse_attn, + ) from tensorrt_llm._torch.modules.mla import MLA - init_src = inspect.getsource(MLA.__init__) - helper_src = inspect.getsource(MLA._deepseek_v4_q_b_layernorm) - forward_src = inspect.getsource(MLA.forward_impl_with_deepseek_v4) + mla_init_src = inspect.getsource(MLA.__init__) + init_src = inspect.getsource(initialize_sparse_attn) + forward_src = inspect.getsource(forward_sparse_attn) + forward_src_dedented = textwrap.dedent(forward_src) + q_norm_node = next( + node + for node in ast.walk(ast.parse(forward_src_dedented)) + if isinstance(node, ast.FunctionDef) and node.name == "_q_b_layernorm" + ) + fused_q_norm_node = next( + node + for node in ast.walk(ast.parse(forward_src_dedented)) + if isinstance(node, ast.FunctionDef) and node.name == "_q_b_layernorm_fused_fp8" + ) + helper_src = ast.get_source_segment(forward_src_dedented, q_norm_node) + fused_helper_src = ast.get_source_segment(forward_src_dedented, fused_q_norm_node) + assert helper_src is not None + assert fused_helper_src is not None init_src_no_ws = "".join(init_src.split()) + fused_helper_src_no_ws = "".join(fused_helper_src.split()) assert "self.q_b_layernorm=RMSNorm(hidden_size=self.qk_head_dim" in init_src_no_ws assert "has_weights=False" in init_src - assert "kv_a_layernorm_hidden_size = (" in init_src assert "self.kv_lora_rank + self.qk_rope_head_dim" in init_src - assert "self.kv_a_layernorm=RMSNorm(hidden_size=kv_a_layernorm_hidden_size" in init_src_no_ws + assert "self.kv_a_layernorm=RMSNorm(" in init_src_no_ws + assert "initialize_sparse_attn" in mla_init_src + assert "deepseek_v4" not in mla_init_src assert "q.dim() == 2" in helper_src assert "self.num_heads_tp * self.qk_head_dim" in helper_src assert "torch.ops.trtllm.deepseek_v4_q_norm" in helper_src @@ -391,7 +410,11 @@ def test_deepseek_v4_mla_q_b_layernorm_init_and_forward_shape(): assert "q.dtype" not in helper_src assert "total_rows" not in helper_src assert "self.q_b_layernorm(" not in helper_src - assert "self._deepseek_v4_q_b_layernorm(q_proj)" in _source_calls(forward_src) + assert "_q_b_layernorm(q_proj)" in _source_calls(forward_src) + assert "q_pe=q_proj.new_empty((num_q_tokens,self.num_heads_tp,rope_dim))" in ( + fused_helper_src_no_ws + ) + assert "torch.ops.trtllm.deepseek_v4_q_norm_fused_fp8(" in fused_helper_src def test_deepseek_v4_compressor_rotate_and_indexer_rope_contracts(): diff --git a/tests/unittest/_torch/modules/test_mla_helix.py b/tests/unittest/_torch/modules/test_mla_helix.py index 114b3c045ff6..e0b8029ccb9a 100644 --- a/tests/unittest/_torch/modules/test_mla_helix.py +++ b/tests/unittest/_torch/modules/test_mla_helix.py @@ -385,7 +385,7 @@ def _run_mla_distributed( ctx_output = input_ctx_rank.new_empty( [input_ctx_rank.shape[0], mla.num_heads_tp * mla.v_head_dim], dtype=input_ctx_rank.dtype ) - mla.forward_impl(position_ids_ctx_rank, input_ctx_rank, attn_metadata, output=ctx_output) + mla.forward_impl(position_ids_ctx_rank, input_ctx_rank, attn_metadata, attn_output=[ctx_output]) # For non-last rank, generate the right latent cache for generation. input_ctx_bs = input_ctx.view(scenario.batch, scenario.ctx_len, scenario.hidden_size) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py index 993ab9af0c5a..fe42add01052 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py @@ -70,7 +70,7 @@ def _prefill_threshold( timestep: Optional[float] = None, ) -> float: assert isinstance(sparse_params, SkipSoftmaxParams) - return sparse_params.scheduler.get_kernel_params( + return sparse_params.scheduler.get_runtime_params( timestep=timestep ).threshold_scale_factor_prefill diff --git a/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py b/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py index cd055f17b1f3..20f98d92b56c 100644 --- a/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py +++ b/tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py @@ -36,7 +36,7 @@ from tensorrt_llm import Mapping from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import DeepseekV4CacheManager -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.params import ( DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO, DEEPSEEK_V4_SPARSE_RATIO, DeepseekV4AttentionType, diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 558a57c70208..f031a962bcfc 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2752,7 +2752,7 @@ def _checkpoint_config(cls) -> dict: @staticmethod def _kernel_params(config: SkipSoftmaxAttentionConfig, **kwargs): sparse_params = config.to_sparse_params(**kwargs) - return sparse_params.scheduler.get_kernel_params() + return sparse_params.scheduler.get_runtime_params() def test_python_api_parses_skip_softmax_config(self): args = TorchLlmArgs(