Function shard_dense_bf16 

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)>

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