Struct DsaTpPlan 

pub struct DsaTpPlan {
    pub tp_rank: usize,
    pub tp_size: usize,
    pub full_heads: usize,
    pub local_heads: usize,
    pub tensors: Vec<DsaTensorPlan>,
}

Per-rank shard plan for one DSA block.

Fields

tp_rank: usize

tp_size: usize

full_heads: usize

local_heads: usize

tensors: Vec<DsaTensorPlan>

Implementations

impl DsaTpPlan

pub fn new( tp_rank: usize, tp_size: usize, cfg: &Glm5NextDsaConfig, ) -> Result<Self>

cfg.local_heads is already per-rank (topology divides before loaders run), so the full count is reconstructed as local * tp_size — the same convention TpGdnDims::from_config uses.

pub fn get(&self, name: &str) -> Option<&DsaTensorPlan>

pub fn local_bytes(&self) -> usize

pub fn full_bytes(&self) -> usize

pub fn needs_output_all_reduce(&self) -> bool

o_proj is row-parallel, so its partial output needs the reduce.

pub fn replicated_bytes(&self) -> usize

Bytes replicated on every rank — the part TP cannot remove.

Trait Implementations

impl Clone for DsaTpPlan

fn clone(&self) -> DsaTpPlan

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for DsaTpPlan

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.