shard_gdn_value_vector

spark_model::tp_shard

Function shard_gdn_value_vector 

Source 
pub fn shard_gdn_value_vector(
    src: DevicePtr,
    dims: &TpGdnDims,
    unit: usize,
    elem_bytes: usize,
    gpu: &dyn GpuBackend,
) -> Result<(DevicePtr, usize)>
Expand description

Shard a per-value-head 1D vector on the value-head axis. Handles BF16 (norm, [full_nv*vd] → [local_nv*vd] with elem_bytes = 2, unit = vd) and FP32 (a_log / dt_bias, [full_nv] → [local_nv] with elem_bytes = 4, unit = 1). unit is the number of elements per value head. Returns (ptr, local_len_elems).