Struct KdaTpPlan 

pub struct KdaTpPlan {
    pub tp_rank: usize,
    pub tp_size: usize,
    pub hidden: usize,
    pub head_dim: usize,
    pub full_heads: usize,
    pub local_heads: usize,
    pub conv_kernel: usize,
    pub gate_rank: usize,
    pub tensors: Vec<KdaTensorPlan>,
}

The complete per-rank shard plan for one KDA block.

Fields

tp_rank: usize

tp_size: usize

hidden: usize

head_dim: usize

full_heads: usize

Pre-shard head count (all ranks combined).

local_heads: usize

Heads this rank owns.

conv_kernel: usize

gate_rank: usize

Low-rank width of the f/g gate bottleneck.

tensors: Vec<KdaTensorPlan>

Implementations

impl KdaTpPlan

pub fn from_config(config: &ModelConfig, gate_rank: usize) -> Result<Self>

Build from a ModelConfig whose linear-head counts are already per-rank local — serve_phases::topology divides them before any loader runs, and TpGdnDims::from_config reconstructs full = local * tp_size the same way.

gate_rank is not a config key; it is the f_a_proj row count read from the checkpoint (128 for GLM-5.3).

pub fn new( tp_rank: usize, tp_size: usize, hidden: usize, head_dim: usize, full_heads: usize, conv_kernel: usize, gate_rank: usize, ) -> Result<Self>

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

pub fn local_bytes(&self) -> usize

Total bytes this rank stores for one KDA block.

pub fn full_bytes(&self) -> usize

Total bytes one KDA block occupies on disk.

pub fn needs_output_all_reduce(&self) -> bool

Whether the layer must all-reduce after o_proj. False at tp_size == 1, where the row-parallel slice is the whole tensor and the reduce is a no-op.

Trait Implementations

impl Clone for KdaTpPlan

fn clone(&self) -> KdaTpPlan

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for KdaTpPlan

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

Formats the value using the given formatter.