metrale_model_arch/kimi_k3/
tp.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The Kimi K3 tensor-parallel plan (`metrale_model_weights::kimi_k3_host::tp::tensor_plan`) in [`TpShardKind`] terms.
4//!
5//! Head counts on `config` are per rank; the returned sizes are full, local
6//! heads times `tp_world_size`.
7//!
8//! Owner: model-arch, Kimi K3.
9//! Invariants: none beyond the types.
10
11use metrale_config::ModelConfig;
12use metrale_model_weights::kimi_k3_host::{MixerKind, MlpKind};
13
14use crate::tp_shard::TpShardKind;
15
16/// 2026-09-25: Always true: the Kimi K3 loader binds tensor-parallel shards.
17pub fn supports_tp() -> bool {
18    true
19}
20
21/// 2026-09-25: `(kind, full_out, full_in)` for one checkpoint key. A key the
22/// plan does not name is `Replicated`.
23pub fn tensor_plan(
24    name: &str,
25    mixer: MixerKind,
26    _mlp: MlpKind,
27    config: &ModelConfig,
28) -> (TpShardKind, usize, usize) {
29    use metrale_model_weights::kimi_k3_host::tp::{TpAxis, tensor_plan};
30    let (axis, n, k) = tensor_plan(name, mixer, config);
31    let kind = match axis {
32        TpAxis::Replicated => TpShardKind::Replicated,
33        TpAxis::Rows => TpShardKind::ColumnParallel,
34        TpAxis::Columns => TpShardKind::RowParallel,
35    };
36    (kind, n, k)
37}
38
39#[cfg(test)]
40mod tests {
41    use super::*;
42    use metrale_config::parse_config;
43
44    const TWIN: &str = include_str!("../../../../docs/k3/fixtures/Kimi-K3-0.40B-config.json");
45
46    fn twin() -> ModelConfig {
47        parse_config(TWIN).expect("0.40B twin")
48    }
49
50    fn divide_heads_for_tp(config: &mut ModelConfig, rank: usize, size: usize) {
51        config.tp_rank = rank;
52        config.tp_world_size = size;
53        if size > 1 {
54            config.num_attention_heads /= size;
55            config.num_key_value_heads /= size;
56            config.linear_num_key_heads /= size;
57            config.linear_num_value_heads /= size;
58        }
59    }
60
61    #[test]
62    fn kimi_k3_supports_tp() {
63        assert!(supports_tp(), "loader must not refuse --tp-size 2");
64    }
65
66    #[test]
67    fn kda_q_is_column_o_is_row() {
68        let mut c = twin();
69        divide_heads_for_tp(&mut c, 0, 2);
70        let q = tensor_plan(
71            "language_model.model.layers.0.self_attn.q_proj.weight",
72            MixerKind::Kda,
73            MlpKind::Dense,
74            &c,
75        );
76        let o = tensor_plan(
77            "language_model.model.layers.0.self_attn.o_proj.weight",
78            MixerKind::Kda,
79            MlpKind::Dense,
80            &c,
81        );
82        assert_eq!(q, (TpShardKind::ColumnParallel, 8 * 32, 1024));
83        assert_eq!(o, (TpShardKind::RowParallel, 1024, 8 * 32));
84        let embed = tensor_plan(
85            "language_model.model.embed_tokens.weight",
86            MixerKind::Kda,
87            MlpKind::Dense,
88            &c,
89        );
90        assert_eq!(embed.0, TpShardKind::Replicated);
91    }
92
93    #[test]
94    fn rank0_and_rank1_q_proj_full_out_match() {
95        let mut c0 = twin();
96        divide_heads_for_tp(&mut c0, 0, 2);
97        let mut c1 = twin();
98        divide_heads_for_tp(&mut c1, 1, 2);
99        let q0 = tensor_plan(
100            "language_model.model.layers.0.self_attn.q_proj.weight",
101            MixerKind::Kda,
102            MlpKind::Dense,
103            &c0,
104        );
105        let q1 = tensor_plan(
106            "language_model.model.layers.0.self_attn.q_proj.weight",
107            MixerKind::Kda,
108            MlpKind::Dense,
109            &c1,
110        );
111        assert_eq!(q0, q1, "full sizes reconstruct from local*tp on both ranks");
112        assert_eq!(c0.tp_rank, 0);
113        assert_eq!(c1.tp_rank, 1);
114        assert_ne!(c0.tp_rank, c1.tp_rank);
115    }
116
117    fn round_bf16(x: f32) -> f32 {
118        let bits = metrale_core::numeric::f32_to_bf16(x);
119        metrale_core::numeric::bf16_bytes_to_f32(bits.to_le_bytes())
120    }
121
122    #[test]
123    fn production_7168_bf16_allreduce_rounding_is_measured() {
124        // 2026-09-25: `tp_allreduce` rounds each rank's FP32 partial to BF16 before
125        // the sum. This measures that rounding at the production hidden width,
126        // 7168 (`docs/k3/fixtures/moonshotai-Kimi-K3-config.json`).
127        let n = 7168;
128        let a = vec![1.0f32 / 3.0; n];
129        let b = vec![2.0f32 / 3.0; n];
130        let mut max = 0.0f32;
131        for i in 0..n {
132            let f32s = a[i] + b[i];
133            let bfs = round_bf16(a[i]) + round_bf16(b[i]);
134            max = max.max((f32s - bfs).abs());
135        }
136        assert!(max > 0.0, "1/3 is not exact in BF16");
137        assert!(
138            max < 0.01,
139            "7168-wide double BF16 round before NCCL max_abs={max}"
140        );
141        eprintln!("K3 TP BF16 allreduce max_abs @7168 (1/3+2/3) = {max}");
142    }
143}
metrale_model_arch/kimi_k3/tp.rs

metrale_model_arch/kimi_k3/
tp.rs
