I want to be upfront about something: this whole project runs on free Kaggle T4 notebooks, an AWS EC2 t3.micro relay that costs almost nothing, and public internet. No A100s. No private datacenter network. No budget.
And yet, ShardFlow v2.1 hits 28.10 TPS peak on Qwen2.5-7B across two separate cloud regions over WAN.
This is the story of how that happened, and specifically the one fix in v2.1 that I did not see coming.
The Problem: Running a 7B Model When You Have No Money
A 7B parameter model in FP16 needs roughly 15 GB of VRAM. A single Kaggle T4 has 16 GB. Technically it fits, barely, with nothing left over for a KV cache.
The solution is tensor parallelism: split the model across two machines. Node 0 (Iowa) handles layers 0 to 14. Node 1 (Oregon) handles layers 14 to 28, plus the LM head and final verification. They talk to each other through a TCP relay running on an EC2 t3.micro in Ohio.
The baseline throughput with this setup and no tricks: 4.92 TPS. Usable, but not fast.
Speculative Decoding: The Idea
LLM inference is slow because it's sequential. You generate one token, wait, generate another, wait. Each round trip across WAN costs you ~86ms RTT. At 1 token per round trip, you're fighting the network the whole time.
Speculative decoding flips this. Instead of sending one token at a time, you run a tiny draft model locally to guess the next K tokens ahead. Then you send all K guesses to the verifier in one shot. If the big model agrees with M of them, you've committed M tokens in a single round trip instead of one.
ShardFlow uses Qwen2.5-0.5B as the draft model, running on cuda:1 of Node 0 while the 7B target slice runs on cuda:0. Zero VRAM contention. The drafter proposes 8 candidates, Node 1 verifies them all in parallel, and you get an average of 4.07 tokens per round trip instead of 1.
With speculative decoding in eager mode: 14.3 TPS peak. 3x better.
The Wall I Hit
I thought 14.3 was the ceiling. The network was the obvious bottleneck: two Kaggle instances in different states, an EC2 relay in between, public internet routing. What else could you do?
Then I looked more carefully at what the draft model was actually doing.
Every round, generating 8 candidate tokens meant running 8 separate forward passes through the 0.5B model. Each forward pass launched roughly 1,500 CUDA kernels, one by one, from a Python loop.
Here's the problem: each CUDA kernel executes in 2 to 5 microseconds on the GPU. But Python needs 8 to 10 microseconds just to issue the launch call. The GPU was sitting idle for more time than it was actually computing. Draft generation per round: 112ms. The GPU idle rate: 65%.
Python was quietly murdering GPU utilization and I had no idea.
CUDA Graphs: What They Are and Why They Helped
A CUDA Graph is a way to capture a sequence of GPU operations once and replay them as a single driver call.
Normally, every time your model does a forward pass, Python issues hundreds or thousands of individual kernel launches. Each one is a separate call to the CUDA driver. That overhead adds up fast, especially when you're doing it in a loop.
With CUDA Graphs, you capture the entire forward pass of the 0.5B draft model: all 24 transformer layers, the LM head, the argmax for the next token. You do this once. After that, replaying the whole thing costs one driver call. No Python in the hot path at all.
Draft generation: 112ms to 25ms. 4.5x faster.
Why It Kept Breaking (and How I Fixed It)
Every time I tried CUDA Graphs, the model started looping: "the the the the the". Clearly something was wrong.
CUDA Graphs capture exact GPU memory addresses at record time. If any tensor gets reallocated during replay, the graph reads from a stale address and you get garbage output.
HuggingFace's default KV cache (DynamicCache) calls torch.cat every single token step. That allocates a new buffer every time. The graph had captured the old address. Replay read from it. Output: garbage.
Four changes fixed this:
1. StaticCache instead of DynamicCache. StaticCache pre-allocates fixed-size buffers for the KV cache. No reallocations during generation. The addresses the graph captured stay valid.
2. In-place tensor mutation. Instead of creating new tensors for intermediate values, everything gets written in-place. Same memory, same address, graph stays happy.
3. Explicit position_ids updates. The graph needs to know which position each token is at. With dynamic allocation, this was implicit. With a static graph, you have to update position_ids manually before each replay.
4. In-place KV rewind. When the speculative verifier rejects some draft tokens, the KV cache needs to roll back to the last accepted position. This rewind has to happen in-place, not by creating a new cache object.
Once all four were in place: no more loops. Clean output. 25ms draft generation.
The Numbers
On Qwen2.5-7B across 2 Kaggle T4s over WAN:
| Version | TPS |
|---|---|
| v1.0 (REST relay) | 2.27 |
| v2.0 baseline | 4.92 |
| v2.0 + neural drafter (eager) | 14.3 peak |
| v2.1 + CUDA graphs | 28.10 peak / 20.31 avg |
Also tested on Qwen2.5-14B with 4-bit NF4 quantization, same two T4s: 14.43 TPS average over WAN. A 14.7B parameter model. Free GPUs.
What I Actually Learned
The network was not the bottleneck. I spent a lot of time assuming the WAN latency was the hard ceiling, that there was nothing left to squeeze. The real bottleneck was Python kernel launch overhead, and it was invisible until I looked at GPU idle time.
Profiling matters more than intuition. "The network is slow" is an easy assumption to make. "Python is launching 1,500 kernels from a loop and the GPU is idle 65% of the time" requires actually measuring.
CUDA Graphs are not magic. They are very specific. Captured addresses must stay valid. Any dynamic allocation breaks them. The StaticCache + in-place mutation combination is what makes them work for autoregressive generation.
Try It Yourself
ShardFlow is open source and designed to reproduce on free Kaggle notebooks. You need two Kaggle accounts and an EC2 t3.micro (or any machine with a public IP).
github.com/rautaditya2606/Shardflow
The README has step-by-step instructions for reproducing the exact benchmark. 583 people have already cloned it. I'd love to know if you get different numbers on different hardware.
v3 is whenever someone sponsors me actual GPUs. Until then, free T4s and Ohio relays it is.
Top comments (0)