transport_bench/
transport_bench.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Transport microbenchmark — loopback TCP latency / bandwidth / collectives.
6//!
7//! Spins up a 2-rank mesh on `127.0.0.1` (a thread per rank, but real
8//! kernel TCP sockets) and measures the cost the distribution layer adds:
9//! point-to-point round-trip latency, streaming bandwidth across payload
10//! sizes, and `all_reduce` latency for a hidden-state-sized vector.
11//!
12//! Run: `cargo run -p rlx-distributed --example transport_bench --release`
13//!
14//! Loopback is ~memcpy speed, so these numbers isolate framing/serialization
15//! overhead; a real Ethernet/Thunderbolt link adds its own latency floor.
16
17use rlx_distributed::{NetTransport, ProcessGroup, ReduceKind};
18use std::net::{Ipv4Addr, SocketAddr, TcpListener};
19use std::sync::Arc;
20use std::thread;
21use std::time::Instant;
22
23fn main() {
24    let world = 2u32;
25    let listeners: Vec<TcpListener> = (0..world)
26        .map(|_| TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap())
27        .collect();
28    let addrs: Vec<SocketAddr> = listeners.iter().map(|l| l.local_addr().unwrap()).collect();
29
30    println!("transport: loopback TCP, world_size = {world}");
31    let handles: Vec<_> = listeners
32        .into_iter()
33        .enumerate()
34        .map(|(rank, listener)| {
35            let addrs = addrs.clone();
36            thread::spawn(move || {
37                let t = NetTransport::from_listener(rank as u32, world, listener, addrs, 64 << 20)
38                    .unwrap();
39                bench(&ProcessGroup::new(Arc::new(t)));
40            })
41        })
42        .collect();
43    for h in handles {
44        h.join().unwrap();
45    }
46}
47
48fn bench(g: &ProcessGroup) {
49    let rank = g.rank();
50    g.barrier().unwrap();
51
52    // ---- point-to-point round-trip latency (ping-pong) ----
53    let iters = 5000;
54    if rank == 0 {
55        let p = [1.0f32];
56        for _ in 0..200 {
57            g.send_f32(1, 0, &p).unwrap();
58            g.recv_f32(1, 0).unwrap();
59        }
60        let t0 = Instant::now();
61        for _ in 0..iters {
62            g.send_f32(1, 0, &p).unwrap();
63            g.recv_f32(1, 0).unwrap();
64        }
65        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
66        println!(
67            "[latency]   ping-pong: {us:6.2} µs RTT   ({:.2} µs one-way)   [{iters} iters]",
68            us / 2.0
69        );
70    } else {
71        for _ in 0..(200 + iters) {
72            let x = g.recv_f32(0, 0).unwrap();
73            g.send_f32(0, 0, &x).unwrap();
74        }
75    }
76    g.barrier().unwrap();
77
78    // ---- streaming bandwidth across payload sizes ----
79    for &elems in &[256usize, 16_384, 262_144, 4_194_304] {
80        let bytes = elems * 4;
81        let reps = (512 * 1024 * 1024 / bytes).clamp(8, 200_000);
82        if rank == 0 {
83            let buf = vec![0.5f32; elems];
84            g.send_f32(1, 1, &[0.0]).unwrap(); // sync start
85            let t0 = Instant::now();
86            for _ in 0..reps {
87                g.send_f32(1, 2, &buf).unwrap();
88            }
89            g.recv_f32(1, 3).unwrap(); // ack: all received
90            let secs = t0.elapsed().as_secs_f64();
91            let gbps = (reps * bytes) as f64 / secs / 1e9;
92            println!(
93                "[bandwidth] {:>9} B × {:>6} = {:>5.1} GB  ->  {gbps:6.2} GB/s",
94                bytes,
95                reps,
96                (reps * bytes) as f64 / 1e9
97            );
98        } else {
99            g.recv_f32(0, 1).unwrap();
100            for _ in 0..reps {
101                g.recv_f32(0, 2).unwrap();
102            }
103            g.send_f32(0, 3, &[1.0]).unwrap();
104        }
105        g.barrier().unwrap();
106    }
107
108    // ---- collective latency for a hidden-state-sized vector ----
109    // e.g. batch1 × seq1 × d_model=4096 — the kind of all_reduce a
110    // tensor-parallel layer issues, or a broadcast the pipeline issues.
111    for &elems in &[4096usize, 4096 * 8] {
112        let mut data = vec![1.0f32; elems];
113        let iters = 2000;
114        for _ in 0..50 {
115            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
116        }
117        g.barrier().unwrap();
118        let t0 = Instant::now();
119        for _ in 0..iters {
120            g.all_reduce(&mut data, ReduceKind::Sum).unwrap();
121        }
122        let us = t0.elapsed().as_secs_f64() * 1e6 / iters as f64;
123        if rank == 0 {
124            println!("[collective] all_reduce {:>6} f32: {us:6.2} µs/op", elems);
125        }
126        g.barrier().unwrap();
127    }
128
129    g.barrier().unwrap();
130}
transport_bench/transport_bench.rs

transport_bench/
transport_bench.rs
