shard_gdn_ba_rows

Function shard_gdn_ba_rows 

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

Shard the BA gate BF16 weight [2*full_nv, h] to [2*local_nv, h].

The interleave is per key-head group ([β₀..β_{vpg-1}, α₀..α_{vpg-1}] per group, vpg = nv/nk), but rank r owns key-head groups [r*local_nk, (r+1)*local_nk) which map to the contiguous row range [r*2*local_nv, (r+1)*2*local_nv) — the rank boundary always lands on a group boundary, so a single contiguous slice preserves the interleave.