metrale_gpu_sys/
wire.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: RDMA handshake wire codecs, used by the RDMA clients and
4//! re-exported by the storage peer daemons, so both ends share one codec.
5//!
6//! Owner: metrale-gpu-sys.
7//! Invariants:
8//! - Every integer is little-endian.
9//! - The byte layouts are pinned by `tests/wire_roundtrip.rs` and
10//!   `tests/transcript_golden.rs`, the constant values by `frozen_wire_constants`.
11//! - Only `std::io` and `anyhow`, with no `cfg` gate, so it builds without rdma-core.
12
13use anyhow::{Context, Result, bail};
14
15pub const STATUS_OK: u8 = 0;
16pub const STATUS_ERR: u8 = 1;
17/// 2026-09-26: Transport byte for two-sided TCP record streaming from the expert peer.
18pub const MODE_TCP: u8 = 0;
19/// 2026-09-26: Transport byte for one-sided RDMA READ: the server sends its
20/// MRs as `VerbsServerParams` and the client reads records from them.
21pub const MODE_VERBS: u8 = 1;
22
23/// 2026-09-26: The server's half of the read-only verbs handshake (expert,
24/// weight and LoRA tiers): its QP identity plus one registered MR per MoE layer
25/// (expert tier) or per shard (weight and LoRA tiers).
26/// The expert tier reads at `layers[layer].0 + expert * record_stride`.
27#[derive(Clone, Debug, PartialEq, Eq)]
28pub struct VerbsServerParams {
29    pub qpn: u32,
30    pub psn: u32,
31    pub gid: [u8; 16],
32    /// 2026-09-26: `(mr_base_addr, rkey)` per MoE layer or per shard, by index.
33    pub layers: Vec<(u64, u32)>,
34}
35
36/// 2026-09-26: The client's half of the verbs handshake: its QP identity only.
37#[derive(Clone, Copy, Debug, PartialEq, Eq)]
38pub struct VerbsClientParams {
39    pub qpn: u32,
40    pub psn: u32,
41    pub gid: [u8; 16],
42}
43
44/// 2026-09-26: The peer's half of the read-write handshake (KV paging and SSM
45/// snapshots): its QP identity plus the base address and rkey of its RW MR.
46#[derive(Clone, Copy, Debug, PartialEq, Eq)]
47pub struct CacheServerParams {
48    pub qpn: u32,
49    pub psn: u32,
50    pub gid: [u8; 16],
51    pub base_addr: u64,
52    pub rkey: u32,
53}
54
55/// 2026-09-26: A remote QP identity a client rail can `connect` to. Both
56/// server-param types implement it, so `RailSet::complete` serves both dialects.
57pub trait RemoteQp {
58    fn qp_identity(&self) -> (u32, u32, [u8; 16]);
59}
60
61impl RemoteQp for VerbsServerParams {
62    fn qp_identity(&self) -> (u32, u32, [u8; 16]) {
63        (self.qpn, self.psn, self.gid)
64    }
65}
66
67impl RemoteQp for CacheServerParams {
68    fn qp_identity(&self) -> (u32, u32, [u8; 16]) {
69        (self.qpn, self.psn, self.gid)
70    }
71}
72
73impl VerbsServerParams {
74    /// 2026-09-26: Wire form: `[u32 qpn][u32 psn][16 gid][u32 n_layers]{[u64 base][u32 rkey]}*`.
75    pub fn write_to<W: std::io::Write>(&self, w: &mut W) -> Result<()> {
76        w.write_all(&self.qpn.to_le_bytes())?;
77        w.write_all(&self.psn.to_le_bytes())?;
78        w.write_all(&self.gid)?;
79        w.write_all(&(self.layers.len() as u32).to_le_bytes())?;
80        for (base, rkey) in &self.layers {
81            w.write_all(&base.to_le_bytes())?;
82            w.write_all(&rkey.to_le_bytes())?;
83        }
84        Ok(())
85    }
86
87    pub fn read_from<R: std::io::Read>(r: &mut R) -> Result<Self> {
88        let qpn = read_u32(r)?;
89        let psn = read_u32(r)?;
90        let mut gid = [0u8; 16];
91        r.read_exact(&mut gid).context("read server gid")?;
92        let n = read_u32(r)? as usize;
93        if n == 0 || n > 4096 {
94            bail!("implausible verbs layer count: {n}");
95        }
96        let mut layers = Vec::with_capacity(n);
97        for _ in 0..n {
98            let mut b8 = [0u8; 8];
99            r.read_exact(&mut b8).context("read mr base")?;
100            let base = u64::from_le_bytes(b8);
101            let rkey = read_u32(r)?;
102            layers.push((base, rkey));
103        }
104        Ok(Self {
105            qpn,
106            psn,
107            gid,
108            layers,
109        })
110    }
111}
112
113impl VerbsClientParams {
114    /// 2026-09-26: Wire form: `[u32 qpn][u32 psn][16 gid]`.
115    pub fn write_to<W: std::io::Write>(&self, w: &mut W) -> Result<()> {
116        w.write_all(&self.qpn.to_le_bytes())?;
117        w.write_all(&self.psn.to_le_bytes())?;
118        w.write_all(&self.gid)?;
119        Ok(())
120    }
121
122    pub fn read_from<R: std::io::Read>(r: &mut R) -> Result<Self> {
123        let qpn = read_u32(r)?;
124        let psn = read_u32(r)?;
125        let mut gid = [0u8; 16];
126        r.read_exact(&mut gid).context("read client gid")?;
127        Ok(Self { qpn, psn, gid })
128    }
129}
130
131impl CacheServerParams {
132    /// 2026-09-26: Wire form: `[u32 qpn][u32 psn][16 gid][u64 base_addr][u32 rkey]`.
133    pub fn write_to<W: std::io::Write>(&self, w: &mut W) -> Result<()> {
134        w.write_all(&self.qpn.to_le_bytes())?;
135        w.write_all(&self.psn.to_le_bytes())?;
136        w.write_all(&self.gid)?;
137        w.write_all(&self.base_addr.to_le_bytes())?;
138        w.write_all(&self.rkey.to_le_bytes())?;
139        Ok(())
140    }
141
142    pub fn read_from<R: std::io::Read>(r: &mut R) -> Result<Self> {
143        let mut b4 = [0u8; 4];
144        let mut b8 = [0u8; 8];
145        let mut gid = [0u8; 16];
146        r.read_exact(&mut b4).context("kv qpn")?;
147        let qpn = u32::from_le_bytes(b4);
148        r.read_exact(&mut b4).context("kv psn")?;
149        let psn = u32::from_le_bytes(b4);
150        r.read_exact(&mut gid).context("kv gid")?;
151        r.read_exact(&mut b8).context("kv base")?;
152        let base_addr = u64::from_le_bytes(b8);
153        r.read_exact(&mut b4).context("kv rkey")?;
154        let rkey = u32::from_le_bytes(b4);
155        Ok(Self {
156            qpn,
157            psn,
158            gid,
159            base_addr,
160            rkey,
161        })
162    }
163}
164
165fn read_u32<R: std::io::Read>(r: &mut R) -> Result<u32> {
166    let mut b = [0u8; 4];
167    r.read_exact(&mut b).context("read u32")?;
168    Ok(u32::from_le_bytes(b))
169}
170
171/// 2026-09-26: Write `[u8 n_rails]` followed by each rail's `VerbsServerParams`.
172/// Fails before writing anything when the count is 0 or above 8.
173pub fn write_server_rails<W: std::io::Write>(w: &mut W, rails: &[VerbsServerParams]) -> Result<()> {
174    if rails.is_empty() || rails.len() > 8 {
175        bail!("implausible server rail count: {}", rails.len());
176    }
177    w.write_all(&[rails.len() as u8])?;
178    for sp in rails {
179        sp.write_to(w)?;
180    }
181    Ok(())
182}
183
184/// 2026-09-26: Read `want` rails of `VerbsServerParams` framed by a leading
185/// `[u8 n_rails]`. Fails when the framed count is 0, above 8, or not `want`.
186pub fn read_server_rails<R: std::io::Read>(
187    r: &mut R,
188    want: usize,
189) -> Result<Vec<VerbsServerParams>> {
190    let mut b1 = [0u8; 1];
191    r.read_exact(&mut b1).context("read server rail count")?;
192    let n = b1[0] as usize;
193    if n == 0 || n > 8 {
194        bail!("implausible server rail count: {n}");
195    }
196    if n != want {
197        bail!("server framed {n} rails but client negotiated {want}");
198    }
199    let mut rails = Vec::with_capacity(n);
200    for _ in 0..n {
201        rails.push(VerbsServerParams::read_from(r)?);
202    }
203    Ok(rails)
204}
metrale_gpu_sys/wire.rs

metrale_gpu_sys/
wire.rs
