Sparse Attention#

Overview
Supported Sparse Attentions
Supported Algorithms
- Capability Comparison
- Algorithm Details
Usage with trtllm-bench and trtllm-serve
Further Reading

Overview#

Sparse attention reduces long-context inference cost by avoiding attention work on KV entries that an algorithm considers unimportant. TensorRT LLM separates two parts of that process:

An algorithm selects tokens or blocks, or decides which kernel tiles can be skipped.
An attention implementation consumes that sparse pattern and computes the output.

This distinction matters for support. A kernel that can compute sparse MQA/GQA does not by itself define how a model selects tokens, and therefore is not a standalone user-facing algorithm.

The user-facing sparse_attention_config API is currently prototype and is supported by the PyTorch execution backend. Each public algorithm has a config class selected by its algorithm field. Model-native algorithms usually read their geometry from the checkpoint; avoid overriding those values unless the model-specific guide says they are tunable.

Supported Sparse Attentions#

TensorRT LLM supports sparse computation for MLA, MQA/GQA, and MHA. This section describes the attention and kernel contracts independently of the algorithm that produces the sparse pattern. The public algorithms that connect selectors, cache management, and these attention implementations are listed in Supported Algorithms.

Sparse MLA#

Sparse MLA consumes token-level selections against a model-specific shared KV representation. DeepSeek Sparse Attention selects entries from a low-rank latent KV cache, while DeepSeek-V4 combines compressed full-head non-RoPE K with its corresponding RoPE K. Both prefill and generation are supported, including mixed batches.

Parameter	Support
GPU architecture	SM90, SM100, SM103, SM120, and SM121
Compute phase	Packed prefill and generation, including mixed batches
Attention type	MLA
Head counts	Checkpoint-defined
Q heads per KV head	Not applicable; the model uses a shared KV representation
Head dimensions	DeepSeek-V3.2: QK `192`, V `128`; DeepSeek-V4: QK/V `512`
Input dtype	BF16
Input layout	Model-native MLA inputs
Output dtype	BF16
KV-cache dtype	BF16 or model- and architecture-specific FP8
KV-cache layout	Paged, model-specific shared KV representation
Sparse granularity	Token
Attention semantics	Causal self-attention

Input dtype refers to model-native MLA inputs, which remain BF16. The FP8 KV-cache entry and any internal FP8 staging do not indicate raw FP8 model-input support.

See test_sparse_mla_forward.py for executable sparse MLA examples.

Sparse MQA/GQA#

The table below compares token-sparse and 128-token block-sparse MQA/GQA. The token-sparse path accepts a precomputed token list for each KV head and query token; query heads in the same KV group share that list. The block-sparse path accepts request-local KV-block selections from a paged HND cache. The shared page-sparse generation path described under Sparse MHA also supports MQA and GQA.

These are attention capabilities, not standalone public SparseAttentionConfig algorithms. A user-facing algorithm must also provide the selector, metadata, cache management, and backend integration.

Parameter	Token-sparse	Block-sparse
GPU architecture	SM100 and SM103	SM100 and SM103
Compute phase	Packed prefill and generation, including linear draft tokens	Packed prefill and generation, including linear multi-query and mixed batches
Attention type	MQA and GQA	MQA and GQA
Head counts	Q heads must be divisible by KV heads; no other discrete limit	Q heads must be divisible by KV heads; no other discrete limit
Q heads per KV head	At most 32	`2`, `4`, `8`, or `16`
Head dimensions	Q/K/V: `64`, `80`, `128`, or `256`	Q/K/V: `128`
Input dtype	BF16 or FP16	BF16 or E4M3 FP8
Input layout	Fused QKV	Q `[tokens, q_heads, 128]`; paged K/V `[pages, kv_heads, 128, 128]`
Output dtype	BF16 or FP16 for every supported head dimension; E4M3 FP8 for head dimensions `64`, `128`, and `256`	BF16
KV-cache dtype	BF16 or FP16 for every supported head dimension; E4M3 FP8 for head dimensions `64`, `128`, and `256`	BF16 or E4M3 FP8
KV-cache layout	Paged cache; page size is a power of two and at least 8 tokens	Paged HND cache with page size `128`; supports shuffled physical pages and strided outer-page storage
Sparse granularity	Token	Block (`128` tokens)
Attention semantics	Causal self-attention	Causal self-attention with bottom-right or explicit per-request query offsets

The token-sparse path is JIT-compiled with NVRTC. During linear draft-token generation, each query has its own causal sparse list, including K/V written earlier in the same speculative forward.

For an FP8 KV cache, token-sparse Q is quantized to E4M3 during QKV preprocessing while the model input remains BF16 or FP16. The path supports both BF16 output with an FP8 KV cache and E4M3 FP8 output.

This is distinct from the block-sparse column: its E4M3 input row is a raw Q/K/V contract of that dedicated backend. Raw E4M3 fused QKV is not a token-sparse MQA/GQA input contract.

Backend developers can use test_sparse_mqa_gqa.py as an executable integration example.

Sparse MHA#

The shared page-sparse MHA path consumes block indices and per-request offsets produced by a sparse selector. Sparse MHA computation starts during generation; prefill attention computation remains dense. An algorithm can still compact the retained KV cache after prefill to reduce cache size and later decode work.

Parameter	Support
GPU architecture	SM100 and SM103
Compute phase	Generation, including single-token and linear draft-token inputs
Attention type	MHA
Head counts	Positive and `num_q_heads == num_kv_heads`; no other discrete limit
Q heads per KV head	`1`
Head dimensions	Q/K/V: `64`, `80`, `128`, or `256`
Input dtype	BF16 or FP16
Input layout	Fused QKV
Output dtype	Model dtype for head dimensions `64`, `80`, `128`, and `256`; E4M3 FP8 for head dimensions `64`, `128`, and `256` with an FP8 KV cache
KV-cache dtype	Model dtype for head dimensions `64`, `80`, `128`, and `256`; E4M3 FP8 for head dimensions `64`, `128`, and `256`
KV-cache layout	Paged KV cache; page size is a power of two and at least 8 tokens
Sparse granularity	Positive-size blocks expanded to KV-cache pages
Attention semantics	Causal self-attention

The E4M3 entries above describe FP8 KV-cache and output paths. The fused model QKV input remains BF16 or FP16; raw E4M3 fused QKV is not supported by the page-sparse MHA path.

Backend developers can use test_sparse_mha.py as an executable integration example.

Block-sparse MHA/MQA/GQA#

The generic block-sparse path executes attention over KV blocks that a sparse algorithm selects for each KV head. The algorithm hands its routes to the core forward through the block_sparse_attn_predict hook as BlockSparseForwardInputs: canonical BSR (block_indptr plus block_indices) or a packed block bitmask, optionally with K/V block summaries so unselected blocks contribute a proxy instead of being dropped. Only the block-sparse FMHA library declares supports_block_sparse_inputs, so a request that carries routes is never served by a dense kernel. Both the contiguous prefill path, used by diffusion models that keep no KV cache, and the paged generation path are provided by the vendored PrimTS kernels.

This is an attention capability, not a standalone public SparseAttentionConfig algorithm. A user-facing algorithm must also provide the selector, metadata, and backend integration.

Parameter	Contiguous prefill	Paged generation
GPU architecture	SM100 and SM103	SM100 and SM103
Compute phase	Prefill with separate Q/K/V and no KV cache	Generation with a fixed per-request query length
Attention type	MHA, MQA, and GQA	MHA, MQA, and GQA
Head counts	Q heads must be divisible by KV heads; no other discrete limit	Q heads must be divisible by KV heads; no other discrete limit
Q heads per KV head	Any divisor of the Q head count	Any divisor of the Q head count
Head dimensions	Q/K/V: `128`	Q/K/V: `128`
Input dtype	BF16 or FP16	BF16 or FP16
Input layout	Separate Q `[tokens, q_heads, 128]` and K/V `[tokens, kv_heads, 128]`	Fused QKV
Output dtype	Model dtype	Model dtype
KV-cache dtype	No KV cache	Model dtype
KV-cache layout	No KV cache	Paged HND cache; page size `64` or `128`
Sparse granularity	Q blocks of `q_block_size` tokens by KV blocks of `8`, `16`, `32`, or a positive multiple of `64` tokens, selected per KV head	KV blocks of a positive multiple of `64` tokens, selected per KV head
Sparse routes	BSR or packed bitmask; optional K/V block summaries (proxy routes); optional packed `kv_valid_bits` for ragged KV tails	BSR; live per-request KV lengths and page tables
Attention semantics	Dense or causal self-attention; proxy routes require dense	Causal self-attention

Routes are validated against the kernel’s static profile on every call, and an unsupported request raises instead of degrading to dense attention. A paged request needs a page that holds at least one 64-token route fragment, which is why page sizes below 64 are rejected.

Backend developers can use test_prims_ts_block_sparse.py as an executable integration example; it covers MHA, GQA, and MQA head topologies, both model dtypes, KV block sizes 64 and 128, page sizes 64 and 128, proxy routes, token-validity masks, and CUDA Graph replay.

Supported Algorithms#

The public sparse_attention_config API connects a sparse algorithm to its selector, runtime metadata, cache management, and attention implementation.

`algorithm`	Config class	Sparse mechanism	Attention implementation	Typical use
`rocket`	`RocketSparseAttentionConfig`	Prompt KV eviction, then page-level Top-K selection during decode	TRTLLM or Vanilla	Training-free sparsity for MHA/MQA/GQA models
`dsa`	`DeepSeekSparseAttentionConfig`	Learned token-level indexer followed by sparse MLA	TRTLLM	DeepSeek-V3.2 and compatible model-native DSA architectures
`deepseek_v4`	`DeepSeekV4SparseAttentionConfig`	Sliding-window attention plus compressed sparse or compressed dense history	TRTLLM	DeepSeek-V4 hybrid attention
`minimax_m3`	`MiniMaxM3SparseAttentionConfig`	Learned block selection followed by sparse GQA	Dedicated Triton or packaged block-sparse implementation	MiniMax-M3 sparse layers
`skip_softmax`	`SkipSoftmaxAttentionConfig`	Dynamically skips eligible softmax work inside the FMHA kernel	TRTLLM	Existing full-attention models with calibrated or direct thresholds

All five configs are supported only by the PyTorch execution backend. The “attention implementation” column refers to the attention kernel/backend used inside that execution backend.

Capability Comparison#

Capability	RocketKV	DSA	DeepSeek-V4	MiniMax-M3	Skip Softmax
Sparse prefill computation	No	Yes	Yes	Yes	Yes
Sparse decode computation	Yes	Yes	Yes	Yes	Yes
Reduces retained main KV history	Yes	No	Yes, through model-native compression	No	No
Requires a model-trained selector	No	Yes	Yes	Yes	No
Selection granularity	Token eviction and pages	Tokens	Compressed entries	Blocks	Kernel tiles

“No” for RocketKV prefill means that prompt attention is still computed densely. RocketKV selects which prompt KV entries to retain, so it reduces cache size and later decode work.

Algorithm Details#

RocketKV#

RocketKV is a training-free, two-stage algorithm for MHA, MQA, and GQA architectures. During prefill, it computes dense attention and permanently evicts prompt KV entries beyond a prompt budget. During decode, it scores retained pages and attends to the selected Top-K pages.

RocketKV currently requires CUDA compute capability 10.0 or newer. KV-cache block reuse and chunked prefill must be disabled, and disaggregated serving is not supported.

from tensorrt_llm import LLM, SamplingParams
from tensorrt_llm.llmapi import KvCacheConfig, RocketSparseAttentionConfig

llm = LLM(
    model="<path_or_hf_id>",
    sparse_attention_config=RocketSparseAttentionConfig(
        prompt_budget=2048,
        kt_cache_dtype="float8_e5m2",
    ),
    kv_cache_config=KvCacheConfig(enable_block_reuse=False),
    enable_chunked_prefill=False,
)
outputs = llm.generate(
    ["To be or not to be..."],
    SamplingParams(max_tokens=128),
)

sparse_attention_config:
  algorithm: rocket
  prompt_budget: 2048
  kt_cache_dtype: float8_e5m2
kv_cache_config:
  enable_block_reuse: false
enable_chunked_prefill: false

The TRTLLM and Vanilla attention implementations support RocketKV. The Vanilla implementation requires a BF16 KT cache.

DeepSeek Sparse Attention#

DeepSeek Sparse Attention (DSA) is a model-native mechanism introduced by DeepSeek V3.2. A learned MQA indexer scores the KV history, Top-K selects token indices, and sparse MLA consumes them. Checkpoint fields define the indexer head count, index head dimension, and Top-K; the safest configuration is to let TensorRT LLM load them from the model.

from tensorrt_llm import LLM
from tensorrt_llm.llmapi import DeepSeekSparseAttentionConfig

llm = LLM(
    model="deepseek-ai/DeepSeek-V3.2",
    sparse_attention_config=DeepSeekSparseAttentionConfig(),
)

On supported Blackwell configurations, Guess-Verify-Refine (GVR) can replace the regular decode Top-K dispatcher. The current implementation accepts index_topk values 512, 1024, and 2048, and indexer compression ratios 1 and 4. Unsupported combinations fall back to the production insertion/radix Top-K path.

sparse_attention_config:
  algorithm: dsa
  index_topk: 2048
  enable_heuristic_topk: true

See the DeepSeek V3/V3.2 example for model precision, hardware, parallelism, MTP, chunked-prefill, cache-reuse, and disaggregated-serving support.

DeepSeek-V4 Hybrid Sparse Attention#

DeepSeek-V4 interleaves three model-native attention modes:

sliding-window attention over recent raw tokens;
compressed sparse attention over 4x-compressed history selected by an indexer;
compressed dense attention over 128x-compressed history.

TensorRT LLM normally constructs DeepSeekV4SparseAttentionConfig from the checkpoint. An explicit config overrides matching fields; it must preserve the model’s attention layout. The current implementation requires window_size=128, compression ratios from {1, 4, 128}, Hopper (SM90) or Blackwell (SM100+) GPUs, KV-cache blocks of 128 or 256 tokens, and beam width 1. Hopper requires kv_cache_config.dtype=fp8_ds_mla; on SM120 and SM121, that cache layout requires 256-token blocks.

sparse_attention_config:
  algorithm: deepseek_v4
  window_size: 128
  index_topk: 512

See the DeepSeek-V4 example for checkpoint-derived configuration and deployment constraints.

MiniMax-M3 Block-Sparse GQA#

MiniMax-M3 uses model-native block-sparse GQA in its sparse layers. An index branch scores main KV-cache blocks, forces configured initial/local blocks into the selection, and chooses the remaining Top-K blocks before sparse GQA. Defaults such as four index heads, index dimension 128, block size 128, and 16 selected blocks come from the checkpoint-compatible config.

sparse_attention_config:
  algorithm: minimax_m3

Two implementations are available:

triton is the default reference implementation.
msa uses fmha_sm100 kernels and requires an SM100-family GPU (SM100 or SM103), the fmha_sm100 package, and sparse_block_size=128.

sparse_attention_config:
  algorithm: minimax_m3
  implementation: msa

The sparse path currently has no dense fallback and does not support KV-cache reuse or MTP. See the MiniMax-M3 deployment guide for supported checkpoints and parallel deployment settings.

Skip Softmax Attention#

Skip Softmax Attention, also known as BLASST, dynamically skips eligible work inside a FlashAttention-style kernel. It does not select tokens, alter the model architecture, or reduce KV-cache storage.

The kernel consumes threshold_scale_factor and combines it with sequence length at runtime. You can provide that value directly:

from tensorrt_llm import LLM
from tensorrt_llm.llmapi import SkipSoftmaxAttentionConfig

llm = LLM(
    model="<path_or_hf_id>",
    sparse_attention_config=SkipSoftmaxAttentionConfig(
        threshold_scale_factor={"prefill": 1000.0, "decode": 500.0},
    ),
)

sparse_attention_config:
  algorithm: skip_softmax
  threshold_scale_factor:
    prefill: 1000.0
    decode: 500.0

Alternatively, provide target_sparsity. This path requires the checkpoint to contain a calibration formula that maps the requested target to the kernel’s threshold scale factor. target_sparsity is calibration guidance rather than a runtime guarantee; the achieved sparsity depends on the model inputs and workload.

sparse_attention_config:
  algorithm: skip_softmax
  target_sparsity:
    prefill: 0.5
    decode: 0.3

Both fields accept a scalar for both phases or a dictionary with prefill and decode values. If both are present, threshold_scale_factor takes precedence. User-provided target_sparsity overrides a checkpoint default.

Model Optimizer can store calibration metadata in the checkpoint’s config.json:

{
  "sparse_attention_config": {
    "config_groups": {
      "group_0": {
        "algorithm": "skip_softmax",
        "threshold_scale_factor": {
          "formula": "a * exp(b * target_sparsity)",
          "prefill": {"a": 100.0, "b": 5.0},
          "decode": {"a": 0.05, "b": 10.0}
        },
        "target_sparsity": {"prefill": 0.5, "decode": 0.3},
        "ignore": ["model.layers.0.self_attn"]
      }
    }
  }
}

The formula is a numexpr expression over target_sparsity and named coefficients. The optional ignore list uses fnmatch layer patterns. At most one checkpoint config group may use the skip_softmax algorithm.

Skip Softmax Attention requires the TRTLLM attention backend. Other attention backends do not apply it.

Usage with trtllm-bench and trtllm-serve#

Sparse attention is configured through sparse_attention_config on the PyTorch backend. DeepSeek-V3.2 provides a mature end-to-end example: its checkpoint defines the DSA indexer geometry and Top-K, so the minimal YAML only needs to select the dsa algorithm.

# config.yml
sparse_attention_config:
  algorithm: dsa

Start an OpenAI-compatible server with the same config file used for other PyTorch backend options:

trtllm-serve deepseek-ai/DeepSeek-V3.2 \
  --backend pytorch \
  --tp_size 8 \
  --ep_size 8 \
  --custom_tokenizer deepseek_v32 \
  --config ./config.yml

For a throughput benchmark, first prepare or supply a tokenized dataset, then pass the same config to trtllm-bench:

trtllm-bench --model deepseek-ai/DeepSeek-V3.2 \
  prepare-dataset \
  --output ./deepseek-v3.2-dataset.json \
  token-norm-dist \
  --input-mean 4096 \
  --output-mean 512 \
  --input-stdev 0 \
  --output-stdev 0 \
  --num-requests 16

trtllm-bench --model deepseek-ai/DeepSeek-V3.2 throughput \
  --backend pytorch \
  --tp 8 \
  --ep 8 \
  --dataset ./deepseek-v3.2-dataset.json \
  --max_batch_size 16 \
  --max_num_tokens 8192 \
  --config ./config.yml

Use a local checkpoint path in place of the Hugging Face model ID when needed. Other sparse algorithms use the same YAML entry point with their own algorithm discriminator and settings. See the DeepSeek V3/V3.2 example for model precision, hardware, parallelism, MTP, chunked-prefill, cache-reuse, and disaggregated-serving configurations.