28 TPS on Qwen2.5-7B across two separate cloud regions over public WAN using speculative decoding + CUDA Graphs [P]
Mirrored from r/MachineLearning for archival readability. Support the source by reading on the original site.
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]
More from r/MachineLearning
-
Archival vs non archival workshop [R]
Aug 23
-
When an AI agent says “done” how do you know it actually happened? [P]
Aug 23
-
I gave a repair tool exactly one worked example. It inferred the rest and fixed 5/5 bugs in code it had never seen — 4 CPU instructions, zero tokens.[p]
Aug 23
-
[N] EACL 2027 Industry Track - Deadline 11 September [N]
Aug 23
Discussion (0)
Sign in to join the discussion. Free account, 30 seconds — email code or GitHub.
Sign in →No comments yet. Sign in and be the first to say something.