FastTree: Optimizing Attention Kernel and Runtime for Tree-Structured LLM Inference

Zaifeng Pan, Yitong Ding, Yue Guan, Zheng Wang, Yida Wang, Yufei Ding

Conference on Machine Learning and Systems 2025 · Day 2 · Session 1: LLM and Diffusion Model Serving

Overview

This article delves into DeepServe, an innovative system designed to significantly enhance the efficiency and quality of serving text-to-image diffusion models. Presented by Shirong Yang, a PhD student from UMass Amherst, and developed in collaboration with UMass and Adobe Research, DeepServe tackles the inherent computational intensity of these generative AI models. The core challenge lies in the multi-step iterative denoising process central to diffusion models, which contributes substantially to high serving costs and latency.

Watch on SlidesLive · Slides

Visual summary for FastTree: Optimizing Attention Kernel and Runtime for Tree-Structured LLM Inference by Zaifeng Pan, Yitong Ding, Yue Guan, Zheng Wang, Yida Wang, Yufei Ding
Visual summary for FastTree: Optimizing Attention Kernel and Runtime for Tree-Structured LLM Inference by Zaifeng Pan, Yitong Ding, Yue Guan, Zheng Wang, Yida Wang, Yufei Ding

Key moments

  1. 0:00 Introduction: High Cost of Diffusion Model Serving
  2. 2:00 Quality vs. Speed Trade-off & Query Difficulty
  3. 4:30 DeepServe's Core Idea: Query-Aware Model Scaling
  4. 5:50 How DeepServe's Cascading Pipeline Works
  5. 6:50 Discriminator Design for Output Quality Estimation
  6. 10:00 Optimizing System Parameters with MILP Resource Allocator
  7. 12:00 DeepServe System Architecture Overview

DeepServe: Efficiently Serving Text-to-Image Diffusion Models with Query-Aware Model Scaling

Speakers: Shirong Yang (PhD Student, UMass Amherst), and collaborators from UMass and Adobe Research

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

This article delves into DeepServe, an innovative system designed to significantly enhance the efficiency and quality of serving text-to-image diffusion models. Presented by Shirong Yang, a PhD student from UMass Amherst, and developed in collaboration with UMass and Adobe Research, DeepServe tackles the inherent computational intensity of these generative AI models. The core challenge lies in the multi-step iterative denoising process central to diffusion models, which contributes substantially to high serving costs and latency.

DeepServe introduces a novel query-aware model scaling approach, moving beyond static model deployments or random dynamic resource allocation. The system intelligently assesses the difficulty of incoming user queries and dynamically routes them to either a lightweight or a heavy-weight diffusion model, leveraging a specially trained discriminator. This adaptive strategy ensures that computationally expensive heavy-weight models are only invoked when genuinely necessary, striking an optimal balance between image quality, generation speed, and operational cost.

The importance of DeepServe cannot be overstated in an era where text-to-image models are becoming pervasive. By optimizing the serving infrastructure, DeepServe makes these powerful generative AI tools more accessible, cost-effective, and responsive for a wide range of applications. It represents a crucial step towards building more efficient and intelligent AI serving systems that can adapt to the dynamic and diverse demands of real-world workloads, pushing the boundaries of what's possible in large-scale AI deployment.

Background

▶ Watch: Introduction: High Cost of Diffusion Model Serving (0:00)

The landscape of text-to-image diffusion models has rapidly evolved, offering users the ability to generate highly realistic images from simple text prompts. The underlying mechanism, however, is computationally demanding. A typical diffusion model starts with pure noise and progressively refines it over multiple denoising steps, which can range from a few to as many as 50, to produce a coherent image. This iterative process is the primary driver of the high computational cost associated with serving these models, directly impacting latency and operational expenses.

Currently, a diverse array of diffusion models exists, each presenting a unique trade-off between image quality and generation time. Models like SD Turbo or SDXS are optimized for speed, delivering quick results but potentially at the expense of intricate detail or overall quality. Conversely, larger, more complex models, such as SD 1.5 or other heavier variants, excel at generating high-fidelity images but demand significantly more computational resources and time. The prevailing practice of committing to a single model for all serving scenarios—whether lightweight for speed or heavy-weight for quality—is inherently inefficient. This static approach fails to account for the variability in user queries and system conditions, leading to either excessive resource consumption for simple requests or compromised quality for complex ones.

A key observation that underpins DeepServe's motivation is that not all queries are equally difficult. Experiments reveal that a significant proportion, estimated between 20% to 40%, of user prompts are "easy queries." For these simpler requests (e.g., "generate an image of a cat"), lightweight models can produce images of comparable, or even superior, quality to their heavier counterparts, with minimal discernible difference to the end-user. In contrast, "hard queries" (e.g., "generate an image of the downtown of a city") necessitate the greater capacity of heavy-weight models to capture complexity and maintain visual fidelity. Prior attempts to address this challenge, such as using existing perceptual metrics like CLIP score or Pick score for routing, have proven ineffective, often performing worse than random routing strategies due to their computational overhead and inability to accurately distinguish query difficulty in the context of generative outputs. Even more advanced systems like Proteus, which offer dynamic resource allocation, still resort to random routing among model variants, missing the crucial opportunity for query-level optimization. This gap highlights the need for an intelligent, query-aware system that can dynamically adapt model selection to the specific demands of each prompt, thereby optimizing resource utilization without compromising the user experience.

Key Findings

▶ Watch: DeepServe's Core Idea: Query-Aware Model Scaling (4:30)

DeepServe presents several pivotal findings that redefine the approach to serving text-to-image diffusion models efficiently and effectively:

  • Superiority of Query-Aware Model Scaling: The most significant finding is the profound impact of dynamically selecting models based on query difficulty. DeepServe demonstrates that a query-aware model cascading strategy, which routes requests to either lightweight or heavy-weight models based on an intelligent assessment, dramatically improves overall system efficiency and response quality compared to static model deployments or methods relying on random routing. This approach capitalizes on the observation that a substantial fraction of queries are "easy," making lightweight models sufficient.
  • Effectiveness of a Learned Discriminator: A critical innovation is the design and implementation of a specialized discriminator. Unlike generic perceptual metrics (e.g., CLIP score, Pick score) that proved ineffective and even detrimental due to added overhead, DeepServe's discriminator is a binary classifier trained specifically with a real vs. fake classification objective. This discriminator accurately assesses whether the output from a lightweight model is "good enough," guiding the routing decision. Its superior performance, consistently outperforming baselines, validates the necessity of a purpose-built mechanism for evaluating generative model outputs in a cascading pipeline.
  • Importance of Co-optimized Resource Allocation: DeepServe highlights that the optimal performance of such a dynamic system hinges on the co-optimization of various system-level parameters. These include the discriminator's confidence threshold, the number of workers allocated to lightweight and heavy-weight models, and batch sizes. The research shows that tuning these parameters independently leads to suboptimal configurations. DeepServe addresses this by employing a Mixed Integer Linear Programming (MILP)-based resource allocator to identify Pareto optimal configurations, ensuring the best possible response quality for a given throughput while adhering to latency constraints.
  • Quantifiable Performance Gains: Experimental results powerfully validate DeepServe's approach. The system demonstrated an improvement in image quality, indicated by a lower FID (Frechet Inception Distance) score, of up to 24%. Simultaneously, it achieved a substantial reduction in SLO violation ratio (missed latency deadlines) by 19% to 70% when compared against state-of-the-art baselines. These numbers underscore DeepServe's ability to deliver both high quality and high efficiency concurrently.
  • Robustness Across Workloads: DeepServe's adaptive nature allows it to maintain optimal performance across a spectrum of workloads, from low to high demand. By dynamically adjusting the confidence threshold and resource allocation, the system consistently provides the best trade-off between quality and latency, demonstrating its practical applicability in volatile real-world serving environments.

Technical Deep Dive

▶ Watch: How DeepServe's Cascading Pipeline Works (5:50)

DeepServe's architecture is meticulously designed around the principle of query-aware model scaling through a model cascading pipeline. The system's intelligence lies in its ability to dynamically route queries based on their perceived difficulty, thereby optimizing resource utilization.

The core data path begins when a client sends a query (a text prompt) to a load balancer. Instead of immediately routing to a fixed model, the load balancer first directs the query to a worker instance running a lightweight diffusion model (e.g., SD Turbo). Crucially, this worker also hosts a specialized discriminator. The lightweight model quickly generates an initial image, which the discriminator then evaluates.

The discriminator is a key innovation. It's a binary classifier specifically trained using a real vs. fake classification objective. During its training phase, the discriminator learns to distinguish between high-quality "real" images (serving as positive samples) and generated images from both lightweight and heavy-weight diffusion models (serving as negative samples). This training enables it to develop a robust understanding of what constitutes a "realistic" or "good enough" image in the context of diffusion model outputs. During inference, the discriminator no longer requires real images; it simply assesses the output from the lightweight diffusion model and produces a confidence score. This score quantifies how realistic the generated image appears, essentially estimating its quality.

If the discriminator's confidence score exceeds a pre-defined threshold, it indicates that the lightweight model's output is sufficient for the query. In this scenario, the result is immediately returned to the user, saving significant computational resources and latency. However, if the confidence score falls below the threshold, signifying that the lightweight model's output is suboptimal or the query is particularly challenging, the query is then rerouted to a worker operating a heavy-weight diffusion model (e.g., SD 1.5). This second stage generates a higher-quality image, which is then returned to the user. This cascading mechanism ensures that the more expensive heavy-weight models are only engaged when truly necessary, maximizing efficiency without compromising quality.

Beyond the data path, DeepServe incorporates a sophisticated control path for dynamic resource management. Workers continuously send status updates (e.g., load, latency) to a central controller. The controller leverages this real-time information to make informed resource allocation decisions. These decisions involve dynamically adjusting parameters such as the confidence threshold for the discriminator, the number of workers assigned to lightweight and heavy-weight models, and even the batch sizes for inference. The controller then updates the load balancer, ensuring that queries are routed to appropriate workers with the correct model configurations. This closed-loop feedback system allows DeepServe to adapt effectively to fluctuating workloads and maintain optimal performance.

A critical technical challenge addressed by DeepServe is the co-optimization of these interacting system parameters. The design space for configurations (e.g., threshold, worker counts, batch sizes) is vast, making independent tuning impractical and suboptimal. To navigate this complexity, DeepServe employs a Mixed Integer Linear Programming (MILP)-based resource allocator. The MILP solver is designed to search for Pareto optimal configurations, which represent the best possible trade-offs between serving throughput and response quality, all while adhering to system constraints such as latency deadlines (e.g., 5-second SLO). By formulating resource allocation as an MILP problem, DeepServe can systematically explore this large design space to identify configurations that deliver the optimal response quality for a given serving throughput, ensuring robust performance under varying workload conditions.

In essence, DeepServe's technical prowess stems from the intelligent integration of a learned, domain-specific discriminator with a dynamic, MILP-optimized resource allocation strategy, all orchestrated within a cascading model serving pipeline.

Experimental Setup & Results

▶ Watch: Optimizing System Parameters with MILP Resource Allocator (10:00)

DeepServe's performance was rigorously evaluated against several baselines using a realistic experimental setup designed to simulate real-world serving conditions for text-to-image diffusion models.

Evaluation Setup:

  • Datasets: Text prompts were sourced from the widely used MS Coco dataset, providing a diverse range of image generation requests.
  • Service Level Objective (SLO): A strict latency deadline of 5 seconds was set for each query, representing a common requirement for interactive AI applications.
  • Model Cascade:
  • Lightweight Model: SD Turbo was chosen for its speed, achieving a latency of approximately 0.1 seconds per inference on an NVIDIA A100 GPU.
  • Heavy-weight Model: SD 1.5 served as the high-quality, higher-latency model, with an inference time of around 1.8 seconds on the same A100 GPU.

Baselines for Comparison:

  • Clipper Light: A static baseline that exclusively uses the lightweight model (SD Turbo) for all queries. This setup prioritizes speed but often compromises image quality.
  • Clipper Heavy: Another static baseline that exclusively employs the heavy-weight model (SD 1.5) for all queries. This prioritizes quality but incurs higher latency and computational costs, making it prone to SLO violations.
  • Proteus: Representing the state-of-the-art in dynamic resource allocation, Proteus adapts resources based on demand but routes queries randomly among available model variants. It lacks query-level intelligence.

Key Metrics:

  • FID (Frechet Inception Distance): A widely accepted metric for evaluating the quality of generated images, where a lower FID score indicates better response quality and greater fidelity to real images.
  • SLO Violation Ratio: Measures the percentage of queries that fail to meet the 5-second latency deadline, indicating the system's ability to maintain responsiveness.

Headline Results:

The experimental results emphatically demonstrated DeepServe's superior performance across various workloads:

  1. Static Baselines' Limitations: As expected, Clipper Light struggled significantly, missing SLOs frequently while also delivering poor image quality (high FID). Clipper Heavy, while producing better image quality, suffered from severe SLO violations due to the high latency of the heavy-weight model. These static approaches highlighted the inherent trade-off that DeepServe aims to overcome.
  1. Proteus's Suboptimality: Proteus offered a middle-ground performance, better than the static extremes by dynamically adjusting resources. However, its random routing strategy prevented it from achieving optimal efficiency or quality, as it did not account for the varying difficulty of individual queries.
  1. DeepServe's Dominance: DeepServe consistently occupied the "bottom-left corner" of performance graphs, signifying simultaneously low FID (high image quality) and low SLO violation ratio. This outstanding performance is directly attributed to its query-aware nature, which intelligently selects the appropriate model for each query based on its content and difficulty.
  • Image Quality Improvement: DeepServe was shown to improve image quality by up to 24% (as measured by a reduction in FID) compared to the best performing baselines.
  • Latency Violation Reduction: The system dramatically reduced latency violations by 19% to 70%, ensuring a much more responsive user experience.
  1. Workload Adaptability: DeepServe proved highly adaptive, maintaining its optimal trade-off between quality and latency under both low and high workload conditions. This adaptability is achieved through the dynamic adjustment of the discriminator's confidence threshold and intelligent resource allocation by the MILP-based resource allocator.
  1. Discriminator Efficacy: A dedicated experiment compared DeepServe's custom-trained discriminator against random routing and existing perceptual metrics like CLIP score and Pick score. The results clearly showed that DeepServe's discriminator (represented by a "red curve" in the presentation) consistently performed better, leading to more effective and efficient routing decisions with superior control over output quality. This validates the design choice of a specialized discriminator over generic metrics.

In summary, DeepServe's experimental validation strongly supports its claim of achieving both high efficiency and quality in text-to-image diffusion model serving, significantly outperforming current state-of-the-art methods.

Practical Implications

▶ Watch: DeepServe System Architecture Overview (12:00)

DeepServe offers profound practical implications for anyone involved in deploying, managing, or utilizing text-to-image diffusion models at scale. Its query-aware model scaling paradigm addresses critical bottlenecks and introduces new efficiencies across the AI serving ecosystem.

For Practitioners and Infrastructure Teams:

  • Significant Cost Savings: By intelligently routing easy queries to lightweight models, DeepServe drastically reduces the computational resources required for serving. This translates directly into lower GPU utilization costs and energy consumption, making large-scale diffusion model deployments more economically viable.
  • Improved User Experience: The system's ability to minimize SLO violations (reducing them by 19-70%) means users experience faster response times and fewer timeouts, leading to a smoother and more satisfying interaction with generative AI applications.
  • Dynamic Resource Management: DeepServe's MILP-based resource allocator provides a robust framework for dynamically adjusting worker allocations and other system parameters. This eliminates the need for manual, reactive scaling, allowing infrastructure teams to maintain optimal performance under fluctuating and unpredictable workloads without constant intervention.
  • Avoiding Over- and Under-provisioning: The system intelligently matches model capacity to query difficulty, preventing the common pitfalls of either over-provisioning expensive resources for simple tasks or under-provisioning, which leads to poor quality or missed deadlines for complex ones.

For Model Builders and Deployers:

  • Flexible Model Deployment: DeepServe provides a flexible framework for deploying multiple model variants (e.g., distilled models like SD Turbo alongside full-sized models like SD 1.5). This allows model builders to develop a spectrum of models optimized for different quality-latency trade-offs, knowing that DeepServe can intelligently leverage them. As discussed in the Q&A, even a single model run with fewer steps (lightweight) versus more steps (heavy-weight) can be integrated.
  • Quality-of-Service Guarantees: By co-optimizing system parameters, DeepServe allows deployers to meet specific quality-of-service targets (e.g., a certain FID score) while maximizing throughput and adhering to latency SLOs. This is crucial for applications with stringent performance requirements.
  • Enhanced Research and Development: The insights gained from identifying "easy" versus "hard" queries could inform future model distillation techniques or prompt engineering strategies, guiding the development of more efficient models from the outset.

Tradeoffs and Limitations:

  • Discriminator Training Overhead: While highly effective, the custom discriminator requires a dedicated training phase using a real vs. fake classification objective. This adds an initial development and training cost, though it pays off in long-term serving efficiency.
  • Initial Latency for Lightweight Model: The cascading pipeline introduces a minimal initial latency as every query first passes through the lightweight model and the discriminator. While this is typically negligible (0.1s for SD Turbo), it's an added step compared to direct routing to a heavy-weight model.
  • Complexity of MILP Solver: Implementing and maintaining an MILP-based resource allocator can be complex, requiring expertise in optimization techniques. The solver needs to be robust and efficient enough to make real-time decisions in a dynamic environment.
  • Difficulty of Prompt-based Pre-filtering: The talk acknowledges that directly estimating image output quality from a text prompt before any image generation is still an open and challenging problem. This reinforces the necessity of the cascaded approach where an initial image is generated for the discriminator to evaluate.

In conclusion, DeepServe offers a sophisticated yet practical solution to the challenges of serving complex generative AI models. It empowers organizations to deploy diffusion models more efficiently, cost-effectively, and with a significantly improved user experience, making advanced AI capabilities more accessible and sustainable.

Key Takeaways

  • Diffusion model serving is computationally expensive: The iterative denoising process of text-to-image diffusion models leads to high serving costs and latency, necessitating optimization strategies.
  • Query difficulty varies significantly: A substantial portion (20-40%) of queries are "easy," meaning lightweight models can produce comparable quality outputs, while "hard" queries require more powerful models.
  • DeepServe employs query-aware model cascading: The system dynamically routes queries through a lightweight model and a specialized discriminator; if the output is not "good enough," the query is rerouted to a heavy-weight model.
  • A learned discriminator is crucial for effective routing: DeepServe's binary classifier, trained with a real vs. fake objective, accurately assesses lightweight model outputs, outperforming generic perceptual metrics and random routing.
  • MILP-based co-optimization is key for system-wide efficiency: DeepServe uses a Mixed Integer Linear Programming (MILP) solver to find Pareto optimal configurations by co-optimizing parameters like confidence threshold, worker counts, and batch sizes, ensuring the best quality for a given throughput and latency SLO.
  • DeepServe delivers significant performance gains: Experiments show DeepServe improves image quality (lower FID) by up to 24% and reduces SLO violation ratio by 19-70% compared to baselines, while adapting effectively to varying workloads.

About the Speaker(s)

Shirong Yang is a PhD student from UMass Amherst. He presented the work on DeepServe, a project conducted jointly with collaborators from UMass Amherst and Adobe Research. His research focuses on optimizing the efficiency and serving capabilities of advanced machine learning models, particularly in the domain of text-to-image diffusion models.

Reviews

Simon Wisk (Open Source Developer & AI Tooling Expert) — SOLID

DeepServe is a competent systems paper on diffusion model serving that solves a real problem — not all prompts are equally hard, so why treat them the same? The cascade architecture plus a learned discriminator is a sensible idea, the MILP-based resource allocator is an interesting wrinkle, and the FID/SLO results are at least directionally credible. But the write-up reads like a conference proceedings summary padded with PR copy, the discriminator training procedure is underspecified enough that you couldn't reproduce it, and there's no code or artifact to point at. Worth knowing about if you're in the diffusion serving space, not a must-watch for everyone else.

Jensen Hitch (AI Compute Platform CEO) — SOLID

DeepServe is a competent, well-scoped systems paper that addresses a real serving bottleneck in text-to-image diffusion models. The core insight — that query difficulty is non-uniform and a learned discriminator can route accordingly — is sound and practically motivated. The MILP-based co-optimization of threshold, worker allocation, and batch size shows genuine systems thinking at the serving layer. But the work stays inside the serving system boundary and doesn't reason outward to the full infrastructure stack, cost-per-image economics at production scale, or what this unlocks for the next generation of generative AI deployment. It's a solid session contribution, not a platform insight.

→ Top-rated talks at Conference on Machine Learning and Systems 2025

All talks from Conference on Machine Learning and Systems 2025