Source-linked AI summary
FlashVector: Agent for Hierarchical Model Serving Stack Optimization
Qi Wu, Lohan Lemire, Kai Meng, Zhongmou Cai, Raphael Bargues, Petr Zhitnikov, Zeyuan Cao, Yao Wang, Shujun Bian, Wei Chen, Sean Sheng
TL;DR
Model-serving optimization spans heterogeneous layers, but automated tuning beyond standalone GPU kernels remains underexplored. FlashVector applies an agentic loop across the serving stack, and its deployment at Unity’s Vector platform produced throughput and latency improvements across model-server and feature-store workloads.
Problem
Each serving layer demands specialized expertise, while automated optimization beyond standalone GPU kernels remains largely unexplored.
Source: §1 · p. 1Method
FlashVector uses a layer-agent framework to optimize stack components independently, verify changes against production traffic, and retain accepted end-to-end gains.
Source: §1 · p. 2Results
Up to 2× throughput increase and up to 1.98× latency speedup on the model server, and up to 1.6× throughput increase on the feature store.
Source: §6 · p. 8Takeaways & Limitations
The same closed loop discovered optimizations across GPU kernels, computation graphs, model-server code, feature processing, and serving parameters.
Source: §6 · p. 8
Abstract
from arXiv · showhide
Model serving is one of the largest cost drivers in production recommender systems. Maximizing its throughput requires navigating a deeply layered hierarchy: GPU kernels, the ML framework computation graph, the model server, and on-demand feature processing -- each demanding specialized domain expertise. Such cross-layer expertise is inherently difficult to acquire, and does not scale with a workload that continuously grows and evolves, leaving significant cost efficiency gains unrealized. While recent AI agents have demonstrated human expert level efficiency in standalone GPU kernel optimization, automated tuning and optimization for the rest of the serving stack remain largely unexplored. We present FlashVector, an agentic system that optimizes performance across all layers of the model serving stack. The key contribution is an extensible framework to generalize the single kernel optimization agent paradigm to heterogeneous technical stacks, and to deliver performance improvements holistically. After deployment in Unity's Vector advertising platform, FlashVector achieved up to 2x throughput increase and up to 1.98x latency speedup on model server, and up to 1.6x throughput increase on feature store. These optimizations were discovered not only at the GPU kernel and computation graph levels, but also across the other components of the model serving stack, such as the model server (NVIDIA Triton's C++ codebase) and the on-demand feature transformation service (Python codebase), demonstrating the extensibility of the framework to more complex system architectures.
1 Introduction
Production recommender serving spans GPU kernels, framework graphs, model servers, and feature processing, each with distinct optimization demands. FlashVector applies an agentic optimization loop across those layers and accepts changes only when production evidence shows an end-to-end gain.
1 Introduction
- Cross-layer optimization is difficult because each serving layer requires distinct expertise, leaving cost-efficiency gains unrealized.The stack includes client-side request handling, model-server batching, framework graph configuration, and GPU-kernel efficiency.Source: §1 · p. 1
- Kernel agents motivate extending LLM-driven optimization to the full serving stack, where code discovery, parameter tuning, and co-design remain challenges.The paper notes that non-kernel layers contain substantial optimization opportunities.Source: §1 · p. 1
- FlashVector uses a shared layer-agent interface and accepts independently proposed layer changes only after replayed production traffic confirms an end-to-end gain above measurement noise.The loop also re-triggers automatically as models, traffic, and infrastructure change.Source: §1 · p. 2
2 Related Work
Related work addresses model-serving systems, serving-time configuration search, and LLM-driven kernel optimization. FlashVector extends these strands by applying a profiling-guided agent loop across heterogeneous serving layers.
Model Serving Systems
- Serving-system research targets routing, batching, and resource allocation, while FlashVector uses an agent to discover fixes across evolving stack bottlenecks.Examples include Clipper [1], TensorFlow Serving [22], and systems for generative models such as Orca, vLLM, and AlpaServe [16].Source: §2 · p. 2
Serving-time Autotuning
- Serving-time autotuners search configuration spaces for throughput or cost goals, whereas FlashVector ranges over those levers alongside code changes.MArk and Cocktail jointly select hardware, models, and autoscaling; Morphling searches resource and batch settings.Source: §2 · p. 2
Kernel Design Agent
- Kernel-agent work applies profiling-guided LLM optimization to GPU kernels; FlashVector extends that loop to model-server code, feature-processing runtimes, and configuration.Those layers require distinct profiling tools, diagnosis context, and correctness criteria.Source: §2 · p. 2
3 Challenges
Extending automated optimization beyond standalone kernels creates four challenges: fragmented tooling, joint code-and-parameter search, production-wide validation, and optimization that must keep pace with change.
3 Challenges
- Distinct languages, tools, and performance signals across layers make bottlenecks difficult to localize without cross-layer expertise.Examples include eBPF, PyTorch Kineto, and Nsight Systems across different stack ecosystems.Source: §3 · p. 2
- Non-kernel throughput optimization jointly searches code changes and multidimensional settings such as batching thresholds, thread pools, and GPU queue depths.The search combines discrete code edits with continuous parameters.Source: §3 · pp. 2–3
- Multi-repository changes must preserve production SLAs, including tight p99 latency budgets, rather than being evaluated as isolated microbenchmarks.The paper notes the risk of cascading downstream failures.Source: §3 · p. 3
- Optimizations can become suboptimal or detrimental as models, pipelines, traffic, and infrastructure change, making performance a continuous target.Manual re-evaluation of the four-layer stack carries an operational burden that teams cannot sustain indefinitely.Source: §3 · p. 3
4 Core System
FlashVector combines layer-specific context with an iterative profile–diagnose–optimize–verify workflow. It refines its knowledge after a run and checks candidate changes for correctness, measurable local improvement, and full-stack performance.
4 Core System
- A one-shot LLM approach lacks precise, layer-specific signals and rarely identifies the correct optimization in a single attempt.Useful inputs can include exact tensor layouts, domain-specific profiles, and business-logic context.Source: §4 · p. 3
- FlashVector loads layer-specific knowledge, iteratively profiles, diagnoses, optimizes, and verifies candidates, then updates its reference documents.These are the system’s three stages: knowledge-base loading, iterative optimization, and knowledge refinement.Source: §4 · p. 3
4.1 Knowledge Base
- The knowledge base covers serving-stack components and includes an expert-curated optimization cookbook to guide decisions within a limited context window.It also records implementation information relevant to model efficiency.Source: §4.1 · p. 3
4.2 Optimization Loop
- Each iteration samples one layer, profiles it across the stack, and applies a layer-specific agent and tooling rather than sweeping all layers at once.Profiles range from local GPU-kernel environments to production-like load tests for model servers and feature processing.Source: §4.2 · p. 4
- Diagnosis tailors context to the sampled layer, ranks bottlenecks by time share and expected gain, and skips candidates whose best gain falls below measurement noise.The top-ranked candidate proceeds to optimization.Source: §4.2 · p. 4
- Candidates must meet change-specific correctness criteria and clear the measured noise floor, then pass a full-stack load test before commitment.A full-stack regression discards a candidate even if it passed local verification.Source: §4.2 · p. 5
- After a run, refinement records changes and outcomes in the knowledge base; the loop also triggers on new releases or stale configurations.In-place knowledge updates bound context size and are reported to mitigate catastrophic forgetting and hallucinations.Source: §4.2 · pp. 5–6
5 Evaluation & Case Studies
FlashVector delivered production-model serving-stack optimizations, including GPU-kernel improvements and bottlenecks beyond those commonly examined automatically. Table 1 summarizes the resulting speedups, with baseline and optimized profiling performed in the same environment.
5 Evaluation & Case Studies
- The production case studies cover full-stack optimizations, with Table 1 summarizing their resulting speedups.The baseline measurements and profiling of optimized variants used the same environment for consistency and fair comparison.Source: §5 · p. 6, Table 1 · p. 6
- The case studies include unusual GPU-kernel optimizations and bottlenecks not commonly found automatically and systematically.Source: §5 · p. 6
- Baseline measurements and optimized-variant profiling used the same environment for a consistent, fair comparison.Source: §5 · p. 6
5.1 Case 1: Optimize Model Server
For an early-stage two-tower retrieval model, FlashVector located a bottleneck in Triton’s input conversion to its Python backend and proposed rewriting the logic in C++.
5.1 Case 1: Optimize Model Server
- 30× speedup followed rewriting Triton’s Python input-deserialization logic in C++.The fix was contributed back to NVIDIA Triton.Source: §5.1 · p. 6, Figure 4 · p. 6
- String-heavy inputs made Python deserialization of Triton’s client-supplied inputs a major bottleneck in the server-to-backend handoff.Source: §5.1 · p. 6
- Figure 4 identifies input serialization between Triton Server and its Python backend as the bottleneck location.Source: Figure 4 · p. 6
5.2 Case 2: Improve Model Inference Efficiency
On a multi-task recommendation model, FlashVector optimized attention and embedding execution, reducing total GPU kernel time and forward-pass latency. The baseline used AOTInductor, FP16 quantization, an RTX PRO 6000 Blackwell GPU, and a batch size of 2000 per request.
5.2 Case 2: Improve Model Inference Efficiency
- The evaluated model combines dense and categorical features, target-aware attention, feature-interaction layers, and task-specific output heads.Source: §5.2 · p. 6
5.2.1 Baseline setup.
- The baseline used PyTorch AOTInductor, FP16 quantization, an NVIDIA RTX PRO 6000 Blackwell GPU, and batch size 2000 per request.Source: §5.2.1 · p. 6
5.2.2 Optimization Highlight.
- Precomputing candidate-only and position-only feature blocks avoids wide per-pair matrices and enables fusing six attention kernels into one.The fused kernel keeps intermediate values in GPU registers.Source: §5.2.2 · p. 6
- 3.56× attention speedup cut per-module latency from 996.5 µs to 279.8 µs.The fused kernel reduced memory traffic to roughly one-third while preserving bit-identical attention weights.Source: §5.2.2 · p. 6, Figure 5 · p. 7
- Decomposing masked-mean pooling into PyTorch primitives exposed embedding lookup and reduction for fusion into one Triton kernel.The embedding bag kernel had consumed 28.1% of GPU time, and the opaque ATen dispatcher prevented fusion with adjacent reductions.Source: §5.2.2 · pp. 6–7, Figure 6 · p. 7
- 33.4% lower total GPU kernel time and 30.1% lower forward-pass latency followed the embedding-layer changes.Dispatcher overhead fell from 44.8% to 28.5%, shifting the primary bottleneck from embedding lookups to GEMM kernels.Source: §5.2.2 · p. 7
5.3 Case 3: Optimize On-demand Feature Processing
FlashVector optimizes on-demand feature processing from request input through feature computation and tensor assembly. Across successive production-measured iterations, the three optimizations increased one pod’s throughput by 1.6%.
5.3.2 Optimization highlight.
- Diagnosis used production eBPF profiles across the Python runtime and Triton threads to identify costs along input receipt, feature computation, and tensor assembly.Source: §5.3 · p. 7, §5.3.2 · p. 7
5.3.1 Baseline setup.
- The baseline handled requests sequentially, with one context row and up to 2,000 candidate rows per request, across 12 Python backend instances on 14 vCPUs.Source: §5.3.1 · p. 7
5.3.2 Optimization highlight.
- For input receipt, native views avoid copying numeric tensors, while native mapping lets C++ look up eligible vocabulary inputs without reading their text into Python.Routes are resolved at initialization; inputs that cannot be proven safe fall back to decoding.Source: §5.3.2 · p. 7
- For irregular per-element work, FlashVector replaces NumPy/pandas operations with hand-written Cython kernels.Table 2 identifies the replaced operations and their call-site speedups.Source: §5.3.2 · p. 7, Table 2 · p. 7
- For tensor assembly, destinations are allocated once and features write directly to their slices, removing both concatenations and a redundant copy.The original two-pass assembly accounted for roughly 20% of serving CPU.Source: §5.3.2 · p. 7
5.3.3 Results.
- 1.6× throughput improvement: one preprocessor pod rose from 940 requests per second at baseline to roughly 1.5k with all three optimizations.The iterations were measured on top of one another, so their double-digit gains compounded.Source: §5.3.3 · p. 8, Table 3 · p. 8
5.4 Case 4: Tune Serving Stack Parameters
FlashVector searches serving runtime configurations alongside code optimization, measuring candidates under held load at a fixed latency SLO. On an already code-optimized model, sustained throughput varied by more than 2× across measured configurations.
5.4 Case 4: Tune Serving Stack Parameters
- More than 2×: sustained throughput varied across configurations for a single already code-optimized model at the same latency SLO.The agent retunes runtime settings as the surrounding stack changes, including after its own code optimizations.Source: §5.4 · p. 8
5.4.1 Baseline setup.
- The sweep currently varies model-instance count and dynamic batch size; queue policy is derived from the SLO rather than searched.A dynamic batch of 6 corresponds to several thousand candidate rows because each request already packs roughly a thousand candidates.Source: §5.4.1 · p. 8
5.4.2 The loop.
- Candidates run under held load and are ranked by sustained throughput at a fixed latency SLO; changes are applied through pull requests and verified with a production shadow-traffic canary.The loop repeats on the next model release, using replayed production requests; it declines a change that does not beat measurement noise.Source: §5.4.2 · p. 8
5.4.3 Results.
- The agent tuned four production models, with each run launched, monitored, interpreted, and delivered as a reviewed pull request.Table 4 lists the selected configurations; its caption defines RPS as throughput sustained at the SLO.Source: §5.4.3 · p. 8, Table 4 · p. 8
6 Discussion and Conclusion
FlashVector generalizes LLM-driven optimization from GPU kernels to the full hierarchical model-serving stack through a shared, production-verified agent loop. Its case studies also illustrate how agents can bridge reference implementations and performance-oriented production code.
6 Discussion and Conclusion
- FlashVector extends LLM-driven GPU kernel optimization to the full hierarchical model-serving stack using a layer-agent abstraction, local proposals with global verification, and an always-on loop.The loop records accepted and rejected changes and re-triggers after retraining, traffic shifts, or hardware upgrades.Source: §6 · p. 8
- Deployment on Unity’s Vector platform delivered up to 2× throughput increase and up to 1.98× latency speedup on the model server, and up to 1.6× throughput increase on the feature store.The optimizations spanned GPU kernels, model-server input serialization, on-demand feature processing, and online parameter search.Source: §6 · p. 8
- The case studies bridge reference and production code by translating implementations into CUDA, C/C++, or lower-level Python and Cython without requiring teams to maintain two implementations by hand.The paper frames this as a broader role for agentic systems between fast-moving reference code and performance-critical systems.Source: §6 · p. 9