spark_model/kimi_k3/
tp.rs

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

spark_model/kimi_k3/
tp.rs
