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 kimi_k3;
24pub mod layer;
25pub mod layers;
26pub mod lora;
27pub mod mistral_loader;
28pub mod model;
29pub mod mtp_layout;
30pub mod precision_schedule;
31pub mod preflight;
32pub mod quant_format;
33mod rank_agree;
34pub mod seq_state_reserve;
35pub mod speculative;
36pub mod ssm_reserve;
37pub mod tp_shard;
38pub mod traits;
39pub mod video_decode_ffmpeg;
40pub mod video_preprocess;
41pub mod vision_item;
42pub mod vision_preprocess;
43pub mod vision_preprocess_glm;
44/// Marconi snapshot-restore threshold: the shipped default, and the setter
45/// `spark serve` uses to pin it from `--marconi-min-tokens` before anything
46/// reads it. Re-exported rather than making `model::mtp_carry` public — the
47/// rest of that module is internal.
48pub use model::mtp_carry::{DEFAULT_MARCONI_MIN_TOKENS, set_marconi_min_tokens};
49pub use vision_item::VisionItem;
50
51pub mod weight_loader;
52pub mod weight_map;
53
54/// True when the checkpoint ships **HF-vanilla** RMSNorm weights — i.e. the norm
55/// weight is used as `out = x * w / rms`, not Qwen3-Next's offset-from-1
56/// `out = x * (1 + w) / rms`.
57///
58/// Such a model must load its norm weights **exactly** and dispatch
59/// `rms_norm_vanilla`. The alternative — pre-subtracting 1.0 and storing
60/// `bf16(w - 1)` for the offset kernel — is only lossless when `w ≈ 1`.
61/// DeepSeek-V4's norm weights are ≈ 0.03, so `w - 1 ≈ -0.97`, and BF16's
62/// rounding error there (~1.9e-3 absolute) becomes a **1.8-3.4 % relative error
63/// on the weight itself** once 1 is added back — catastrophic cancellation.
64/// Measured over all 249 V4 norm tensors: up to 19 % on `q_norm`, and 100 %
65/// with sign flips on the compressor norms.
66///
67/// This is an explicit model dispatch, NOT an inference from weight statistics.
68pub fn ships_vanilla_norm_weights(config: &avarok_core::config::ModelConfig) -> bool {
69    model_type_ships_vanilla_norm_weights(&config.model_type)
70}
71
72/// The dispatch predicate itself, on the bare `model_type`, so it is unit-testable
73/// without constructing a full `ModelConfig`.
74pub fn model_type_ships_vanilla_norm_weights(model_type: &str) -> bool {
75    // 🪤 `glm5_next` added 2026-08-27. GLM-5.3's norms are PLAIN `x * rms * w` — the same
76    // trap `glm5next_layer` documents for its per-layer norms. This predicate additionally
77    // picks the kernel for the MODEL-LEVEL final norm (`model/impl_a1.rs`), which is applied
78    // outside any layer, so omitting GLM here silently normalises the final hidden state with
79    // the `(1 + w)` offset and corrupts every token's logits. Nothing about the shapes says so.
80    matches!(
81        model_type,
82        "deepseek_v4" | "deepseek_v41" | "laguna" | "glm5_next" | "kimi_k3"
83    )
84}
85
86/// Must chunked prefill run as a SINGLE chunk for this model?
87///
88/// True only for models that reach the chunk-LOCAL MLA prefill in
89/// `qwen3_attention/prefill.rs`, which attends over the current chunk's K/V alone —
90/// multi-chunk there silently corrupts attention output (Mistral-Small-4, 2026-05-01: 8 K
91/// collapses to "The\nThe…").
92///
93/// 🔴 `kv_lora_rank > 0` is a PROXY for that kernel and `glm5_next` breaks it: GLM-5.3 is
94/// MLA (rank 512) but prefills through `Glm5NextLayer::prefill`, a per-token walk that
95/// attends the whole paged prefix at each absolute position — chunk boundaries are
96/// invisible to it. Answering true capped every GLM prompt at `2 × --max-prefill-tokens`,
97/// because `prefill_a_step` splits the FIRST chunk at the cap regardless and this gate then
98/// made the remainder one unsplit chunk the buffer arena refused. ANOMALIES A61.
99/// K3 also uses the default per-token prefill walk: `K3BoundLayer::decode` advances
100/// its sequence-owned KDA/MLA cache at each absolute position. Its MLA decode
101/// attends that whole retained cache, so it must honor the scheduler chunk budget.
102pub fn requires_single_chunk_prefill(model_type: &str, kv_lora_rank: usize) -> bool {
103    kv_lora_rank > 0 && !matches!(model_type, "glm5_next" | "kimi_k3")
104}
105
106#[cfg(test)]
107mod single_chunk_prefill_tests {
108    use super::requires_single_chunk_prefill as single;
109
110    #[test]
111    fn kimi_k3_keeps_prefill_bounded_beyond_two_chunks() {
112        // Live regression: a 130-token prompt with a 32-token budget first
113        // advances 33 tokens, then the old MLA gate dispatches all 97 remaining
114        // tokens into a 33-token arena. K3's per-token MLA retains its full cache.
115        assert!(!single("kimi_k3", 512));
116        assert!(!single("kimi_k3", 64));
117        for model in ["deepseek_v3", "deepseek_v4", "mistral", "unknown_mla"] {
118            assert!(
119                single(model, 512),
120                "{model} must retain its correctness gate"
121            );
122        }
123    }
124
125    /// GLM-5.3 is MLA and must still be chunked — that is the whole of A61.
126    #[test]
127    fn glm5_next_is_mla_but_chunks_fine() {
128        assert!(!single("glm5_next", 512));
129        // Every other MLA family keeps the single-chunk guard.
130        assert!(single("deepseek_v4", 512));
131        assert!(single("mistral", 512));
132        // Non-MLA models were never gated.
133        assert!(!single("qwen3_5_moe", 0));
134    }
135}
136
137#[cfg(test)]
138mod norm_convention_tests {
139    use super::model_type_ships_vanilla_norm_weights as vanilla;
140
141    /// Only explicitly listed model families take the vanilla path. Every
142    /// other family keeps the offset-from-1 convention it was validated under.
143    #[test]
144    fn vanilla_norm_models_are_explicit() {
145        assert!(vanilla("deepseek_v4"));
146        assert!(vanilla("laguna"));
147        // GLM-5.3's norms are plain; the final norm is applied outside any layer.
148        assert!(vanilla("glm5_next"));
149        assert!(vanilla("kimi_k3"));
150        for other in [
151            "qwen3_next",
152            "qwen3_5_moe",
153            "qwen3_moe",
154            "deepseek_v3",
155            "llama",
156            "mistral",
157            "nemotron",
158            "",
159        ] {
160            assert!(!vanilla(other), "{other} must keep offset-from-1 semantics");
161        }
162    }
163}
spark_model/lib.rs

spark_model/
lib.rs
