Struct TpAttentionDims 

pub struct TpAttentionDims {
    pub tp_rank: usize,
    pub tp_size: usize,
    pub h: usize,
    pub head_dim: usize,
    pub full_q_n: usize,
    pub full_o_in: usize,
    pub full_kv_n: usize,
    pub gated: bool,
}

Pre-TP-shard attention dimensions reconstructed from config.

main.rs divides num_attention_heads and num_key_value_heads by tp_world_size at startup, so by the time a loader runs, config holds per-rank-local head counts. The full_* fields multiply back up to the pre-shard sizes that slice_for_rank and friends expect.

When config.attn_gated is true (Qwen3-Next), the Q projection output dim is doubled — the second half is the per-token gate applied after attention. full_q_n includes the gate; full_o_in does NOT (O proj’s input dim matches the un-gated attention output, since the gate is applied before O proj).

Fields

tp_rank: usize

tp_size: usize

tp_world_size clamped to >= 1. Loaders should treat tp_size == 1 as the no-shard fast path.

h: usize

Hidden size (model embed dim) — never sharded.

head_dim: usize

full_q_n: usize

Q-projection output dim. For gated attention this is doubled (the second half is the gate).

full_o_in: usize

O-projection input dim. Equals the un-gated attention output — num_attention_heads * tp_size * head_dim, NOT doubled.

full_kv_n: usize

num_key_value_heads_local * tp_size * head_dim — full K/V pre-shard.

gated: bool

Whether the loader is operating on a gated-attention config.

Implementations

impl TpAttentionDims

pub fn from_config(config: &ModelConfig) -> Self

pub fn proj_shape(&self, name: &str) -> Option<(usize, usize, TpShardKind)>

(out_dim, in_dim, kind) for a given QKVO projection.

Trait Implementations

impl Clone for TpAttentionDims

fn clone(&self) -> TpAttentionDims

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for TpAttentionDims

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

Formats the value using the given formatter.

impl Copy for TpAttentionDims