Disaggregating the Disaggregated Model

Introduction

Speculative decoding is a powerful optimization: it can speed up token generation by 2x or more. The idea is simple: a draft model races ahead and proposes tokens, and the large target model verifies them in parallel, yielding many tokens for every forward pass of the target model. We’ve previously shown that moving the draft model onto a faster accelerator can significantly improve overall system performance.

In Gimlet Cloud, we’ve improved speculative decoding performance by disaggregating the draft model itself: running its prefill and decode phases on separate devices.

In this post, we’ll start by showing why moving the drafter to faster hardware is the biggest lever for interactivity. Then we’ll explore how disaggregating the draft model itself can preserve those gains while improving other important metrics, including time to first token and energy efficiency.

The power of drafting faster

Let’s start with the basics: how does speculative decoding work? We use a small and fast draft model to quickly generate a potentially valid sequence, and then we employ the target model to verify whether it would have output the same tokens. The diagram below graphically demonstrates the benefit; in the time it would take the target model to decode two tokens (”captial of”), we’ve already generated 5 tokens (“capital of France is Paris”).

How far ahead should we speculate; in other words, what draft length should we choose? Longer drafts can yield more accepted tokens after verification, but going too far can slow end-to-end performance: acceptance rates fall, and more draft work is discarded. Let’s model how draft length affects interactivity:

Interactivity peaks at a draft length of 9 tokens when draft token generation takes 5 ms (blue line).

The graph also shows the effect of running the same model on faster hardware (purple line); the latency improvement yields higher interactivity at every draft length and also moves the peak forwards to 15 tokens.

Moving just the draft model onto faster hardware increases the number of tokens that we can verify in parallel on the target model and thus speeds up the experience for users.

What is this faster hardware?

We keep talking about faster accelerators, but you might be wondering: what are these chips? What makes them so special?

GPUs are great at handling prefill: tokens get computed in parallel and saturate the processor. When it comes to decode, they become less efficient. As shown in the roofline below, GPUs end up “memory-bound”: they are bottlenecked on loading weights into the chip itself.

A new generation of chipmakers are designing hardware to address this memory-bound region gap. SRAM-based chips are one example (read our previous post to learn more).

Because these chips can then access weights at much higher bandwidths, workloads that are memory-bound on GPU (ie decode) speed-up significantly on these chips. This comes at a cost, as these chips tend to have less memory capacity than GPUs so loading the same model requires more chips. More chips = more power = lower throughput per kW of the overall system; it’s a tradeoff of energy efficiency for interactivity.

Now this is where we can leverage the power of disaggregation: we can start with a GPU only system and move any memory-bound workloads onto high-memory bandwidth accelerators and improve the overall system performance.

Modeling a disaggregated drafter

For this post, we want to demonstrate the benefits of the different proposed designs using tangible numbers. We use a high-level, utilization-optimizing model with a fixed workload (ISL=4096, OSL=512, no cached prefix). It selects fractional device ratios to match capacity to each phase, assumes fixed batching, and does not model prefill–decode interference.

Below are the hardware parameters that we used for our modeling. The GPU numbers are inspired from real workloads, and the accelerator figures are set to reflect the speedup and energy costs typical of purpose-built inference hardware. The idea is to keep the analysis focused on the underlying tradeoffs rather than on one vendor implementation.

See the appendix for more modeling details.

Moving the draft model to faster hardware

Let’s see what happens when we move the draft model to the faster accelerator in our model, as shown in the figure:

As plotted below, interactivity improves significantly across all draft lengths, and the optimal draft length moves from 8 to 15. Peak interactivity rises from about 61 to 137 output tokens/s/request, a 2.25x improvement. This is an impressive speedup for a change that only affects the small draft model.

This gain doesn’t come for free: the cost of generating a draft token is higher for the accelerator, which is reflected in the throughput/kw energy efficiency plot below.

We see the energy tradeoff because draft-decode takes a larger share of the total energy on the accelerator: at draft length 10, draft decode accounts for about 5.8% of forward energy on GPU and 22.4% on the accelerator. Each drafted token costs 30 mJ on GPU versus 150 mJ on the accelerator, reducing energy efficiency despite faster decoding, as shown in the two charts below.

This is a tradeoff we make: we get faster tokens at higher energy costs, but this is a significant step forward on the interactivity–efficiency frontier.

How we’ve leveraged draft disaggregation in Gimlet Cloud

We found several discrepancies from the model when bringing disaggregated speculative-decoding deployments into Gimlet Cloud. Draft decode performance matched our models across deployments, but draft prefill varied significantly. On one accelerator, draft prefill latency was high enough to increase system TTFT compared with running the draft model on a GPU. On another, draft prefill consumed substantially more energy, reducing throughput per kW.

This led us to the final architecture shown below: disaggregating the draft model by moving draft prefill back to a GPU.

We modeled this configuration before integration and found no clear advantage or disadvantage in either interactivity or throughput per kW. Interactivity remained unchanged whether draft prefill ran on the GPU or the accelerator, even after we updated the model parameters to reflect what we observed in the field.

In the first deployment, moving draft prefill to GPU brought the TTFT to match or be below GPU latency. We modeled this below by setting draft prefill on the accelerator to 150 ms (15x originally modeled): moving it to the GPU reduces TTFT from 212.5 ms to 190 ms at a draft length of 15, matching the GPU-only configuration’s TTFT while retaining the accelerator’s interactivity advantage, as shown above.

In the other, energy inefficient deployment, separating the stages pushed the efficiency up significantly. We modeled this below by setting accelerator draft prefill energy usage to 100,000 mJ per prompt (2.5x), compared with 20,000 mJ on the GPU. Moving draft prefill to the GPU improves peak system throughput/kW by approximately 17%; and once again does not affect interactivity.

Keeping draft prefill on a GPU while moving draft decode to an accelerator gives us design flexibility. We can optimize GPU prefill for the draft model, then focus on developing fast decode kernels that play to each accelerator’s strengths. We also retain the option to optimize accelerator prefill later and move the work back if its performance matches or exceeds the GPU’s.

Conclusion

Running draft models on fast accelerators significantly improves interactivity while also allowing us to speculate much further ahead. While testing these heterogenous systems in Gimlet Cloud, we discovered that real accelerators tend to vary in how well they handle prefill performance. In such cases, we found that moving draft prefill to GPU could deliver gains in TTFT or energy efficiency, depending on the accelerator’s prefill performance.

We’ve focused on linear speculative decoding, but faster drafting could benefit other approaches too. Tree-based methods such as SpecInfer could benefit from more breadth of paths or deeper drafting on particular paths. DSpark is a method that might also see benefits from faster decoding, but the need to transfer model hidden state between the target and draft model would add a different constraint to the disaggregated design. We think these would be valuable future investigations.

This is the broader opportunity in heterogeneous inference: design systems to target the unique capabilities of each chip. At Gimlet, we’re always tinkering with new configurations to push the Pareto frontier of inference. If you want to help design and build these systems, we’d love to talk.

Appendix

Modeling Assumptions continued.

We modeled expected accepted draft length a saturating exponential, which matches the diminishing return of accepted draft tokens as the draft length goes up.

The table below lists out the modeling assumptions in more detail: