avarok_core/kimi_k3/
binding_memory.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! Additional GPU weight storage created by the marked K3 model binder.
3//! Dense/packed layer binding aliases resident allocations. Only the three
4//! engine-facing FP32 tensors allocate BF16 copies; source allocations remain.
5//! This excludes engine workspace, KV, CUDA/NCCL and lazy inference allocations.
6use anyhow::{Context, Result};
7
8pub fn engine_bf16_weight(name: &str) -> bool {
9    [
10        "model.embed_tokens.weight",
11        "model.norm.weight",
12        "lm_head.weight",
13    ]
14    .iter()
15    .any(|suffix| name == *suffix || name.strip_suffix(suffix).is_some_and(|p| p.ends_with('.')))
16}
17
18pub fn extra_gpu_bytes(name: &str, is_fp32: bool, elements: usize) -> Result<usize> {
19    if is_fp32 && engine_bf16_weight(name) {
20        elements
21            .checked_mul(2)
22            .context("K3 BF16 binding allocation overflow")
23    } else {
24        Ok(0)
25    }
26}
27
28#[cfg(test)]
29mod tests {
30    use super::*;
31    #[test]
32    fn fp32_engine_copies_are_counted_and_other_weights_alias() {
33        for key in [
34            "model.embed_tokens.weight",
35            "model.norm.weight",
36            "lm_head.weight",
37        ] {
38            for name in [key.to_string(), format!("language_model.{key}")] {
39                assert_eq!(extra_gpu_bytes(&name, true, 64).unwrap(), 128);
40                assert_eq!(extra_gpu_bytes(&name, false, 64).unwrap(), 0);
41            }
42        }
43        assert_eq!(
44            extra_gpu_bytes("model.layers.0.self_attn.q_proj.weight", true, 64).unwrap(),
45            0
46        );
47        assert!(!engine_bf16_weight("not_lm_head.weight"));
48        assert!(extra_gpu_bytes("lm_head.weight", true, usize::MAX).is_err());
49    }
50}
avarok_core/kimi_k3/binding_memory.rs

avarok_core/kimi_k3/
binding_memory.rs
