r/MachineLearning · · 1 min read

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.

submitted by /u/katua_bkl
[link] [comments]

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.

More from r/MachineLearning