metrale_comm/nccl_backend/
recv_buffer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Sizing of the 2-rank all-reduce receive buffer and the check
4//! that a payload fits it. At `world_size == 2` the all-reduce receives the
5//! partner's data into this buffer, so a payload larger than it would write
6//! past the allocation.
7//!
8//! Owner: metrale-comm.
9//! Invariants:
10//! - The size functions never wrap: an overflow is an error.
11
12use anyhow::{Context, Result};
13
14/// 2026-09-26: Bytes per element of the all-reduce: BF16. `NcclBackend`
15/// divides byte counts by it to get NCCL element counts, and
16/// [`required_model_recv_bytes`] sizes the buffer with it.
17pub const ALL_REDUCE_DTYPE_BYTES: usize = 2;
18
19/// 2026-09-26: `max_batch_tokens × hidden_size × dtype_bytes`, in bytes.
20///
21/// # Errors
22/// When the product overflows `usize`.
23pub fn required_recv_bytes(
24    max_batch_tokens: usize,
25    hidden_size: usize,
26    dtype_bytes: usize,
27) -> Result<usize> {
28    max_batch_tokens
29        .checked_mul(hidden_size)
30        .and_then(|elems| elems.checked_mul(dtype_bytes))
31        .with_context(|| {
32            format!(
33                "receive-buffer size overflows usize: \
34                 max_batch_tokens={max_batch_tokens} × hidden_size={hidden_size} \
35                 × dtype_bytes={dtype_bytes}"
36            )
37        })
38}
39
40/// 2026-09-26: Receive-buffer bytes for `tokens` rows of BF16 activations or
41/// logits: `tokens × max(hidden, vocab) × 2`. Serve passes `max_batch_tokens`,
42/// `hidden_size` and `vocab_size`.
43///
44/// # Errors
45/// When the product overflows `usize`.
46pub fn required_model_recv_bytes(tokens: usize, hidden: usize, vocab: usize) -> Result<usize> {
47    required_recv_bytes(tokens, hidden.max(vocab), ALL_REDUCE_DTYPE_BYTES)
48}
49
50/// 2026-09-26: Refuse a payload of more than `capacity` bytes. A free
51/// function, so the tests below run it without a communicator or a GPU.
52pub(crate) fn ensure_payload_fits(
53    bytes: usize,
54    capacity: usize,
55    rank: usize,
56    world_size: usize,
57) -> Result<()> {
58    if bytes > capacity {
59        anyhow::bail!(
60            "2-rank all-reduce payload exceeds receive-buffer capacity: \
61             requested {bytes} bytes ({} elements × {ALL_REDUCE_DTYPE_BYTES} B/elem), \
62             capacity {capacity} bytes, rank {rank}, world_size {world_size}. \
63             The receive buffer is sized from the configured maximum transfer \
64             (max_batch_tokens × hidden_size × {ALL_REDUCE_DTYPE_BYTES} B); a larger \
65             payload would write past the allocation. Refusing to send.",
66            bytes / ALL_REDUCE_DTYPE_BYTES,
67        );
68    }
69    Ok(())
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    const BF16: usize = 2;
77    const FP32: usize = 4;
78
79    #[test]
80    fn short_prefill_still_covers_vocab_parallel_logits() {
81        let capacity = required_model_recv_bytes(33, 1024, 163840).unwrap();
82        assert!(ensure_payload_fits(163584 * 2, capacity, 0, 2).is_ok());
83        assert!(ensure_payload_fits(33 * 163840 * 2, capacity, 0, 2).is_ok());
84        assert!(required_model_recv_bytes(usize::MAX, 1024, 163840).is_err());
85        assert_eq!(required_model_recv_bytes(2, 4096, 1024).unwrap(), 16384);
86    }
87
88    /// 2026-09-26: 64 MiB, the boundary the tests below compare against.
89    const OLD_FIXED_BUFFER: usize = 64 * 1024 * 1024;
90
91    /// 2026-09-26: The all-reduce element is two bytes (BF16).
92    #[test]
93    fn all_reduce_dtype_width_is_bf16() {
94        assert_eq!(ALL_REDUCE_DTYPE_BYTES, BF16);
95    }
96
97    /// 2026-09-26: 8192 × 4096 × 2 is exactly 64 MiB, and a payload equal to
98    /// the capacity is accepted.
99    #[test]
100    fn exact_boundary_bf16() {
101        let need = required_recv_bytes(8192, 4096, BF16).unwrap();
102        assert_eq!(need, 67_108_864);
103        assert_eq!(
104            need, OLD_FIXED_BUFFER,
105            "DS4F default sat exactly on the cap"
106        );
107        assert!(ensure_payload_fits(need, need, 0, 2).is_ok());
108    }
109
110    /// 2026-09-26: One token more than 64 MiB fits a capacity sized for it and
111    /// is refused against 64 MiB.
112    #[test]
113    fn over_boundary_bf16() {
114        let need = required_recv_bytes(8193, 4096, BF16).unwrap();
115        assert_eq!(need, 67_117_056);
116        assert!(need > OLD_FIXED_BUFFER);
117        assert!(ensure_payload_fits(need, need, 0, 2).is_ok());
118        assert!(ensure_payload_fits(need, OLD_FIXED_BUFFER, 0, 2).is_err());
119    }
120
121    /// 2026-09-26: `hidden_size = 6144` at 8192 tokens needs 32 MiB more than
122    /// 64 MiB, and is refused against 64 MiB.
123    #[test]
124    fn wide_model_bf16() {
125        let need = required_recv_bytes(8192, 6144, BF16).unwrap();
126        assert_eq!(need, 100_663_296);
127        assert_eq!(
128            need - OLD_FIXED_BUFFER,
129            33_554_432,
130            "the overrun this fixes"
131        );
132        assert!(ensure_payload_fits(need, need, 0, 2).is_ok());
133        assert!(ensure_payload_fits(need, OLD_FIXED_BUFFER, 0, 2).is_err());
134    }
135
136    /// 2026-09-26: A 4-byte element doubles the size.
137    #[test]
138    fn fp32_is_representable() {
139        let need = required_recv_bytes(8192, 4096, FP32).unwrap();
140        assert_eq!(need, 134_217_728);
141        assert_eq!(need, 2 * required_recv_bytes(8192, 4096, BF16).unwrap());
142        assert_eq!(
143            need,
144            2 * OLD_FIXED_BUFFER,
145            "FP32 would have been a 2× overrun"
146        );
147        assert!(ensure_payload_fits(need, need, 0, 2).is_ok());
148    }
149
150    /// 2026-09-26: One row (`hidden × 2`) fits a capacity sized for 8192 rows.
151    #[test]
152    fn decode_b1_bf16() {
153        let cap = required_recv_bytes(8192, 4096, BF16).unwrap();
154        let decode = 4096 * BF16;
155        assert_eq!(decode, 8_192);
156        assert!(ensure_payload_fits(decode, cap, 0, 2).is_ok());
157    }
158
159    /// 2026-09-26: Zero tokens size to zero, and a zero payload fits any
160    /// capacity, including zero.
161    #[test]
162    fn zero_elements() {
163        let cap = required_recv_bytes(8192, 4096, BF16).unwrap();
164        assert_eq!(required_recv_bytes(0, 4096, BF16).unwrap(), 0);
165        assert!(ensure_payload_fits(0, cap, 0, 2).is_ok());
166        assert!(ensure_payload_fits(0, 0, 0, 2).is_ok());
167    }
168
169    /// 2026-09-26: An overflowing product is an error, not a wrapped value.
170    #[test]
171    fn overflow_is_rejected() {
172        assert!(required_recv_bytes(usize::MAX, 4096, BF16).is_err());
173        assert!(required_recv_bytes(usize::MAX, 1, 2).is_err());
174        assert!(required_recv_bytes(usize::MAX / 2 + 1, 2, 1).is_err());
175        // 2026-09-26: A wrapping multiply would give 0 here.
176        assert!(required_recv_bytes(1 << 62, 4, 1).is_err());
177    }
178
179    /// 2026-09-26: One byte short is refused, and the message names the
180    /// payload, the capacity and the world size.
181    #[test]
182    fn one_byte_short_is_rejected() {
183        let need = required_recv_bytes(8192, 4096, BF16).unwrap();
184        let err = ensure_payload_fits(need, need - 1, 1, 2)
185            .unwrap_err()
186            .to_string();
187        assert!(
188            err.contains(&need.to_string()),
189            "must report requested bytes: {err}"
190        );
191        assert!(
192            err.contains(&(need - 1).to_string()),
193            "must report capacity: {err}"
194        );
195        assert!(err.contains("world_size 2"), "must report the path: {err}");
196    }
197}
metrale_comm/nccl_backend/recv_buffer.rs

metrale_comm/nccl_backend/
recv_buffer.rs
