Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

98$\times$ Faster LLM Routing Without a Dedicated GPU: Flash Attention, Prompt Compression, and Near-Streaming for the vLLM Semantic Router

Xunzhuo Liu, Bowei He, Xue Liu, Andy Luo, Haichen Zhang, Huamin Chen

Abstract

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.

98$\times$ Faster LLM Routing Without a Dedicated GPU: Flash Attention, Prompt Compression, and Near-Streaming for the vLLM Semantic Router

Abstract

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 memory makes long-context classification (8K--32K tokens) impossible: at 8K tokens, three concurrent classifiers need 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 to and end-to-end (E2E) latency from 4{,}918\,ms to 127\,ms (\textbf{38.7}), 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 512 tokens without neural inference, capping both latency and GPU memory at a constant regardless of original prompt length (E2E 12762\,ms, \textbf{2.0}). \emph{Stage~3}: near-streaming body processing with adaptive chunking and zero-copy JSON eliminates serialization overhead (E2E 6250\,ms, \textbf{1.2}). Cumulatively: \textbf{98} 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.
Paper Structure (55 sections, 4 equations, 4 figures, 14 tables)

This paper contains 55 sections, 4 equations, 4 figures, 14 tables.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Envoy External Processing
  4. Attention Mechanisms and Flash Attention
  5. Prompt Compression
  6. ModernBERT
  7. System Architecture
  8. System Design
  9. CK Flash Attention for ONNX Runtime on ROCm
  10. ONNX Runtime with ROCm Execution Provider
  11. Why Flash Attention Is Unavailable in ONNX Runtime for ROCm
  12. Custom CK Flash Attention Operator
  13. ONNX Graph Rewrite.
  14. ModernBERT Attention Mapping.
  15. HIP Kernel and ORT Registration.
  16. ...and 40 more sections

Figures (4)

  • Figure 1: Co-located deployment: the router's ONNX classifiers share the MI300X GPU with vLLM inference, consuming only ${\sim}$0.8 GB vs. vLLM's ${\sim}$190 GB---no dedicated accelerator needed.
  • Figure 2: Near-streaming body handler state machine. The first chunk determines the path: specified models bypass body inspection entirely; "auto" requests accumulate for classification with incremental preprocessing.
  • Figure 3: Attention memory per classifier session vs. sequence length. SDPA exceeds the ${\sim}$718 MB available (dashed line) at ${\sim}$6K tokens. FA stays below but grows linearly. With compression, all inputs are reduced to 512 tokens, yielding a constant footprint regardless of original prompt length.
  • Figure 4: E2E latency comparison (log scale). GPU+FA with streaming and compression achieves 50 ms at 8K tokens. At 16K+, only GPU+FA operates (108 ms at 16K; Table \ref{['tab:combined']}).