pub fn shard_dense_bf16(
src: DevicePtr,
out_dim: usize,
in_dim: usize,
kind: TpShardKind,
tp_rank: usize,
tp_size: usize,
gpu: &dyn GpuBackend,
) -> Result<(DevicePtr, usize, usize)>Expand description
2026-09-25: Shard a BF16 dense weight [out_dim, in_dim] according to kind.
Returns (sharded_ptr, sharded_out, sharded_in). When tp_size <= 1
or kind == Replicated, returns the source pointer untouched; the caller
must not free the source separately (no shard happened).
Otherwise allocates a new device buffer holding the local rank’s slice,
copies into it, and returns the new pointer. The caller owns the source
and must gpu.free it after the shard is built. Errors when
tp_rank >= tp_size or the split dim is not divisible by tp_size.