avarok_core/kimi_k3/
tp.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! K3 tensor-parallel storage plan, shared by host loading and device binding.
4//! Config head counts are rank-local; expert and dense widths remain full.
5
6use super::MixerKind;
7use crate::config::ModelConfig;
8use crate::mxfp4_e8m0::GROUP_SIZE;
9use anyhow::{Result, ensure};
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum TpAxis {
13    Replicated,
14    Rows,
15    Columns,
16}
17
18pub fn tensor_plan(name: &str, mixer: MixerKind, config: &ModelConfig) -> (TpAxis, usize, usize) {
19    // Packed weights and E8M0 scales share the logical projection plan.
20    let canonical;
21    let name = if let Some(stem) = name
22        .strip_suffix(".weight_packed")
23        .or_else(|| name.strip_suffix(".weight_scale"))
24    {
25        canonical = format!("{stem}.weight");
26        canonical.as_str()
27    } else {
28        name
29    };
30    let tp = config.tp_world_size.max(1);
31    let h = config.hidden_size;
32    let kda_heads = config.linear_num_key_heads * tp;
33    let kda_d = config.linear_key_head_dim;
34    let kda_q = kda_heads * kda_d;
35    let conv_k = config.linear_conv_kernel_dim.max(1);
36    let mla_heads = config.num_attention_heads * tp;
37    let qk = config.qk_nope_head_dim + config.qk_rope_head_dim;
38    let dv = mla_heads * config.v_head_dim;
39    let kv_b = mla_heads * (config.qk_nope_head_dim + config.v_head_dim);
40    let inter = config.intermediate_size;
41    let eh = config.moe_intermediate_size;
42    let lat = config.moe_latent_size;
43
44    if name.ends_with(".self_attn.q_proj.weight")
45        || name.ends_with(".self_attn.k_proj.weight")
46        || name.ends_with(".self_attn.v_proj.weight")
47    {
48        return (TpAxis::Rows, kda_q, h);
49    }
50    if name.ends_with(".self_attn.q_conv1d.weight")
51        || name.ends_with(".self_attn.k_conv1d.weight")
52        || name.ends_with(".self_attn.v_conv1d.weight")
53    {
54        return (TpAxis::Rows, kda_q, conv_k);
55    }
56    if name.ends_with(".self_attn.g_proj.weight") {
57        let n = match mixer {
58            MixerKind::Kda => kda_q,
59            MixerKind::Mla => dv,
60        };
61        return (TpAxis::Rows, n, h);
62    }
63    if name.ends_with(".self_attn.o_proj.weight") {
64        let inn = match mixer {
65            MixerKind::Kda => kda_q,
66            MixerKind::Mla => dv,
67        };
68        return (TpAxis::Columns, h, inn);
69    }
70    if name.ends_with(".self_attn.b_proj.weight") {
71        return (TpAxis::Rows, kda_heads, h);
72    }
73    if name.ends_with(".self_attn.A_log") {
74        return (TpAxis::Rows, kda_heads, 1);
75    }
76    if name.ends_with(".self_attn.dt_bias") {
77        return (TpAxis::Rows, kda_q, 1);
78    }
79    if name.ends_with(".self_attn.f_b_proj.weight") {
80        return (TpAxis::Rows, kda_q, kda_d);
81    }
82    if name.ends_with(".self_attn.q_b_proj.weight") {
83        return (TpAxis::Rows, mla_heads * qk, config.q_lora_rank);
84    }
85    if name.ends_with(".self_attn.kv_b_proj.weight") {
86        return (TpAxis::Rows, 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 (TpAxis::Rows, inter, h);
90    }
91    if name.ends_with(".mlp.down_proj.weight") {
92        return (TpAxis::Columns, 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 (TpAxis::Rows, eh, lat);
97        }
98        if name.ends_with(".w2.weight") {
99            return (TpAxis::Columns, lat, eh);
100        }
101    }
102    (TpAxis::Replicated, 1, 1)
103}
104
105/// Byte ranges from one full checkpoint tensor that form this rank's tensor.
106/// Ranges preserve row-major packed nibbles and their matching E8M0 groups.
107#[derive(Debug, Clone)]
108pub struct TpBytePlan {
109    pub axis: TpAxis,
110    pub local_shape: Vec<usize>,
111    pub source_bytes: usize,
112    offset: usize,
113    rows: usize,
114    stride: usize,
115    row_bytes: usize,
116    zero_f32_tail_from: Option<usize>,
117}
118
119impl TpBytePlan {
120    pub fn local_bytes(&self) -> usize {
121        self.rows * self.row_bytes
122    }
123
124    pub fn byte_ranges(&self) -> impl Iterator<Item = std::ops::Range<usize>> + '_ {
125        (0..self.rows).map(|row| {
126            let start = self.offset + row * self.stride;
127            start..start + self.row_bytes
128        })
129    }
130
131    /// Check storage-only padding before any device allocations. The official
132    /// K3 export stores 96 trained A_log heads followed by 32 zero F32 values.
133    /// Keep the logical head count; never truncate an unverified nonzero tail.
134    pub fn validate_source(&self, source: &[u8]) -> Result<()> {
135        ensure!(
136            source.len() == self.source_bytes,
137            "K3 TP source bytes {} != {}",
138            source.len(),
139            self.source_bytes
140        );
141        if let Some(start) = self.zero_f32_tail_from {
142            ensure!(
143                source[start..]
144                    .chunks_exact(4)
145                    .all(|v| f32::from_le_bytes(v.try_into().unwrap()) == 0.0),
146                "K3 padded A_log has a nonzero/nonfinite tail; refusing to discard trained values"
147            );
148        }
149        Ok(())
150    }
151
152    /// Gather only the local bytes from an mmap/read-only checkpoint view.
153    pub fn copy_shard(&self, source: &[u8]) -> Result<Vec<u8>> {
154        self.validate_source(source)?;
155        let mut local = Vec::with_capacity(self.local_bytes());
156        for range in self.byte_ranges() {
157            local.extend_from_slice(&source[range]);
158        }
159        Ok(local)
160    }
161}
162
163fn product(values: &[usize]) -> Result<usize> {
164    values.iter().try_fold(1usize, |n, &v| {
165        ensure!(v > 0, "K3 TP zero dimension/element width");
166        n.checked_mul(v)
167            .ok_or_else(|| anyhow::anyhow!("K3 TP byte-size overflow"))
168    })
169}
170
171/// Plan checkpoint storage BEFORE allocating a device buffer.
172///
173/// BF16/F32 keep their dtype. Packed experts require `[N, K/2]` bytes plus
174/// `[N, K/32]` E8M0 scales; a column split cannot bisect a 32-value group.
175/// Unknown non-packed tensors remain replicated. Runtime must separately
176/// enforce the model's tensor inventory and exclude unsupported modalities.
177pub fn plan_tensor_bytes(
178    name: &str,
179    shape: &[usize],
180    element_bytes: usize,
181    mixer: MixerKind,
182    config: &ModelConfig,
183) -> Result<TpBytePlan> {
184    let world = config.tp_world_size;
185    let rank = config.tp_rank;
186    ensure!(
187        world > 0 && rank < world,
188        "K3 TP rank {rank} invalid for world {world}"
189    );
190    ensure!(
191        !shape.is_empty(),
192        "{name}: K3 TP requires a non-scalar tensor"
193    );
194    let elements = product(shape)?;
195    let source_bytes = product(&[elements, element_bytes])?;
196    let (axis, n, logical_k) = tensor_plan(name, mixer, config);
197    // Verified export exception, not a general shape relaxation. See HF
198    // moonshotai/Kimi-K3 discussion #150 and NVIDIA Megatron Bridge K3 docs.
199    if name.ends_with(".self_attn.A_log")
200        && mixer == MixerKind::Kda
201        && n == 96
202        && config.linear_key_head_dim == 128
203        && shape == [128]
204        && element_bytes == 4
205    {
206        ensure!(
207            n.is_multiple_of(world),
208            "K3 A_log heads must divide TP world"
209        );
210        let row_bytes = n / world * element_bytes;
211        return Ok(TpBytePlan {
212            axis,
213            local_shape: vec![n / world],
214            source_bytes,
215            offset: rank * row_bytes,
216            rows: 1,
217            stride: row_bytes,
218            row_bytes,
219            zero_f32_tail_from: Some(n * element_bytes),
220        });
221    }
222    let packed = name.ends_with(".weight_packed");
223    let scale = name.ends_with(".weight_scale");
224    let storage_k = if packed || scale {
225        ensure!(
226            name.contains(".block_sparse_moe.experts.") && axis != TpAxis::Replicated,
227            "{name}: unsupported K3 packed projection"
228        );
229        ensure!(
230            element_bytes == 1,
231            "{name}: packed MXFP4/E8M0 storage must be one byte"
232        );
233        ensure!(
234            logical_k.is_multiple_of(GROUP_SIZE),
235            "{name}: K {logical_k} must align to E8M0 groups of 32"
236        );
237        if axis == TpAxis::Columns {
238            ensure!(
239                logical_k.is_multiple_of(world) && (logical_k / world).is_multiple_of(GROUP_SIZE),
240                "{name}: TP{world} column partition splits an E8M0 group of 32"
241            );
242        }
243        let divisor = if packed { 2 } else { GROUP_SIZE };
244        let k = logical_k / divisor;
245        ensure!(
246            shape == [n, k],
247            "{name}: packed shape {shape:?} != [{n}, {k}]"
248        );
249        k
250    } else {
251        if axis != TpAxis::Replicated {
252            ensure!(
253                matches!(element_bytes, 2 | 4),
254                "{name}: K3 dense TP requires BF16 or FP32"
255            );
256            ensure!(
257                shape[0] == n && elements == product(&[n, logical_k])?,
258                "{name}: dense shape {shape:?} != logical [{n}, {logical_k}]"
259            );
260        }
261        logical_k
262    };
263    let mut local_shape = shape.to_vec();
264    let (offset, rows, stride, row_bytes) = match axis {
265        TpAxis::Replicated => (0, 1, source_bytes, source_bytes),
266        TpAxis::Rows => {
267            ensure!(
268                n.is_multiple_of(world),
269                "{name}: rows {n} not divisible by TP{world}"
270            );
271            local_shape[0] /= world;
272            let row_bytes = source_bytes / world;
273            (rank * row_bytes, 1, row_bytes, row_bytes)
274        }
275        TpAxis::Columns => {
276            ensure!(
277                shape.len() == 2 && storage_k.is_multiple_of(world),
278                "{name}: columns {storage_k} not divisible by TP{world}"
279            );
280            local_shape[1] /= world;
281            let stride = product(&[storage_k, element_bytes])?;
282            let row_bytes = stride / world;
283            (rank * row_bytes, n, stride, row_bytes)
284        }
285    };
286    ensure!(
287        product(&local_shape)? * element_bytes == rows * row_bytes,
288        "{name}: local TP byte accounting mismatch"
289    );
290    Ok(TpBytePlan {
291        axis,
292        local_shape,
293        source_bytes,
294        offset,
295        rows,
296        stride,
297        row_bytes,
298        zero_f32_tail_from: None,
299    })
300}
301
302#[cfg(test)]
303mod tests;
avarok_core/kimi_k3/tp.rs

avarok_core/kimi_k3/
tp.rs
