shard_dense_bf16

Function shard_dense_bf16 

Source
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.