Struct TensorParallelConfig 

pub struct TensorParallelConfig {
    pub tp_rank: usize,
    pub tp_size: usize,
    pub hidden_size: usize,
    pub intermediate_size: usize,
    pub num_heads: usize,
    pub num_kv_heads: usize,
    pub head_dim: usize,
}

Tensor parallel configuration for a single GPU.

Fields

tp_rank: usize

This GPU’s rank within the TP group

tp_size: usize

Total GPUs in the TP group

hidden_size: usize

Full hidden size (before sharding)

intermediate_size: usize

Full intermediate size (before sharding)

num_heads: usize

Number of attention heads (must be divisible by tp_size)

num_kv_heads: usize

Number of KV heads (must be divisible by tp_size)

head_dim: usize

Head dimension

Implementations

impl TensorParallelConfig

pub fn new( tp_rank: usize, tp_size: usize, hidden_size: usize, intermediate_size: usize, num_heads: usize, num_kv_heads: usize, ) -> Self

Create a new TP config.

Panics

Panics if heads or intermediate size not divisible by tp_size.

pub fn local_num_heads(&self) -> usize

Number of Q heads on this GPU.

pub fn local_num_kv_heads(&self) -> usize

Number of KV heads on this GPU.

pub fn local_q_size(&self) -> usize

Local Q projection size: local_heads × head_dim.

pub fn local_kv_size(&self) -> usize

Local KV projection size: local_kv_heads × head_dim.

pub fn local_intermediate_size(&self) -> usize

Local intermediate (FFN) size.

pub fn weight_memory_fraction(&self) -> f64

Memory savings from tensor parallelism (ratio).

TP shards attention + FFN weights. Embedding and norms are replicated.

Trait Implementations

impl Clone for TensorParallelConfig

fn clone(&self) -> TensorParallelConfig

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for TensorParallelConfig

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

Formats the value using the given formatter.