28 TPS on Qwen2.5-7B across two separate cloud regions over public WAN using speculative decoding + CUDA Graphs [P]
been building ShardFlow for the past few months, a distributed LLM inference
framework that splits any HuggingFace transformer across N GPU machines and uses
neural speculative decoding to deal with WAN latency.
the setup for the benchmark: two T4 nodes in separate GCP regions (Iowa + Oregon)
talking through an AWS EC2 TCP relay in Ohio. ~86ms RTT on public internet.
the key insight with speculative decoding here is that WAN latency stops being a
per-token cost and becomes a per-round cost. with K=8 drafting you're committing
4.07 tokens per round trip instead of 1. at 86ms RTT that's a big deal.
numbers on Qwen2.5-7B:
non-speculative baseline: 4.92 TPS
neural drafter (eager): 14.3 TPS peak
+ CUDA Graphs on drafter: 28.10 TPS peak / 20.31 TPS avg
also ran Qwen2.5-14B with NF4 4-bit quant, same two nodes: 14.43 TPS avg.
the v2.1 fix that surprised me most: draft generation was launching ~1500 CUDA
kernels per round from a Python loop. each kernel 2-5us, Python launch overhead
8-10us. GPU sitting idle 65% of the time. capturing the full 0.5B forward pass
as a CUDA Graph and replaying with one driver call dropped draft latency from
112ms to 25ms.
other things in the stack: zero-copy Rust TCP relay, StaticCache + in-place KV
rewind for graph compatibility, meta-device model slicing to avoid loading 15GB
into CPU RAM.
repo: https://github.com/rautaditya2606/Shardflow
happy to answer questions on the speculative decoding implementation or the CUDA
graphs stuff specifically.
[link] [comments]
Want to read more?
Check out the full article on the original site