metrale_model_layers/
lib.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Model layers: the per-layer forward implementations, LoRA
4//! adapters, weight maps, the SSM reserve, speculative helpers and the
5//! vision/video preprocessing, plus the per-model-type dispatch predicates
6//! below.
7//!
8//! Owner: model-layers.
9//! Invariants: none beyond the types.
10
11#![deny(warnings)]
12#![deny(clippy::all)]
13#![allow(clippy::too_many_arguments)]
14#![allow(clippy::needless_range_loop)]
15#![allow(clippy::if_same_then_else)]
16#![allow(clippy::ptr_arg)]
17#![allow(clippy::type_complexity)]
18
19pub mod circuit_exec;
20pub mod forward;
21pub mod layer;
22pub mod layers;
23pub mod lora;
24pub mod quant_format;
25pub mod speculative;
26pub mod ssm_reserve;
27pub mod video_decode_ffmpeg;
28pub mod video_preprocess;
29pub mod vision_item;
30pub mod vision_preprocess;
31pub mod vision_preprocess_glm;
32/// 2026-09-25: Marconi snapshot-restore threshold: the default, and the
33/// setter `met serve` calls with `--marconi-min-tokens` before it builds the
34/// model.
35pub use crate::mtp_carry::{DEFAULT_MARCONI_MIN_TOKENS, set_marconi_min_tokens};
36pub use vision_item::VisionItem;
37
38pub mod weight_map;
39
40/// 2026-09-26: True when the checkpoint's RMSNorm weights are plain
41/// multipliers, `out = x * w / rms` (the `rms_norm_vanilla` kernel), rather
42/// than offsets from 1, `out = x * (1 + w) / rms` (the `rms_norm` kernel).
43/// It selects the final-norm kernel (model-engine `model/impl_a1.rs`) and the
44/// attention norm kernels (`layers/qwen3_attention/init_proj_kernels.rs` and
45/// `init_prefill_kernels.rs`). The answer comes from `model_type` alone, not
46/// from the weight values.
47pub fn ships_vanilla_norm_weights(config: &metrale_config::ModelConfig) -> bool {
48    model_type_ships_vanilla_norm_weights(&config.model_type)
49}
50
51/// 2026-09-25: The predicate behind `ships_vanilla_norm_weights`, on the bare
52/// `model_type`, so a test needs no `ModelConfig`.
53pub fn model_type_ships_vanilla_norm_weights(model_type: &str) -> bool {
54    // 2026-09-25: The final norm runs outside every layer, so a plain-weight
55    // model missing from this list gets the `(1 + w)` offset on its final hidden
56    // state. Nothing in the weight shapes reveals the convention.
57    matches!(
58        model_type,
59        "deepseek_v4" | "deepseek_v41" | "laguna" | "glm5_next" | "kimi_k3"
60    )
61}
62
63/// 2026-09-25: Must chunked prefill run as one chunk for this model?
64///
65/// True for every MLA model (`kv_lora_rank > 0`) except `glm5_next` and
66/// `kimi_k3`. The Qwen3 attention layer's paged MLA prefill
67/// (`layers/qwen3_attention/prefill/paged_mla.rs`) attends over the K/V of the
68/// call's own tokens only and refuses `seq_len_start != 0`, so a second chunk
69/// cannot run there.
70///
71/// `glm5_next` and `kimi_k3` do not prefill through that layer:
72/// `Glm5NextLayer::prefill` (model-arch `glm5next_layer/mod.rs`) takes
73/// `seq_len_start`, and `K3BoundLayer` (model-arch `kimi_k3/bound.rs`) keeps the
74/// trait's default prefill, which calls `decode` once per token. Both keep the
75/// scheduler's chunk budget.
76pub fn requires_single_chunk_prefill(model_type: &str, kv_lora_rank: usize) -> bool {
77    kv_lora_rank > 0 && !matches!(model_type, "glm5_next" | "kimi_k3")
78}
79
80#[cfg(test)]
81mod single_chunk_prefill_tests {
82    use super::requires_single_chunk_prefill as single;
83
84    #[test]
85    fn kimi_k3_keeps_prefill_bounded_beyond_two_chunks() {
86        assert!(!single("kimi_k3", 512));
87        assert!(!single("kimi_k3", 64));
88        for model in ["deepseek_v3", "deepseek_v4", "mistral", "unknown_mla"] {
89            assert!(
90                single(model, 512),
91                "{model} must retain its correctness gate"
92            );
93        }
94    }
95
96    #[test]
97    fn glm5_next_is_mla_but_chunks_fine() {
98        assert!(!single("glm5_next", 512));
99        assert!(single("deepseek_v4", 512));
100        assert!(single("mistral", 512));
101        assert!(!single("qwen3_5_moe", 0));
102    }
103}
104
105#[cfg(test)]
106mod norm_convention_tests {
107    use super::model_type_ships_vanilla_norm_weights as vanilla;
108
109    /// 2026-09-25: Only the listed model types take the plain-weight norm.
110    #[test]
111    fn vanilla_norm_models_are_explicit() {
112        assert!(vanilla("deepseek_v4"));
113        assert!(vanilla("laguna"));
114        assert!(vanilla("glm5_next"));
115        assert!(vanilla("kimi_k3"));
116        for other in [
117            "qwen3_next",
118            "qwen3_5_moe",
119            "qwen3_moe",
120            "deepseek_v3",
121            "llama",
122            "mistral",
123            "nemotron",
124            "",
125        ] {
126            assert!(!vanilla(other), "{other} must keep offset-from-1 semantics");
127        }
128    }
129}
130
131pub mod drafter_context;
132pub mod mtp_carry;
metrale_model_layers/lib.rs

metrale_model_layers/
lib.rs
