Function required_recv_bytes 

pub fn required_recv_bytes(
    max_batch_tokens: usize,
    hidden_size: usize,
    dtype_bytes: usize,
) -> Result<usize>

Bytes required for the 2-rank all-reduce receive buffer.

The largest payload any caller can hand a collective is one full arena buffer — max_batch_tokens × hidden_size × dtype — which is exactly how moe_output is sized (spark-runtime/src/buffers/sizes.rs). The tensor-parallel attention and SSM reduces produce the same [num_tokens, hidden_size] BF16 shape, and num_tokens is capped by max_batch_tokens, so this bound covers every caller of all_reduce / all_reduce_async.

Arithmetic is checked: a configuration whose buffer does not fit in a usize is rejected at startup rather than wrapping into a small allocation, which is how a sizing bug becomes an out-of-bounds write.