spark_runtime/weights/
name_utils.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Tensor-NAME parsing: the numeric trailing-index sort, the n-gram table
4//! predicate, and the expert-index extractor. Pure string work, which is
5//! why it is unit-tested here rather than behind a loader.
6
7pub(crate) fn split_trailing_index(name: &str) -> (String, u64) {
8    let mut segs: Vec<&str> = name.split('.').collect();
9    for i in (0..segs.len()).rev() {
10        if let Ok(n) = segs[i].parse::<u64>() {
11            segs.remove(i);
12            return (segs.join("."), n);
13        }
14    }
15    (name.to_string(), u64::MAX)
16}
17
18/// Whether a tensor is an n-gram embedding TABLE — the huge blobs that are
19/// served off NVMe instead of being uploaded with the rest of the checkpoint.
20///
21/// TWO NAMING FAMILIES, because two model families ship this mechanism:
22///
23///   LongCat-Flash-Lite  `*.ngram_embeddings.embedders.{i}.weight`
24///                       12 tables, ~5.2 GB each
25///   Qwen3.8-Flash-Next  `*.ple.ple_embedding.ngram_embedding.shard_{i}.weight`
26///                       128 shards of one logical table, 47.7 GB (FP8) to
27///                       95.4 GB (BF16) in total
28///
29/// Matching is by TENSOR NAME rather than by file, deliberately: RadixArk
30/// isolates the Qwen shards in dedicated `model-plefp8-*` files, but Inferact
31/// buries all 128 inside one 95.4 GB `model-00001-of-00004.safetensors`. A
32/// filename rule would work for one release and silently fail the other,
33/// where "silently" means a 221 GB OOM pre-flight on a 121 GB box.
34///
35/// What must NOT match, because the loader needs these resident:
36///   - LongCat's small `post_projs`
37///   - Qwen's `ngram_embedding.weight_scale` (one BF16 scalar)
38///   - Qwen's `ngram_heads_offsets` / `ngram_heads_vocab_sizes` (I64, 16 each)
39pub fn is_ngram_table(name: &str) -> bool {
40    if !name.ends_with(".weight") {
41        return false;
42    }
43    name.contains("ngram_embeddings.embedders.") || name.contains("ngram_embedding.shard_")
44}
45
46#[cfg(test)]
47mod ngram_defer_tests {
48    use super::*;
49    use crate::weights::{DeferredTensor, WeightDtype, WeightStore};
50
51    #[test]
52    fn ngram_table_predicate_matches_only_the_big_tables() {
53        assert!(is_ngram_table("model.ngram_embeddings.embedders.0.weight"));
54        assert!(is_ngram_table("model.ngram_embeddings.embedders.11.weight"));
55        // The small projections are ordinary tensors and must still load.
56        assert!(!is_ngram_table(
57            "model.ngram_embeddings.post_projs.0.weight"
58        ));
59        assert!(!is_ngram_table("model.embed_tokens.weight"));
60        assert!(!is_ngram_table(
61            "model.layers.0.mlp.experts.3.gate_proj.weight"
62        ));
63    }
64
65    /// Qwen3.8-Flash-Next stores ONE logical table as 128 `shard_{i}` tensors
66    /// under a PLE block, rather than LongCat's 12 separate embedders.
67    #[test]
68    fn ngram_table_predicate_matches_the_qwen_ple_shards() {
69        let base = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding";
70        assert!(is_ngram_table(&format!("{base}.shard_0.weight")));
71        assert!(is_ngram_table(&format!("{base}.shard_127.weight")));
72
73        // The per-table scalar scale (RadixArk's FP8 release) must stay
74        // resident — the row cache needs it to dequantize.
75        assert!(!is_ngram_table(&format!("{base}.weight_scale")));
76        // The head range tables are 16 I64 values the loader reads directly.
77        assert!(!is_ngram_table(
78            "model.language_model.layers.1.ple.ple_embedding.ngram_heads_offsets"
79        ));
80        assert!(!is_ngram_table(
81            "model.language_model.layers.1.ple.ple_embedding.ngram_heads_vocab_sizes"
82        ));
83        // The PLE block's own projections are ordinary tensors.
84        assert!(!is_ngram_table(
85            "model.language_model.layers.1.ple.key_proj.weight"
86        ));
87        assert!(!is_ngram_table(
88            "model.language_model.layers.1.ple.value_proj.weight"
89        ));
90        assert!(!is_ngram_table(
91            "model.language_model.layers.1.ple.conv1d.weight"
92        ));
93    }
94
95    #[test]
96    fn deferred_sorts_numerically_not_lexicographically() {
97        let mut st = WeightStore::empty();
98        for i in [0usize, 2, 10, 11, 9] {
99            st.defer(
100                format!("model.ngram_embeddings.embedders.{i}.weight"),
101                DeferredTensor {
102                    path: std::path::PathBuf::from("x"),
103                    offset: i as u64,
104                    shape: vec![1, 1],
105                    dtype: WeightDtype::BF16,
106                },
107            );
108        }
109        let got: Vec<u64> = st.deferred_sorted().iter().map(|(_, d)| d.offset).collect();
110        assert_eq!(got, vec![0, 2, 9, 10, 11], "table order must be numeric");
111    }
112}
113
114/// Parse expert index from tensor name (e.g. "model.layers.3.mlp.experts.42.gate_proj.weight" → 42).
115pub fn parse_expert_index(name: &str) -> Option<usize> {
116    let parts: Vec<&str> = name.split('.').collect();
117    for (i, part) in parts.iter().enumerate() {
118        if *part == "experts" && i + 1 < parts.len() {
119            return parts[i + 1].parse().ok();
120        }
121    }
122    None
123}
124
125#[cfg(test)]
126mod from_str_tests {
127    use crate::weights::WeightDtype;
128
129    #[test]
130    fn from_safetensors_str_matches_disk_mapping() {
131        // The RDMA weight peer publishes these raw header strings; the client
132        // must resolve them to the exact WeightDtype the disk loaders use, else
133        // byte_size/shape diverge and logits break. Locks the closed mapping.
134        use WeightDtype::*;
135        for (s, want) in [
136            ("F32", FP32),
137            ("BF16", BF16),
138            ("U8", UInt8),
139            ("I8", UInt8), // packed NVFP4 raw container
140            ("F8_E4M3", FP8E4M3),
141            ("F8_E8M0", FP8E8M0),
142            ("I64", Int64),
143        ] {
144            assert_eq!(
145                WeightDtype::from_safetensors_str(s).unwrap(),
146                want,
147                "dtype {s}"
148            );
149        }
150        // F16 is converted to BF16 at disk-load; a store (and therefore a
151        // peer manifest) can never contain it, so the wire mapping rejects it.
152        assert!(WeightDtype::from_safetensors_str("F16").is_err());
153        assert!(WeightDtype::from_safetensors_str("bogus").is_err());
154    }
155
156    #[test]
157    fn f16_bytes_convert_to_bf16_via_f32() {
158        use half::{bf16, f16};
159        // Cover sign, exact powers of two, a value needing mantissa rounding
160        // (f16 has 10 mantissa bits, bf16 only 7), f16 max, and a subnormal.
161        let vals = [0.0f32, 1.0, -1.5, 0.1, 65504.0, -6.1035156e-5];
162        let src: Vec<u8> = vals
163            .iter()
164            .flat_map(|v| f16::from_f32(*v).to_le_bytes())
165            .collect();
166        let out = crate::weights::f16_to_bf16_bytes(&src);
167        assert_eq!(out.len(), src.len());
168        for (i, v) in vals.iter().enumerate() {
169            let got = bf16::from_le_bytes([out[2 * i], out[2 * i + 1]]);
170            let want = bf16::from_f32(f16::from_f32(*v).to_f32());
171            assert_eq!(got, want, "value {v}");
172        }
173    }
174}
spark_runtime/weights/name_utils.rs

spark_runtime/weights/
name_utils.rs
