spark_runtime/weights/
k3.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! Rank-aware K3 checkpoint I/O. Plans every tensor before GPU allocation;
3//! uploads only rank-local bytes. Other model loaders are unchanged.
4use super::{WeightLoader, WeightStore, WeightTensor, evict_page_cache};
5use crate::gpu::GpuBackend;
6use anyhow::{Context, Result, ensure};
7use avarok_core::config::ModelConfig;
8use avarok_core::kimi_k3::tp::TpAxis;
9use std::path::Path;
10
11mod index;
12mod inventory;
13#[cfg(test)]
14mod tests;
15
16pub struct K3SafetensorsLoader {
17    config: ModelConfig,
18}
19impl K3SafetensorsLoader {
20    pub fn new(config: ModelConfig) -> Result<Self> {
21        ensure!(
22            matches!(
23                config.model_type.as_str(),
24                "kimi_k3" | "kimi_linear" | "Kimi-K3"
25            ),
26            "K3 rank loader requires a K3 configuration"
27        );
28        ensure!(
29            config.tp_world_size > 0 && config.tp_rank < config.tp_world_size,
30            "invalid K3 TP topology"
31        );
32        Ok(Self { config })
33    }
34}
35impl WeightStore {
36    /// Set only after the rank-aware loader has validated and sliced all tensors.
37    pub fn prepartitioned_tp(&self) -> Option<(usize, usize)> {
38        self.prepartitioned_tp
39    }
40}
41
42impl WeightLoader for K3SafetensorsLoader {
43    fn load(&self, model_dir: &Path, gpu: &dyn GpuBackend, reserve: usize) -> Result<WeightStore> {
44        let inventory = inventory::scan(model_dir, &self.config)?;
45        let mut local_bytes = 0usize;
46        let mut source_bytes = 0usize;
47        let mut replicated_bytes = 0usize;
48        let mut staging_peak = 0usize;
49        let mut binding_extra_bytes = 0usize;
50        for file in &inventory {
51            for tensor in &file.tensors {
52                local_bytes = local_bytes
53                    .checked_add(tensor.plan.local_bytes())
54                    .context("K3 local byte overflow")?;
55                source_bytes = source_bytes
56                    .checked_add(tensor.plan.source_bytes)
57                    .context("K3 source byte overflow")?;
58                if tensor.plan.axis == TpAxis::Replicated {
59                    replicated_bytes = replicated_bytes
60                        .checked_add(tensor.plan.local_bytes())
61                        .context("K3 replicated byte overflow")?;
62                }
63                staging_peak = staging_peak.max(tensor.plan.local_bytes());
64                binding_extra_bytes = binding_extra_bytes
65                    .checked_add(avarok_core::kimi_k3::binding_memory::extra_gpu_bytes(
66                        &tensor.name,
67                        tensor.dtype == super::WeightDtype::FP32,
68                        tensor.plan.local_bytes() / tensor.dtype.byte_size(),
69                    )?)
70                    .context("K3 binding byte overflow")?;
71            }
72        }
73        // Marked layer binding aliases the uploaded pointers; only engine-facing
74        // FP32 embed/head/norm create BF16 device copies. Keep staging headroom
75        // and the explicit engine/KV/runtime reserve separate from this bound.
76        let admitted = local_bytes
77            .checked_add(binding_extra_bytes)
78            .and_then(|n| n.checked_add(staging_peak))
79            .and_then(|n| n.checked_add(reserve))
80            .context("K3 admission byte overflow")?;
81        let initial_free = gpu.free_memory()?;
82        tracing::info!(
83            event = "k3_weight_plan",
84            rank = self.config.tp_rank,
85            world = self.config.tp_world_size,
86            source_bytes,
87            local_bytes,
88            replicated_bytes,
89            staging_peak,
90            binding_extra_bytes,
91            reserve,
92            admitted,
93            free_bytes = initial_free,
94            "K3 pre-upload memory plan; engine workspace/KV/runtime require explicit reserve"
95        );
96        ensure!(
97            admitted <= initial_free,
98            "K3 OOM preflight: rank {} needs {} bytes (local {}, replicated {}, staging {}, binding copies {}, reserve {}), free {}",
99            self.config.tp_rank,
100            admitted,
101            local_bytes,
102            replicated_bytes,
103            staging_peak,
104            binding_extra_bytes,
105            reserve,
106            initial_free
107        );
108        let mut store = WeightStore::empty();
109        let mut uploaded = 0usize;
110        let result: Result<()> = (|| {
111            for item in inventory {
112                let file = std::fs::File::open(&item.path)?;
113                let mmap = unsafe { memmap2::MmapOptions::new().map(&file)? };
114                let tensors = safetensors::SafeTensors::deserialize(&mmap)?;
115                for spec in item.tensors {
116                    let tensor = tensors.tensor(&spec.name)?;
117                    ensure!(
118                        tensor.shape() == spec.source_shape && tensor.dtype() == spec.source_dtype,
119                        "{}: tensor metadata changed after preflight",
120                        spec.name
121                    );
122                    let local = spec.plan.copy_shard(tensor.data())?;
123                    // No UVM fallback: admission or allocation failures are failures.
124                    let ptr = gpu.alloc(local.len()).with_context(|| {
125                        format!(
126                            "K3 rank {} allocating {} local bytes for {}",
127                            self.config.tp_rank,
128                            local.len(),
129                            spec.name
130                        )
131                    })?;
132                    if let Err(e) = gpu.copy_h2d(&local, ptr) {
133                        let _ = gpu.free(ptr);
134                        return Err(e);
135                    }
136                    uploaded += local.len();
137                    store.weights.insert(
138                        spec.name,
139                        WeightTensor {
140                            ptr,
141                            shape: spec.plan.local_shape,
142                            dtype: spec.dtype,
143                        },
144                    );
145                }
146                drop(tensors);
147                drop(mmap);
148                evict_page_cache(&file);
149                let free = gpu.free_memory()?;
150                tracing::info!(
151                    event = "k3_weight_shard_loaded",
152                    rank = self.config.tp_rank,
153                    uploaded_bytes = uploaded,
154                    free_bytes = free,
155                    sampled_device_delta_bytes = initial_free.saturating_sub(free),
156                    "K3 rank-local shard uploaded"
157                );
158                ensure!(
159                    free >= reserve,
160                    "K3 reserve exhausted during rank-local upload"
161                );
162            }
163            Ok(())
164        })();
165        if let Err(err) = result {
166            for (name, tensor) in store.weights.drain() {
167                if let Err(cleanup) = gpu.free(tensor.ptr) {
168                    tracing::error!(%cleanup,%name,"K3 allocation cleanup failed");
169                }
170            }
171            return Err(err);
172        }
173        ensure!(uploaded == local_bytes, "K3 upload accounting mismatch");
174        store.prepartitioned_tp = Some((self.config.tp_rank, self.config.tp_world_size));
175        Ok(store)
176    }
177}
spark_runtime/weights/k3.rs

spark_runtime/weights/
k3.rs
