优化Transformer长上下文推理,减少显存访问提升速度
Deep Kernel Fusion for Transformers
- 深度融合SwiGLU模块计算与内存操作,降低显存带宽压力
- 在H100上实现最高13.2%加速,A100上达9.7%
- 兼容多种模型与硬件,适合长文本生成场景
长上下文代理大模型推理正受制于显存带宽而非算力。在此背景下,因权重过大超出缓存容量而成为主要瓶颈的SwiGLU MLP模块,长期未获有效优化。我们提出DeepFusionKernel,一种深度融合的内核设计,显著减少HBM流量并提升缓存复用率,在H100上相比SGLang最高提速13.2%,A100上达9.7%。该内核与SGLang集成,并搭配内核调度器,确保生成长度变化时仍保持稳定加速,且可适配多样模型、推理配置与硬件平台。
原文摘要 · Abstract (English)
Agentic LLM inference with long contexts is increasingly limited by memory bandwidth rather than compute. In this setting, SwiGLU MLP blocks, whose large weights exceed cache capacity, become a major yet under-optimized bottleneck. We propose DeepFusionKernel, a deeply fused kernel that cuts HBM traffic and boosts cache reuse, delivering up to 13.2% speedup on H100 and 9.7% on A100 over SGLang. Integrated with SGLang and paired with a kernel scheduler, DeepFusionKernel ensures consistent accelerations over generation lengths, while remaining adaptable to diverse models, inference configurations, and hardware platforms.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。