spark_runtime/fast_weights/
skip.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Which tensors the fast loader does NOT upload.
4//!
5//! Three independent rules, worth reading together because each one withholds
6//! bytes a downstream loader might expect:
7//!
8//!   1. **EP sharding** — remote experts belong to another rank.
9//!   2. **`skip_activation_scales`** — W4A4 `*.input_scale`, opt-in.
10//!   3. **`skip_mtp`** — `mtp.*` for a loader that builds no MTP head, opt-in.
11//!   4. **`skip_vision`** — the vision tower, for a text-only port that binds
12//!      no vision encoder. Derived from the loader, not opt-in per model.
13//!
14//! Rules 2 and 3 default OFF and are allow-listed per model, because
15//! withholding a tensor a loader DOES read is invisible until the output is
16//! subtly wrong. Rule 1 is structural and always active under EP.
17//!
18//! **DEFER is a fifth rule and a different kind.** A skipped tensor is gone; a
19//! deferred one is recorded with its on-disk location because the model's own
20//! loader will read it from the host. [`FastSafetensorsLoader::is_deferred`]
21//! lives here so the two questions are answered side by side, but it is asked
22//! AFTER `should_skip_tensor` — see [`crate::weights::deferred`].
23
24use super::FastSafetensorsLoader;
25use crate::weights::{WeightDtype, parse_expert_index};
26
27impl FastSafetensorsLoader {
28    /// Does the model's loader claim this tensor? `false` when no hook is set,
29    /// which is every model but the ones that opt in.
30    ///
31    /// 🪤 Keyed on the STORE dtype — the width the tensor would have had in the
32    /// store, which for an F16 export is BF16. A predicate that says "BF16
33    /// routed expert" therefore also catches the F16 spelling of the same
34    /// export, which is the honest answer.
35    pub fn is_deferred(&self, name: &str, dtype: WeightDtype) -> bool {
36        self.defer.as_ref().is_some_and(|f| f(name, dtype))
37    }
38}
39
40impl FastSafetensorsLoader {
41    pub(super) fn should_skip_tensor(&self, name: &str) -> bool {
42        // Checked before the EP short-circuit: a text-only port skips the
43        // vision tower at tp/ep 1 too.
44        if self.skip_vision && super::is_vision_tensor(name) {
45            return true;
46        }
47        // MTP head weights for a model whose loader does not build one.
48        if self.skip_mtp && name.starts_with("mtp.") {
49            return true;
50        }
51        // W4A4 activation scales: never read on the w4a16 path (the NVFP4
52        // loader falls back to `DevicePtr::NULL`), and 4-byte allocations are
53        // almost pure granule padding at expert scale.
54        if self.skip_activation_scales && name.ends_with(".input_scale") {
55            return true;
56        }
57        if self.ep_world_size <= 1 {
58            return false;
59        }
60        if name.starts_with("mtp.") {
61            return false;
62        }
63        if let Some(idx) = parse_expert_index(name) {
64            let per_rank = self.num_experts / self.ep_world_size;
65            let local_start = self.ep_rank * per_rank;
66            let local_end = if self.ep_rank == self.ep_world_size - 1 {
67                self.num_experts
68            } else {
69                local_start + per_rank
70            };
71            idx < local_start || idx >= local_end
72        } else {
73            false
74        }
75    }
76}
77
78#[cfg(test)]
79mod activation_scale_tests {
80    use super::FastSafetensorsLoader;
81
82    fn loader() -> FastSafetensorsLoader {
83        let mut l = FastSafetensorsLoader::new();
84        l.skip_activation_scales = true;
85        l
86    }
87
88    /// 🔴 The BLAST RADIUS of the opt-in, pinned. The rule is a `.input_scale`
89    /// suffix and nothing else; every model added to the server's allow-list
90    /// depends on it never reaching `weight_scale` / `weight_scale_2`, which
91    /// ARE the w4a16 operands. A predicate widened to `contains("scale")` or
92    /// to `.starts_with` would pass every other test in this crate and quietly
93    /// unbind every quantised projection in the checkpoint.
94    #[test]
95    fn the_rule_reaches_only_activation_scales() {
96        let l = loader();
97        assert!(l.should_skip_tensor("model.layers.0.mlp.gate_proj.input_scale"));
98        assert!(l.should_skip_tensor("model.layers.0.mlp.experts.7.up_proj.input_scale"));
99
100        for keep in [
101            "model.layers.0.mlp.gate_proj.weight",
102            "model.layers.0.mlp.gate_proj.weight_scale",
103            "model.layers.0.mlp.gate_proj.weight_scale_2",
104            "model.layers.0.self_attn.k_proj.k_scale",
105            "model.layers.0.mlp.gate_proj.input_scale_inv",
106            "input_scale",
107        ] {
108            assert!(!l.should_skip_tensor(keep), "{keep} must still load");
109        }
110    }
111
112    /// OPT-IN: off by default, so a model that never asks keeps every scalar.
113    #[test]
114    fn nothing_is_skipped_unless_the_model_opted_in() {
115        let l = FastSafetensorsLoader::new();
116        assert!(!l.skip_activation_scales);
117        assert!(!l.should_skip_tensor("model.layers.0.mlp.gate_proj.input_scale"));
118    }
119}
120
121#[cfg(test)]
122#[path = "defer_tests.rs"]
123mod defer_tests;
spark_runtime/fast_weights/skip.rs

spark_runtime/fast_weights/
skip.rs
