spark_model/
lib.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3#![deny(warnings)]
4#![deny(clippy::all)]
5// Kernel-launch helpers and trait-impl wide signatures legitimately exceed
6// clippy's 7-argument default. The same goes for the indexing-loop patterns
7// that mirror the kernel grids we dispatch.
8#![allow(clippy::too_many_arguments)]
9#![allow(clippy::needless_range_loop)]
10// Some FP/integer special-case branches return the same value but have
11// distinct semantic meanings (NaN vs zero, etc.). Audit shows these are
12// intentional.
13#![allow(clippy::if_same_then_else)]
14// The HSS / disk-spill plumbing threads `Vec<u32>` through trait methods so
15// callers can grow them in place; converting to slices breaks the contract.
16#![allow(clippy::ptr_arg)]
17// HF safetensors index tuples are wide on purpose.
18#![allow(clippy::type_complexity)]
19
20pub mod engine;
21pub mod factory;
22pub mod forward;
23pub mod layer;
24pub mod layers;
25pub mod lora;
26pub mod mistral_loader;
27pub mod model;
28pub mod mtp_layout;
29pub mod precision_schedule;
30pub mod preflight;
31pub mod quant_format;
32mod rank_agree;
33pub mod seq_state_reserve;
34pub mod speculative;
35pub mod ssm_reserve;
36pub mod tp_shard;
37pub mod traits;
38pub mod video_decode_ffmpeg;
39pub mod video_preprocess;
40pub mod vision_item;
41pub mod vision_preprocess;
42pub use vision_item::VisionItem;
43
44pub mod weight_loader;
45pub mod weight_map;
46
47/// True when the checkpoint ships **HF-vanilla** RMSNorm weights — i.e. the norm
48/// weight is used as `out = x * w / rms`, not Qwen3-Next's offset-from-1
49/// `out = x * (1 + w) / rms`.
50///
51/// Such a model must load its norm weights **exactly** and dispatch
52/// `rms_norm_vanilla`. The alternative — pre-subtracting 1.0 and storing
53/// `bf16(w - 1)` for the offset kernel — is only lossless when `w ≈ 1`.
54/// DeepSeek-V4's norm weights are ≈ 0.03, so `w - 1 ≈ -0.97`, and BF16's
55/// rounding error there (~1.9e-3 absolute) becomes a **1.8-3.4 % relative error
56/// on the weight itself** once 1 is added back — catastrophic cancellation.
57/// Measured over all 249 V4 norm tensors: up to 19 % on `q_norm`, and 100 %
58/// with sign flips on the compressor norms.
59///
60/// This is an explicit model dispatch, NOT an inference from weight statistics.
61pub fn ships_vanilla_norm_weights(config: &atlas_core::config::ModelConfig) -> bool {
62    model_type_ships_vanilla_norm_weights(&config.model_type)
63}
64
65/// The dispatch predicate itself, on the bare `model_type`, so it is unit-testable
66/// without constructing a full `ModelConfig`.
67pub fn model_type_ships_vanilla_norm_weights(model_type: &str) -> bool {
68    // 🪤 `glm5_next` added 2026-08-27. GLM-5.3's norms are PLAIN `x * rms * w` — the same
69    // trap `glm5next_layer` documents for its per-layer norms. This predicate additionally
70    // picks the kernel for the MODEL-LEVEL final norm (`model/impl_a1.rs`), which is applied
71    // outside any layer, so omitting GLM here silently normalises the final hidden state with
72    // the `(1 + w)` offset and corrupts every token's logits. Nothing about the shapes says so.
73    matches!(model_type, "deepseek_v4" | "laguna" | "glm5_next")
74}
75
76/// Must chunked prefill run as a SINGLE chunk for this model?
77///
78/// True only for models that reach the chunk-LOCAL MLA prefill in
79/// `qwen3_attention/prefill.rs`, which attends over the current chunk's K/V alone —
80/// multi-chunk there silently corrupts attention output (Mistral-Small-4, 2026-05-01: 8 K
81/// collapses to "The\nThe…").
82///
83/// 🔴 `kv_lora_rank > 0` is a PROXY for that kernel and `glm5_next` breaks it: GLM-5.3 is
84/// MLA (rank 512) but prefills through `Glm5NextLayer::prefill`, a per-token walk that
85/// attends the whole paged prefix at each absolute position — chunk boundaries are
86/// invisible to it. Answering true capped every GLM prompt at `2 × --max-prefill-tokens`,
87/// because `prefill_a_step` splits the FIRST chunk at the cap regardless and this gate then
88/// made the remainder one unsplit chunk the buffer arena refused. ANOMALIES A61.
89pub fn requires_single_chunk_prefill(model_type: &str, kv_lora_rank: usize) -> bool {
90    kv_lora_rank > 0 && model_type != "glm5_next"
91}
92
93#[cfg(test)]
94mod single_chunk_prefill_tests {
95    use super::requires_single_chunk_prefill as single;
96
97    /// GLM-5.3 is MLA and must still be chunked — that is the whole of A61.
98    #[test]
99    fn glm5_next_is_mla_but_chunks_fine() {
100        assert!(!single("glm5_next", 512));
101        // Every other MLA family keeps the single-chunk guard.
102        assert!(single("deepseek_v4", 512));
103        assert!(single("mistral", 512));
104        // Non-MLA models were never gated.
105        assert!(!single("qwen3_5_moe", 0));
106    }
107}
108
109#[cfg(test)]
110mod norm_convention_tests {
111    use super::model_type_ships_vanilla_norm_weights as vanilla;
112
113    /// Only explicitly listed model families take the vanilla path. Every
114    /// other family keeps the offset-from-1 convention it was validated under.
115    #[test]
116    fn vanilla_norm_models_are_explicit() {
117        assert!(vanilla("deepseek_v4"));
118        assert!(vanilla("laguna"));
119        // GLM-5.3's norms are plain; the final norm is applied outside any layer.
120        assert!(vanilla("glm5_next"));
121        for other in [
122            "qwen3_next",
123            "qwen3_5_moe",
124            "qwen3_moe",
125            "deepseek_v3",
126            "llama",
127            "mistral",
128            "nemotron",
129            "",
130        ] {
131            assert!(!vanilla(other), "{other} must keep offset-from-1 semantics");
132        }
133    }
134}
spark_model/lib.rs

spark_model/
lib.rs
