PaperlayerView exact revision on arXiv

Source-linked AI summary

Request-Adaptive Execution for Efficient Hybrid LLM Inference

Fahao Chen, Xuzhe Yan, Qi Zhang, Peng Li, Dongxiao Yu, Song Guo

arXiv:2610.37432v1cs.DC

TL;DR

Static hybrid layouts do not fit workloads that rely on contextual information differently, risking wasted computation or lower quality. Flex combines request-adaptive Attention/SSM execution with runtime support for batching and context reuse. Across six datasets it improves the quality-efficiency trade-off, while system evaluation reports throughput gains of up to 3.87× over vLLM.

  • Problem

    Static hybrid layouts assume requests with different contextual demands benefit from the same structure, although a layout can waste computation on simple requests or sacrifice quality on context-sensitive requests.

    Source: §1 · p. 2

  • Method

    Flex pairs Attention and SSM branches in adaptive blocks, uses prefill-oriented gating to fix each request’s branch choices through decoding, and supports divergent execution with selective batching and hybrid-aware context caching.

    Source: §1 · p. 2

  • Results

    Across six datasets, Flex achieves a better quality-efficiency trade-off than homogeneous and static hybrid baselines, with the best results on 13 of 18 quality metrics and throughput up to 3.87× that of vLLM.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Flex combines request-adaptive modeling with runtime optimizations to improve serving throughput and reduce inference latency relative to baselines.

    Source: §10 · p. 13

Abstract

from arXiv · show

Hybrid large language models (LLMs) have emerged as a promising way to combine the modeling capability of the Attention mechanism with the efficiency of state space models (SSMs). However, most existing hybrid LLMs still inherit a static design philosophy: once the Attention-to-SSM layout is chosen, every request must follow the same hybrid execution. This one-layout-fits-all design is inefficient for real-world serving workloads, where different requests rely on contextual information in different ways. This paper presents Morph, a serving system that co-designs the model architecture and runtime for request-adaptive hybrid LLM inference. Its model architecture pairs Attention and SSM layers within each block and uses prefill-oriented gating to select a request-specific layout that remains fixed throughout decoding. Its serving runtime preserves batching efficiency under divergent layer choices through selective batching and enables correct prefix reuse through hybrid-aware context caching. Experiments on six datasets show that Morph achieves a better quality-efficiency trade-off than existing static hybrid baselines. System-level evaluation further shows that Morph improves the serving throughput by up to 3.87x over baselines while substantially reducing inference latency.

1 Introduction

Static hybrid layouts assume every request benefits from the same Attention-to-SSM arrangement, although workloads rely on context differently. Flex adapts each request’s layout and co-designs the model and runtime to support efficient execution and context reuse.

1 Introduction

  • Because workloads rely on contextual information differently, a static layout can waste computation on simple requests or sacrifice quality on context-sensitive ones.The passage contrasts reasoning tasks such as ARC-Challenge with summarization such as PubMed.Source: §1 · p. 2
  • Flex pairs Attention and SSM branches in adaptive blocks, where gates choose one branch per block to form a request-specific layout.Shared components include the FFN; the layout balances modeling capability and inference cost.Source: §1 · p. 2
  • Prefill-oriented gating uses all input tokens to choose each request’s branch per block, then fixes those choices through decoding so prior-token caches remain reusable.Token-level choices could require recomputation when later tokens cannot reuse earlier Attention KVs or SSM states.Source: §1 · p. 2
  • Selective batching keeps branch-independent operations in the packed batch, splits requests for Attention or SSM execution, and merges outputs back into the original layout.Flex also coordinates branch resources according to relative workload to reduce synchronization stalls.Source: §1 · p. 2
  • Hybrid-aware caching separates textual-prefix matching from layer-cache matching, allowing compatible layer-wise caches to be reused under different request layouts.Priority-based layer-wise eviction accounts for context sharing, reuse likelihood, and recomputation cost.Source: §1 · pp. 2–3
  • Across six datasets, Flex achieves the best results on 13 of 18 quality metrics, reduces estimated TFLOPs/token by up to 67.2% relative to LLaMA, and boosts throughput by up to 3.87× over vLLM.It also substantially reduces time-to-first-token and time-per-output-token.Source: §1 · p. 3

2 Background and Motivation

Attention offers global context but incurs growing inference costs, while SSMs use compact recurrent states but can be less effective at precise long-range recall. Static hybrid layouts trade quality against efficiency, yet preferred layouts vary across workloads and even across requests within one benchmark.

2.1 Background

  • Attention aggregates information from previous tokens but has quadratic prefill computation and growing KV-cache costs, whereas SSMs use compact recurrent states for efficient long-sequence inference.SSM state compression may limit precise long-range context recall; hybrid layouts can achieve higher accuracy than Attention-only or SSM-only configurations at moderate computation cost.Source: §2.1 · p. 3

2.2 Motivation

  • Throughput rises from 23.83 tokens/s under 1:1 to 87.11 tokens/s under 1:12; quality varies by workload, with 1:6 best on ARC-Challenge and 1:1 best on PIQA and PubMed.Source: §2.2 · p. 4
  • The best layout also varies across six ARC-Challenge subsets: Sets 1 and 3 favor more Attention, Sets 4 and 6 favor 1:6, and Set 5 favors 1:12.Each disjoint subset contains 100 requests.Source: Figure 4 · p. 4, §2.2 · p. 4
  • No single static Attention-to-SSM layout is best across workloads or across requests within the same workload.These observations motivate adapting the hybrid structure to request-specific contextual requirements.Source: §2.2 · p. 4

3 System Overview

Flex combines a request-adaptive hybrid model with a serving runtime for divergent layer choices and prefix reuse. Gates choose a branch during prefill, and those choices remain fixed throughout decoding.

3 System Overview

  • Each adaptive hybrid block pairs a lightweight gate with Attention and SSM branches; each request selects one branch per block to form its layout.Shared components such as FFNs are also included in each block.Source: §3 · p. 4
  • Selective batching executes shared operations on the packed batch and splits requests into branch-specific sub-batches; context caching matches textual prefixes and layer-selection histories.The cache identifies reusable Attention KV caches and SSM states.Source: §3 · p. 4
  • Gates choose branches during prefill, and the runtime records those choices, reuses compatible caches, and keeps the choices fixed while updating KV caches or SSM states through decoding.The figure marks activated layers with filled colors and inactive layers unfilled.Source: §3 · pp. 4–5, Figure 6 · p. 5

4 Prefill-Oriented Gating

Flex makes hybrid layer selection request-specific while keeping each request’s choices fixed during decoding. Its training combines soft routing with a cost-aware objective to balance generation quality and inference efficiency.

4.1 Gating Design

  • Prefill-oriented gating selects Attention or SSM once per request and block, then reuses those choices throughout decoding to preserve cached states.The gate aggregates prefill hidden states into a request representation and uses top-1 selection; token-level choices could force recomputation when later tokens cannot reuse earlier caches.Source: §4.1 · p. 5

4.2 Hybrid Model Training

  • The hybrid model pairs pretrained Attention and SSM layers into blocks, trains them with differentiable weighted soft routing, and switches to hard routing at inference.Structured pruning aligns hidden dimensions, and KL-divergence-based importance scores select layers when source depths differ.Source: §4.2 · p. 5
  • A training-only cost-aware regularizer biases gate choices toward efficient layers alongside the supervised fine-tuning objective.Its normalized layer-cost weights represent Attention and SSM costs, while λ controls the quality-efficiency trade-off.Source: §4.2 · pp. 5–6

5 Selective Batching Engine

Flex preserves batching for requests that choose different Attention and SSM branches by splitting only branch-dependent computation and coordinating branch resources. It merges branch outputs back into the packed batch for shared operations.

5.1 Split-and-Merge Execution

  • Flex packs requests together for shared operations, compacts them into Attention and SSM sub-batches only for branch-specific computation, then merges outputs into the original layout.During decoding, recorded prefill choices form the sub-batches, and each branch reads and updates its corresponding cache.Source: §5.1 · p. 6

5.2 Load-Aware Branch Coordination

  • Flex predicts relative branch costs from request counts and input or KV-cache lengths, then assigns more GPU SM resources to the heavier branch to reduce synchronization stalls.Prediction error is within 6.32% for Attention and 5.18% for SSM; profiling is a one-time effort per model-GPU pair.Source: §5.2 · pp. 6–7
  • Separate Green Contexts and CUDA streams run the two sub-batches concurrently on assigned SM partitions, synchronizing before merging outputs.Existing FlashAttention and mamba_ssm kernels run without modification.Source: §5.2 · p. 7

6 Hybrid-Aware Context Caching

Flex’s hybrid-aware cache reuses prefixes only along compatible layer-selection paths, allowing partial reuse when a request’s layout diverges. Its eviction policy ranks individual cache entries rather than discarding all paths for a context together.

6.1 Context Caching Mechanism

  • A textual prefix alone cannot guarantee cache compatibility in Flex because requests with the same prefix may select different Attention and SSM layouts.Reusing mismatched states can degrade inference quality, while rejecting the entry requires recomputing the shared prefix.Source: §6.1 · p. 7
  • A two-level radix tree indexes textual prefixes outside and layer-selection paths with their Attention KV and SSM caches inside, reusing only compatible paths.Flex selects a compatible prefix and cache path by estimated recomputation cost; at the first route mismatch, valid preceding hidden states support recomputation from that block onward.Source: §6.1 · p. 7, Figure 9 · p. 8

6.2 Priority-based Eviction Policy

  • Flex evicts individual layer-wise cache entries by retention priority, removing an outer context node only after its caches are gone and it is no longer needed for shared prefixes.Priority reflects recent access, saved recomputation cost, memory footprint, and a depth factor favoring earlier blocks.Source: §6.2 · p. 8

7 Implementation

Flex’s prototype combines a PyTorch model implementation with a vLLM-style serving runtime and CUDA/C++ components.

7 Implementation

  • The approximately 6.5K-line prototype uses Python, PyTorch, HuggingFace Transformers, and CUDA/C++, with a vLLM-style execution loop extended for selective batching and hybrid-aware caching.It uses NVIDIA GreenContext for concurrent branch execution and a unified paging-based memory manager for Attention KV and SSM caches.Source: §7 · p. 8

8 Evaluation

Across six datasets, Flex improves the quality-efficiency trade-off over static hybrid models, and its serving runtime outperforms the evaluated systems in throughput, TTFT, and TPOT.

8.1 Evaluation Setup

  • Evaluation uses three model scales, six datasets spanning reasoning, QA, and summarization, and one H100 GPU for system-level tests.The 61B system evaluation uses FP8 quantization.Source: §8.1 · p. 8, §8.3 · p. 9

8.2 Overall Model Performance

  • Flex achieves the best results on 13 of 18 quality metrics while reducing computation cost versus Attention-only Llama models across all three scales.The reported TFLOPs/token reductions are 40.5%, 33.2%, and 67.2% at the small, medium, and large scales, respectively.Source: §8.2 · p. 9, Table 2 · p. 10

8.3 Overall System Performance

  • Flex reaches 3.79×, 3.87×, and 3.44× the throughput of vLLM for the 4.4B, 14.5B, and 61B models, respectively.It also reduces TTFT by an average of 80.9%, 78.9%, and 75.2% versus vLLM, SGLang, and Jenga, respectively, and TPOT by 71.6%, 63.1%, and 56.5% versus those baselines.Source: Figure 10 · p. 10, §8.3 · p. 11

8.2 Overall Model Performance

  • Adaptive routing outperforms both fixed-layout strategies across all 18 model-dataset combinations, gaining 4.86 percentage points in accuracy and 2.55 points in ROUGE-L on average over D-Optimal.D-Optimal selects a fixed layout per dataset, while Adaptive chooses a layout for each request.Source: §8.2 · p. 9, Table 3 · p. 11

8.4 Ablation Study

  • Selective batching lowers end-to-end latency versus vLLM by 54.6%, 66.0%, and 59.9% on average across the 4.4B, 14.5B, and 61B models.Load-aware branch coordination further cuts latency versus Flex w/o Coord. by 21.7%, 28.2%, and 26.2% on average at those scales.Source: §8.4 · p. 11, Figure 11 · p. 12
  • Hybrid-aware caching raises average effective hit rate to 43.4% from 18.7% and reduces average TTFT from 10.68s to 5.92s.At 20 articles, hit rates are 24% versus 7%, and TTFT is 9.36s versus 16.18s for the baseline.Source: Figure 12 · p. 12, §8.4 · p. 12
  • Attention selection varies by workload and block depth; on ARC-Challenge, the six most frequent of 35 layouts cover over 90% of 500 requests.Attention selection rises toward deeper blocks on ARC-Challenge, peaks in early-to-middle blocks on TruthfulQA, and is highest in middle blocks on CNN/DailyMail.Source: Figure 13 · p. 12, §8.4 · p. 12

9 Related Work

Related work combines Attention and SSMs, improves LLM batching, and reuses prefix caches; Flex differs through request-adaptive layouts and runtime support for divergent branches and layer-wise caches.

Hybrid LLMs

  • Prior hybrid LLMs combine Attention and SSM layers but use static layouts, whereas Flex enables request-adaptive layer selection.Examples include Jamba, Zamba [14], Hymba [11], Mamba-in-Llama [38], and Nemotron-H [3].Source: §9 · pp. 12–13

Efficient batching

  • Serving systems have advanced scheduling and batching, while Flex batches requests whose Attention and SSM layer choices diverge.Related systems include Orca, vLLM [21], SGLang [51], Sarathi-Serve [1], Splitwise [27], and others.Source: §9 · p. 13

Context caching

  • Prefix-caching methods reuse repeated contexts, while Flex's two-level radix tree distinguishes layer-wise Attention KV and SSM caches for the same textual prefix.Related approaches include RadixAttention [51], Parrot [24], CacheBlend [41], and Mooncake.Source: §9 · p. 13

10 Conclusion

Flex co-designs a request-adaptive hybrid model and serving runtime, reporting a better quality-efficiency trade-off than static hybrid baselines and improved serving throughput with lower latency.

10 Conclusion

  • Flex combines a request-adaptive hybrid model with an optimized serving runtime through algorithm-system co-design.The model uses prefill-oriented gating to choose Attention or SSM branches per request, then keeps that layout fixed during decoding.Source: §10 · p. 13
  • Selective batching and hybrid-aware context caching support efficient mixed-branch execution and correct prefix reuse.Source: §10 · p. 13
  • Experiments report a better quality-efficiency trade-off than static hybrid baselines, alongside higher serving throughput and lower inference latency.Source: §10 · p. 13

Related papers

Loading 2610.37432v1…