让大模型路由快98倍,不占专用显卡,支持超长文本
98$\times$ Faster LLM Routing Without a Dedicated GPU: Flash Attention, Prompt Compression, and Near-Streaming for the vLLM Semantic Router
- 用Flash Attention降低内存开销,实现长文本快速分类
- 通过提示压缩将输入固定在512词元内,稳定延迟与显存占用
- 近流式处理消除序列化开销,适合部署在共享显卡的生产环境
系统级路由需对LLM请求进行安全分类、领域分发和隐私信息检测,既要低延迟又要轻量运行,不宜占用专用GPU——该资源应留给LLM推理。当路由与vLLM共用同一块显卡时,传统注意力机制的$O(n^2)$内存开销使长文本分类(8K–32K tokens)不可行:8K长度下三个并发分类器仅注意力掩码就需约4.5GB内存,远超vLLM剩余空间。本文针对vLLM语义路由提出三阶段优化,在AMD Instinct MI300X上验证:第1阶段,基于ROCm的ONNX Runtime自定义CK Flash Attention算子,将注意力内存从$O(n^2)$降至$O(n)$,端到端延迟由4,918毫秒降至127毫秒(38.7×加速),实现8K–32K token分类;第2阶段,采用TextRank、位置加权、TF-IDF与新颖性评分等经典NLP提示压缩方法,将所有输入压缩至约512词元,无需神经推理,使延迟与显存占用恒定不变(端到端127→62毫秒,2.0×);第3阶段,采用自适应分块与零拷贝JSON的近流式处理,消除序列化开销(端到端62→50毫秒,1.2×)。累计提升达98×(4,918毫秒→50毫秒),16K词元路由仅耗时108毫秒,路由器总显存占用低于800MB,足以与LLM服务共享显卡,无需专用加速器。第1阶段适配AMD ROCm(NVIDIA已有cuDNN FlashAttention);第2、3阶段为硬件无关设计。
原文摘要 · Abstract (English)
System-level routers that intercept LLM requests for safety classification, domain routing, and PII detection must be both fast and operationally lightweight: they should add minimal latency to every request, yet not require a dedicated GPU -- an expensive resource better used for LLM inference itself. When the router co-locates on the same GPU as vLLM serving instances, standard attention's $O(n^2)$ memory makes long-context classification (8K--32K tokens) impossible: at 8K tokens, three concurrent classifiers need ${\sim}$4.5\,GB for attention masks alone, far exceeding the memory left by vLLM. We present three staged optimizations for the vLLM Semantic Router, benchmarked on AMD Instinct MI300X, that solve both the latency and the memory problem. \emph{Stage~1}: a custom CK Flash Attention operator for ONNX Runtime on ROCm reduces attention memory from $O(n^2)$ to $O(n)$ and end-to-end (E2E) latency from 4{,}918\,ms to 127\,ms (\textbf{38.7$\times$}), enabling 8K--32K tokens where SDPA OOMs. \emph{Stage~2}: classical NLP prompt compression (TextRank, position weighting, TF-IDF, and novelty scoring) reduces all inputs to ${\sim}$512 tokens without neural inference, capping both latency and GPU memory at a constant regardless of original prompt length (E2E 127$\to$62\,ms, \textbf{2.0$\times$}). \emph{Stage~3}: near-streaming body processing with adaptive chunking and zero-copy JSON eliminates serialization overhead (E2E 62$\to$50\,ms, \textbf{1.2$\times$}). Cumulatively: \textbf{98$\times$} improvement (4{,}918\,ms to 50\,ms), 16K-token routing in 108\,ms, and a total router GPU footprint under 800\,MB -- small enough to share a GPU with LLM serving and removing the need for a dedicated accelerator. Stage~1 targets AMD ROCm (NVIDIA GPUs already have FlashAttention via cuDNN); Stages~2 and~3 are hardware-agnostic.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。