Sparse Attention#

Source NVIDIA/TensorRT-LLM.

  1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  2# SPDX-License-Identifier: Apache-2.0
  3#
  4# Licensed under the Apache License, Version 2.0 (the "License");
  5# you may not use this file except in compliance with the License.
  6# You may obtain a copy of the License at
  7#
  8# http://www.apache.org/licenses/LICENSE-2.0
  9#
 10# Unless required by applicable law or agreed to in writing, software
 11# distributed under the License is distributed on an "AS IS" BASIS,
 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13# See the License for the specific language governing permissions and
 14# limitations under the License.
 15r"""
 16This example demonstrates how to use sparse attention with TensorRT-LLM.
 17
 18Supported sparse attention algorithms:
 19- RocketKV
 20- DSA
 21
 22Usage:
 23```bash
 24python llm_sparse_attention.py \
 25    --model_path nvidia/Llama-3.1-8B-Instruct-FP8 \
 26    --algo ROCKETKV \
 27    --attention_backend TRTLLM \
 28    --window_size 32 \
 29    --kernel_size 63 \
 30    --prompt_budget 2048
 31```
 32
 33When ``--input_file`` is omitted, the example uses a built-in
 34needle-in-a-haystack prompt that exceeds the default ``prompt_budget`` and
 35exercises the sparse attention path.
 36"""
 37import argparse
 38import json
 39
 40from tensorrt_llm import LLM, SamplingParams
 41from tensorrt_llm.llmapi import (CudaGraphConfig, DeepSeekSparseAttentionConfig,
 42                                 KvCacheConfig, MoeConfig,
 43                                 RocketSparseAttentionConfig)
 44
 45# The built-in prompt follows a self-contained needle-in-a-haystack layout:
 46# 1. Cycle routine expedition log templates to form a deterministic haystack.
 47# 2. Replace one numbered log entry with a unique access-code needle.
 48# 3. Append a question that asks the model to retrieve the needle.
 49# With the default model, 128 entries produce about 2.8K tokens, exceeding the
 50# default 2,048-token prompt budget so that sparse attention is exercised.
 51_DEFAULT_LOG_TEMPLATES = (
 52    "The survey team checked the northern weather station and recorded normal "
 53    "temperature and pressure readings.",
 54    "Technicians inspected backup batteries, radio transmitters, and emergency "
 55    "lighting; all systems passed routine checks.",
 56    "Researchers cataloged soil samples, labeled storage containers, and "
 57    "updated the expedition inventory.",
 58    "The navigation group reviewed trail maps, satellite images, and the next "
 59    "day's travel schedule.",
 60    "At sunset, the field crew secured scientific instruments and uploaded the "
 61    "day's measurements.",
 62    "The medical officer reviewed first-aid supplies and confirmed the "
 63    "evacuation plan with base camp.",
 64    "Engineers calibrated wind sensors and verified timestamps against the "
 65    "observatory master clock.",
 66    "The logistics team counted food, water, fuel, and spare parts before "
 67    "closing the storage area.",
 68)
 69_DEFAULT_NUM_LOG_ENTRIES = 128
 70_DEFAULT_NEEDLE_INDEX = 87
 71_DEFAULT_ACCESS_CODE = "314159"
 72
 73
 74def _build_default_prompt():
 75    log_entries = []
 76    for index in range(_DEFAULT_NUM_LOG_ENTRIES):
 77        if index == _DEFAULT_NEEDLE_INDEX:
 78            message = (f"The expedition access code is {_DEFAULT_ACCESS_CODE}. "
 79                       "Remember this code for the final question.")
 80        else:
 81            message = _DEFAULT_LOG_TEMPLATES[index %
 82                                             len(_DEFAULT_LOG_TEMPLATES)]
 83        log_entries.append(f"Log entry {index}: {message}")
 84
 85    return (
 86        "Read the following expedition field log carefully. One entry contains "
 87        "an access code that you will need to recall.\n\n" +
 88        "\n".join(log_entries) + "\n\nWhat is the expedition access code?")
 89
 90
 91DEFAULT_PROMPTS = [_build_default_prompt()]
 92
 93
 94def read_input(input_file):
 95    results = []
 96    with open(input_file, 'r') as f:
 97        for line in f:
 98            ret = json.loads(line)
 99            results.append(ret)
100    return results
101
102
103def parse_arguments():
104    parser = argparse.ArgumentParser()
105    parser.add_argument('--model_path',
106                        type=str,
107                        default="nvidia/Llama-3.1-8B-Instruct-FP8",
108                        help="The local path or Hugging Face ID of the model.")
109    parser.add_argument(
110        '--input_file',
111        type=str,
112        default=None,
113        help="Optional path to a JSONL input file. The built-in "
114        "long prompt is used when omitted.")
115
116    # Build config
117    parser.add_argument('--algo',
118                        type=str,
119                        default='ROCKETKV',
120                        choices=['ROCKETKV', 'DSA'])
121    parser.add_argument('--attention_backend',
122                        type=str,
123                        default='TRTLLM',
124                        choices=['VANILLA', 'TRTLLM'])
125
126    # RocketKV config
127    parser.add_argument('--window_size',
128                        type=int,
129                        default=32,
130                        help="The window size for RocketKV.")
131    parser.add_argument('--kernel_size',
132                        type=int,
133                        default=63,
134                        help="The kernel size for RocketKV.")
135    parser.add_argument('--prompt_budget',
136                        type=int,
137                        default=2048,
138                        help="The prompt budget for RocketKV.")
139    parser.add_argument('--topk',
140                        type=int,
141                        default=64,
142                        help='Top-k for RocketKV')
143    parser.add_argument('--kt_cache_dtype',
144                        type=str,
145                        default='float8_e5m2',
146                        choices=['bfloat16', 'float8_e5m2'])
147    parser.add_argument('--index_max_chunk_size',
148                        type=int,
149                        default=32768,
150                        help="The maximum chunk size for the indexer.")
151    parser.add_argument("--max_seq_len",
152                        type=int,
153                        default=10240,
154                        help="The maximum sequence length.")
155    parser.add_argument("--max_batch_size",
156                        type=int,
157                        default=256,
158                        help="The maximum batch size.")
159    parser.add_argument("--max_new_tokens",
160                        type=int,
161                        default=128,
162                        help="The maximum new tokens.")
163    parser.add_argument(
164        "--max_num_tokens",
165        type=int,
166        default=81920,
167        help=
168        "The maximum total tokens (context + generation) across all sequences in a batch."
169    )
170
171    # Parallelism
172    parser.add_argument('--moe_backend',
173                        type=str,
174                        default='CUTLASS',
175                        choices=[
176                            'CUTLASS', 'TRTLLM', 'VANILLA', 'DEEPGEMM',
177                            'CUTEDSL', 'TRITON'
178                        ])
179    parser.add_argument('--tp_size', type=int, default=1)
180    parser.add_argument('--moe_ep_size', type=int, default=-1)
181    parser.add_argument('--enable_attention_dp',
182                        default=False,
183                        action='store_true')
184
185    # KV cache
186    parser.add_argument('--kv_cache_dtype', type=str, default='auto')
187    parser.add_argument("--kv_cache_fraction", type=float, default=0.7)
188    parser.add_argument('--tokens_per_block', type=int, default=32)
189    parser.add_argument('--num_samples', type=int, default=10)
190
191    # Runtime
192    parser.add_argument('--print_iter_log',
193                        default=False,
194                        action='store_true',
195                        help='Print iteration logs during execution')
196    parser.add_argument('--use_cuda_graph', default=False, action='store_true')
197    parser.add_argument('--cuda_graph_padding_enabled',
198                        default=False,
199                        action='store_true')
200    parser.add_argument('--cuda_graph_batch_sizes',
201                        nargs='+',
202                        type=int,
203                        default=None)
204    parser.add_argument('--enable_chunked_prefill',
205                        default=False,
206                        action='store_true',
207                        help='Enable chunked prefill')
208    args = parser.parse_args()
209    return args
210
211
212def run_llm(args, sparse_attention_config):
213    if args.input_file is None:
214        prompts = DEFAULT_PROMPTS
215        reference = [None] * len(prompts)
216    else:
217        data = read_input(args.input_file)
218        num_samples = args.num_samples if args.num_samples is not None else len(
219            data)
220        data = data[:num_samples]
221        prompts = [{
222            'prompt': sample['input_context'] + sample['input_query']
223        } for sample in data]
224        reference = [sample['outputs'] for sample in data]
225
226    kv_cache_config = KvCacheConfig(
227        enable_block_reuse=
228        False,  # sparse attention does not support kv cache reuse now
229        free_gpu_memory_fraction=args.kv_cache_fraction,
230        tokens_per_block=args.tokens_per_block,
231        dtype=args.kv_cache_dtype,
232    )
233
234    cuda_graph_config = CudaGraphConfig(
235        batch_sizes=args.cuda_graph_batch_sizes,
236        enable_padding=args.cuda_graph_padding_enabled,
237    ) if args.use_cuda_graph else None
238
239    llm = LLM(
240        model=args.model_path,
241        backend='pytorch',
242        kv_cache_config=kv_cache_config,
243        attn_backend=args.attention_backend,
244        sparse_attention_config=sparse_attention_config,
245        max_batch_size=args.max_batch_size,
246        max_seq_len=args.max_seq_len,
247        max_num_tokens=args.max_num_tokens,
248        tensor_parallel_size=args.tp_size,
249        moe_expert_parallel_size=args.moe_ep_size,
250        enable_attention_dp=args.enable_attention_dp,
251        cuda_graph_config=cuda_graph_config,
252        print_iter_log=args.print_iter_log,
253        enable_iter_perf_stats=args.print_iter_log,
254        moe_config=MoeConfig(backend=args.moe_backend),
255        enable_chunked_prefill=args.enable_chunked_prefill,
256    )
257
258    sampling_params = SamplingParams(add_special_tokens=False,
259                                     max_tokens=args.max_new_tokens,
260                                     temperature=0.8,
261                                     top_p=0.95)
262
263    outputs = llm.generate(prompts, sampling_params)
264    for idx, output in enumerate(outputs):
265        result = f'Generated text: {output.outputs[0].text!r}'
266        if reference[idx] is not None:
267            result += f', ref: {reference[idx]}'
268        print(result)
269
270
271def run_RocketKV(args):
272    sparse_attention_config = RocketSparseAttentionConfig(
273        window_size=args.window_size,
274        kernel_size=args.kernel_size,
275        prompt_budget=args.prompt_budget,
276        topk=args.topk,
277        kt_cache_dtype=args.kt_cache_dtype,
278    )
279    run_llm(args, sparse_attention_config)
280
281
282def run_DSA(args):
283    sparse_attention_config = DeepSeekSparseAttentionConfig(
284        indexer_max_chunk_size=args.index_max_chunk_size, )
285    run_llm(args, sparse_attention_config)
286
287
288def main():
289    args = parse_arguments()
290    if args.algo == 'ROCKETKV':
291        run_RocketKV(args)
292    elif args.algo == 'DSA':
293        run_DSA(args)
294    else:
295        raise ValueError(f"Invalid algorithm: {args.algo}")
296
297
298if __name__ == "__main__":
299    main()