spark_model/kimi_k3/
tp.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! K3 Megatron TP plan. `supports_tp` is true — the umbrella loader used to
4//! refuse `--tp-size 2`. Full `slice_for_rank` bind lives with the weight
5//! loader; this slice publishes the plan + the fail-fast flag.
6//!
7//! Head counts on `config` are already per-rank. Full sizes = local * tp.
8
9use atlas_core::config::ModelConfig;
10use atlas_core::kimi_k3::{MixerKind, MlpKind};
11
12use crate::tp_shard::TpShardKind;
13
14/// K3 does tensor-parallel. Callers must not fail-fast `supports_tp` false.
15pub fn supports_tp() -> bool {
16    true
17}
18
19/// `(kind, full_out, full_in)` for one checkpoint key.
20pub fn tensor_plan(
21    name: &str,
22    mixer: MixerKind,
23    _mlp: MlpKind,
24    config: &ModelConfig,
25) -> (TpShardKind, usize, usize) {
26    let tp = config.tp_world_size.max(1);
27    let h = config.hidden_size;
28    let kda_heads = config.linear_num_key_heads * tp;
29    let kda_d = config.linear_key_head_dim;
30    let kda_q = kda_heads * kda_d;
31    let conv_k = config.linear_conv_kernel_dim.max(1);
32    let mla_heads = config.num_attention_heads * tp;
33    let qk = config.qk_nope_head_dim + config.qk_rope_head_dim;
34    let dv = mla_heads * config.v_head_dim;
35    let kv_b = mla_heads * (config.qk_nope_head_dim + config.v_head_dim);
36    let inter = config.intermediate_size;
37    let eh = config.moe_intermediate_size;
38    let lat = config.moe_latent_size;
39
40    if name.ends_with(".self_attn.q_proj.weight")
41        || name.ends_with(".self_attn.k_proj.weight")
42        || name.ends_with(".self_attn.v_proj.weight")
43    {
44        return (TpShardKind::ColumnParallel, kda_q, h);
45    }
46    if name.ends_with(".self_attn.q_conv1d.weight")
47        || name.ends_with(".self_attn.k_conv1d.weight")
48        || name.ends_with(".self_attn.v_conv1d.weight")
49    {
50        return (TpShardKind::ColumnParallel, kda_q, conv_k);
51    }
52    if name.ends_with(".self_attn.g_proj.weight") {
53        let n = match mixer {
54            MixerKind::Kda => kda_q,
55            MixerKind::Mla => dv,
56        };
57        return (TpShardKind::ColumnParallel, n, h);
58    }
59    if name.ends_with(".self_attn.o_proj.weight") {
60        let inn = match mixer {
61            MixerKind::Kda => kda_q,
62            MixerKind::Mla => dv,
63        };
64        return (TpShardKind::RowParallel, h, inn);
65    }
66    if name.ends_with(".self_attn.b_proj.weight") {
67        return (TpShardKind::ColumnParallel, kda_heads, h);
68    }
69    if name.ends_with(".self_attn.A_log") {
70        return (TpShardKind::ColumnParallel, kda_heads, 1);
71    }
72    if name.ends_with(".self_attn.dt_bias") {
73        return (TpShardKind::ColumnParallel, kda_q, 1);
74    }
75    if name.ends_with(".self_attn.f_b_proj.weight") {
76        return (TpShardKind::ColumnParallel, kda_q, kda_d);
77    }
78    if name.ends_with(".self_attn.q_b_proj.weight") {
79        return (
80            TpShardKind::ColumnParallel,
81            mla_heads * qk,
82            config.q_lora_rank,
83        );
84    }
85    if name.ends_with(".self_attn.kv_b_proj.weight") {
86        return (TpShardKind::ColumnParallel, kv_b, config.kv_lora_rank);
87    }
88    if name.ends_with(".mlp.gate_proj.weight") || name.ends_with(".mlp.up_proj.weight") {
89        return (TpShardKind::ColumnParallel, inter, h);
90    }
91    if name.ends_with(".mlp.down_proj.weight") {
92        return (TpShardKind::RowParallel, h, inter);
93    }
94    if name.contains(".block_sparse_moe.experts.") {
95        if name.ends_with(".w1.weight") || name.ends_with(".w3.weight") {
96            return (TpShardKind::ColumnParallel, eh, lat);
97        }
98        if name.ends_with(".w2.weight") {
99            return (TpShardKind::RowParallel, lat, eh);
100        }
101    }
102    (TpShardKind::Replicated, 1, 1)
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108    use atlas_core::config::parse_config;
109
110    const TWIN: &str = include_str!("../../../../docs/k3/fixtures/Kimi-K3-0.40B-config.json");
111
112    fn twin() -> ModelConfig {
113        parse_config(TWIN).expect("0.40B twin")
114    }
115
116    fn divide_heads_for_tp(config: &mut ModelConfig, rank: usize, size: usize) {
117        config.tp_rank = rank;
118        config.tp_world_size = size;
119        if size > 1 {
120            config.num_attention_heads /= size;
121            config.num_key_value_heads /= size;
122            config.linear_num_key_heads /= size;
123            config.linear_num_value_heads /= size;
124        }
125    }
126
127    #[test]
128    fn kimi_k3_supports_tp() {
129        assert!(supports_tp(), "loader must not refuse --tp-size 2");
130    }
131
132    #[test]
133    fn kda_q_is_column_o_is_row() {
134        let mut c = twin();
135        divide_heads_for_tp(&mut c, 0, 2);
136        let q = tensor_plan(
137            "language_model.model.layers.0.self_attn.q_proj.weight",
138            MixerKind::Kda,
139            MlpKind::Dense,
140            &c,
141        );
142        let o = tensor_plan(
143            "language_model.model.layers.0.self_attn.o_proj.weight",
144            MixerKind::Kda,
145            MlpKind::Dense,
146            &c,
147        );
148        assert_eq!(q, (TpShardKind::ColumnParallel, 8 * 32, 1024));
149        assert_eq!(o, (TpShardKind::RowParallel, 1024, 8 * 32));
150        let embed = tensor_plan(
151            "language_model.model.embed_tokens.weight",
152            MixerKind::Kda,
153            MlpKind::Dense,
154            &c,
155        );
156        assert_eq!(embed.0, TpShardKind::Replicated);
157    }
158
159    #[test]
160    fn rank0_and_rank1_q_proj_full_out_match() {
161        let mut c0 = twin();
162        divide_heads_for_tp(&mut c0, 0, 2);
163        let mut c1 = twin();
164        divide_heads_for_tp(&mut c1, 1, 2);
165        let q0 = tensor_plan(
166            "language_model.model.layers.0.self_attn.q_proj.weight",
167            MixerKind::Kda,
168            MlpKind::Dense,
169            &c0,
170        );
171        let q1 = tensor_plan(
172            "language_model.model.layers.0.self_attn.q_proj.weight",
173            MixerKind::Kda,
174            MlpKind::Dense,
175            &c1,
176        );
177        assert_eq!(q0, q1, "full sizes reconstruct from local*tp on both ranks");
178        assert_eq!(c0.tp_rank, 0);
179        assert_eq!(c1.tp_rank, 1);
180        assert_ne!(c0.tp_rank, c1.tp_rank);
181    }
182
183    fn round_bf16(x: f32) -> f32 {
184        let bits = atlas_core::numeric::f32_to_bf16(x);
185        atlas_core::numeric::bf16_bytes_to_f32(bits.to_le_bytes())
186    }
187
188    #[test]
189    fn production_7168_bf16_allreduce_rounding_is_measured() {
190        // tp_allreduce stores f32 partials as BF16, NCCL-sums, widens back.
191        // Twin 16/16 held. This is the production-width error, not an assumption.
192        let n = 7168;
193        let a = vec![1.0f32 / 3.0; n];
194        let b = vec![2.0f32 / 3.0; n];
195        let mut max = 0.0f32;
196        for i in 0..n {
197            let f32s = a[i] + b[i];
198            let bfs = round_bf16(a[i]) + round_bf16(b[i]);
199            max = max.max((f32s - bfs).abs());
200        }
201        assert!(max > 0.0, "1/3 is not exact in BF16");
202        assert!(
203            max < 0.01,
204            "7168-wide double BF16 round before NCCL max_abs={max}"
205        );
206        eprintln!("K3 TP BF16 allreduce max_abs @7168 (1/3+2/3) = {max}");
207    }
208}
spark_model/kimi_k3/tp.rs

spark_model/kimi_k3/
tp.rs
