sync: whisper support in base-convert (open-core source)

base-convert/Cargo.toml

@@ -12,7 +12,7 @@ members = [
 ]
 
 [workspace.package]
-version = "0.2.0"
+version = "0.2.1"
 edition = "2021"
 license = "Apache-2.0"
 repository = "https://github.com/basecompute/baseRT"

base-convert/crates/base-arch/src/gemma.rs

@@ -430,7 +430,49 @@ impl GgufMapper for Gemma3Mapper {
         } else {
             "gemma"
         };
-        extract_config(m, prefix)
+        let mut config = extract_config(m, prefix)?;
+
+        // Gemma 3 interleaves 5 sliding-window layers per 1 global layer, and
+        // the SWA layers use their OWN rope theta (10000) while global layers
+        // use `rope.freq_base` (1e6). `extract_config` is shared with gemma/
+        // gemma2 and reads neither, so a GGUF-sourced bundle came out with
+        // sliding_window = 0 and rope_local_theta = 0 — i.e. FULL attention at
+        // the global theta on every layer.
+        //
+        // That is silent and severe: greedy generation still looks fine
+        // (argmax is robust and recent context dominates) but teacher-forced
+        // NLL was +2.15 nats vs the HF reference, PPL 168 where converting the
+        // same weights from safetensors gives 19.4 against HF's 19.6.
+        //
+        // gemma3-only: gemma2's window/pattern differ and it has no local
+        // theta, so the fallback prefixes must keep the old behaviour.
+        if prefix == "gemma3" {
+            let u32_key = |k: &str| match m.get(k) {
+                Some(KvValue::U32(v)) => Some(*v),
+                Some(KvValue::U64(v)) => Some(*v as u32),
+                _ => None,
+            };
+            let f32_key = |k: &str| m.get(k).and_then(|v| v.as_f32());
+            // Window size is model-dependent (512 on 1B, 1024 on 4B/12B/27B)
+            // and IS exported — read it rather than assuming.
+            config.sliding_window = u32_key(&format!("{prefix}.attention.sliding_window"))
+                .unwrap_or(config.sliding_window);
+            // Pattern and local theta are Gemma-3 architecture constants that
+            // llama.cpp bakes in and does not export. Prefer a key if a future
+            // exporter adds one; otherwise fall back to the constants, which
+            // match every published gemma-3 config.json.
+            if config.sliding_window > 0 {
+                if config.sliding_window_pattern == 0 {
+                    config.sliding_window_pattern =
+                        u32_key(&format!("{prefix}.attention.sliding_window_pattern")).unwrap_or(6);
+                }
+                if config.rope_local_theta == 0.0 {
+                    config.rope_local_theta =
+                        f32_key(&format!("{prefix}.rope.local.freq_base")).unwrap_or(10_000.0);
+                }
+            }
+        }
+        Ok(config)
     }
 
     fn map_tensor_name(&self, n: &str) -> Option<String> {
@@ -905,6 +947,65 @@ mod tests {
     /// Mirrors the suffix predicate in `convert_hf_to_gguf.py`'s
     /// `Gemma3Model.norm_shift`. Embeddings, projections, and biases
     /// (1-D non-norm) are unaffected.
+    /// A GGUF-sourced gemma-3 bundle must carry the sliding-window config.
+    ///
+    /// `extract_config` is shared with gemma/gemma2 and reads none of these, so
+    /// the mapper used to emit sliding_window = 0 / rope_local_theta = 0 —
+    /// FULL attention at the global theta on all 26 layers. It is silent
+    /// (greedy output still reads fine) but measured +2.15 nats of
+    /// teacher-forced NLL vs the HF reference: PPL 168 from GGUF where the same
+    /// weights converted from safetensors give 19.4 against HF's 19.6.
+    ///
+    /// Window size IS exported and varies by size (512 on 1B, 1024 on 4B+), so
+    /// it must be read, not assumed. Pattern (6) and local theta (10000) are
+    /// architecture constants llama.cpp bakes in and does not export.
+    #[test]
+    fn gemma3_gguf_carries_sliding_window_config() {
+        let mut g: BTreeMap<String, KvValue> = BTreeMap::new();
+        // Exactly the keys gemma-3-1b-it-Q8_0.gguf exports.
+        g.insert("gemma3.embedding_length".into(), KvValue::U32(1152));
+        g.insert("gemma3.block_count".into(), KvValue::U32(26));
+        g.insert("gemma3.attention.head_count".into(), KvValue::U32(4));
+        g.insert("gemma3.attention.head_count_kv".into(), KvValue::U32(1));
+        g.insert("gemma3.attention.key_length".into(), KvValue::U32(256));
+        g.insert("gemma3.attention.value_length".into(), KvValue::U32(256));
+        g.insert("gemma3.feed_forward_length".into(), KvValue::U32(6912));
+        g.insert("gemma3.attention.layer_norm_rms_epsilon".into(), KvValue::F32(1e-6));
+        g.insert("gemma3.rope.freq_base".into(), KvValue::F32(1_000_000.0));
+        g.insert("gemma3.vocab_size".into(), KvValue::U32(262144));
+        g.insert("gemma3.attention.sliding_window".into(), KvValue::U32(512));
+
+        let c = Gemma3Mapper.config_from_gguf(&g).unwrap();
+        assert_eq!(c.sliding_window, 512, "read from the GGUF, not assumed");
+        assert_eq!(c.sliding_window_pattern, 6, "5 SWA : 1 global");
+        assert_eq!(c.rope_local_theta, 10_000.0, "SWA layers use their own theta");
+        assert_eq!(c.rope_theta, 1_000_000.0, "global layers keep freq_base");
+
+        // 4B/12B/27B ship a 1024 window — the value must track the file.
+        g.insert("gemma3.attention.sliding_window".into(), KvValue::U32(1024));
+        assert_eq!(Gemma3Mapper.config_from_gguf(&g).unwrap().sliding_window, 1024);
+    }
+
+    /// The gemma3 defaults must NOT leak onto the gemma2/gemma fallback
+    /// prefixes the same mapper accepts — their window and pattern differ and
+    /// they have no local theta at all.
+    #[test]
+    fn gemma2_gguf_keeps_legacy_behaviour() {
+        let mut g: BTreeMap<String, KvValue> = BTreeMap::new();
+        g.insert("gemma2.embedding_length".into(), KvValue::U32(2304));
+        g.insert("gemma2.block_count".into(), KvValue::U32(26));
+        g.insert("gemma2.attention.head_count".into(), KvValue::U32(8));
+        g.insert("gemma2.attention.head_count_kv".into(), KvValue::U32(4));
+        g.insert("gemma2.feed_forward_length".into(), KvValue::U32(9216));
+        g.insert("gemma2.attention.layer_norm_rms_epsilon".into(), KvValue::F32(1e-6));
+        g.insert("gemma2.vocab_size".into(), KvValue::U32(256000));
+        g.insert("gemma2.attention.sliding_window".into(), KvValue::U32(4096));
+
+        let c = Gemma3Mapper.config_from_gguf(&g).unwrap();
+        assert_eq!(c.sliding_window_pattern, 0, "no gemma3 pattern on gemma2");
+        assert_eq!(c.rope_local_theta, 0.0, "no gemma3 local theta on gemma2");
+    }
+
     #[test]
     fn gemma3_norm_shift_predicate() {
         let m = Gemma3HfMapper;

base-convert/crates/base-arch/src/lib.rs

@@ -11,6 +11,7 @@ pub mod gemma;
 pub mod llama;
 pub mod qwen;
 pub mod tokenizer;
+pub mod whisper;
 
 /// What a converter needs to produce one layer's worth of canonical
 /// tensor names + fuse/route them into the `.base` writer.
@@ -119,6 +120,13 @@ pub fn hf_mapper_for_model_type(model_type: &str) -> Option<&'static dyn HfMappe
         // text conversion skips like any other non-text tower. Config is
         // gemma4-shaped (uniform head_dim, rope_parameters, layer_types).
         "gemma4" | "gemma4_text" | "gemma4_unified" => Some(&gemma::Gemma4HfMapper),
+        // Whisper encoder-decoder speech models (openai/whisper-*). HF
+        // safetensors only — whisper GGML files are not GGUF, so the GGUF
+        // dispatch table stays untouched. Tensor renaming is
+        // whisper-specific (`whisper::map_hf_tensor_name`); the convert
+        // path emits everything as f16 (the engine's fused whisper
+        // kernels are f16-only in v1).
+        "whisper" => Some(&whisper::WhisperHfMapper),
         _ => None,
     }
 }
@@ -148,6 +156,7 @@ pub const SUPPORTED_HF_MODEL_TYPES: &[&str] = &[
     "gemma4",
     "gemma4_text",
     "gemma4_unified",
+    "whisper",
 ];
 
 pub trait GgufMapper: Sync {
@@ -299,6 +308,29 @@ pub struct ArchConfig {
     /// mRoPE interleaves the section frequencies rather than
     /// concatenating them (Qwen3.5 = true).
     pub mrope_interleaved: bool,
+
+    // ── Whisper encoder-decoder fields (zero/empty for other archs) ──
+    // The decoder half reuses the standard fields above (hidden_size /
+    // num_hidden_layers / num_attention_heads / intermediate_size /
+    // max_position_embeddings); the encoder half is described here.
+    /// Encoder transformer depth (`encoder_layers`). 0 = not an
+    /// encoder-decoder model.
+    pub encoder_layers: u32,
+    /// Encoder attention head count (`encoder_attention_heads`).
+    pub encoder_attention_heads: u32,
+    /// Encoder FFN width (`encoder_ffn_dim`).
+    pub encoder_ffn_dim: u32,
+    /// Decoder FFN width (`decoder_ffn_dim`). Mirrors
+    /// `intermediate_size`; emitted under its HF name so the runtime's
+    /// whisper config parser reads the contract key directly.
+    pub decoder_ffn_dim: u32,
+    /// Mel filterbank bin count (`num_mel_bins`, 80 or 128). The engine
+    /// computes the Slaney filterbank from this — it is not stored.
+    pub num_mel_bins: u32,
+    /// Encoder positional-embedding length (`max_source_positions`, 1500).
+    pub max_source_positions: u32,
+    /// Decoder positional-embedding length (`max_target_positions`, 448).
+    pub max_target_positions: u32,
 }
 
 impl ArchConfig {
@@ -471,6 +503,34 @@ impl ArchConfig {
             m.insert("mrope_section".into(), json!(self.mrope_section));
             m.insert("mrope_interleaved".into(), json!(self.mrope_interleaved));
         }
+        // Whisper encoder-decoder fields — only emit when the encoder
+        // half is populated so decoder-only archs' headers stay tidy.
+        if self.encoder_layers > 0 {
+            m.insert("model_type".into(), json!("whisper"));
+            m.insert("encoder_layers".into(), json!(self.encoder_layers));
+            m.insert(
+                "encoder_attention_heads".into(),
+                json!(self.encoder_attention_heads),
+            );
+            m.insert("encoder_ffn_dim".into(), json!(self.encoder_ffn_dim));
+            m.insert("decoder_ffn_dim".into(), json!(self.decoder_ffn_dim));
+            // Whisper's encoder width equals the decoder width; emit the
+            // HF `d_model` key the runtime's whisper config parser reads.
+            m.insert("d_model".into(), json!(self.hidden_size));
+            m.insert("num_mel_bins".into(), json!(self.num_mel_bins));
+            m.insert(
+                "max_source_positions".into(),
+                json!(self.max_source_positions),
+            );
+            m.insert(
+                "max_target_positions".into(),
+                json!(self.max_target_positions),
+            );
+            // Whisper uses LayerNorm, not RMSNorm; the contract key is
+            // `norm_eps` (rms_norm_eps above carries the same value for
+            // struct-level uniformity).
+            m.insert("norm_eps".into(), json!(self.rms_norm_eps));
+        }
         m
     }
 }