最快MLA内核背后的秘密：TokenSpeed如何在Blackwell上征服Agentic推理-北京尧图网络

最快MLA内核背后的秘密TokenSpeed如何在Blackwell上征服Agentic推理【免费下载链接】tokenspeedTokenSpeed is a speed-of-light LLM inference engine.项目地址: https://gitcode.com/gh_mirrors/to/tokenspeedTokenSpeed是一个面向Agentic 推理工作负载的光速 LLM 推理引擎其核心亮点之一是在 NVIDIA Blackwell 上实现了业界最快的MLAMulti-head Latent Attention多头潜在注意力内核。本文带你揭秘这套 MLA 内核凭什么在 Agentic 场景高并发、短步长解码、极低的 next-token 延迟要求中拉开差距以及它背后的三个关键优化技巧。一、为什么MLA内核是Agentic推理的核心战场Agentic智能体负载有两个鲜明特征高频、短解码编码类 Agent 不断发起短 prompt、短输出的请求每一步 decode 的延迟都直接影响思考速度超长 KV 缓存Agent 对话动辄积累 64K~80K token 的历史上下文decode 阶段要反复扫描巨大的 KV cache。MLA 通过把 K/V 压进低秩潜在空间大幅降低了 KV cache 的显存占用——这正是 DeepSeek、Kimi 这类旗舰模型普遍采用它的原因。但省显存之后MLA 注意力内核本身的速度就成了瓶颈。TokenSpeed 的 tokenspeed-mla/ 子项目专为 BlackwellSM100/SM103/SM107即 B200/GB300 一代打磨了完整的 MLA 内核家族内核作用源码位置tokenspeed_mla_prefill变长 ragged prefill无 paddingtokenspeed-mla/python/tokenspeed_mla/mla_prefill.pytokenspeed_mla_decodeFP16/BF16/FP8 三路 decodeSplit-KV 自动 workspacetokenspeed-mla/python/tokenspeed_mla/mla_decode.pymla_kv_pack_quantize_fp8融合的 K/V 打包 FP8 量化 Triton 内核tokenspeed-mla/python/tokenspeed_mla/mla_kv_pack_quantize_fp8.py二、秘密 ①两内核 Split-KV 架构——把长序列拆给整个 GPU主流方案的 MLA decode 常用单内核 swap-AB策略当 KV 序列很长时每个线程块要串行扫完全部 80K 个 keyGPU 大量 SM 空闲等待。TokenSpeed 的 decode 内核采用双内核流水线第一个内核按split_kv把 KV 序列切成多段各线程块并行计算局部注意力partial output LSE第二个内核对分段部分结果做归约合并。split_kv与 workspace 大小由运行时根据batch_size、q_len、max_seq_len自动推导并缓存见 mla_decode.py 中的_get_split_kv_and_workspace_size用户无需手动调参。效果立竿见影测试条件q_len4KV 长 80KFP8以num_heads32为例batch4 时 TensorRT-LLM 需要 176.8μsTokenSpeed 仅 49.5μs约3.6×到 batch16 时差距扩大到 652.9μs vs 148.1μs约4.4×。橙色薄片是归约内核的开销——占比极小说明双内核设计几乎没有额外代价。三、秘密 ②query-token 折叠——消灭小 batch 的 tile 浪费Agentic 场景的 decode 常常是q_len4、num_heads16/32/64这样又小又碎的形状。矩阵乘法单元MMA偏好 M 维度接近 128 的 tile但小 head 数会严重浪费算力。TokenSpeed 的解法是fold_sq_factor折叠实现见 mla_helpers.py运行时找出最大的因子F满足q_len % F 0且num_heads × F ≤ 128把F个 query token 折叠进 head 维度等效形状变为H_eff num_heads × Fq_len_eff q_len / F。举个例子num_heads64, q_len4时取F2——两个 query token 被折叠进 M 维H_eff128tile 刚好填满剩余 query 组交给调度器第二维处理。仅这一个变换就让 M 维度利用率接近 100%正是 token 级 agent 流量的甜点。四、秘密 ③Blackwell 微架构级调优清单除了算法层设计内核还针对 Blackwell 硬件做了一批显微镜下的优化见 tokenspeed-mla/README.md 的 Kernel Capability Summary2-CTA UTCMMA 指令两个线程块协同执行一条 MMA直接砍半 shared memory 占用最少 mbarrier降低同步屏障数量减少 warp 间等待K/V 加载 warp 拆分K 加载完成后 V 已在 L2 缓存下一个 tile 的 K 加载不必等 V 加载完天然隐藏访存延迟epilogue 多阶段 STG 子切片把输出写回切成多个 sub-tile 流水执行FP8 部分结果归约reducer 按 M128/M64 路径选择 1/2/4 个 D512 输出波段避免单元素加载这类隐形开销编译缓存内核以静态配置dtype、d_qk、d_v、causal、LSE 等为 key 缓存 JIT 编译结果避免重复编译拖慢服务启动。Prefill 侧同样可观——开源版 CuTe DSL JIT 实现无 padding 的 ragged varlen在中等序列长度上与 TensorRT-LLM 打平甚至反超长序列场景差距在 3% 以内五、从内核到吞吐Agentic负载下的系统级收益MLA 内核再快也要落到端到端吞吐上才算数。TokenSpeed 把内核、C 有限状态机调度和 KV cache 复用整合成完整引擎在 Kimi K2.5B200上的 Agentic 工作负载基准中在 Agent 实际运行的高吞吐区间持续领先 TensorRT-LLM四张图分别对应 Attn TP4/TP8 × MoE TP/EP 四种并行配置。注意曲线交点右侧——正是高并发 Agent 流量所在的区域TokenSpeed 的千 token/分钟·GPU吞吐全程压过对手。六、上手指南如何跑起来环境要求SM100/SM103/SM107 的 Blackwell GPUSM107 需 CuTe DSL ≥ 4.8.0 及匹配的 CUDA 工具链调用 decode传入query [B, q_len, H, D_qk]、kv_cache [num_pages, page_size, D_total]与block_tables即可split_kv/workspace 全自动完整签名见 README.md 的 Minimal Usage 一节复现性能数据仓库内置了基准脚本例如python ./tokenspeed-mla/python/tokenspeed_mla/mla_decode_fp8.py --batch_size 4 --num_heads 16 ...回归测试PYTHONPATHtokenspeed-mla/python python -m pytest -q tokenspeed-mla/tests/test_mla_decode.py覆盖 packed-query 几何、split-KV 尺寸、FP8/FP16/BF16 输出与 LSE、CUDA Graph 回放、滑动窗口等test_mla_decode.py。总结TokenSpeed 征服 Blackwell 上 Agentic 推理的秘诀并不神秘而是三层优化的叠加️架构层双内核 Split-KV 把 80K 长上下文喂饱整块 GPU调度层query-token 折叠让 MMA tile 在小 batch 下也不浪费⚙️微架构层2-CTA MMA、拆分加载 warp、L2 命中调序等 Blackwell 专属打磨。最终成绩decode 延迟 2~4.4× 领先prefill 与最强闭源方案打平端到端 Agent 吞吐在高并发区间持续占优。对于正在生产环境跑 DeepSeek/Kimi 系 MLA 模型的团队这套内核值得重点关注。【免费下载链接】tokenspeedTokenSpeed is a speed-of-light LLM inference engine.项目地址: https://gitcode.com/gh_mirrors/to/tokenspeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

最快MLA内核背后的秘密：TokenSpeed如何在Blackwell上征服Agentic推理

相关推荐

Agent开发：2026高薪风口，小白也能收藏入局，实现技术逆袭！

2026 精选 Java 八股文及答案，全面覆盖，直击面试

浮点数内存分布：C语言

更多资讯

Infection Monkey 部署完全指南：从 AppImage、Docker、Windows 到云市场的 Monkey Island 搭建手册

Agora Flat v1.6.0 版本全解析：课件无感切换、Agora 登录与回放交互升级

ALS1-Camera组件化相机系统：解决复杂动画下的镜头抖动与穿模

AI编程智能体实战：从零搭建自动化开发流程与避坑指南

UE5 Coop联机开发指南：服务器权威与网络同步核心实践

工业级电源路径保护：TPS259483AYWPR+STM32L152ZD硬核方案

今日更新

Agent Skills 完全指南：原理、写法、安装与实战避坑

Agent Skills 实战：从 Genkit 定义到 GKE 部署与排查

Agent Skills 实战：从设计到调试的完整指南

本周精选

CMOS设计三大物理效应：体效应、沟道长度调制与亚阈值导电

微信小程序在线阅读毕设全解析：Java后端+MySQL+避坑实战

Android音频问题五层分析法：从App到Codec的系统化诊断

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？