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,
}Expand description
2026-09-25: Pre-TP-shard attention dimensions reconstructed from config.
The server’s serve_phases/topology.rs divides num_attention_heads and
num_key_value_heads by the TP size before any loader runs, so config
holds per-rank head counts. The full_* fields multiply back up to the
pre-shard sizes the shard functions take.
When config.attn_gated is true the Q projection’s output includes the
gate, doubling it: full_q_n includes the gate and full_o_in does not.
Fields§
§tp_rank: usize§tp_size: usize2026-09-25: tp_world_size clamped to >= 1.
h: usize2026-09-25: Hidden size, not divided per rank.
head_dim: usize§full_q_n: usize2026-09-25: Q-projection output dim, doubled for gated attention.
full_o_in: usize2026-09-25: O-projection input dim, the un-gated attention output:
num_attention_heads * tp_size * head_dim, not doubled.
full_kv_n: usize2026-09-25: num_key_value_heads_local * tp_size * head_dim: full K/V pre-shard.
gated: bool2026-09-25: config.attn_gated.
Implementations§
Source§impl TpAttentionDims
impl TpAttentionDims
pub fn from_config(config: &ModelConfig) -> Self
Sourcepub fn proj_shape(&self, name: &str) -> Option<(usize, usize, TpShardKind)>
pub fn proj_shape(&self, name: &str) -> Option<(usize, usize, TpShardKind)>
2026-09-25: (out_dim, in_dim, kind) for a given QKVO projection; None for
any other name.
Trait Implementations§
Source§impl Clone for TpAttentionDims
impl Clone for TpAttentionDims
Source§fn clone(&self) -> TpAttentionDims
fn clone(&self) -> TpAttentionDims
1.0.0 · Source§fn clone_from(&mut self, source: &Self)
fn clone_from(&mut self, source: &Self)
source. Read more