TensorRT-LLM 如何用 sparse_attention_config 配置 Sparse Attention 降低长上下文推理开销【免费下载链接】TensorRT-LLMTensorRT LLM provides users with an easy-to-use Python API to define Large Language Models (LLMs) and supports state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs. TensorRT LLM also contains components to create Python and C runtimes that orchestrate the inference execution in a performant way.项目地址: https://gitcode.com/GitHub_Trending/te/TensorRT-LLM长上下文推理时KV cache 随序列长度线性增长attention 要遍历全部历史 KV成为显存和算力的主要瓶颈。TensorRT-LLM 通过用户可见的sparse_attention_configAPI 让稀疏 attention 算法跳过算法判定为不重要的 KV 条目从而降低这部分开销。该 API 目前是 prototype 状态仅支持 PyTorch execution backend即--backend pytorch或backendpytorch。官方特性文档 Sparse Attention 将稀疏 attention 拆成两部分算法负责挑选 token/block 或决定可跳过的 kernel tileattention 实现负责按稀疏模式计算输出。当前有五个用户可见算法配置类都从 tensorrt_llm/llmapi/llm_args.py 导出algorithm配置类稀疏机制适用对象rocketRocketSparseAttentionConfigprompt KV 驱逐 decode 阶段 page 级 Top-K无需训练的 MHA/MQA/GQA 模型dsaDeepSeekSparseAttentionConfig学习型 token 级 indexer 稀疏 MLADeepSeek-V3.2 等模型原生 DSA 架构deepseek_v4DeepSeekV4SparseAttentionConfig滑窗 压缩稀疏/压缩 dense 历史DeepSeek-V4 混合 attentionminimax_m3MiniMaxM3SparseAttentionConfig学习型 block 选择 稀疏 GQAMiniMax-M3 稀疏层skip_softmaxSkipSoftmaxAttentionConfigFMHA kernel 内动态跳过 softmax 工作现有全 attention 模型需校准或直接阈值各算法的关键差异引自文档的 Capability ComparisonRocketKV 不做稀疏 prefillprompt attention 仍密集计算但它选择保留哪些 prompt KV直接减小 cache 和后续 decode 工作量DeepSeek-V4 通过模型原生压缩减少保留的主 KV 历史skip_softmax不选 token、不改架构、不减少 KV cache 存储只在 kernel 内跳过工作只有模型原生 selectorDSA、DeepSeek-V4、MiniMax-M3要求 checkpoint 自带训练好的选择器。本文按选算法 → 写配置 → 启动/评测 → 验证的路径展开以 RocketKV 和 DSA 为主路径Skip Softmax 作为对普通 MHA/GQA 模型的替代分支。准备条件不同算法对硬件和运行方式的要求不同配置前先确认所有算法只能运行在 PyTorch backend 上。RocketKV要求 CUDA compute capability ≥ 10.0Blackwell 或更新 GPU支持 FP16/BF16/FP8、Paged KV Cache、Tensor Parallel、CUDA Graph必须enable_block_reuseFalse不支持 chunked prefill也不支持 disagg-serving因为它依赖的 Python KT cache manager 无法完成 prefill 到 decode 引擎的 KV cache 传输。来源examples/sparse_attention/RocketKV.md。DSADeepSeek-V3.2checkpoint 自带 DSA indexer 几何indexer 头数、头维度、Top-K最安全的做法是让 TensorRT LLM 从模型加载不要覆盖这些值。模型整体要求参见 DeepSeek-V3/V3.2 示例如 FP8 需 8×H200/B200 或 16×H100SM80/SM86 不支持。Skip Softmax要求 TRTLLM attention backend其他 attention backend 不会应用它。DeepSeek-V4要求window_size128、压缩比属于{1, 4, 128}、SM90 或 SM100 GPU、KV cache block 为128或256tokens、beam width 为1Hopper 上要求kv_cache_config.dtypefp8_ds_mlaSM120/SM121 上该 cache 布局要求 256-token blocks。在 Python API 中配置LLM构造函数接收sparse_attention_config参数。示例中的path_to_model需要替换为你的本地 checkpoint 路径或 Hugging Face 模型 ID文档原有占位符。RocketKV可选分支适用于普通 MHA/MQA/GQA 模型from tensorrt_llm import LLM, SamplingParams from tensorrt_llm.llmapi import RocketSparseAttentionConfig, KvCacheConfig rocket_config RocketSparseAttentionConfig( window_size32, # 始终保留的近期窗口大小 kernel_size63, # 重要性打分的池化核大小 prompt_budget2048, # 从 prompt 中保留的 token 数Stage 1 topk64, # 生成阶段动态选取的 KT page 数Stage 2 topr128, # 打分时保留的 query 通道数 kt_cache_dtypefloat8_e5m2, # 辅助 KT cache 的数据类型 ) kv_config KvCacheConfig(enable_block_reuseFalse) # RocketKV 必须禁用 block reuse llm LLM( modelpath_to_model, backendpytorch, # RocketKV 目前要求 PyTorch backend sparse_attention_configrocket_config, kv_cache_configkv_config, ) outputs llm.generate( [To be or not to be, that is the question.], SamplingParams(max_tokens128), )参数含义与文档说明prompt_budget默认 2048是 prompt 压缩后的保留 token 数topk默认 64是生成阶段动态选取的 KT page 数注意选取粒度是 KT cache page而实际 attention kernel 按 KV cache page size 取数topr默认 128用于计算 Query 与 KT Cache 相关性的 query 特征维度文档建议将其设为等于head_dim以跳过topr_filter计算、获得更好的性能和精度window_size默认 32内最近的 token 始终保留另外保留prompt_budget - window_size个前缀中的重要 tokenkernel_size默认 63是 context 阶段一维 max-pooling 的核大小kt_cache_dtype可选float8_e5m2推荐节省显存并加速预测 kernel或bfloat16page_size默认 4是稀疏 token 选取粒度Triton kernel 限制只支持 2 的幂文档指出page_size 4时精度通常保持良好。DSADeepSeek-V3.2checkpoint 已定义 indexer 几何与 Top-K配置可以留空from tensorrt_llm import LLM from tensorrt_llm.llmapi import DeepSeekSparseAttentionConfig llm LLM( modeldeepseek-ai/DeepSeek-V3.2, sparse_attention_configDeepSeekSparseAttentionConfig(), )在受支持的 Blackwell 配置上可用 Guess-Verify-RefineGVR替换常规 decode Top-K dispatcher当前实现接受index_topk为512、1024、2048indexer 压缩比1和4不满足条件时回退到生产的 insertion/radix Top-K 路径sparse_attention_config: algorithm: dsa index_topk: 2048 enable_heuristic_topk: trueSkip Softmax可选分支不选 token、不减少 KV cache 存储而是让 FMHA kernel 动态跳过符合条件的 softmax 工作。可以直接给threshold_scale_factor标量或含prefill/decode的字典from tensorrt_llm import LLM from tensorrt_llm.llmapi import SkipSoftmaxAttentionConfig llm LLM( modelpath_to_model, sparse_attention_configSkipSoftmaxAttentionConfig( threshold_scale_factor{prefill: 1000.0, decode: 500.0}, ), )也可以给target_sparsity但这条路径要求 checkpoint 的config.json中包含把目标稀疏度映射到 kernel 阈值系数的校准公式由 Model Optimizer 写入sparse_attention_config.config_groups。target_sparsity是校准指引而不是运行时保证实际达到的稀疏度取决于模型输入和工作负载两个字段同时给出时threshold_scale_factor优先用户提供的target_sparsity会覆盖 checkpoint 默认值。用 YAML 配置文件驱动 trtllm-serve / trtllm-bench / trtllm-eval用命令行工具时稀疏 attention 选项通过--config指向的 YAML 文件指定参数名与对应配置类一致来自 RocketKV 文档backend: pytorch attn_backend: TRTLLM sparse_attention_config: algorithm: rocket kt_cache_dtype: float8_e5m2 window_size: 32 prompt_budget: 2048 kv_cache_config: enable_block_reuse: false enable_chunked_prefill: falseDSA 场景的最小 YAML 只需选择算法其余几何由 checkpoint 提供# config.yml sparse_attention_config: algorithm: dsa启动 OpenAI 兼容服务文档示例path_or_hf_id位置按需替换为本地 checkpoint 路径trtllm-serve deepseek-ai/DeepSeek-V3.2 \ --backend pytorch \ --tp_size 8 \ --ep_size 8 \ --custom_tokenizer deepseek_v32 \ --config ./config.yml运行示例脚本并验证效果仓库提供了参考脚本 examples/llm-api/llm_sparse_attention.py支持--algo ROCKETKV和--algo DSA并带--window_size、--kernel_size、--prompt_budget、--topk、--kt_cache_dtype、--max_seq_len、--max_new_tokens等参数。按 RocketKV 文档 中的示例命令运行把--model_path调整为你本地的 Llama checkpoint 路径python3 ../llm-api/llm_sparse_attention.py \ --model_path path_to_model \ --algo ROCKETKV \ --attention_backend TRTLLM \ --window_size 32 \ --kernel_size 63 \ --prompt_budget 2048 \ --topk 64 \ --topr 128 \ --kt_cache_dtype float8_e5m2 \ --max_seq_len 10240 \ --max_num_tokens 10240 \ --max_new_tokens 128验证配置是否生效文档给出的两条路径生成请求能正常返回。服务启动后用 curl 查询引自 DeepSeek-V3/V3.2 示例 的验证方式curl http://localhost:8000/v1/completions \ -H Content-Type: application/json \ -d { model: deepseek-ai/DeepSeek-V3.2, prompt: Where is New York?, max_tokens: 16, temperature: 0 }跑吞吐基准对比。用trtllm-bench准备 tokenized 数据集并执行 throughput 测试把同一份 config 传给基准命令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精度评估可选。trtllm-eval也接受同一份 config文档给出的示例trtllm-eval --model path_to_model --config extra_config.yaml longbench_v2 --max_output_length 1024限制与排查RocketKV 的三个硬性限制enable_block_reuse必须为False、chunked prefill 不支持、disaggregated serving 不支持在配置阶段就会触发问题写 YAML 时直接按上文的kv_cache_config/enable_chunked_prefill字段固定不要与常规长文本服务的缓存复用配置混用。MiniMax-M3 的稀疏路径目前没有 dense fallback也不支持 KV cache 复用和 MTPimplementation: msa分支要求 SM100/SM103 GPU、fmha_sm100包和sparse_block_size128其余情况使用默认的triton参考实现。支持的 checkpoint 与并行部署设置见 MiniMax-M3 部署指南。DeepSeek-V4 显式覆盖 config 时必须保持模型的 attention 布局官方构造逻辑默认从 checkpoint 推导DeepSeekV4SparseAttentionConfig部署约束见 DeepSeek-V4 示例。模型原生算法DSA、DeepSeek-V4、MiniMax-M3的几何参数从 checkpoint 读取除非模型专属指南明确说明可调否则不要覆盖这些值。文档没有给出固定性能收益数值是否需要 sparse attention 及其参数取舍如prompt_budget、topk应由你用自己的工作负载数据经trtllm-bench对比得出。进一步阅读Sparse Attention 特性文档 的 Further Reading 一节指向了 KV Cache Compression 文档、稀疏 attention 博客框架设计与逐算法实现、Skip Softmax 的 kernel 细节与端到端基准以及 Sparse Attention 开发指南如何新增一个稀疏 attention 算法。【免费下载链接】TensorRT-LLMTensorRT LLM provides users with an easy-to-use Python API to define Large Language Models (LLMs) and supports state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs. TensorRT LLM also contains components to create Python and C runtimes that orchestrate the inference execution in a performant way.项目地址: https://gitcode.com/GitHub_Trending/te/TensorRT-LLM创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考