inferencelayer 0.2.3

//! Model weight loading: safetensors (BF16 or F32 on disk) → Q4/f32 GPU buffers, organised per
//! layer. Two architectures share the loader: **LFM2** (conv/attention hybrid) and **Gemma-3 text**
//! (uniform sliding/full attention with sandwich norms, GeGLU, and the `(1+w)` RMSNorm convention —
//! folded into the buffers at load so the kernels stay convention-free).

use crate::GpuCtx;
use anyhow::{Context, Result, bail};
use half::bf16;
use std::path::Path;

/// Which transformer architecture the checkpoint is (drives plan construction + tensor names).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Arch {
    Lfm2,
    Gemma3,
    /// Gemma-4 text: Gemma-3's attention/norm stack + sparse MoE feed-forward (softmax router →
    /// top-k → renormalize → × per-expert scale; routed experts SUMMED with the dense MLP).
    Gemma4,
    /// Qwen-3 dense: standard pre-norm Llama-style stack + per-head qk-RMSNorm. Plain `w` norm
    /// convention (no `(1+w)`), SiLU GLU, single RoPE theta, full attention on every layer, no
    /// sandwich norms (residuals fold into the GEMVs like LFM2), no embedding scale.
    Qwen3,
    /// Llama / Mistral: the SAME pre-norm SwiGLU stack as `Qwen3` with ONE difference — **no
    /// qk-norm** (Llama/Mistral never normalize q/k). Plain `w` RMSNorm, single RoPE theta, NeoX
    /// half-split rope, GQA, no attention bias, no sandwich norms. Covers `LlamaForCausalLM` and
    /// `MistralForCausalLM` (structurally identical here); Mistral's `sliding_window`, when the
    /// config sets it, applies uniformly to every layer with the SAME (global) rope base. Rides
    /// the GLM-OCR no-qk-norm kernel variant (which reads ones q/k-norm buffers but skips the norm).
    Llama,
    /// Qwen2 / Qwen2.5 (`Qwen2ForCausalLM`): `Llama` PLUS q/k/v projection **bias** (the o_proj
    /// stays bias-free). No qk-norm (that is Qwen3's addition). The bias is applied to the fused
    /// qkv output right before qk-norm/RoPE by a small standalone `qkv_bias_add` pass, inserted
    /// into the plan ONLY for this arch (`Op::Attn::qkv_bias = Some`), so every bias-free arch is
    /// byte-identical. `use_sliding_window` + `sliding_window` honored (Qwen2.5 leaves it off).
    Qwen2,
    /// Qwen2-MoE (`Qwen2MoeForCausalLM`, e.g. Qwen1.5-MoE-A2.7B): [`Qwen2`] attention (q/k/v **bias**,
    /// no qk-norm) + a shared-gated sparse MoE structurally IDENTICAL to [`Arch::Qwen35`]'s — routed
    /// experts (`mlp.experts.N.{gate,up,down}_proj`, separate like Mixtral) PLUS an always-computed
    /// `shared_expert` (into the dense GLU slots) scaled by `sigmoid(shared_expert_gate·x)`. So it
    /// reuses BOTH proven pieces: the Qwen2 bias attention and the Qwen3.5 shared-gate MoE forward
    /// (the fused-down path keys off `Moe.shared_gate.is_some()`, not the arch). `intermediate =
    /// shared_expert_intermediate_size` (the shared expert). MoE on every layer (`decoder_sparse_step
    /// = 1`, `mlp_only_layers = []`) — mixed dense/MoE layers are refused.
    Qwen2Moe,
    /// OLMo 1 (`OlmoForCausalLM`, AllenAI OLMo-1B/7B): the `Llama` compute stack (pre-norm SwiGLU,
    /// NeoX rope, GQA, no qk-norm, no biases) EXCEPT the norm is a **non-parametric LayerNorm**
    /// (`elementwise_affine = False`, no weight, no bias) instead of RMSNorm. The engine's first
    /// LayerNorm arch — exploited cheaply: non-parametric LayerNorm(x) = RMSNorm(x − mean(x)), so
    /// the decode plan MEAN-CENTERS the norm input (the `CENTER` kernel) into a scratch and feeds it
    /// to the SAME RMSNorm-fused kernels with ones weights (loaded as ones). Only the M=1 decode
    /// path is wired; batched/spec refuse. `clip_qkv` (present on some OLMo checkpoints) is refused
    /// until ported. See [`arch_is_layernorm`].
    Olmo,
    /// StableLM-2 (`StableLmForCausalLM`, e.g. StableLM-2-1.6B): the Llama stack with **AFFINE
    /// LayerNorm** (weight+bias), q/k/v projection **bias** (Qwen2 mechanism), and **partial rotary**
    /// (`partial_rotary_factor`, the Phi-3 `rotary_dim`), no qk-norm (the 1.6B). The affine LN can't
    /// fold its bias into the RMSNorm-fused kernels, so it runs the `LAYERNORM_AFFINE` kernel (norm
    /// buffers packed `[weight | bias]`) into PLAIN projections + the standalone `MLP_ACT_MUL` — so
    /// no struct change and backend-agnostic. M=1 decode only; batched/spec refuse. See
    /// [`arch_layernorm_affine`].
    StableLm,
    /// Falcon-7B (`FalconForCausalLM`, `new_decoder_architecture = false`, `multi_query = true`):
    /// PARALLEL attn/MLP — one `input_layernorm(x)` feeds both, `x = x + attn(ln) + mlp(ln)` — with
    /// a NON-GATED GELU MLP (`dense_4h_to_h(gelu(dense_h_to_4h(x)))`, exact-erf gelu), MULTIQUERY
    /// (1 kv head), NO biases, and the fused `query_key_value` `[nh·hd | hd | hd]` loaded straight
    /// into the qkv slot (same [q|k|v] layout as `q4cat`, n_kv = 1). Reuses affine LayerNorm
    /// ([`arch_layernorm_affine`]), the parallel path ([`arch_is_parallel`]) and the non-gated MLP
    /// ([`arch_mlp_nogate`]). Falcon-40B's two-LN / grouped-qkv `new_decoder_architecture` is
    /// refused. M=1 decode only; batched/spec refuse.
    Falcon,
    /// Cohere / Command-R (`CohereForCausalLM`): PARALLEL attn/MLP (one `input_layernorm(x)` feeds
    /// both) like [`Arch::Falcon`], but a GATED SwiGLU MLP and NON-parametric-BIAS affine LayerNorm
    /// (weight, NO bias — packed `[weight | zeros]` for LAYERNORM_AFFINE). No biases anywhere; tied
    /// head; full rope; GQA. The one wrinkle — `logit_scale` (logits ×= scale) — is FOLDED into the
    /// LM head at load (Q4-exact under positive scaling), so no runtime step. `use_qk_norm=true`
    /// (Command-R7B's per-head LayerNorm qk-norm) is refused. M=1 decode only; batched/spec refuse.
    Cohere,
    /// Nemotron / Minitron (`NemotronForCausalLM`): Llama-ish but with (a) **LayerNorm1P** — affine
    /// LayerNorm applied as `(1 + weight)` (folded at load, like Gemma3's RMSNorm), packed
    /// `[(weight+1) | bias]`; (b) a **squared-ReLU MLP** `down(relu(up·x)²)` with NO gate — realised
    /// as the gated FFN with w1=w3=up_proj + a ReLU activation (`relu(u)·u = relu(u)²`); (c) **partial
    /// rotary**. No biases (attention_bias / mlp_bias = false). Sequential (pre-norm). M=1 decode
    /// only; batched/spec refuse. See [`arch_act_relu`].
    Nemotron,
    /// Phi-2 (`PhiForCausalLM`): PARALLEL attn/MLP (one `input_layernorm(x)` feeds both) + a
    /// NON-gated `fc2(gelu_new(fc1·x))` MLP (tanh GELU) + partial rotary + AFFINE LayerNorm + BIASES
    /// EVERYWHERE (q/k/v, o_proj/dense, fc1, fc2, untied lm_head). The FFN + lm_head biases are
    /// PACKED into the existing `qkv_bias` / `embedding_norm` buffers (offset slices → BIAS_ADD), so
    /// there's no per-layer struct field. See [`arch_has_biases`]. M=1 decode only; batched/spec refuse.
    Phi2,
    /// OLMo2 (`Olmo2ForCausalLM`): RMSNorm (NOT LayerNorm) but POST-norm — `h = x +
    /// post_attn_norm(attn(x))`, `h = h + post_ff_norm(mlp(h))`, NO pre-norm — plus a FULL-VECTOR
    /// qk-norm (RMSNorm over the whole q `[nh·hd]` and k `[nkv·hd]`, not per-head). Reuses the
    /// sandwich `post_op_norm` / `post_ffn_norm`; `arch_skip_prenorm` makes the projections read the
    /// raw residual and inserts `QK_FULLNORM`. SwiGLU, NeoX rope, GQA, no biases. M=1 decode only;
    /// batched/spec refuse. See [`arch_skip_prenorm`].
    Olmo2,
    /// Mamba-1 (`MambaForCausalLM`): a state-space model, NOT a transformer. Every layer is an
    /// `Op::Mamba` selective-scan mixer (pre-norm RMSNorm + residual); no attention, no KV cache.
    /// `d_inner = expand·hidden`, `dt_rank = ceil(hidden/16)` unless set. M=1 decode only.
    Mamba,
    /// RWKV-4 (`RwkvForCausalLM`): an RNN, NOT a transformer — no attention, no KV cache. Each block
    /// is two token-shift-mixed sub-layers with their OWN affine LayerNorm: a TIME-MIX (`Op::Rwkv`
    /// runs the WKV linear-attention recurrence — a numerically-stable exponential-decay running
    /// weighted sum over `time_decay`/`time_first`, gated by `sigmoid(receptance)`, then `output`)
    /// and a CHANNEL-MIX (`sigmoid(receptance) · value(relu(key(x))²)`). Both mix the current token
    /// with the previous one (`x·mix + x_prev·(1−mix)`); the token-shift state and the WKV state
    /// (num/den/max) persist across tokens. Block 0 has an extra `pre_ln` (ln0) on the embeddings.
    /// UNTIED `head`; final norm is affine LayerNorm. M=1 decode only; batched/spec refuse.
    Rwkv,
    /// Phi-3 (`Phi3ForCausalLM`): the `Llama` compute stack (pre-norm SwiGLU, RMSNorm, NeoX rope,
    /// GQA, no qk-norm, no bias) reached through Phi-3's FUSED projections — `self_attn.qkv_proj`
    /// `[(nh+2nkv)·hd, hidden]` loads straight into the fused qkv slot (already q4cat's layout, and
    /// Phi-3 uses plain rotate_half rope so no de-interleave), and `mlp.gate_up_proj` `[2·inter,
    /// hidden]` splits gate=first-half/up=second-half (the GLM-OCR trick). Untied head is
    /// auto-detected (`lm_head.weight`). `partial_rotary_factor` → `rotary_dim` (full for Phi-3).
    Phi3,
    /// DeepSeek-V2 MLA (`DeepseekV2ForCausalLM`) — Multi-head Latent Attention: q via optional
    /// low-rank (`q_a→q_a_layernorm→q_b`), KV compressed to a `kv_lora_rank` latent + a shared
    /// `qk_rope`-dim `k_pe`, expanded per head by `kv_b`; DECOUPLED rope (GPT-J interleaved-pair) on
    /// the pe dims only. Attention reuses `ATTN_K` (v zero-padded to `qk_head_dim`, o_proj
    /// zero-column-padded). FFN is a dense SwiGLU, OR — when `n_routed_experts > 0` — DeepSeek's
    /// fine-grained MoE: the `mlp.experts` ride the Mixtral pure-routed machinery (softmax→top-k→
    /// renorm, `per_expert_scale = routed_scaling_factor`), and the always-on ungated
    /// `shared_experts` (one pre-concatenated MLP) folds into the dense GLU slots so the
    /// always-added dense path IS their sum. DeepSeek-**V3** (`DeepseekV3ForCausalLM`) shares this
    /// arch and adds only the aux-loss-free router (sigmoid scoring + additive selection bias +
    /// group-limited top-k) — detected from `scoring_func`/`n_group` and served by `MOE_ROUTER_V3`
    /// (see [`V3Router`]). `first_k_dense_replace > 0` (mixed dense/MoE layer sizes) and
    /// `norm_topk_prob = false` (the engine router always renorms) are refused/assumed. See
    /// HANDOFF_MLA.md.
    DeepseekV2,
    /// Qwen3-MoE (`Qwen3MoeForCausalLM`, e.g. Qwen3-30B-A3B / 235B): `Qwen3` attention (per-head
    /// qk-norm, GQA, no bias, single-θ NeoX rope) + the SAME pure sparse MoE as [`Arch::Mixtral`]
    /// (softmax→top-k→renormalize, `norm_topk_prob`, no shared expert, no dense MLP; disabled via a
    /// tiny ZERO dense GLU). Experts ship as separate `mlp.experts.N.{gate,up,down}_proj` tensors.
    Qwen3Moe,
    /// Mixtral (`MixtralForCausalLM`): Mistral attention (no bias, no qk-norm, GQA, sliding
    /// window) + a PURE sparse-MoE feed-forward — 8 experts, top-2, softmax→top-k→renormalize
    /// (the engine's existing router order, `per_expert_scale = 1`), NO shared expert and NO dense
    /// MLP. The engine's MoE always accumulates the routed sum onto a dense/shared slot, so Mixtral
    /// zeroes a TINY (intermediate=32) dense GLU → the dense contribution is exactly 0 and the
    /// block output is purely routed. Experts ship as SEPARATE `block_sparse_moe.experts.N.w{1,2,3}`
    /// tensors (w1=gate, w2=down, w3=up), concatenated into the [`Moe`] slots at load.
    Mixtral,
    /// Granite 3.x (`GraniteForCausalLM`): the `Llama` compute stack + FOUR scalar multipliers,
    /// each FOLDED into weights at load so NO kernel/plan changes are needed —
    /// `embedding_multiplier` × the input embedding rows; `attention_multiplier` × q_proj rows
    /// (·√head_dim, so the kernel's 1/√hd scale yields `(q·k)·attention_multiplier`);
    /// `residual_multiplier` × o_proj & down_proj rows (the residual add then carries the scaled
    /// output); `logits_scaling` ÷ the FINAL norm weight (`logits = head·norm_out/logits_scaling`).
    /// No qk-norm, no bias.
    Granite,
    /// Qwen3.5/3.6 hybrid: Gated DeltaNet linear attention + gated full attention. Covers BOTH
    /// flavors of the family, which differ ONLY in the feed-forward:
    /// - **MoE** (`Qwen3_5Moe*`): routed experts + a shared expert → `num_experts > 0`,
    ///   `Layer::moe = Some(..)`, and the dense GLU slots hold the SHARED expert.
    /// - **Dense** (`Qwen3_5ForConditionalGeneration`, e.g. the claim-extractor 4B): a plain SwiGLU
    ///   MLP → `num_experts == 0`, `Layer::moe = None`, dense GLU slots hold `mlp.{gate,up,down}`.
    ///
    /// Everything else — the hybrid `layer_types` split, `attn_output_gate`, partial rotary,
    /// qk-norm, the DeltaNet geometry — is identical, so the two share every code path here.
    Qwen35,
    /// GLM-OCR's text decoder (`GlmOcrForConditionalGeneration`): a Qwen3-style dense GQA stack
    /// with GLM's differences — sandwich norms on BOTH halves (plain `w`, unlike Gemma's `(1+w)`),
    /// NO decoder qk-norm, a fused `gate_up_proj` (gate = first half), an untied head, and M-RoPE
    /// with CONTIGUOUS `[t|h|w]` sections applied GPT-J interleaved-pair style — converted to the
    /// NeoX kernels by de-interleaving the q/k projection rows at load (scores are invariant when
    /// q and k share the permutation).
    GlmOcr,
    /// Moshi 7B's TEMPORAL transformer (kyutai moshiko; the depformer stays outside): a bias-free
    /// MHA (n_kv == n_heads) SwiGLU stack with FUSED q|k|v (`in_projs.0`, GPT-J interleaved rope
    /// de-interleaved at load like GLM-OCR), gate-FIRST fused `gating.linear_in`, rms_norm_f32
    /// eps 1e-8, sliding window 3000, untied `text_linear` head (logits = the TEXT stream), and
    /// kyutai-native tensor names. Input embeddings are a 17-table sum computed by the CALLER —
    /// serve it via `Lfm2Gpu::forward_from_embeds`, never `forward` (the `Step::Embed` table is
    /// `text_emb` only, a placeholder). CPU twin + trace gates: `moshi_lm.rs`.
    Moshi,
}

/// Whether a checkpoint is a decoder (generation, `Lfm2Gpu`) or an encoder (embedding,
/// `EncoderGpu`) — the top-level dispatch between the two runtimes.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ModelKind {
    /// Autoregressive decoder — flows to [`Lfm2Config::from_json`] / the generation plan.
    Decoder(Arch),
    /// Embedding encoder — flows to
    /// [`EncoderConfig::from_json`](crate::encoder_weights::EncoderConfig::from_json).
    Encoder(crate::encoder_weights::EncArch),
}

/// Whether an architecture normalizes with **LayerNorm** (mean-subtract) rather than RMSNorm.
/// NON-parametric ones (OLMo1) mean-center the norm input so the RMSNorm-fused kernels compute
/// LayerNorm (LayerNorm(x) = RMSNorm(x − mean(x))); AFFINE ones (see [`arch_layernorm_affine`])
/// take the dedicated `LAYERNORM_AFFINE` kernel instead.
pub fn arch_is_layernorm(arch: Arch) -> bool {
    matches!(
        arch,
        Arch::Olmo
            | Arch::StableLm
            | Arch::Falcon
            | Arch::Cohere
            | Arch::Nemotron
            | Arch::Phi2
            | Arch::Rwkv
    )
}

/// Whether a LayerNorm arch has AFFINE norms (weight AND bias). These take the `LAYERNORM_AFFINE`
/// kernel (which computes `(x−mean)/std·w + b` in one pass, reading a norm buffer that packs
/// `[weight | bias]`) feeding PLAIN projections — the `+bias` cannot be folded into the
/// RMSNorm-fused kernels. `false` for OLMo1 (non-parametric, handled by centering).
pub fn arch_layernorm_affine(arch: Arch) -> bool {
    matches!(
        arch,
        Arch::StableLm | Arch::Falcon | Arch::Cohere | Arch::Nemotron | Arch::Phi2 | Arch::Rwkv
    )
}

/// PARALLEL attn/MLP arches: ONE LayerNorm(cur) feeds both the attention and the MLP, which BOTH
/// accumulate into the same residual (no intermediate norm). The decode plan hoists the shared
/// norm to the top of the layer. Falcon-7B (single `input_layernorm`).
pub fn arch_is_parallel(arch: Arch) -> bool {
    matches!(arch, Arch::Falcon | Arch::Cohere | Arch::Phi2)
}

/// NON-GATED MLP arches: `fc2(act(fc1·x))` — no gate / silu·mul (a different MLP than SwiGLU).
/// Realised by loading `fc1→w3`, `fc2→w2` and multiplying the activation by ones. Falcon.
pub fn arch_mlp_nogate(arch: Arch) -> bool {
    matches!(arch, Arch::Falcon | Arch::Phi2)
}

/// Arches whose GELU is the EXACT (erf) form rather than the tanh approximation. Falcon's `gelu`.
pub fn arch_gelu_exact(arch: Arch) -> bool {
    matches!(arch, Arch::Falcon)
}

/// Arches whose MLP activation is ReLU. Nemotron's squared-ReLU (`relu(up·x)²`) is realised on the
/// gated FFN path with both streams bound to the up-projection: `relu(u)·u = relu(u)²`.
pub fn arch_act_relu(arch: Arch) -> bool {
    matches!(arch, Arch::Nemotron)
}

/// Arches with pervasive projection BIASES (Phi-2: q/k/v + o_proj + fc1 + fc2 + untied lm_head).
/// The FFN + lm_head biases are packed into the layer's `qkv_bias` (after [q|k|v]: fc1, then the
/// parallel-combined o_proj+fc2) and into `embedding_norm` (after [weight|bias]: lm_head bias), so
/// the BIAS_ADD kernel reaches them by offset — no per-layer struct field.
pub fn arch_has_biases(arch: Arch) -> bool {
    matches!(arch, Arch::Phi2)
}

/// POST-norm arches (OLMo2): the attn/mlp project the RAW residual (NO pre-norm — RMSNorm weight=1
/// is not identity, so the norm is genuinely skipped), a FULL-VECTOR qk-norm runs on the qkv
/// (`QK_FULLNORM`), and each sub-layer's OUTPUT is RMSNorm'd before the residual add (the sandwich
/// `post_op_norm` / `post_ffn_norm`, which the loader sets).
pub fn arch_skip_prenorm(arch: Arch) -> bool {
    matches!(arch, Arch::Olmo2)
}

/// Classify `config.json` bytes as encoder vs decoder. Encoder families are matched by
/// `architectures[0]`; decoder families dispatch through [`decoder_arch`], which REFUSES
/// unrecognized architectures instead of defaulting (a Llama checkpoint must never silently
/// load as LFM2 and decode garbage).
pub fn detect_model_kind(bytes: &[u8]) -> Result<ModelKind> {
    let v: serde_json::Value = serde_json::from_slice(bytes)?;
    let name = v
        .get("architectures")
        .and_then(|a| a.get(0))
        .and_then(|x| x.as_str());
    if let Some(enc) = name.and_then(crate::encoder_weights::encoder_arch) {
        return Ok(ModelKind::Encoder(enc));
    }
    Ok(ModelKind::Decoder(decoder_arch(name)?))
}

/// The decoder-architecture dispatch shared by [`detect_model_kind`] and
/// [`Lfm2Config::from_json`]. Explicit arms only; an unrecognized `architectures[0]` is a loud
/// error naming the supported set. A MISSING `architectures` key (synthetic/legacy fixtures)
/// keeps the historical LFM2 default — real HF checkpoints always carry the field.
fn decoder_arch(name: Option<&str>) -> Result<Arch> {
    Ok(match name {
        Some("Gemma3ForCausalLM") => Arch::Gemma3,
        Some(a) if a.starts_with("Gemma4") => Arch::Gemma4,
        Some("Qwen3ForCausalLM") => Arch::Qwen3,
        // Llama / Mistral: pre-norm SwiGLU, no qk-norm (see `Arch::Llama`). Both map here; they
        // differ only in `sliding_window` (Mistral sets it), read from config below.
        Some("LlamaForCausalLM" | "MistralForCausalLM") => Arch::Llama,
        // Qwen2 / Qwen2.5: Llama + q/k/v bias (see `Arch::Qwen2`).
        Some("Qwen2ForCausalLM") => Arch::Qwen2,
        // ARK-ASR-3B (`arkasr`): a Whisper-large-v3 encoder + an MLP adapter bolted onto a
        // STOCK Qwen2.5 decoder — `model.layers.*` is Qwen2 tensor-for-tensor (q/k/v bias,
        // biasless o_proj, SwiGLU, two RMSNorms). Only the decoder half loads here; the audio
        // tower is driven separately by `asr_llm.rs`, which splices its output in as embeddings.
        Some("ArkasrForConditionalGeneration") => Arch::Qwen2,
        // Qwen2-MoE: Qwen2 bias attention + Qwen3.5-style shared-gated MoE (see `Arch::Qwen2Moe`).
        Some("Qwen2MoeForCausalLM") => Arch::Qwen2Moe,
        // OLMo 1: Llama stack with non-parametric LayerNorm (see `Arch::Olmo`).
        Some("OlmoForCausalLM") => Arch::Olmo,
        // StableLM-2: affine LayerNorm + qkv bias + partial rotary (see `Arch::StableLm`).
        Some("StableLmForCausalLM") => Arch::StableLm,
        // Falcon-7B: parallel attn/MLP + non-gated GELU MLP + multiquery (see `Arch::Falcon`).
        Some("FalconForCausalLM") => Arch::Falcon,
        // Cohere / Command-R: parallel attn/MLP + gated SwiGLU + logit_scale (see `Arch::Cohere`).
        Some("CohereForCausalLM") => Arch::Cohere,
        // Nemotron / Minitron: LayerNorm1P + squared-ReLU MLP + partial rotary (see `Arch::Nemotron`).
        Some("NemotronForCausalLM") => Arch::Nemotron,
        // Phi-2: parallel attn/MLP + non-gated GELU + biases everywhere (see `Arch::Phi2`).
        Some("PhiForCausalLM") => Arch::Phi2,
        // OLMo2: post-norm RMSNorm + full-vector qk-norm (see `Arch::Olmo2`).
        Some("Olmo2ForCausalLM") => Arch::Olmo2,
        // Mamba-1: selective state-space model, not a transformer (see `Arch::Mamba`).
        Some("MambaForCausalLM") => Arch::Mamba,
        // RWKV-4: an RNN (token-shift + WKV recurrence), not a transformer (see `Arch::Rwkv`).
        Some("RwkvForCausalLM") => Arch::Rwkv,
        // Phi-3: Llama compute via fused qkv_proj + gate_up_proj (see `Arch::Phi3`).
        Some("Phi3ForCausalLM") => Arch::Phi3,
        // Granite 3.x: Llama + four scalar multipliers folded at load (see `Arch::Granite`).
        Some("GraniteForCausalLM") => Arch::Granite,
        // Mixtral: Mistral attention + pure sparse MoE (see `Arch::Mixtral`).
        Some("MixtralForCausalLM") => Arch::Mixtral,
        // Qwen3-MoE: Qwen3 (qk-norm) attention + the Mixtral pure-MoE (see `Arch::Qwen3Moe`).
        Some("Qwen3MoeForCausalLM") => Arch::Qwen3Moe,
        // DeepSeek-V2/V3 MLA (see `Arch::DeepseekV2`). Same attention; V3 differs only in the MoE
        // router (sigmoid + group-limited), detected from `scoring_func`/`n_group` in the loader.
        Some("DeepseekV2ForCausalLM") | Some("DeepseekV3ForCausalLM") => Arch::DeepseekV2,
        // Both flavors of the Qwen3.5 hybrid family: the MoE checkpoints (`Qwen3_5Moe*`) and the
        // DENSE ones (`Qwen3_5ForCausalLM` / `Qwen3_5ForConditionalGeneration` — the multimodal
        // wrapper the claim-extractor 4B ships as). They differ only in the feed-forward; see
        // `Arch::Qwen35`.
        Some(a) if a.starts_with("Qwen3_5") => Arch::Qwen35,
        Some("GlmOcrForConditionalGeneration") => Arch::GlmOcr,
        // Moshi temporal (the config.json is engine-written; kyutai ships no HF architectures key).
        Some("MoshiLmModel") => Arch::Moshi,
        // LFM2 dense decoders (+ the bare backbone: tied head decodes through embed_tokens).
        // Lfm2Moe* is deliberately NOT matched — no MoE-LFM2 arm exists.
        Some("Lfm2ForCausalLM" | "LFM2ForCausalLM" | "Lfm2Model") => Arch::Lfm2,
        None => Arch::Lfm2,
        Some(other) => bail!(
            "unsupported decoder architecture {other:?}; supported: Lfm2ForCausalLM, \
             DeepseekV2ForCausalLM, Gemma3ForCausalLM, Gemma4*, GraniteForCausalLM, \
             MixtralForCausalLM, Phi3ForCausalLM, Qwen2ForCausalLM, Qwen3ForCausalLM, \
             Qwen3MoeForCausalLM, Qwen3_5* (MoE and dense), LlamaForCausalLM, \
             MistralForCausalLM, GlmOcrForConditionalGeneration"
        ),
    })
}

/// Model hyper-parameters needed to drive the kernels (LFM2 or Gemma-3 text).
#[derive(Clone, Debug)]
pub struct Lfm2Config {
    pub arch: Arch,
    pub hidden: usize,
    pub vocab: usize,
    pub n_layers: usize,
    pub intermediate: usize,
    pub n_heads: usize,
    pub n_kv_heads: usize,
    pub head_dim: usize,
    pub conv_l: usize,
    pub eps: f32,
    pub rope_theta: f32,
    /// Gemma-3: RoPE base for the sliding-window layers (global layers use `rope_theta`).
    pub rope_local_theta: f32,
    /// Gemma-3: attention window of the sliding layers (keys `max(0, T-window)..T`).
    pub sliding_window: usize,
    /// Per-layer operator kind: `true` = GQA attention, `false` = gated short-conv (LFM2 only).
    pub layer_is_attn: Vec<bool>,
    /// Per-layer sliding flag (Gemma-3/4; all-false for LFM2).
    pub layer_is_sliding: Vec<bool>,
    /// Per-layer K=V sharing (Gemma-4 `attention_k_eq_v` on GLOBAL layers): no v_proj — V is the
    /// k-projection output (pre-norm), passed through a weightless RMSNorm, NOT roped.
    pub layer_k_eq_v: Vec<bool>,
    /// MLP gate activation is gelu_pytorch_tanh (Gemma-3/4) instead of SiLU (LFM2).
    pub act_gelu: bool,
    /// MoE (Gemma-4): routed experts per MoE layer (0 = dense model).
    pub num_experts: usize,
    /// MoE: experts activated per token.
    pub top_k_experts: usize,
    /// MoE: per-expert feed-forward intermediate width.
    pub moe_intermediate: usize,
    /// Pipeline sharding: this engine owns the embedding (stage 0). Full models: true.
    pub stage_first: bool,
    /// Pipeline sharding: this engine owns the final norm + lm_head. Full models: true.
    pub stage_last: bool,
    /// RoPE rotary dimension (partial RoPE): rotation applies to the FIRST `rotary_dim` dims of
    /// each head, pairs `(j, j + rotary_dim/2)`; the rest pass through. `0` = full RoPE.
    pub rotary_dim: usize,
    /// Interleaved M-RoPE section widths `(t, h, w)` summing to `rotary_dim / 2` — multimodal
    /// checkpoints only. `None` = ordinary 1-D RoPE. Text is unaffected either way (all three axes
    /// carry the same position), so this only bites once an image is in the prompt.
    pub mrope_section: Option<[usize; 3]>,
    /// Gated DeltaNet geometry (Qwen3.5 hybrid layers; zeros on other architectures).
    pub dn_nk: usize,
    pub dn_nv: usize,
    pub dn_dk: usize,
    pub dn_dv: usize,
    pub dn_kernel: usize,
    /// How the BODY projection/MLP matrices are STORED — the speed/fidelity trade, resolved at
    /// load. It lives on the config (rather than on each [`Q4`] container) for two reasons: the
    /// loader's per-architecture `q4` closures all already receive `cfg`, and the config travels
    /// into `Weights.cfg`, so the plan builder can pick the matching kernel family without a
    /// parallel lookup. `Q4` is the default and the historical behaviour, byte-for-byte.
    pub wdtype: WDtype,
    /// How the LM HEAD is stored, independently of the body. The head decodes through a separate
    /// kernel path ([`crate::Q4LmHead`]), so it can be upgraded (e.g. `q4,head=f16` — Unsloth's
    /// canonical recovery lever) without the body plan needing a per-site bind-group choice.
    /// Defaults equal to `wdtype`.
    pub head_wdtype: WDtype,
    /// PER-LAYER body dtype overrides (`blk:` policy terms). EMPTY = uniform `wdtype` everywhere —
    /// the fast path, and the reason the default stays byte-identical. Non-empty: entry `i` is
    /// layer `i`'s dtype; the plan builder picks each layer's kernel family from
    /// [`Self::wdtype_at`] instead of the global.
    pub layer_wdtype: Vec<WDtype>,
    /// Per-layer CONTAINER kinds for native-GGUF serving (empty = derived from the wdtype
    /// policy: Q4 → engine Q4_0 planar, F16 → f16). Non-empty only when `gguf::load_gguf_native`
    /// kept an artifact's own block formats — one layer can then mix e.g. a Q8_0N qkv with a
    /// Q6K down (llama-quantize's ne0 % 256 fallback), which is why the granularity is per SITE,
    /// not per layer.
    pub layer_kinds: Vec<LayerKinds>,
    /// The LM head's container kind when it differs from what `head_wdtype` derives (native GGUF).
    pub head_kind: Option<WKind>,
    /// Gemma-4 EDGE geometry (E2B/E4B) — per-layer head widths, KV sharing, per-layer-embedding
    /// dims, softcap. `EdgeCfg::default()` (all empty/zero) on every other architecture, and the
    /// accessors below collapse to the uniform scalars in that case, so non-edge plans are
    /// byte-for-byte unchanged.
    pub edge: EdgeCfg,
}

/// Gemma-4 edge (E2B/E4B) checkpoint geometry. One struct so non-Gemma construction sites carry a
/// single `EdgeCfg::default()` instead of six dead fields.
#[derive(Clone, Debug, Default)]
pub struct EdgeCfg {
    /// Per-layer head_dim: E2B full-attention layers run 512-wide heads while sliding layers run
    /// 256 (`global_head_dim` vs `head_dim`). EMPTY = uniform `head_dim` everywhere.
    pub layer_head_dim: Vec<usize>,
    /// Per-layer MLP intermediate. `use_double_wide_mlp` doubles it on EXACTLY the KV-shared
    /// layers (HF ties the two: `Gemma4TextMLP.__init__`). EMPTY = uniform `intermediate`.
    pub layer_intermediate: Vec<usize>,
    /// Per-layer KV-cache source: entry `i` is the layer whose K/V cache layer `i` READS.
    /// Own index = the layer computes and stores its own K/V. A consumer (`entry != i`) runs
    /// q_proj/q_norm/q-rope/o_proj ONLY — its on-disk k_proj/v_proj/k_norm tensors are DEAD
    /// (HF `_keys_to_ignore_on_load_unexpected`; presence is not evidence of use). Donors are
    /// the LAST layer of each attention type below `n_layers - num_kv_shared_layers`.
    /// EMPTY = every layer owns its cache.
    pub layer_kv_src: Vec<usize>,
    /// Per-layer-embedding (PLE) width per layer (`hidden_size_per_layer_input`; 256 on E2B).
    /// 0 = no PLE. The table row is `[n_layers · ple_dim]`, layer-major.
    pub ple_dim: usize,
    /// Final-logit softcapping `s·tanh(z/s)` (0 = off). Monotonic, so GREEDY argmax is exempt;
    /// applied on the sampling/logprob paths only.
    pub final_logit_softcapping: f32,
    /// "Proportional" RoPE (full-attention layers): number of ROTATED frequency pairs; the
    /// remaining `head_dim/2 − pairs` pairs get ZERO frequency (cos=1/sin=0 = identity), so the
    /// standard full-width rope kernel needs no partial branch — the TABLE does the masking.
    /// The frequency exponent denominator stays the FULL head_dim (512), not the rotated width.
    /// 0 = ordinary rope.
    pub rope_prop_pairs: usize,
}

impl Lfm2Config {
    /// Layer `i`'s attention head width ([`EdgeCfg::layer_head_dim`], else uniform).
    pub fn head_dim_at(&self, i: usize) -> usize {
        self.edge.layer_head_dim.get(i).copied().unwrap_or(self.head_dim)
    }
    /// The widest head in the model — scratch/table sizing.
    pub fn max_head_dim(&self) -> usize {
        self.edge
            .layer_head_dim
            .iter()
            .copied()
            .max()
            .unwrap_or(self.head_dim)
    }
    /// Layer `i`'s MLP intermediate width ([`EdgeCfg::layer_intermediate`], else uniform).
    pub fn intermediate_at(&self, i: usize) -> usize {
        self.edge
            .layer_intermediate
            .get(i)
            .copied()
            .unwrap_or(self.intermediate)
    }
    /// The widest MLP in the model — scratch sizing.
    pub fn max_intermediate(&self) -> usize {
        self.edge
            .layer_intermediate
            .iter()
            .copied()
            .max()
            .unwrap_or(self.intermediate)
    }
    /// The layer whose KV cache layer `i` reads ([`EdgeCfg::layer_kv_src`], else itself).
    pub fn kv_src_at(&self, i: usize) -> usize {
        self.edge.layer_kv_src.get(i).copied().unwrap_or(i)
    }
    /// Does layer `i` compute and store its own K/V (false = shared-KV consumer, Q-only)?
    pub fn owns_kv(&self, i: usize) -> bool {
        self.kv_src_at(i) == i
    }
    /// Widest FULL-attention head — the global-rope table width (uniform models: `head_dim`).
    pub fn head_dim_full(&self) -> usize {
        (0..self.n_layers)
            .filter(|&i| !self.layer_is_sliding.get(i).copied().unwrap_or(false))
            .map(|i| self.head_dim_at(i))
            .max()
            .unwrap_or(self.head_dim)
    }
    /// Widest SLIDING head — the local-rope table width (uniform models: `head_dim`).
    pub fn head_dim_sliding(&self) -> usize {
        (0..self.n_layers)
            .filter(|&i| self.layer_is_sliding.get(i).copied().unwrap_or(false))
            .map(|i| self.head_dim_at(i))
            .max()
            .unwrap_or(self.head_dim)
    }
    /// Any Gemma-4 edge feature active? The batched/spec plan builders refuse on this until
    /// they are wired for per-layer geometry (the solo M=1 path serves edge models).
    pub fn is_edge(&self) -> bool {
        !self.edge.layer_head_dim.is_empty()
            || !self.edge.layer_kv_src.is_empty()
            || !self.edge.layer_intermediate.is_empty()
            || self.edge.ple_dim > 0
    }
}

/// How ONE tensor's bytes are laid out on the GPU — the CONTAINER format, distinct from the
/// `WDtype` POLICY. `Q4_0`/`F16` are the engine's own packings; the rest are llama.cpp block
/// formats served as-is (padded to word alignment at load, values untouched):
/// [`crate::gemv_q4k_k_lcpp_src`] and friends read them directly.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum WKind {
    /// Engine planar Q4_0: separate f16-scales + nibble buffers (5-binding kernels).
    Q4_0,
    /// Engine f16: plain row-major halves in `quants` (4-binding kernels).
    F16,
    /// Native Q4_K superblocks, 36 words each.
    Q4K,
    /// Native Q6_K superblocks padded to 53 words.
    Q6K,
    /// Native Q5_K superblocks (44 words, naturally aligned).
    Q5K,
    /// Native Q5_0 blocks padded to 6 words.
    Q5_0N,
    /// Native Q8_0 blocks padded to 9 words (also the lossless upcast target for mixed
    /// legacy concatenations: `d8 = d5, q8 = q5 − 16` represents every Q5_0/Q4_0 value exactly).
    Q8_0N,
    /// Native IQ4_NL blocks padded to 5 words (16-entry nonlinear codebook).
    Iq4Nl,
    /// Native IQ4_XS superblocks (34 words, naturally aligned; same codebook, 6-bit sub-scales).
    Iq4Xs,
    /// Native IQ2_XXS superblocks padded to 17 words; the 256x8 grid rides the scales slot.
    Iq2Xxs,
    /// Native IQ2_XS superblocks padded to 19 words; 512x8 grid in the scales slot.
    Iq2Xs,
    /// Native IQ2_S superblocks padded to 21 words; 1024x8 grid in the scales slot.
    Iq2S,
    /// Native IQ3_XXS superblocks padded to 25 words; 256x4 grid in the scales slot.
    Iq3Xxs,
    /// Native IQ3_S superblocks padded to 28 words; 512x4 grid in the scales slot.
    Iq3S,
    /// Native IQ1_S superblocks padded to 13 words; the shared 2048x8 IQ1 grid in the scales slot.
    Iq1S,
    /// Native IQ1_M superblocks (14 words, naturally aligned; same IQ1 grid).
    Iq1M,
}

impl WKind {
    /// Codebook kinds whose GRID buffer rides the `scales` binding — they take the 5-binding
    /// (Q4_0-shaped) bind group, not the 4-binding no-scale-table shape.
    pub fn is_iq_grid(self) -> bool {
        matches!(
            self,
            WKind::Iq2Xxs
                | WKind::Iq2Xs
                | WKind::Iq2S
                | WKind::Iq3Xxs
                | WKind::Iq3S
                | WKind::Iq1S
                | WKind::Iq1M
        )
    }
}

/// One layer's per-site container kinds (the sites the M=1 attention plan binds).
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct LayerKinds {
    pub qkv: WKind,
    pub o: WKind,
    pub gate_up: WKind,
    pub down: WKind,
}

impl LayerKinds {
    pub fn uniform(k: WKind) -> Self {
        Self { qkv: k, o: k, gate_up: k, down: k }
    }
}

impl WKind {
    /// The PACKING dtype that produces this container from f32 weights (safetensors loaders).
    /// Native block kinds only ever come from a GGUF's own bytes — packing to them from f32 is
    /// a different operation (a quantizer), so they panic here rather than silently mis-pack.
    pub fn pack_dtype(self) -> WDtype {
        match self {
            WKind::Q4_0 => WDtype::Q4,
            WKind::F16 => WDtype::F16,
            WKind::Q8_0N => WDtype::Q8,
            other => panic!("{other:?} is a GGUF-native container, not a packing target"),
        }
    }
}

impl Lfm2Config {
    /// Layer `li`'s body weight dtype: the per-layer override when one exists, else the uniform
    /// `wdtype`. The plan builder consults this per layer to pick the kernel family + bind-group
    /// shape for that layer's sites.
    pub fn wdtype_at(&self, li: usize) -> WDtype {
        self.layer_wdtype.get(li).copied().unwrap_or(self.wdtype)
    }

    /// Does ANY body layer store f16? (Gates building the f16 pipelines, the lcpp-family assert,
    /// and the not-yet-wired batch/prefill refusals.)
    pub fn any_f16_body(&self) -> bool {
        self.wdtype == WDtype::F16
            || self.layer_wdtype.iter().any(|d| *d == WDtype::F16)
            || self
                .layer_kinds
                .iter()
                .any(|k| [k.qkv, k.o, k.gate_up, k.down].contains(&WKind::F16))
    }

    /// The container kind the wdtype POLICY derives (used wherever no explicit kind is set).
    pub fn base_kind(d: WDtype) -> WKind {
        match d {
            WDtype::F16 => WKind::F16,
            WDtype::Q8 => WKind::Q8_0N,
            WDtype::Q4 => WKind::Q4_0,
        }
    }

    /// Layer `li`'s per-site container kinds: the explicit native-GGUF kinds when present, else
    /// uniform from the layer's wdtype.
    pub fn kinds_at(&self, li: usize) -> LayerKinds {
        self.layer_kinds
            .get(li)
            .copied()
            .unwrap_or_else(|| LayerKinds::uniform(Self::base_kind(self.wdtype_at(li))))
    }

    /// The LM head's container kind.
    pub fn head_kind(&self) -> WKind {
        self.head_kind.unwrap_or_else(|| Self::base_kind(self.head_wdtype))
    }

    /// Does any body site use a NON-Q4_0 container (f16 or native blocks)? Gates the batch/
    /// prefill refusals and the lcpp-family assert — every non-Q4_0 kernel is lcpp-shaped.
    pub fn any_nonq4_body(&self) -> bool {
        (0..self.n_layers).any(|li| {
            let k = self.kinds_at(li);
            [k.qkv, k.o, k.gate_up, k.down]
                .iter()
                .any(|x| *x != WKind::Q4_0)
        })
    }


    /// Parse `config.json` bytes into the hyper-parameters (incl. block_auto_adjust ff-dim).
    pub fn from_json(bytes: &[u8]) -> Result<Self> {
        let mut c = Self::from_json_inner(bytes)?;
        if c.arch == Arch::Moshi {
            // kyutai's temporal stack is RING attention: every layer attends only the last
            // `context` (3000) steps. The generic parse left window=0 (full attention), which
            // runs the model out of its training distribution the moment a live session
            // crosses 3000 frames — observed as speech degenerating into stuttering loops at
            // exactly the 4-minute mark. Window applies to ALL 32 layers.
            let v: serde_json::Value = serde_json::from_slice(bytes)?;
            let w = v
                .get("context")
                .or_else(|| v.get("sliding_window"))
                .and_then(|x| x.as_u64())
                .unwrap_or(3000) as usize;
            c.sliding_window = w;
            c.layer_is_sliding = vec![true; c.n_layers];
        }
        Ok(c)
    }

    fn from_json_inner(bytes: &[u8]) -> Result<Self> {
        let v: serde_json::Value = serde_json::from_slice(bytes)?;
        let g = |k: &str| v.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
        let arch = decoder_arch(
            v.get("architectures")
                .and_then(|a| a.get(0))
                .and_then(|x| x.as_str()),
        )?;
        if arch == Arch::Qwen35 {
            // Multimodal wrapper: the language model's hyper-parameters live under `text_config`.
            let t = v.get("text_config").unwrap_or(&v);
            let tg = |k: &str| t.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
            let tf = |k: &str| t.get(k).and_then(|x| x.as_f64());
            let hidden = tg("hidden_size").context("hidden_size")?;
            let n_layers = tg("num_hidden_layers").context("num_hidden_layers")?;
            let head_dim = tg("head_dim").context("head_dim")?;
            // layer_types drives the hybrid split: full_attention → true, linear_attention → false.
            let layer_is_attn: Vec<bool> = t
                .get("layer_types")
                .and_then(|x| x.as_array())
                .context("layer_types")?
                .iter()
                .map(|x| x.as_str() == Some("full_attention"))
                .collect();
            anyhow::ensure!(layer_is_attn.len() == n_layers, "layer_types length");
            anyhow::ensure!(
                t.get("attn_output_gate").and_then(|x| x.as_bool()) == Some(true),
                "qwen3.5 without attn_output_gate is untested — refuse loudly"
            );
            let rp = t.get("rope_parameters").unwrap_or(t);
            // The multimodal checkpoints ship M-RoPE (`mrope_section`/`mrope_interleaved`):
            // frequencies are PARTITIONED across the (t, h, w) position axes. On TEXT all three axes
            // carry the same position id, so every frequency sees the same angle p·inv_freq[i]
            // regardless of its section, and M-RoPE degenerates to exactly the 1-D partial rotary —
            // which is why this engine served Qwen3.5-VL text correctly for a long time while
            // reading none of these keys.
            //
            // That equivalence ends the moment an IMAGE is in the prompt: its tokens share one `t`
            // while `h`/`w` walk the merged grid, so the axes disagree and the section decides which
            // frequency rotates by which. The section is now read and honoured (see
            // `forward::mrope_axis_map`). `None` keeps the 1-D path byte-for-byte.
            let mrope_section = rp
                .get("mrope_section")
                .and_then(|x| x.as_array())
                .map(|a| -> Result<[usize; 3]> {
                    anyhow::ensure!(
                        a.len() == 3,
                        "mrope_section must be (t, h, w); got {} entries",
                        a.len()
                    );
                    let mut s = [0usize; 3];
                    for (i, v) in a.iter().enumerate() {
                        s[i] = v.as_u64().context("mrope_section entry")? as usize;
                    }
                    Ok(s)
                })
                .transpose()?;
            // Only the INTERLEAVED layout is implemented. A chunked-M-RoPE checkpoint would need a
            // different axis map, and silently applying the interleaved one would rotate the right
            // frequencies by the wrong axis — fluent output, skewed geometry. Refuse instead.
            if mrope_section.is_some() {
                anyhow::ensure!(
                    rp.get("mrope_interleaved").and_then(|x| x.as_bool()) == Some(true),
                    "chunked (non-interleaved) M-RoPE is unimplemented — refuse loudly"
                );
            }
            let partial = rp
                .get("partial_rotary_factor")
                .and_then(|x| x.as_f64())
                .or_else(|| tf("partial_rotary_factor"))
                .unwrap_or(1.0);
            // MoE vs dense is decided by the config's own keys, not by the architecture string:
            // the MoE checkpoints carry `num_experts` (+ a shared expert), the dense ones carry a
            // plain `intermediate_size` and nothing else. `num_experts == 0` is the engine's
            // established "dense" signal (Gemma-4 already uses it), and `Layer::moe` stays `None`.
            let num_experts = tg("num_experts").unwrap_or(0);
            let is_moe = num_experts > 0;
            let intermediate = if is_moe {
                // The dense GLU slots hold the SHARED expert (512 wide).
                tg("shared_expert_intermediate_size").context("shared_expert")?
            } else {
                tg("intermediate_size").context("intermediate_size")?
            };
            return Ok(Self {
                arch,
                hidden,
                vocab: tg("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate,
                n_heads: tg("num_attention_heads").context("num_attention_heads")?,
                n_kv_heads: tg("num_key_value_heads").context("num_key_value_heads")?,
                head_dim,
                conv_l: 0,
                eps: tf("rms_norm_eps").unwrap_or(1e-6) as f32,
                rope_theta: rp
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .or_else(|| tf("rope_theta"))
                    .unwrap_or(10_000_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn,
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts,
                top_k_experts: if is_moe {
                    tg("num_experts_per_tok").context("num_experts_per_tok")?
                } else {
                    0
                },
                moe_intermediate: if is_moe {
                    tg("moe_intermediate_size").context("moe_intermediate_size")?
                } else {
                    0
                },
                stage_first: true,
                stage_last: true,
                rotary_dim: (head_dim as f64 * partial) as usize,
                mrope_section,
                dn_nk: tg("linear_num_key_heads").context("linear_num_key_heads")?,
                dn_nv: tg("linear_num_value_heads").context("linear_num_value_heads")?,
                dn_dk: tg("linear_key_head_dim").context("linear_key_head_dim")?,
                dn_dv: tg("linear_value_head_dim").context("linear_value_head_dim")?,
                dn_kernel: tg("linear_conv_kernel_dim").context("linear_conv_kernel_dim")?,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::GlmOcr {
            // Multimodal wrapper: the text stack's hyper-parameters live under `text_config`,
            // rope under `rope_parameters` (theta 10_000, mrope_section [16, 24, 24]).
            let t = v.get("text_config").unwrap_or(&v);
            let tg = |k: &str| t.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
            let hidden = tg("hidden_size").context("hidden_size")?;
            let n_heads = tg("num_attention_heads").context("num_attention_heads")?;
            let n_layers = tg("num_hidden_layers").context("num_hidden_layers")?;
            let rp = t.get("rope_parameters").unwrap_or(t);
            let mrope_section = rp
                .get("mrope_section")
                .and_then(|x| x.as_array())
                .map(|a| -> Result<[usize; 3]> {
                    anyhow::ensure!(a.len() == 3, "mrope_section must be (t, h, w)");
                    let mut sct = [0usize; 3];
                    for (i, x) in a.iter().enumerate() {
                        sct[i] = x.as_u64().context("mrope_section entry")? as usize;
                    }
                    Ok(sct)
                })
                .transpose()?;
            return Ok(Self {
                arch,
                hidden,
                vocab: tg("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: tg("intermediate_size").context("intermediate_size")?,
                n_heads,
                n_kv_heads: tg("num_key_value_heads").unwrap_or(n_heads),
                head_dim: tg("head_dim").unwrap_or(hidden / n_heads),
                conv_l: 0,
                eps: t
                    .get("rms_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-5) as f32,
                rope_theta: rp
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(10_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Qwen2Moe {
            // Qwen2 attention (bias, no qk-norm) + Qwen3.5-style shared-gated MoE. The shared expert
            // fills the dense GLU slots, so `intermediate = shared_expert_intermediate_size`. MoE is
            // on EVERY layer (decoder_sparse_step = 1, mlp_only_layers = []) — anything else (mixed
            // dense/MoE) would need per-layer slot sizing, refused here.
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            anyhow::ensure!(
                g("decoder_sparse_step").unwrap_or(1) == 1,
                "Qwen2-MoE decoder_sparse_step != 1 (mixed dense/MoE layers) not yet supported"
            );
            anyhow::ensure!(
                v.get("mlp_only_layers")
                    .and_then(|x| x.as_array())
                    .is_none_or(|a| a.is_empty()),
                "Qwen2-MoE mlp_only_layers (mixed dense/MoE layers) not yet supported"
            );
            anyhow::ensure!(
                v.get("norm_topk_prob")
                    .and_then(|x| x.as_bool())
                    .unwrap_or(true),
                "Qwen2-MoE norm_topk_prob = false (engine router always renormalizes) not supported"
            );
            let use_sliding = v
                .get("use_sliding_window")
                .and_then(|x| x.as_bool())
                .unwrap_or(false);
            let sliding = if use_sliding {
                v.get("sliding_window")
                    .and_then(|x| x.as_u64())
                    .unwrap_or(0) as usize
            } else {
                0
            };
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: g("shared_expert_intermediate_size")
                    .context("shared_expert_intermediate_size")?,
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim: g("head_dim").unwrap_or(hidden / n_heads),
                conv_l: 0,
                eps: v
                    .get("rms_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-6) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1_000_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: sliding,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![sliding > 0; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts: g("num_experts").context("num_experts")?,
                top_k_experts: g("num_experts_per_tok").context("num_experts_per_tok")?,
                moe_intermediate: g("moe_intermediate_size").context("moe_intermediate_size")?,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Olmo {
            // OLMo 1: Llama stack with NON-PARAMETRIC LayerNorm (no norm weights in the checkpoint;
            // the loader feeds ones). `clip_qkv` (qkv clamp) not yet ported.
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            anyhow::ensure!(
                v.get("clip_qkv").is_none_or(|x| x.is_null()),
                "OLMo clip_qkv (qkv clamp) not yet supported"
            );
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: g("intermediate_size").context("intermediate_size")?,
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim: hidden / n_heads,
                conv_l: 0,
                // Non-parametric LayerNorm eps: OLMo defaults 1e-5; read either key if present.
                eps: v
                    .get("layer_norm_eps")
                    .or_else(|| v.get("rms_norm_eps"))
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-5) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(10_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::StableLm {
            // StableLM-2: affine LayerNorm + q/k/v bias + partial rotary. `qk_layernorm` (the 12B)
            // is a separate qk-norm variant, not yet ported.
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            let head_dim = g("head_dim").unwrap_or(hidden / n_heads);
            anyhow::ensure!(
                !v.get("qk_layernorm")
                    .and_then(|x| x.as_bool())
                    .unwrap_or(false),
                "StableLM qk_layernorm (the 12B qk-norm variant) not yet supported"
            );
            // partial_rotary_factor → rotary_dim (Phi-3's convention; full ⇒ 0).
            let rotary_dim = {
                let f = v
                    .get("partial_rotary_factor")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1.0);
                let rd = (head_dim as f64 * f) as usize;
                if rd >= head_dim { 0 } else { rd }
            };
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: g("intermediate_size").context("intermediate_size")?,
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim,
                conv_l: 0,
                eps: v
                    .get("layer_norm_eps")
                    .or_else(|| v.get("rms_norm_eps"))
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-5) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(10_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Falcon {
            // Falcon-7B: parallel attn/MLP, multiquery (1 kv head), non-gated GELU MLP, no biases.
            // The 40B `new_decoder_architecture` (two LNs, grouped qkv) and biased variants are
            // refused. MLP is 4·hidden (Falcon has no `intermediate_size`).
            anyhow::ensure!(
                !v.get("new_decoder_architecture")
                    .and_then(|x| x.as_bool())
                    .unwrap_or(false),
                "Falcon new_decoder_architecture (40B two-LN / grouped qkv) not yet supported"
            );
            anyhow::ensure!(
                !v.get("bias").and_then(|x| x.as_bool()).unwrap_or(false),
                "Falcon bias=true not yet supported"
            );
            anyhow::ensure!(
                v.get("parallel_attn")
                    .and_then(|x| x.as_bool())
                    .unwrap_or(true),
                "Falcon parallel_attn=false (sequential) not yet supported"
            );
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            let multi_query = v
                .get("multi_query")
                .and_then(|x| x.as_bool())
                .unwrap_or(true);
            let n_kv = if multi_query {
                1
            } else {
                g("num_kv_heads")
                    .or_else(|| g("n_head_kv"))
                    .unwrap_or(n_heads)
            };
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: g("intermediate_size").unwrap_or(4 * hidden),
                n_heads,
                n_kv_heads: n_kv,
                head_dim: hidden / n_heads,
                conv_l: 0,
                eps: v
                    .get("layer_norm_epsilon")
                    .or_else(|| v.get("layer_norm_eps"))
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-5) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(10_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false, // exact-erf gelu via arch_gelu_exact, not the tanh act_gelu flag
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Cohere {
            // Cohere / Command-R: parallel attn/MLP + gated SwiGLU + affine (weight-only) LayerNorm.
            // `logit_scale` is folded into the LM head at load. `use_qk_norm` (Command-R7B) refused.
            anyhow::ensure!(
                !v.get("use_qk_norm")
                    .and_then(|x| x.as_bool())
                    .unwrap_or(false),
                "Cohere use_qk_norm (Command-R7B per-head LayerNorm qk-norm) not yet supported"
            );
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: g("intermediate_size").context("intermediate_size")?,
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim: g("head_dim").unwrap_or(hidden / n_heads),
                conv_l: 0,
                eps: v
                    .get("layer_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-5) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(10_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false, // gated SwiGLU (silu)
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Nemotron {
            // Nemotron / Minitron: LayerNorm1P + squared-ReLU MLP + partial rotary, no biases.
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            let head_dim = g("head_dim").unwrap_or(hidden / n_heads);
            let rotary_dim = {
                let f = v
                    .get("partial_rotary_factor")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(0.5);
                let rd = (head_dim as f64 * f) as usize;
                if rd >= head_dim { 0 } else { rd }
            };
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: g("intermediate_size").context("intermediate_size")?,
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim,
                conv_l: 0,
                eps: v
                    .get("norm_eps")
                    .or_else(|| v.get("layer_norm_eps"))
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-5) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(10_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false, // ReLU via arch_act_relu
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Phi2 {
            // Phi-2: parallel attn/MLP + non-gated tanh-GELU MLP + partial rotary + biases.
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            let head_dim = g("head_dim").unwrap_or(hidden / n_heads);
            let rotary_dim = {
                let f = v
                    .get("partial_rotary_factor")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(0.4);
                let rd = (head_dim as f64 * f) as usize;
                if rd >= head_dim { 0 } else { rd }
            };
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: g("intermediate_size").context("intermediate_size")?,
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim,
                conv_l: 0,
                eps: v
                    .get("layer_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-5) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(10_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: true, // gelu_new (tanh) → MLP_ACT_MUL mode 1
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Olmo2 {
            // OLMo2: Llama config shape; the post-norm + full-vector qk-norm live in the loader/plan.
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: g("intermediate_size").context("intermediate_size")?,
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim: g("head_dim").unwrap_or(hidden / n_heads),
                conv_l: 0,
                eps: v
                    .get("rms_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-6) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(500_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false, // SwiGLU (silu)
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Mamba {
            // Mamba: no attention (dummy head dims); the mixer dims are re-read in the loader. The
            // per-layer RMSNorm rides `operator_norm`; there is no rope / KV cache.
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: 32, // tiny ZERO dense GLU → the post-mixer FFN slot adds 0 (no FFN)
                n_heads: 1,
                n_kv_heads: 1,
                head_dim: hidden, // dummy: keeps attention scratch safely sized, unused on Mamba
                conv_l: 0,
                eps: v
                    .get("layer_norm_epsilon")
                    .or_else(|| v.get("rms_norm_eps"))
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-5) as f32,
                rope_theta: 10_000.0,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![false; n_layers], // SSM layers, no KV cache
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Rwkv {
            // RWKV-4: an RNN. No attention (dummy head dims); the per-block affine LayerNorms ride
            // operator_norm/ffn_norm and the mixer dims are re-read in the loader. The trailing GLU
            // FFN slot is a tiny ZERO GLU (adds 0 — the channel-mix runs inside Op::Rwkv).
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: 32, // tiny ZERO dense GLU → the trailing FFN slot adds 0
                n_heads: 1,
                n_kv_heads: 1,
                head_dim: hidden, // dummy: keeps attention scratch safely sized, unused on RWKV
                conv_l: 0,
                eps: v
                    .get("layer_norm_epsilon")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-5) as f32,
                rope_theta: 10_000.0,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![false; n_layers], // RNN, no KV cache
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if matches!(
            arch,
            Arch::Qwen3 | Arch::Llama | Arch::Qwen2 | Arch::Phi3 | Arch::Granite
        ) {
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            let head_dim = g("head_dim").unwrap_or(hidden / n_heads);
            // Mistral sets `sliding_window` (Llama leaves it null); when present it applies to
            // EVERY layer with the global rope base (no separate local theta, unlike Gemma).
            // Qwen2 gates it behind `use_sliding_window` (Qwen2.5 ships it false).
            let use_sliding = arch != Arch::Qwen2
                || v.get("use_sliding_window")
                    .and_then(|x| x.as_bool())
                    .unwrap_or(false);
            let sliding = if use_sliding {
                v.get("sliding_window")
                    .and_then(|x| x.as_u64())
                    .unwrap_or(0) as usize
            } else {
                0
            };
            // Phi-3 `partial_rotary_factor` → rotary_dim (full → 0, matching the kernel's "0 =
            // full" convention). Other archs in this arm are always full rotary.
            let rotary_dim = if arch == Arch::Phi3 {
                let f = v
                    .get("partial_rotary_factor")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1.0);
                let rd = (head_dim as f64 * f) as usize;
                if rd >= head_dim { 0 } else { rd }
            } else {
                0
            };
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: g("intermediate_size").context("intermediate_size")?,
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim,
                conv_l: 0,
                eps: v
                    .get("rms_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-6) as f32,
                // Qwen3 defaults 1e6; Llama-3 = 5e5, Llama-2/Mistral/Phi-3 vary — real checkpoints
                // always carry `rope_theta`, so the default only affects synthetic fixtures.
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1_000_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: sliding,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![sliding > 0; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim,
                mrope_section: None, // text-only architecture: 1-D rope, unchanged
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Mixtral {
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            let sliding = v
                .get("sliding_window")
                .and_then(|x| x.as_u64())
                .unwrap_or(0) as usize;
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                // TINY dense GLU: Mixtral has NO dense MLP — the loader zeroes these slots so the
                // (always-computed) dense path contributes exactly 0 and the block is pure-routed.
                intermediate: 32,
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim: g("head_dim").unwrap_or(hidden / n_heads),
                conv_l: 0,
                eps: v
                    .get("rms_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-5) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1_000_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: sliding,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![sliding > 0; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts: g("num_local_experts").context("num_local_experts")?,
                top_k_experts: g("num_experts_per_tok").context("num_experts_per_tok")?,
                moe_intermediate: g("intermediate_size").context("intermediate_size")?,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Qwen3Moe {
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: 32, // tiny zero dense GLU → pure routed (see `Arch::Mixtral`)
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim: g("head_dim").unwrap_or(hidden / n_heads),
                conv_l: 0,
                eps: v
                    .get("rms_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-6) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1_000_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts: g("num_experts").context("num_experts")?,
                top_k_experts: g("num_experts_per_tok").context("num_experts_per_tok")?,
                moe_intermediate: g("moe_intermediate_size").context("moe_intermediate_size")?,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::DeepseekV2 {
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            // MLA expands to full MHA: head_dim = qk_nope + qk_rope, n_kv_heads = n_heads. The
            // MLA-specific dims live in `Op::Mla` (read in the loader), not here.
            let head_dim = g("qk_nope_head_dim").context("qk_nope_head_dim")?
                + g("qk_rope_head_dim").context("qk_rope_head_dim")?;
            // MoE: n_routed_experts > 0. The always-on shared experts fold into the DENSE GLU slots
            // (their concatenation = their sum), so `intermediate = n_shared·moe_intermediate`; the
            // routed experts ride the `Moe` struct. `first_k_dense_replace > 0` (mixed dense/MoE
            // layers with DIFFERENT dense sizes) is not yet handled — refused loudly.
            let n_routed = g("n_routed_experts").unwrap_or(0);
            let is_moe = n_routed > 0;
            if is_moe {
                anyhow::ensure!(
                    g("first_k_dense_replace").unwrap_or(0) == 0,
                    "DeepSeek first_k_dense_replace > 0 (mixed dense/MoE layer sizes) not yet supported"
                );
                // Both router flavors renormalize the selected top-k weights; a checkpoint that
                // disables this (e.g. DeepSeek-V2-Lite) would be silently mis-scored — refuse it.
                anyhow::ensure!(
                    v.get("norm_topk_prob")
                        .and_then(|x| x.as_bool())
                        .unwrap_or(true),
                    "DeepSeek norm_topk_prob = false (router must renormalize) not yet supported"
                );
                let scoring = v
                    .get("scoring_func")
                    .and_then(|x| x.as_str())
                    .unwrap_or("softmax");
                let method = v
                    .get("topk_method")
                    .and_then(|x| x.as_str())
                    .unwrap_or("greedy");
                if scoring == "sigmoid" {
                    // DeepSeek-V3 aux-loss-free router: sigmoid scoring + additive selection bias +
                    // group-limited top-k (`noaux_tc`). Supported via MOE_ROUTER_V3 — needs the group
                    // params to be present so the loader can build the V3 router.
                    anyhow::ensure!(
                        g("n_group").unwrap_or(0) > 0 && g("topk_group").unwrap_or(0) > 0,
                        "DeepSeek-V3 sigmoid router needs n_group / topk_group"
                    );
                } else {
                    // V2: plain softmax greedy. `group_limited_greedy` (V2 with groups) is a
                    // different selection and stays refused until ported.
                    anyhow::ensure!(
                        scoring == "softmax",
                        "DeepSeek scoring_func = {scoring:?} (expected softmax or sigmoid)"
                    );
                    anyhow::ensure!(
                        method == "greedy",
                        "DeepSeek softmax topk_method = {method:?} (only plain greedy supported)"
                    );
                }
            }
            let moe_inter = g("moe_intermediate_size").unwrap_or(0);
            let n_shared = g("n_shared_experts").unwrap_or(0);
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: if is_moe {
                    n_shared * moe_inter // shared experts, concatenated into the dense GLU
                } else {
                    g("intermediate_size").context("intermediate_size")?
                },
                n_heads,
                n_kv_heads: n_heads,
                head_dim,
                conv_l: 0,
                eps: v
                    .get("rms_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-6) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(10_000.0) as f32,
                rope_local_theta: 10_000.0,
                sliding_window: 0,
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding: vec![false; n_layers],
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: false,
                num_experts: n_routed,
                top_k_experts: if is_moe {
                    g("num_experts_per_tok").context("num_experts_per_tok")?
                } else {
                    0
                },
                moe_intermediate: moe_inter,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None,
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        if arch == Arch::Gemma4 {
            // Gemma-4 = the Gemma-3 stack + MoE feed-forward + the EDGE features (E2B/E4B):
            // per-layer-type head geometry, KV-shared layers with a double-wide MLP on exactly
            // those layers, per-layer embeddings (PLE), "proportional" rope, final-logit softcap.
            // Features no accessible checkpoint exercises are still refused loudly below.
            //
            // Multimodal wrapper (Gemma4ForConditionalGeneration): the text stack lives under
            // `text_config`; bare Gemma4* text checkpoints (and the synthetic fixtures) keep
            // root-level keys.
            let t = v.get("text_config").unwrap_or(&v);
            let g = |k: &str| t.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            let layer_is_sliding = match t.get("layer_types").and_then(|x| x.as_array()) {
                Some(a) => {
                    // A malformed pattern would mis-window and mis-rope every layer; the real
                    // E2B pattern is 4:1 (full at 4,9,…), NOT the library-default 5:1 — always
                    // read the checkpoint's list, never synthesize.
                    anyhow::ensure!(
                        a.len() == n_layers,
                        "layer_types length != num_hidden_layers"
                    );
                    a.iter()
                        .map(|t| t.as_str() == Some("sliding_attention"))
                        .collect::<Vec<_>>()
                }
                None => vec![false; n_layers],
            };
            let moe_on = t
                .get("enable_moe_block")
                .and_then(|x| x.as_bool())
                .unwrap_or(false);
            let head_dim = g("head_dim").unwrap_or(hidden / n_heads);
            // Per-layer-TYPE head width: full-attention layers run `global_head_dim`-wide heads
            // (E2B: 512 vs 256). HF defaults global_head_dim to 512 even when the key is ABSENT
            // (`Gemma4TextConfig.__post_init__` pops it from kwargs — it is not a declared
            // field); this parse deliberately defaults to `head_dim` instead: every real
            // Gemma-4 checkpoint ships the key explicitly, and the HF default would silently
            // widen the synthetic/legacy fixtures' full layers.
            let ghd = g("global_head_dim").filter(|&x| x != 0).unwrap_or(head_dim);
            let layer_head_dim: Vec<usize> = if ghd != head_dim {
                layer_is_sliding
                    .iter()
                    .map(|&s| if s { head_dim } else { ghd })
                    .collect()
            } else {
                Vec::new()
            };
            let n_kv = g("num_key_value_heads").unwrap_or(n_heads);
            let k_eq_v = t
                .get("attention_k_eq_v")
                .and_then(|x| x.as_bool())
                .unwrap_or(false);
            // HF gates the global-KV-head override on `attention_k_eq_v` — false (E2B) means
            // the field is IGNORED. True (31B) makes it real work (4 kv heads on full layers)
            // that no checkpoint on this fleet exercises yet: refuse.
            anyhow::ensure!(
                !k_eq_v || g("num_global_key_value_heads").is_none_or(|gk| gk == 0 || gk == n_kv),
                "Gemma-4 num_global_key_value_heads != num_key_value_heads not yet supported"
            );
            let layer_k_eq_v: Vec<bool> = layer_is_sliding.iter().map(|s| k_eq_v && !s).collect();
            // KV sharing: the LAST `num_kv_shared_layers` layers read the cache of the last
            // layer of their own attention type below the shared region (E2B: sliding→13,
            // full→14) and do NO K/V work of their own. Their on-disk k_proj/v_proj/k_norm are
            // DEAD tensors (HF drops them via `_keys_to_ignore_on_load_unexpected`).
            let nks = g("num_kv_shared_layers").unwrap_or(0);
            anyhow::ensure!(nks < n_layers, "num_kv_shared_layers must leave a donor prefix");
            let first_shared = n_layers - nks;
            let layer_kv_src: Vec<usize> = if nks == 0 {
                Vec::new()
            } else {
                (0..n_layers)
                    .map(|i| -> Result<usize> {
                        if i < first_shared {
                            return Ok(i);
                        }
                        (0..first_shared)
                            .rev()
                            .find(|&d| layer_is_sliding[d] == layer_is_sliding[i])
                            .with_context(|| {
                                format!("KV-shared layer {i} has no donor of its attention type")
                            })
                    })
                    .collect::<Result<_>>()?
            };
            let intermediate = g("intermediate_size").context("intermediate_size")?;
            // `use_double_wide_mlp` doubles the MLP on EXACTLY the KV-shared layers — HF ties
            // the two in `Gemma4TextMLP.__init__` (the layers that skip KV work get the width).
            let dw = t
                .get("use_double_wide_mlp")
                .and_then(|x| x.as_bool())
                .unwrap_or(false);
            let layer_intermediate: Vec<usize> = if dw && nks > 0 {
                (0..n_layers)
                    .map(|i| intermediate * if i >= first_shared { 2 } else { 1 })
                    .collect()
            } else {
                Vec::new()
            };
            // PLE (per-layer embeddings): table row = [n_layers · ple_dim], layer-major.
            let ple_dim = g("hidden_size_per_layer_input").unwrap_or(0);
            let vocab = g("vocab_size").context("vocab_size")?;
            if ple_dim > 0 {
                anyhow::ensure!(
                    g("vocab_size_per_layer_input").unwrap_or(vocab) == vocab,
                    "Gemma-4 PLE vocab != main vocab not yet supported"
                );
                anyhow::ensure!(
                    !moe_on,
                    "Gemma-4 PLE + MoE together are untested — refuse loudly"
                );
            }
            // Rope: real checkpoints ship a per-layer-type `rope_parameters` dict; legacy flat
            // keys (and the synthetic fixtures) fall back. Full-attention "proportional" rope
            // keeps `prf·ghd/2` REAL frequencies of the full-width schedule (the exponent
            // denominator stays ghd, NOT the rotated width) and zero-pads the rest — zero
            // frequency is cos=1/sin=0 = identity, so the standard full-width rope kernel
            // needs no partial branch; the TABLE does the masking.
            let rp_full = t
                .get("rope_parameters")
                .and_then(|r| r.get("full_attention"));
            let rp_slid = t
                .get("rope_parameters")
                .and_then(|r| r.get("sliding_attention"));
            let rope_theta = rp_full
                .and_then(|r| r.get("rope_theta"))
                .or_else(|| t.get("rope_theta"))
                .and_then(|x| x.as_f64())
                .unwrap_or(1_000_000.0) as f32;
            let rope_local_theta = rp_slid
                .and_then(|r| r.get("rope_theta"))
                .or_else(|| t.get("rope_local_base_freq"))
                .and_then(|x| x.as_f64())
                .unwrap_or(10_000.0) as f32;
            let prf = rp_full
                .and_then(|r| r.get("partial_rotary_factor"))
                .and_then(|x| x.as_f64())
                .unwrap_or(0.0);
            let rope_prop_pairs = if prf > 0.0 && prf < 1.0 {
                ((prf * ghd as f64) as usize) / 2
            } else {
                0
            };
            let final_logit_softcapping = t
                .get("final_logit_softcapping")
                .and_then(|x| x.as_f64())
                .unwrap_or(0.0) as f32;
            return Ok(Self {
                arch,
                hidden,
                vocab,
                n_layers,
                intermediate,
                n_heads,
                n_kv_heads: n_kv,
                head_dim,
                conv_l: 0,
                eps: t
                    .get("rms_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-6) as f32,
                rope_theta,
                rope_local_theta,
                sliding_window: g("sliding_window").unwrap_or(512),
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding,
                layer_k_eq_v,
                act_gelu: true,
                num_experts: if moe_on {
                    g("num_experts").unwrap_or(0)
                } else {
                    0
                },
                top_k_experts: if moe_on {
                    g("top_k_experts").unwrap_or(0)
                } else {
                    0
                },
                moe_intermediate: if moe_on {
                    g("moe_intermediate_size").unwrap_or(0)
                } else {
                    0
                },
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None, // text-only architecture: 1-D rope, unchanged
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
                head_wdtype: WDtype::Q4,
                layer_wdtype: Vec::new(),
                layer_kinds: Vec::new(),
                head_kind: None,
                edge: EdgeCfg {
                    layer_head_dim,
                    layer_intermediate,
                    layer_kv_src,
                    ple_dim,
                    final_logit_softcapping,
                    rope_prop_pairs,
                },
            });
        }
        if arch == Arch::Gemma3 {
            let hidden = g("hidden_size").context("hidden_size")?;
            let n_heads = g("num_attention_heads").context("num_attention_heads")?;
            let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
            // Sliding/full layout drives BOTH the attention window and the per-layer RoPE
            // theta (local vs global), so a wrong layout silently degrades quality at any
            // sequence length. Two checkpoint conventions exist: an explicit `layer_types`
            // array (source of truth when present — gemma-3-270m, the osfql fine-tunes) and
            // `sliding_window_pattern` (HF-vintage gemma-3-1b-it: layer i is sliding unless
            // `(i+1) % pattern == 0`, per transformers' Gemma3DecoderLayer). Neither present
            // ⇒ refuse loudly rather than mis-rope 5/6 of the layers with the global theta.
            let layer_is_sliding = match (
                v.get("layer_types").and_then(|x| x.as_array()),
                v.get("sliding_window_pattern").and_then(|x| x.as_u64()),
            ) {
                (Some(a), _) => {
                    anyhow::ensure!(a.len() == n_layers, "layer_types length");
                    a.iter()
                        .map(|t| t.as_str() == Some("sliding_attention"))
                        .collect::<Vec<_>>()
                }
                (None, Some(p)) if p > 0 => {
                    (0..n_layers).map(|i| (i + 1) % (p as usize) != 0).collect()
                }
                _ => anyhow::bail!(
                    "gemma3 config has neither `layer_types` nor a positive \
                     `sliding_window_pattern`: cannot derive the sliding/full layout (and \
                     with it the per-layer RoPE theta) — refusing instead of defaulting"
                ),
            };
            return Ok(Self {
                arch,
                hidden,
                vocab: g("vocab_size").context("vocab_size")?,
                n_layers,
                intermediate: g("intermediate_size").context("intermediate_size")?,
                n_heads,
                n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
                head_dim: g("head_dim").unwrap_or(hidden / n_heads),
                conv_l: 0,
                eps: v
                    .get("rms_norm_eps")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1e-6) as f32,
                rope_theta: v
                    .get("rope_theta")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(1_000_000.0) as f32,
                rope_local_theta: v
                    .get("rope_local_base_freq")
                    .and_then(|x| x.as_f64())
                    .unwrap_or(10_000.0) as f32,
                sliding_window: g("sliding_window").unwrap_or(512),
                layer_is_attn: vec![true; n_layers],
                layer_is_sliding,
                layer_k_eq_v: vec![false; n_layers],
                act_gelu: true,
                num_experts: 0,
                top_k_experts: 0,
                moe_intermediate: 0,
                stage_first: true,
                stage_last: true,
                rotary_dim: 0,
                mrope_section: None, // text-only architecture: 1-D rope, unchanged
                dn_nk: 0,
                dn_nv: 0,
                dn_dk: 0,
                dn_dv: 0,
                dn_kernel: 0,
                wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
            });
        }
        let hidden = g("hidden_size").context("hidden_size")?;
        let n_heads = g("num_attention_heads").context("num_attention_heads")?;
        // block_auto_adjust_ff_dim: ff = round_up(2/3 * intermediate, block_multiple_of).
        let raw_ff = g("intermediate_size").context("intermediate_size")?;
        let intermediate = if v
            .get("block_auto_adjust_ff_dim")
            .and_then(|x| x.as_bool())
            .unwrap_or(false)
        {
            let m = g("block_multiple_of").unwrap_or(256);
            let ff = (2 * raw_ff) / 3;
            m * ff.div_ceil(m)
        } else {
            raw_ff
        };
        let layer_is_attn = v
            .get("layer_types")
            .and_then(|x| x.as_array())
            .map(|a| {
                a.iter()
                    .map(|t| t.as_str() == Some("full_attention"))
                    .collect::<Vec<_>>()
            })
            .context("layer_types")?;
        let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
        Ok(Self {
            arch,
            hidden,
            vocab: g("vocab_size").context("vocab_size")?,
            n_layers,
            intermediate,
            n_heads,
            n_kv_heads: g("num_key_value_heads").unwrap_or(n_heads),
            head_dim: g("head_dim").unwrap_or(hidden / n_heads),
            conv_l: g("conv_L_cache").unwrap_or(3),
            eps: v.get("norm_eps").and_then(|x| x.as_f64()).unwrap_or(1e-5) as f32,
            rope_theta: v
                .get("rope_theta")
                .and_then(|x| x.as_f64())
                .unwrap_or(1_000_000.0) as f32,
            rope_local_theta: 10_000.0,
            sliding_window: 0,
            layer_is_attn,
            layer_is_sliding: vec![false; n_layers],
            layer_k_eq_v: vec![false; n_layers],
            act_gelu: false,
            num_experts: 0,
            top_k_experts: 0,
            moe_intermediate: 0,
            stage_first: true,
            stage_last: true,
            rotary_dim: 0,
            mrope_section: None, // text-only architecture: 1-D rope, unchanged
            dn_nk: 0,
            dn_nv: 0,
            dn_dk: 0,
            dn_dv: 0,
            dn_kernel: 0,
            wdtype: WDtype::Q4,
            head_wdtype: WDtype::Q4,
            layer_wdtype: Vec::new(),
            layer_kinds: Vec::new(),
            head_kind: None,
            edge: EdgeCfg::default(),
        })
    }
}

/// The per-layer operator weights (projection matrices are Q4_0; conv kernel + norms stay f32).
#[derive(Clone)]
pub enum Op {
    /// Gated short-conv (LIV): `in_proj` → [B|C|x], depthwise causal conv, `out_proj`.
    Conv {
        in_proj: Q4,          // [3*hidden, hidden] (Q4, fused into the conv kernel)
        conv_w: wgpu::Buffer, // [hidden, conv_l]  (depthwise kernel, the [_,1,_] dim squeezed)
        out_proj: Q4,         // [hidden, hidden] (Q4, bandwidth-bound)
    },
    /// GQA attention with per-head QK-RMSNorm + RoPE.
    Attn {
        qkv: Q4, // [(n_heads+2*n_kv_heads)*head_dim, hidden] — q|k|v concatenated (Q4, bandwidth-bound)
        o: Q4,   // [hidden, n_heads*head_dim]
        q_norm: wgpu::Buffer, // [head_dim]
        k_norm: wgpu::Buffer, // [head_dim]
        /// Sliding-attention window (0 = full causal). Gemma-3 sliding layers; 0 for LFM2.
        window: u32,
        /// RoPE uses the LOCAL base frequency (Gemma-3 sliding layers) instead of the global one.
        local_rope: bool,
        /// Qwen3.5 `attn_output_gate`: per-head gate rows de-interleaved from the doubled q_proj
        /// (`[n_heads·head_dim, hidden]`); attention output is `attn · sigmoid(gate)` pre-o_proj.
        attn_gate: Option<Q4>,
        /// Qwen2 q/k/v projection bias, concatenated `[q(qd) | k(kd) | v(kd)]` f32; added to the
        /// fused qkv output before qk-norm/RoPE. `None` for every bias-free arch (Llama/Qwen3/…).
        qkv_bias: Option<wgpu::Buffer>,
    },
    /// Gated DeltaNet linear attention (Qwen3.5/3.6 hybrid layers): fused input projection
    /// `[qkv(2·nk·dk + nv·dv) | z(nv·dv) | b(nv) | a(nv)]`, depthwise causal conv over the qkv
    /// span, the recurrent delta rule per v-head, gated RMSNorm, and the output projection.
    DeltaNet {
        in_proj: Q4,          // [2·nk·dk + nv·dv + nv·dv + 2·nv, hidden]
        conv_w: wgpu::Buffer, // [2·nk·dk + nv·dv, kernel] f32, taps oldest-first
        /// `[A_log(nv) | dt_bias(nv) | norm_w(dv)]` f32 (gate params + gated-norm weight).
        gpar: wgpu::Buffer,
        out_proj: Q4, // [hidden, nv·dv]
    },
    /// DeepSeek MLA (Multi-head Latent Attention). Low-rank compressed KV + decoupled rope. The
    /// projections run as GEMV plan steps; the one novel kernel (`MLA_ASSEMBLE`) does the
    /// interleaved-pair rope on the pe slice, the shared-k_pe broadcast, and the v zero-pad, then
    /// the standard attention + o_proj tail runs unchanged (o_proj is zero-column-padded at load
    /// so it consumes the [nh·qk_head_dim] padded attention output). See HANDOFF_MLA.md.
    Mla {
        q_a: Option<Q4>, // [q_lora_rank, hidden]; None ⇒ q_or_qb is the direct q_proj
        q_a_norm: Option<wgpu::Buffer>, // [q_lora_rank]
        q_or_qb: Q4, // q_b_proj [nh·qk_head_dim, q_lora_rank] or q_proj [nh·qk_head_dim, hidden]
        kv_a: Q4,    // kv_a_proj_with_mqa [kv_lora_rank + qk_rope, hidden]
        kv_a_norm: wgpu::Buffer, // [kv_lora_rank]
        kv_b: Q4,    // [nh·(qk_nope + v_head_dim), kv_lora_rank]
        o: Q4,       // o_proj PADDED to [hidden, nh·qk_head_dim] (zero cols v..qk)
        qk_nope: usize,
        qk_rope: usize,
        v_head_dim: usize,
        kv_lora: usize,
        q_lora: usize, // 0 ⇒ direct q_proj (q_or_qb is q_proj, q_a/q_a_norm are None)
    },
    /// Mamba-1 selective state-space mixer (`MambaForCausalLM`). `in_proj` → [x_ssm | z(gate)];
    /// depthwise causal conv1d over x_ssm (+ SiLU); `x_proj` → [dt | B | C]; `dt_proj` → dt;
    /// the selective scan updates the SSM state `h[d_inner, d_state]` (`h = exp(dt·A)·h + dt·B·x`,
    /// `y = C·h + D·x`), gated by `SiLU(z)`; `out_proj` → hidden. M=1 decode = one recurrence step
    /// per channel; conv1d + SSM state persist across tokens (like the conv/DeltaNet states).
    Mamba {
        in_proj: Q4,          // [2·d_inner, hidden]
        conv_w: wgpu::Buffer, // [d_inner, d_conv] depthwise (the [_,1,_] squeezed)
        conv_b: wgpu::Buffer, // [d_inner]
        x_proj: Q4,           // [dt_rank + 2·d_state, d_inner]
        dt_proj: Q4,          // [d_inner, dt_rank]
        /// `[A_log(d_inner·d_state) | D(d_inner) | dt_bias(d_inner)]` f32.
        ssm_par: wgpu::Buffer,
        out_proj: Q4, // [hidden, d_inner]
        d_inner: usize,
        d_state: usize,
        d_conv: usize,
        dt_rank: usize,
    },
    /// RWKV-4 block (`RwkvForCausalLM`): time-mix (WKV recurrence) + channel-mix, each token-shift
    /// mixed. Time-mix: `xk/xv/xr = x·mix + x_prev·(1−mix)`, then `k=Wk·xk`, `v=Wv·xv`,
    /// `r=sigmoid(Wr·xr)`; the WKV recurrence over `time_decay`(pre-folded `−exp`)/`time_first`
    /// yields `wkv`; `out = Wo·(r·wkv)`. Channel-mix: `xk/xr` mixed, `kv = Wv·relu(Wk·xk)²`,
    /// `out = sigmoid(Wr·xr)·kv`. Both sub-layers use an affine LayerNorm (ln1→`operator_norm`,
    /// ln2→`ffn_norm`) and add into the residual; `ln0` (present only on block 0) is an extra affine
    /// LayerNorm applied to the embeddings first. The token-shift and WKV states persist across
    /// tokens (M=1 decode = one recurrence step). The trailing GLU-FFN slot is a ZERO GLU (adds 0).
    Rwkv {
        /// Affine LN `[weight|bias]` applied to the embeddings before block 0; `None` on i>0.
        ln0: Option<wgpu::Buffer>,
        /// Time-mix token-shift mix `[time_mix_k | time_mix_v | time_mix_r]` (3·hidden).
        att_mix: wgpu::Buffer,
        att_k: Q4,                // [hidden, hidden]
        att_v: Q4,                // [hidden, hidden]
        att_r: Q4,                // [hidden, hidden]
        att_o: Q4,                // [hidden, hidden]
        time_decay: wgpu::Buffer, // [hidden], pre-folded `−exp(time_decay)`
        time_first: wgpu::Buffer, // [hidden]
        /// Channel-mix token-shift mix `[time_mix_k | 0 | time_mix_r]` (3·hidden; middle unused).
        ffn_mix: wgpu::Buffer,
        ffn_k: Q4, // [ffn_dim, hidden]
        ffn_v: Q4, // [hidden, ffn_dim]
        ffn_r: Q4, // [hidden, hidden]
        ffn_dim: usize,
    },
}

/// One transformer layer: pre-op norm, the operator, pre-MLP norm, and the GLU MLP (Q4_0).
/// Gemma-3 additionally carries the sandwich norms (`post_*`, applied to the block output BEFORE
/// the residual add); `None` for LFM2, whose residual adds are fused into the GEMVs.
#[derive(Clone)]
pub struct Layer {
    pub operator_norm: wgpu::Buffer,
    pub ffn_norm: wgpu::Buffer,
    pub op: Op,
    pub w1: Q4, // gate [intermediate, hidden]
    pub w2: Q4, // down [hidden, intermediate]
    pub w3: Q4, // up   [intermediate, hidden]
    pub post_op_norm: Option<wgpu::Buffer>,
    pub post_ffn_norm: Option<wgpu::Buffer>,
    /// Routed-expert weights (Gemma-4 MoE layers); `None` on dense layers/models.
    pub moe: Option<Moe>,
    /// Gemma-4 edge extras (PLE injection + layer scalar); `None` everywhere else.
    pub edge: Option<LayerEdge>,
}

/// Gemma-4 edge per-layer extras: the PLE injection weights and the whole-stream layer scalar.
/// The server checkpoints (31B) carry `layer_scalar` WITHOUT PLE, so the two nest.
#[derive(Clone)]
pub struct LayerEdge {
    /// PLE injection — structurally a THIRD sandwich sub-block appended after the FFN one:
    /// `x += post_norm(proj(gelu(gate(x)) · ple_i))`. `None` on non-PLE checkpoints.
    pub ple: Option<PleLayer>,
    /// `layer_scalar` — multiplies the ENTIRE residual stream at layer end (after the PLE add).
    /// A 1-element checkpoint tensor, read to a host scalar at load; 1.0 when absent.
    pub layer_scalar: f32,
}

/// One layer's PLE injection weights.
#[derive(Clone)]
pub struct PleLayer {
    /// `per_layer_input_gate` `[ple_dim, hidden]` (Q4).
    pub gate: Q4,
    /// `per_layer_projection` `[hidden, ple_dim]` (Q4).
    pub proj: Q4,
    /// `post_per_layer_input_norm` `[hidden]` f32 (plain-w convention, like every Gemma-4 norm).
    pub post_norm: wgpu::Buffer,
}

/// MoE weights for one layer: the (small, f32) router + per-expert scales, and the expert GLU
/// matrices CONCATENATED across experts (expert `e`'s gate rows are `[e·mi, (e+1)·mi)` of `w1`,
/// etc.) so a slot kernel indexes them by the router's chosen expert id at dispatch time.
#[derive(Clone)]
pub struct Moe {
    /// Router projection `[num_experts, hidden]`, f32 (tiny; router accuracy matters).
    pub router: wgpu::Buffer,
    /// Qwen3.5 shared-expert gate `[hidden]` f32: the shared expert's output is scaled by
    /// `sigmoid(w·x_normed)` before joining the routed sum. `None` for Gemma-4.
    pub shared_gate: Option<wgpu::Buffer>,
    /// Learned per-expert scale `[num_experts]`, multiplied onto the renormalized top-k weights.
    pub per_expert_scale: wgpu::Buffer,
    /// Expert gate matrices, concatenated: `[num_experts · moe_intermediate, hidden]` (Q4).
    pub w1: Q4,
    /// Expert up matrices, same layout as `w1`.
    pub w3: Q4,
    /// Expert down matrices, concatenated: `[num_experts · hidden, moe_intermediate]` (Q4).
    pub w2: Q4,
    /// DeepSeek-V3 router (sigmoid scoring + additive selection bias + group-limited top-k).
    /// `Some` swaps the softmax `MOE_ROUTER` for `MOE_ROUTER_V3`; `None` = the V2/Mixtral softmax
    /// router. The per-expert selection weights it emits are the un-biased sigmoid scores,
    /// renormalized and pre-scaled by `routed_scaling_factor` (so `per_expert_scale` is unused).
    pub router_v3: Option<V3Router>,
}

/// DeepSeek-V3 aux-loss-free router parameters (`DeepseekV3TopkRouter`). Selection adds
/// `bias` to the sigmoid scores, groups the experts into `n_group` groups, scores each group by
/// its top-2 corrected-score sum, keeps the `topk_group` best groups, then takes the global top-k
/// over the surviving experts. The routed WEIGHTS are the ORIGINAL (un-biased) sigmoid scores of
/// the chosen experts, renormalized and multiplied by `scaling`.
#[derive(Clone)]
pub struct V3Router {
    /// `e_score_correction_bias` `[num_experts]`, f32 — added to scores for SELECTION only.
    pub bias: wgpu::Buffer,
    pub n_group: u32,
    pub topk_group: u32,
    /// `routed_scaling_factor`, folded into the emitted weights.
    pub scaling: f32,
}

/// All model weights resident on the GPU.
/// Lazily-readable safetensors checkpoint: a single `model.safetensors` or an HF sharded set
/// (`model.safetensors.index.json` + `model-XXXXX-of-YYYYY.safetensors`). Headers are parsed once;
/// tensor bytes are `pread` on demand — a 65 GB 32B checkpoint never touches RAM whole (the eager
/// `fs::read` path needed the full file per process, which forbids multi-worker sharding hosts).
pub struct LazySt {
    backing: Backing,
    /// name → tensor locator.
    map: std::collections::HashMap<String, TensorLoc>,
}

/// Where the tensor bytes physically live. `Files` preads on demand (native, RAM-flat for huge
/// checkpoints); `Bytes` slices in-memory blobs — how a browser worker uses its SHIPPED
/// safetensors, since wasm has no filesystem.
enum Backing {
    Files(Vec<std::fs::File>),
    Bytes(Vec<Vec<u8>>),
}

/// Index one safetensors header (JSON tensor table) into `map`, offsetting ranges by
/// `data_base` (= 8 + header length) and tagging each tensor with its `fi` blob/file index.
fn index_header(
    hdr: &[u8],
    data_base: u64,
    fi: usize,
    map: &mut std::collections::HashMap<String, TensorLoc>,
) -> Result<()> {
    let v: serde_json::Value = serde_json::from_slice(hdr)?;
    let obj = v.as_object().context("safetensors header")?;
    for (name, meta) in obj {
        if name == "__metadata__" {
            continue;
        }
        let dtype = match meta.get("dtype").and_then(|d| d.as_str()) {
            Some("F32") => safetensors::Dtype::F32,
            Some("BF16") => safetensors::Dtype::BF16,
            Some("F16") => safetensors::Dtype::F16,
            // Non-float tensors are never model weights (e.g. BERT checkpoints ship an I64
            // `embeddings.position_ids` arange buffer). Skip them at indexing rather than
            // rejecting the checkpoint; a skipped tensor that IS requested later fails
            // precisely as `missing tensor <name>`.
            other => {
                eprintln!("skipping non-float tensor {name} (dtype {other:?})");
                continue;
            }
        };
        let shape: Vec<usize> = meta
            .get("shape")
            .and_then(|s| s.as_array())
            .context("shape")?
            .iter()
            .filter_map(|x| x.as_u64().map(|x| x as usize))
            .collect();
        let offs = meta
            .get("data_offsets")
            .and_then(|o| o.as_array())
            .context("data_offsets")?;
        let s0 = offs[0].as_u64().context("off0")?;
        let s1 = offs[1].as_u64().context("off1")?;
        map.insert(
            name.clone(),
            TensorLoc {
                file: fi,
                dtype,
                shape,
                range: (data_base + s0, data_base + s1),
            },
        );
    }
    Ok(())
}

/// Where one tensor lives: file index, dtype, shape, absolute byte range in that file.
struct TensorLoc {
    file: usize,
    dtype: safetensors::Dtype,
    shape: Vec<usize>,
    range: (u64, u64),
}

/// Positioned read of `buf.len()` bytes at `offset` — thread-safe, no seek (concurrent shard
/// workers pread the same file). Unix uses `pread`; Windows `seek_read`; other targets (wasm)
/// return unsupported — a browser worker loads from shipped BYTES, never a file, so this is
/// never reached there at runtime, but the engine core must still compile for `wasm32`.
fn pread_exact(file: &std::fs::File, buf: &mut [u8], offset: u64) -> std::io::Result<()> {
    #[cfg(unix)]
    {
        use std::os::unix::fs::FileExt;
        file.read_exact_at(buf, offset)
    }
    #[cfg(windows)]
    {
        use std::os::windows::fs::FileExt;
        let mut done = 0usize;
        while done < buf.len() {
            match file.seek_read(&mut buf[done..], offset + done as u64)? {
                0 => return Err(std::io::ErrorKind::UnexpectedEof.into()),
                n => done += n,
            }
        }
        Ok(())
    }
    #[cfg(not(any(unix, windows)))]
    {
        let _ = (file, buf, offset);
        Err(std::io::Error::new(
            std::io::ErrorKind::Unsupported,
            "positioned file reads are unavailable on this target; use a bytes-based loader",
        ))
    }
}

impl LazySt {
    pub fn open(dir: &Path) -> Result<Self> {
        let single = dir.join("model.safetensors");
        let paths: Vec<std::path::PathBuf> = if single.exists() {
            vec![single]
        } else {
            let idx = dir.join("model.safetensors.index.json");
            let v: serde_json::Value = serde_json::from_slice(
                &std::fs::read(&idx).with_context(|| format!("missing {}", idx.display()))?,
            )?;
            let wm = v
                .get("weight_map")
                .and_then(|m| m.as_object())
                .context("index.json weight_map")?;
            let mut names: Vec<String> = wm
                .values()
                .filter_map(|f| f.as_str().map(str::to_string))
                .collect();
            names.sort();
            names.dedup();
            names.into_iter().map(|f| dir.join(f)).collect()
        };
        let mut files = Vec::with_capacity(paths.len());
        let mut map = std::collections::HashMap::new();
        for (fi, path) in paths.iter().enumerate() {
            use std::io::Read;
            let mut f =
                std::fs::File::open(path).with_context(|| format!("open {}", path.display()))?;
            let mut len8 = [0u8; 8];
            f.read_exact(&mut len8)?;
            let hlen = u64::from_le_bytes(len8);
            let mut hdr = vec![0u8; hlen as usize];
            f.read_exact(&mut hdr)?;
            index_header(&hdr, 8 + hlen, fi, &mut map)?;
            files.push(f);
        }
        Ok(Self {
            backing: Backing::Files(files),
            map,
        })
    }

    /// Build a lazy view over IN-MEMORY safetensors blobs (one per file of a single- or
    /// HF-sharded checkpoint) — the wasm/browser path, where weights arrive as shipped bytes.
    pub fn from_bytes(blobs: Vec<Vec<u8>>) -> Result<Self> {
        anyhow::ensure!(!blobs.is_empty(), "no safetensors blobs");
        let mut map = std::collections::HashMap::new();
        for (fi, blob) in blobs.iter().enumerate() {
            anyhow::ensure!(
                blob.len() >= 8,
                "blob {fi} too small for a safetensors header"
            );
            let hlen = u64::from_le_bytes(blob[..8].try_into().expect("8 bytes"));
            let end = 8 + hlen as usize;
            anyhow::ensure!(
                end <= blob.len(),
                "blob {fi} header length exceeds the blob"
            );
            index_header(&blob[8..end], 8 + hlen, fi, &mut map)?;
        }
        Ok(Self {
            backing: Backing::Bytes(blobs),
            map,
        })
    }

    pub fn names(&self) -> Vec<String> {
        let mut v: Vec<String> = self.map.keys().cloned().collect();
        v.sort();
        v
    }

    pub fn has(&self, name: &str) -> bool {
        self.map.contains_key(name)
    }

    pub fn shape(&self, name: &str) -> Result<&[usize]> {
        Ok(&self
            .map
            .get(name)
            .with_context(|| format!("missing tensor {name}"))?
            .shape)
    }

    /// Read one tensor's RAW bytes + dtype, without converting to f32. Lets a caller convert in
    /// slices (e.g. per-MoE-expert) to avoid materializing the whole f32 buffer — essential on
    /// wasm32, where a 256-expert gate_up as f32 is exactly 2 GB (> isize::MAX) and overflows.
    #[cfg(target_arch = "wasm32")]
    pub fn tensor_raw(&self, name: &str) -> Result<(Vec<u8>, safetensors::Dtype)> {
        let loc = self
            .map
            .get(name)
            .with_context(|| format!("missing tensor {name}"))?;
        let (a, b) = loc.range;
        let raw: Vec<u8> = match &self.backing {
            Backing::Files(files) => {
                let mut raw = vec![0u8; (b - a) as usize];
                pread_exact(&files[loc.file], &mut raw, a)
                    .with_context(|| format!("pread {name}"))?;
                raw
            }
            Backing::Bytes(blobs) => blobs[loc.file]
                .get(a as usize..b as usize)
                .with_context(|| format!("tensor {name} range out of blob"))?
                .to_vec(),
        };
        Ok((raw, loc.dtype))
    }

    /// Read one tensor as f32 (dtype-converting), touching only its byte range.
    pub fn tensor_f32(&self, name: &str) -> Result<Vec<f32>> {
        let loc = self
            .map
            .get(name)
            .with_context(|| format!("missing tensor {name}"))?;
        let (a, b) = loc.range;
        let raw: Vec<u8> = match &self.backing {
            Backing::Files(files) => {
                let mut raw = vec![0u8; (b - a) as usize];
                pread_exact(&files[loc.file], &mut raw, a)
                    .with_context(|| format!("pread {name}"))?;
                raw
            }
            Backing::Bytes(blobs) => blobs[loc.file]
                .get(a as usize..b as usize)
                .with_context(|| format!("tensor {name} range out of blob"))?
                .to_vec(),
        };
        Ok(match loc.dtype {
            safetensors::Dtype::F32 => bytemuck::cast_slice::<u8, f32>(&raw).to_vec(),
            safetensors::Dtype::F16 => raw
                .chunks_exact(2)
                .map(|c| half::f16::from_le_bytes([c[0], c[1]]).to_f32())
                .collect(),
            _ => bf16_to_f32(&raw),
        })
    }
}

/// Load a norm with the `(1+w)` convention (weights stored as deltas from 1 — Gemma-3 and
/// Qwen3.5 layernorms) from a lazily-opened checkpoint.
fn get_p1_at(st: &LazySt, ctx: &GpuCtx, name: &str) -> Result<wgpu::Buffer> {
    let mut vals = st.tensor_f32(name)?;
    for x in &mut vals {
        *x += 1.0;
    }
    Ok(ctx.storage(&vals))
}

/// Bytes per element for a safetensors float dtype (F32=4, F16/BF16=2). Used to slice a raw tensor
/// buffer by element index without converting the whole thing to f32.
#[cfg(target_arch = "wasm32")]
fn dtype_bytes(dt: safetensors::Dtype) -> usize {
    match dt {
        safetensors::Dtype::F32 => 4,
        _ => 2, // F16 / BF16
    }
}

/// Convert a raw float byte slice (F32/F16/BF16) to f32 — the slice form of [`LazySt::tensor_f32`]'s
/// conversion, so a caller can convert a tensor in pieces (per MoE expert) instead of all at once.
#[cfg(target_arch = "wasm32")]
fn raw_to_f32(raw: &[u8], dt: safetensors::Dtype) -> Vec<f32> {
    match dt {
        safetensors::Dtype::F32 => bytemuck::cast_slice::<u8, f32>(raw).to_vec(),
        safetensors::Dtype::F16 => raw
            .chunks_exact(2)
            .map(|c| half::f16::from_le_bytes([c[0], c[1]]).to_f32())
            .collect(),
        _ => bf16_to_f32(raw),
    }
}

/// Build ONE qwen3.5-MoE layer from `prefix` (`{lm}.layers.{i}` or `mtp.layers.0` — the MTP
/// draft layer is a standard full-attention qwen3.5 layer, so the loader is shared). Covers the
/// doubled-q_proj de-interleave (attn_output_gate), the fused DeltaNet projection, and the
/// routed-experts + shared-expert MoE split.
fn build_qwen35_layer(
    st: &LazySt,
    ctx: &GpuCtx,
    cfg: &Lfm2Config,
    prefix: &str,
    is_attn: bool,
) -> Result<Layer> {
    let c = cfg;
    let (nk, nv, dk, dv) = (c.dn_nk, c.dn_nv, c.dn_dk, c.dn_dv);
    let (qd, kd) = (c.n_heads * c.head_dim, c.n_kv_heads * c.head_dim);
    let tens = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(name) };
    let q4 = |name: &str| -> Result<Q4> {
        let shape = st.shape(name)?;
        let (rows, cols) = (shape[0], shape[1]);
        pack_weight(ctx, &tens(name)?, rows, cols, cfg.wdtype)
    };
    let p = prefix;
    let op = if is_attn {
        // Full attention. q_proj is DOUBLED: per head 512 rows = [q(256) | gate(256)]
        // interleaved — de-interleave into q rows and gate rows at load.
        let qp = tens(&format!("{p}.self_attn.q_proj.weight"))?;
        anyhow::ensure!(
            st.shape(&format!("{p}.self_attn.q_proj.weight"))? == [2 * qd, c.hidden],
            "q_proj must be doubled (attn_output_gate)"
        );
        let hd2 = 2 * c.head_dim;
        let mut qrows = vec![0f32; qd * c.hidden];
        let mut grows = vec![0f32; qd * c.hidden];
        for h in 0..c.n_heads {
            let src = h * hd2 * c.hidden;
            let dst = h * c.head_dim * c.hidden;
            let half = c.head_dim * c.hidden;
            qrows[dst..dst + half].copy_from_slice(&qp[src..src + half]);
            grows[dst..dst + half].copy_from_slice(&qp[src + half..src + 2 * half]);
        }
        let kp = tens(&format!("{p}.self_attn.k_proj.weight"))?;
        let vp = tens(&format!("{p}.self_attn.v_proj.weight"))?;
        let mut qkv_f = qrows;
        qkv_f.extend_from_slice(&kp);
        qkv_f.extend_from_slice(&vp);
        Op::Attn {
            // Assembled (de-interleaved q + concatenated k,v) → packed at the body dtype so f16
            // reaches the DeltaNet/attention body, not only the closure-loaded projections.
            qkv: pack_weight(ctx, &qkv_f, qd + 2 * kd, c.hidden, c.wdtype)?,
            o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
            q_norm: get_p1_at(st, ctx, &format!("{p}.self_attn.q_norm.weight"))?,
            k_norm: get_p1_at(st, ctx, &format!("{p}.self_attn.k_norm.weight"))?,
            window: 0,
            local_rope: false,
            attn_gate: Some(pack_weight(ctx, &grows, qd, c.hidden, c.wdtype)?),
            qkv_bias: None,
        }
    } else {
        // Gated DeltaNet: one fused Q4 projection [qkv | z | b | a].
        let la = format!("{p}.linear_attn");
        let mut in_f = tens(&format!("{la}.in_proj_qkv.weight"))?;
        in_f.extend_from_slice(&tens(&format!("{la}.in_proj_z.weight"))?);
        in_f.extend_from_slice(&tens(&format!("{la}.in_proj_b.weight"))?);
        in_f.extend_from_slice(&tens(&format!("{la}.in_proj_a.weight"))?);
        let rows = 2 * nk * dk + nv * dv + nv * dv + 2 * nv;
        // conv1d.weight [conv_dim, 1, K] → squeeze to [conv_dim, K].
        let conv = tens(&format!("{la}.conv1d.weight"))?;
        let gpar: Vec<f32> = tens(&format!("{la}.A_log"))?
            .into_iter()
            .chain(tens(&format!("{la}.dt_bias"))?)
            .chain(tens(&format!("{la}.norm.weight"))?)
            .collect();
        Op::DeltaNet {
            in_proj: pack_weight(ctx, &in_f, rows, c.hidden, c.wdtype)?,
            conv_w: ctx.storage(&conv),
            gpar: ctx.storage(&gpar),
            out_proj: q4(&format!("{la}.out_proj.weight"))?,
        }
    };
    let mp = format!("{p}.mlp");

    // DENSE flavor (`Qwen3_5ForConditionalGeneration`, e.g. claim-extractor-4B): a plain SwiGLU
    // MLP, no router and no experts. The attention/DeltaNet stack above is IDENTICAL to the MoE
    // flavor's — only the feed-forward differs — so we reuse everything and just fill the dense GLU
    // slots from `mlp.{gate,up,down}_proj` with `moe: None` (the forward already branches on that).
    if c.num_experts == 0 {
        return Ok(Layer {
            operator_norm: get_p1_at(st, ctx, &format!("{p}.input_layernorm.weight"))?,
            ffn_norm: get_p1_at(st, ctx, &format!("{p}.post_attention_layernorm.weight"))?,
            op,
            w1: q4(&format!("{mp}.gate_proj.weight"))?,
            w2: q4(&format!("{mp}.down_proj.weight"))?,
            w3: q4(&format!("{mp}.up_proj.weight"))?,
            post_op_norm: None,
            post_ffn_norm: None,
            moe: None,
            edge: None,
        });
    }

    // MoE on every layer: routed experts (3D tensors flattened; gate|up split on the ROW axis
    // per expert) + the shared expert into the dense GLU slots.
    // The MoE expert kernels (moe_gate_grouped / moe_down_grouped) are Q4-only — no f16 twins —
    // so a non-Q4 body dtype cannot be honoured for the experts. Refuse rather than silently
    // serving Q4 experts under an f16 body (attention/DeltaNet above ARE f16-capable, but a
    // half-f16 model is not what `f16` requested). Dense Qwen3.5 (num_experts == 0) never reaches
    // here.
    anyhow::ensure!(
        c.wdtype == WDtype::Q4,
        "weight dtype {:?} is not servable for a Qwen3.5-MoE model: the expert GEMM kernels are \
         Q4-only. Use `q4`, or a dense checkpoint for f16.",
        c.wdtype
    );
    let e_n = c.num_experts;
    let mi = c.moe_intermediate;
    let per = 2 * mi * c.hidden;
    let gup = format!("{mp}.experts.gate_up_proj");
    let dnp = format!("{mp}.experts.down_proj");
    // Routed experts: gate|up split on the ROW axis per expert, then down. q4_0 quantizes ROW-WISE
    // (one scale per 32-col block within a row), so quantizing per-expert and concatenating the
    // blocks is BITWISE-IDENTICAL to quantizing the whole flattened tensor.
    #[cfg(not(target_arch = "wasm32"))]
    let (s1, z1, s3, z3, s2, z2) = {
        // Native (has the RAM): keep the proven whole-tensor path.
        let gu = tens(&gup)?;
        let mut w1f = Vec::with_capacity(e_n * mi * c.hidden);
        let mut w3f = Vec::with_capacity(e_n * mi * c.hidden);
        for e in 0..e_n {
            let base = e * per;
            w1f.extend_from_slice(&gu[base..base + mi * c.hidden]);
            w3f.extend_from_slice(&gu[base + mi * c.hidden..base + per]);
        }
        drop(gu);
        let (s1, z1) = quantize_q4_0(&w1f, e_n * mi, c.hidden);
        drop(w1f);
        let (s3, z3) = quantize_q4_0(&w3f, e_n * mi, c.hidden);
        drop(w3f);
        let dnf = tens(&dnp)?;
        let (s2, z2) = quantize_q4_0(&dnf, e_n * c.hidden, mi);
        (s1, z1, s3, z3, s2, z2)
    };
    #[cfg(target_arch = "wasm32")]
    let (s1, z1, s3, z3, s2, z2) = {
        // wasm32: convert + quantize PER EXPERT off the raw (bf16) buffer — the whole gate_up as f32
        // is num_experts·2·mi·hidden·4 B = 2 GB = isize::MAX+1 and overflows a single wasm `Vec`.
        let (gu, gdt) = st.tensor_raw(&gup)?;
        let el = dtype_bytes(gdt);
        let (mut s1, mut z1, mut s3, mut z3) = (Vec::new(), Vec::new(), Vec::new(), Vec::new());
        for e in 0..e_n {
            let base = e * per * el;
            let gate = raw_to_f32(&gu[base..base + mi * c.hidden * el], gdt);
            let (se, ze) = quantize_q4_0(&gate, mi, c.hidden);
            s1.extend_from_slice(&se);
            z1.extend_from_slice(&ze);
            let up = raw_to_f32(&gu[base + mi * c.hidden * el..base + per * el], gdt);
            let (se, ze) = quantize_q4_0(&up, mi, c.hidden);
            s3.extend_from_slice(&se);
            z3.extend_from_slice(&ze);
        }
        drop(gu);
        let (dn, ddt) = st.tensor_raw(&dnp)?;
        let del = dtype_bytes(ddt);
        let (mut s2, mut z2) = (Vec::new(), Vec::new());
        for e in 0..e_n {
            let base = e * c.hidden * mi * del;
            let d = raw_to_f32(&dn[base..base + c.hidden * mi * del], ddt);
            let (se, ze) = quantize_q4_0(&d, c.hidden, mi);
            s2.extend_from_slice(&se);
            z2.extend_from_slice(&ze);
        }
        (s1, z1, s3, z3, s2, z2)
    };
    let moe = Moe {
        router: ctx.storage(&tens(&format!("{mp}.gate.weight"))?),
        shared_gate: Some(ctx.storage(&tens(&format!("{mp}.shared_expert_gate.weight"))?)),
        per_expert_scale: ctx.storage(&vec![1.0f32; e_n]),
        w1: Q4 {
            scales: ctx.storage_bytes(&f32_to_f16_bytes(&s1)),
            quants: ctx.storage(bytemuck::cast_slice(&z1)),
        },
        w3: Q4 {
            scales: ctx.storage_bytes(&f32_to_f16_bytes(&s3)),
            quants: ctx.storage(bytemuck::cast_slice(&z3)),
        },
        w2: Q4 {
            scales: ctx.storage_bytes(&f32_to_f16_bytes(&s2)),
            quants: ctx.storage(bytemuck::cast_slice(&z2)),
        },
        router_v3: None,
    };
    Ok(Layer {
        operator_norm: get_p1_at(st, ctx, &format!("{p}.input_layernorm.weight"))?,
        ffn_norm: get_p1_at(st, ctx, &format!("{p}.post_attention_layernorm.weight"))?,
        op,
        w1: q4(&format!("{mp}.shared_expert.gate_proj.weight"))?,
        w2: q4(&format!("{mp}.shared_expert.down_proj.weight"))?,
        w3: q4(&format!("{mp}.shared_expert.up_proj.weight"))?,
        post_op_norm: None,
        post_ffn_norm: None,
        moe: Some(moe),
        edge: None,
    })
}

/// Load the `mtp.*` draft head when present (see [`MtpHead`]); flushes upload staging.
fn load_mtp(st: &LazySt, ctx: &GpuCtx, cfg: &Lfm2Config) -> Result<Option<MtpHead>> {
    // The MTP draft head is a speculative-decode OPTIMIZATION, not required for correct output. Load
    // it only when the WHOLE head is present — `mtp.fc.weight` AND the `mtp.layers.0` transformer
    // block. A partial head (e.g. a shipped shard that carried `mtp.fc` but whose `mtp.layers.0.*`
    // was mis-routed to another stage) is skipped with a warning rather than crashing the load.
    if !st.has("mtp.fc.weight") {
        return Ok(None);
    }
    if !st.has("mtp.layers.0.self_attn.q_proj.weight") {
        eprintln!(
            "qwen35 load: partial MTP head (mtp.fc present, mtp.layers.0 absent) — skipping MTP \
             (speculative decode disabled for this stage)"
        );
        return Ok(None);
    }
    let shape = st.shape("mtp.fc.weight")?;
    let h = cfg.hidden;
    anyhow::ensure!(shape == [h, 2 * h], "mtp.fc must be [hidden, 2·hidden]");
    let fc = st.tensor_f32("mtp.fc.weight")?;
    let (mut left, mut right) = (Vec::with_capacity(h * h), Vec::with_capacity(h * h));
    for r in 0..h {
        left.extend_from_slice(&fc[r * 2 * h..r * 2 * h + h]);
        right.extend_from_slice(&fc[r * 2 * h + h..(r + 1) * 2 * h]);
    }
    drop(fc);
    let q4v = |v: &[f32]| -> Q4 {
        let (scales, quants) = quantize_q4_0(v, h, h);
        Q4 {
            scales: ctx.storage_bytes(&f32_to_f16_bytes(&scales)),
            quants: ctx.storage(bytemuck::cast_slice(&quants)),
        }
    };
    // Acceptance probe matrix (OSFKB_MTP_WIRING): 0 = measured-best (fc swapped, norms by
    // name); 1 = fc swapped + norms crossed; 2 = fc unswapped + norms crossed. Real-model
    // acceptance is the only oracle for the reference concat/norm assignment.
    let wiring = std::env::var("OSFKB_MTP_WIRING")
        .ok()
        .and_then(|v| v.parse::<u32>().ok())
        .unwrap_or(0);
    let (fc_h, fc_e) = if wiring == 2 {
        (q4v(&left), q4v(&right))
    } else {
        (q4v(&right), q4v(&left))
    };
    let (nm_emb, nm_hid) = if wiring >= 1 {
        (
            "mtp.pre_fc_norm_hidden.weight",
            "mtp.pre_fc_norm_embedding.weight",
        )
    } else {
        (
            "mtp.pre_fc_norm_embedding.weight",
            "mtp.pre_fc_norm_hidden.weight",
        )
    };
    let head = MtpHead {
        fc_hidden: fc_h,
        fc_emb: fc_e,
        pre_norm_emb: get_p1_at(st, ctx, nm_emb)?,
        pre_norm_hidden: get_p1_at(st, ctx, nm_hid)?,
        norm: get_p1_at(st, ctx, "mtp.norm.weight")?,
        layer: build_qwen35_layer(st, ctx, cfg, "mtp.layers.0", true)?,
        head: None,
        embed: None,
    };
    ctx.queue.submit(std::iter::empty());
    let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
    eprintln!("qwen35 load: mtp head done");
    Ok(Some(head))
}

/// Build ONE GLM-OCR decoder layer from `prefix` (`model.language_model.layers.{i}`; the MTP
/// draft head at index `n_layers` is the same layer shape plus its merge/head tensors —
/// [`load_glm_mtp`]). GPT-J→NeoX q/k de-interleave, fused q|k|v Q4, split fused gate_up, GLM
/// sandwich norms (plain `w`).
fn build_glm_layer(st: &LazySt, ctx: &GpuCtx, cfg: &Lfm2Config, p: &str) -> Result<Layer> {
    let (hd, hidden) = (cfg.head_dim, cfg.hidden);
    // The qk-norm slots are ONES: the GLM attention pipeline skips the norm entirely (this
    // arch has none), so the buffers exist only to satisfy the bind-group layout.
    let ones: Vec<f32> = vec![1.0f32; hd];
    let tens = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(name) };
    let get = |name: &str| -> Result<wgpu::Buffer> { Ok(ctx.storage(&tens(name)?)) };
    let q4 = |name: &str| -> Result<Q4> {
        let shape = st.shape(name)?;
        let (rows, cols) = (shape[0], shape[1]);
        pack_weight(ctx, &tens(name)?, rows, cols, cfg.wdtype)
    };
    // De-interleave q/k rows per head: GPT-J pair rotation → the NeoX-halves kernels.
    // Lane j of a head moves to lane j/2 (even) or hd/2 + j/2 (odd); v/o are untouched,
    // and q·k dot products are invariant because q and k share the permutation.
    let deinter = |name: &str, heads_n: usize| -> Result<Vec<f32>> {
        let w = tens(name)?;
        anyhow::ensure!(w.len() == heads_n * hd * hidden, "{name} shape");
        let mut out = vec![0f32; w.len()];
        let half = hd / 2;
        for h in 0..heads_n {
            for j in 0..hd {
                let dst_lane = if j % 2 == 0 { j / 2 } else { half + j / 2 };
                let src = (h * hd + j) * hidden;
                let dst = (h * hd + dst_lane) * hidden;
                out[dst..dst + hidden].copy_from_slice(&w[src..src + hidden]);
            }
        }
        Ok(out)
    };
    let mut qkv_f = deinter(&format!("{p}.self_attn.q_proj.weight"), cfg.n_heads)?;
    qkv_f.extend_from_slice(&deinter(
        &format!("{p}.self_attn.k_proj.weight"),
        cfg.n_kv_heads,
    )?);
    qkv_f.extend_from_slice(&tens(&format!("{p}.self_attn.v_proj.weight"))?);
    let rows = (cfg.n_heads + 2 * cfg.n_kv_heads) * hd;
    let (sc, qz) = quantize_q4_0(&qkv_f, rows, hidden);
    // Fused gate_up [2·intermediate, hidden] — gate is the FIRST half of the rows.
    let gu = tens(&format!("{p}.mlp.gate_up_proj.weight"))?;
    anyhow::ensure!(gu.len() == 2 * cfg.intermediate * hidden, "gate_up shape");
    let half_rows = cfg.intermediate * hidden;
    let (s1, z1) = quantize_q4_0(&gu[..half_rows], cfg.intermediate, hidden);
    let (s3, z3) = quantize_q4_0(&gu[half_rows..], cfg.intermediate, hidden);
    Ok(Layer {
        operator_norm: get(&format!("{p}.input_layernorm.weight"))?,
        ffn_norm: get(&format!("{p}.post_attention_layernorm.weight"))?,
        op: Op::Attn {
            qkv: Q4 {
                scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                quants: ctx.storage(bytemuck::cast_slice(&qz)),
            },
            o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
            q_norm: ctx.storage(&ones),
            k_norm: ctx.storage(&ones),
            window: 0,
            local_rope: false,
            attn_gate: None,
            qkv_bias: None,
        },
        w1: Q4 {
            scales: ctx.storage_bytes(&f32_to_f16_bytes(&s1)),
            quants: ctx.storage(bytemuck::cast_slice(&z1)),
        },
        w2: q4(&format!("{p}.mlp.down_proj.weight"))?,
        w3: Q4 {
            scales: ctx.storage_bytes(&f32_to_f16_bytes(&s3)),
            quants: ctx.storage(bytemuck::cast_slice(&z3)),
        },
        // GLM sandwich norms, plain `w` (Gemma's arm loads these with `get_p1`).
        post_op_norm: Some(get(&format!("{p}.post_self_attn_layernorm.weight"))?),
        post_ffn_norm: Some(get(&format!("{p}.post_mlp_layernorm.weight"))?),
        moe: None,
        edge: None,
    })
}

/// Load GLM-OCR's MTP draft head — `model.language_model.layers.{n_layers}`, one past the
/// decoder stack (HF ignores it at inference; vLLM serves it as the Glm4Moe MTP). DeepSeek-V3
/// naming maps onto [`MtpHead`] as: `eh_proj` ≡ `mtp.fc` (split-load), `enorm` ≡
/// `pre_fc_norm_embedding`, `hnorm` ≡ `pre_fc_norm_hidden`, `shared_head.norm` ≡ `mtp.norm` —
/// all plain `w` (GLM norms are not `(1+w)`). UNLIKE the qwen3.5 head, this one ships its OWN
/// `shared_head.head` and `embed_tokens` copies (byte-hashes differ from the main lm_head/embed),
/// so nothing is shared with the main model.
///
/// Gated on `OSFKB_SERVE_MTP > 0`: the head costs ~450 MB VRAM (f32 embed + Q4 head + layer) and
/// is dead weight unless the serving layer actually speculates.
fn load_glm_mtp(st: &LazySt, ctx: &GpuCtx, cfg: &Lfm2Config) -> Result<Option<MtpHead>> {
    let want = std::env::var("OSFKB_SERVE_MTP")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(0);
    if want == 0 {
        return Ok(None);
    }
    let p = format!("model.language_model.layers.{}", cfg.n_layers);
    if !st.has(&format!("{p}.eh_proj.weight")) {
        eprintln!("glm load: OSFKB_SERVE_MTP set but the checkpoint ships no MTP head — plain decode");
        return Ok(None);
    }
    if !st.has(&format!("{p}.self_attn.q_proj.weight")) {
        eprintln!("glm load: partial MTP head ({p}.eh_proj present, transformer block absent) — skipping MTP");
        return Ok(None);
    }
    let h = cfg.hidden;
    let shape = st.shape(&format!("{p}.eh_proj.weight"))?;
    anyhow::ensure!(shape == [h, 2 * h], "eh_proj must be [hidden, 2·hidden]");
    let fc = st.tensor_f32(&format!("{p}.eh_proj.weight"))?;
    let (mut left, mut right) = (Vec::with_capacity(h * h), Vec::with_capacity(h * h));
    for r in 0..h {
        left.extend_from_slice(&fc[r * 2 * h..r * 2 * h + h]);
        right.extend_from_slice(&fc[r * 2 * h + h..(r + 1) * 2 * h]);
    }
    drop(fc);
    let q4v = |v: &[f32]| -> Q4 {
        let (scales, quants) = quantize_q4_0(v, h, h);
        Q4 {
            scales: ctx.storage_bytes(&f32_to_f16_bytes(&scales)),
            quants: ctx.storage(bytemuck::cast_slice(&quants)),
        }
    };
    // The reference concat is `eh_proj(cat[enorm(embed); hnorm(hidden)])` (vLLM glm4_moe_mtp) —
    // embed FIRST, so the LEFT half multiplies the embedding. Same acceptance-probe matrix as the
    // qwen3.5 head (`OSFKB_MTP_WIRING`): 0 = reference (left→emb, norms by name); 1 = norms
    // crossed; 2 = fc halves ALSO swapped.
    let wiring = std::env::var("OSFKB_MTP_WIRING")
        .ok()
        .and_then(|v| v.parse::<u32>().ok())
        .unwrap_or(0);
    let (fc_h, fc_e) = if wiring == 2 {
        (q4v(&left), q4v(&right))
    } else {
        (q4v(&right), q4v(&left))
    };
    let (nm_emb, nm_hid) = if wiring >= 1 {
        ("hnorm", "enorm")
    } else {
        ("enorm", "hnorm")
    };
    let getp = |name: &str| -> Result<wgpu::Buffer> {
        Ok(ctx.storage(&st.tensor_f32(&format!("{p}.{name}.weight"))?))
    };
    let head_q4 = {
        let name = format!("{p}.shared_head.head.weight");
        let shape = st.shape(&name)?;
        let (scales, quants) = quantize_q4_0(&st.tensor_f32(&name)?, shape[0], shape[1]);
        Q4 {
            scales: ctx.storage_bytes(&f32_to_f16_bytes(&scales)),
            quants: ctx.storage(bytemuck::cast_slice(&quants)),
        }
    };
    let head = MtpHead {
        fc_hidden: fc_h,
        fc_emb: fc_e,
        pre_norm_emb: getp(nm_emb)?,
        pre_norm_hidden: getp(nm_hid)?,
        norm: getp("shared_head.norm")?,
        layer: build_glm_layer(st, ctx, cfg, &p)?,
        head: Some(head_q4),
        embed: Some(ctx.storage(&st.tensor_f32(&format!("{p}.embed_tokens.weight"))?)),
    };
    ctx.queue.submit(std::iter::empty());
    let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
    eprintln!("glm load: mtp head done (own head + embed)");
    Ok(Some(head))
}

/// The checkpoint's Multi-Token-Prediction draft head (DeepSeek-V3-style, `mtp.*` tensors):
/// one full qwen3.5 transformer layer + an `fc` merging `[norm(hidden); norm(embed(token))]`,
/// sharing the main model's embedding and lm_head. Drives self-speculative decoding — the same
/// draft source the reference vLLM deployment runs with `num_speculative_tokens: 4`.
pub struct MtpHead {
    /// LEFT half of `mtp.fc.weight` [hidden, 2·hidden]: the columns multiplying the NORMED
    /// HIDDEN — split at load so the merge is two standard GEMV-accumulates (`fc·[a;b] =
    /// fc_left·a + fc_right·b`), no concat buffer or kernel. If real-model acceptance lands
    /// near zero the halves' roles are swapped: flip them HERE, nowhere else.
    pub fc_hidden: Q4,
    /// RIGHT half of `mtp.fc.weight`: the columns multiplying the normed token embedding.
    pub fc_emb: Q4,
    /// `mtp.pre_fc_norm_embedding` ((1+w) folded).
    pub pre_norm_emb: wgpu::Buffer,
    /// `mtp.pre_fc_norm_hidden` ((1+w) folded).
    pub pre_norm_hidden: wgpu::Buffer,
    /// `mtp.norm` ((1+w) folded) — final norm before the SHARED lm_head.
    pub norm: wgpu::Buffer,
    /// `mtp.layers.0` — a standard qwen3.5 full-attention + MoE layer.
    pub layer: Layer,
    /// The head's OWN Q4 lm_head (GLM-OCR `shared_head.head` — byte-distinct from the main
    /// lm_head). `None` = share the main model's (qwen3.5).
    pub head: Option<Q4>,
    /// The head's OWN f32 embedding table (GLM-OCR `layers.16.embed_tokens`). `None` = share
    /// the main model's.
    pub embed: Option<wgpu::Buffer>,
}

pub struct Weights {
    pub cfg: Lfm2Config,
    /// [vocab, hidden] f32 — the token-gather row. Gemma-3: pre-scaled by `sqrt(hidden)` at load
    /// (the architecture's embedding normalizer), so the gather kernel stays scale-free.
    pub embed: wgpu::Buffer,
    /// [vocab, hidden] Q4_0 — the tied lm_head (UNscaled).
    pub embed_q4: Q4,
    pub embedding_norm: wgpu::Buffer,
    pub layers: Vec<Layer>,
    /// MTP draft head when the checkpoint ships one (qwen3.5-MoE multimodal checkpoints do).
    pub mtp: Option<MtpHead>,
    /// Q1/GGUF (Bonsai-class) models only: an f16 copy of the gather table. The f32 `embed`
    /// stays the source for the buffer-COPY staging paths (no size limit on copies), but a
    /// storage BINDING caps at 2 GiB-4 on wgpu-Vulkan — the gather kernels bind THIS table as
    /// two row-aligned sub-ranges instead (Metal-only builds could bind the f32 whole; the
    /// split keeps one code path for all backends).
    pub embed_f16: Option<wgpu::Buffer>,
    /// Q1/GGUF models only: the UPPER rows of the f32 staging table + the first row it holds.
    /// Vulkan's max_buffer_size (~4 GiB) cannot hold the 5 GB f32 table in one buffer, and the
    /// staging paths COPY rows out by (buffer, offset) — so the table ships as two halves and
    /// [`Self::embed_row_src`] picks per token. `None` = `embed` holds every row.
    pub embed_split: Option<(wgpu::Buffer, u32)>,
    /// Gemma-4 edge: model-level per-layer-embedding weights; `None` everywhere else.
    pub ple: Option<PleWeights>,
}

/// Gemma-4 edge model-level PLE weights. The token table stays Q4_0 on the GPU: E2B's bf16 table
/// is 4.7 GB — above Vulkan's max_buffer_size — and its information is embedding-like (the same
/// class this engine already serves Q4 as the tied head). One buffer, one gather kernel.
pub struct PleWeights {
    /// `embed_tokens_per_layer` `[vocab, n_layers·ple_dim]` Q4_0, layer-major columns.
    pub table: Q4,
    /// `per_layer_model_projection` `[n_layers·ple_dim, hidden]` (Q4).
    pub model_proj: Q4,
    /// `per_layer_projection_norm` `[ple_dim]` f32 (plain-w).
    pub proj_norm: wgpu::Buffer,
}

impl Weights {
    /// The (buffer, byte offset) holding token `tok`'s f32 embedding row — the staging-copy
    /// sources. Split models (Q1/GGUF, table > max_buffer_size) resolve the half here.
    pub fn embed_row_src(&self, tok: u32) -> (&wgpu::Buffer, u64) {
        let row = (self.cfg.hidden as u64) * 4;
        match &self.embed_split {
            Some((hi, split)) if tok >= *split => (hi, u64::from(tok - *split) * row),
            _ => (&self.embed, u64::from(tok) * row),
        }
    }

    /// Assemble the MTP draft head as a standalone ONE-LAYER MODEL sharing this model's
    /// embedding and lm_head (wgpu buffers are ref-counted — clones share VRAM): `stage_first =
    /// false` because its residual stream is fed by the fc merge instead of a token gather,
    /// `stage_last = true` for `mtp.norm` + the shared head. The draft engine therefore reuses
    /// every architecture arm of the generic plan builder — a future Gemma-MoE (or dense) draft
    /// head takes exactly the same path, no draft-specific kernels or plans.
    pub fn mtp_engine_weights(&self) -> Option<Weights> {
        let m = self.mtp.as_ref()?;
        let mut cfg = self.cfg.clone();
        cfg.n_layers = 1;
        cfg.layer_is_attn = vec![matches!(m.layer.op, Op::Attn { .. })];
        cfg.layer_is_sliding = vec![false];
        cfg.layer_k_eq_v = vec![false];
        cfg.stage_first = false;
        cfg.stage_last = true;
        // The draft head's containers were packed at the MAIN model's base dtype; a per-layer
        // override vector sized for the trunk must not leak into this 1-layer config.
        cfg.layer_wdtype = Vec::new();
        cfg.layer_kinds = Vec::new();
        cfg.head_kind = None;
        // Same rule for the edge geometry vectors: trunk-sized per-layer vectors must not leak
        // into this 1-layer config (no draft head exists on a Gemma-4 edge model anyway).
        cfg.edge = EdgeCfg::default();
        Some(Weights {
            cfg,
            // A head that ships its OWN embed/lm_head (GLM-OCR) decodes through those; otherwise
            // the draft shares the main model's tables (qwen3.5 — ref-counted buffer clones).
            embed: m
                .embed
                .clone()
                .unwrap_or_else(|| self.embed.clone()),
            embed_q4: m.head.clone().unwrap_or_else(|| self.embed_q4.clone()),
            embedding_norm: m.norm.clone(),
            layers: vec![m.layer.clone()],
            mtp: None,
            embed_f16: None,
            embed_split: None,
            ple: None,
        })
    }

    /// Load only layers `start..end` for pipeline-parallel sharding (Qwen3 for now): stage 0
    /// (`start == 0`) owns the f32 embedding gather; the last stage (`end == n_layers`) owns the
    /// final norm + tied Q4 head; middle stages hold neither (1-element dummy buffers — their plan
    /// never binds them). Layer indices stay ABSOLUTE in tensor names but local in `layers`.
    pub fn load_shard(
        ctx: &GpuCtx,
        dir: impl AsRef<Path>,
        start: usize,
        end: usize,
    ) -> Result<Self> {
        let dir = dir.as_ref();
        let cfg = Lfm2Config::from_json(&std::fs::read(dir.join("config.json"))?)?;
        let st = LazySt::open(dir)?;
        Self::load_shard_from(ctx, cfg, st, start, end)
    }

    /// Load a shard from in-memory `config.json` bytes + safetensors blob(s) — the wasm/browser
    /// path, where a worker's weights arrive by shipping (`inferencelayer::shard` OP_SHIP), never from
    /// a filesystem. `blobs` are whole safetensors files in order (single- or HF-sharded).
    pub fn load_shard_bytes(
        ctx: &GpuCtx,
        config_json: &[u8],
        blobs: Vec<Vec<u8>>,
        start: usize,
        end: usize,
    ) -> Result<Self> {
        let cfg = Lfm2Config::from_json(config_json)?;
        let st = LazySt::from_bytes(blobs)?;
        Self::load_shard_from(ctx, cfg, st, start, end)
    }

    fn load_shard_from(
        ctx: &GpuCtx,
        cfg: Lfm2Config,
        st: LazySt,
        start: usize,
        end: usize,
    ) -> Result<Self> {
        anyhow::ensure!(
            matches!(cfg.arch, Arch::Qwen3 | Arch::Qwen35),
            "load_shard: Qwen3 / Qwen3.5-MoE only"
        );
        anyhow::ensure!(start < end && end <= cfg.n_layers, "bad shard range");
        let tens = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(name) };
        let get = |name: &str| -> Result<wgpu::Buffer> { Ok(ctx.storage(&tens(name)?)) };
        let q4 = |name: &str| -> Result<Q4> {
            let shape = st.shape(name)?;
            let (rows, cols) = (shape[0], shape[1]);
            pack_weight(ctx, &tens(name)?, rows, cols, cfg.wdtype)
        };
        let q4cat = |names: &[String]| -> Result<Q4> {
            // Appends PACKED bytes per projection so all three are never held in f32 at once;
            // valid because both dtypes are strictly per-row (see `pack_weight_bytes`).
            let (mut scales, mut quants) = (Vec::new(), Vec::new());
            for name in names {
                let shape = st.shape(name)?;
                let (rows, cols) = (shape[0], shape[1]);
                let (sv, q) = pack_weight_bytes(&tens(name)?, rows, cols, cfg.wdtype)?;
                scales.extend_from_slice(&sv);
                quants.extend_from_slice(&q);
            }
            Ok(Q4 {
                scales: ctx.storage_bytes(if scales.is_empty() { &[0u8; 4] } else { &scales }),
                quants: ctx.storage_bytes(&quants),
            })
        };
        let (stage_first, stage_last) = (start == 0, end == cfg.n_layers);
        if cfg.arch == Arch::Qwen35 {
            let lm = "model.language_model";
            // The LAST stage also needs the f32 embedding when the checkpoint ships an MTP head:
            // the draft chain gathers embed(token) on the stage that owns norm + head.
            let embed = if stage_first || (stage_last && st.has("mtp.fc.weight")) {
                get(&format!("{lm}.embed_tokens.weight"))?
            } else {
                ctx.empty(1)
            };
            let (embed_q4, embedding_norm) = if stage_last {
                // TIED EMBEDDINGS: with no `lm_head.weight`, the head decodes through the embedding
                // table — and on a multimodal checkpoint that table lives under the language-model
                // prefix, not at `model.embed_tokens.weight`. The unsharded loader already does
                // this; the staged one hardcoded `lm_head.weight`, so a sharded serve of ANY tied
                // model (the claim-extractor 4B among them) died at load with `missing tensor
                // lm_head.weight` — the profiler is how this surfaced.
                let head_name = if st.has("lm_head.weight") {
                    "lm_head.weight".to_string()
                } else {
                    format!("{lm}.embed_tokens.weight")
                };
                (
                    q4(&head_name)?,
                    get_p1_at(&st, ctx, &format!("{lm}.norm.weight"))?,
                )
            } else {
                (
                    Q4 {
                        scales: ctx.storage_bytes(&[0u8; 4]),
                        quants: ctx.storage(&[0.0f32]),
                    },
                    ctx.empty(1),
                )
            };
            let c = &cfg;
            let mut layers = Vec::with_capacity(end - start);
            for i in start..end {
                layers.push(build_qwen35_layer(
                    &st,
                    ctx,
                    c,
                    &format!("{lm}.layers.{i}"),
                    c.layer_is_attn[i],
                )?);
                ctx.queue.submit(std::iter::empty());
                let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
                eprintln!("qwen35 shard load: layer {i} done");
            }
            // The MTP draft head rides the LAST stage (it needs the final norm + shared head);
            // other stages skip it.
            let mtp = if stage_last {
                load_mtp(&st, ctx, c)?
            } else {
                None
            };
            let mut cfg = cfg;
            cfg.n_layers = end - start;
            cfg.layer_is_attn = cfg.layer_is_attn[start..end].to_vec();
            cfg.layer_is_sliding = vec![false; end - start];
            cfg.layer_k_eq_v = vec![false; end - start];
            cfg.stage_first = stage_first;
            cfg.stage_last = stage_last;
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        let embed = if stage_first {
            get("model.embed_tokens.weight")?
        } else {
            ctx.empty(1)
        };
        let head_name = if st.has("lm_head.weight") {
            "lm_head.weight"
        } else if cfg.arch == Arch::Moshi {
            "text_linear.weight" // untied text head; kyutai naming
        } else {
            "model.embed_tokens.weight"
        };
        let (embed_q4, embedding_norm) = if stage_last {
            (q4(head_name)?, get("model.norm.weight")?)
        } else {
            (
                Q4 {
                    scales: ctx.storage_bytes(&[0u8; 4]),
                    quants: ctx.storage(&[0.0f32]),
                },
                ctx.empty(1),
            )
        };
        let mut layers = Vec::with_capacity(end - start);
        for i in start..end {
            let p = format!("model.layers.{i}");
            let op = Op::Attn {
                qkv: q4cat(&[
                    format!("{p}.self_attn.q_proj.weight"),
                    format!("{p}.self_attn.k_proj.weight"),
                    format!("{p}.self_attn.v_proj.weight"),
                ])?,
                o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                q_norm: get(&format!("{p}.self_attn.q_norm.weight"))?,
                k_norm: get(&format!("{p}.self_attn.k_norm.weight"))?,
                window: 0,
                local_rope: false,
                attn_gate: None,
                qkv_bias: None,
            };
            layers.push(Layer {
                operator_norm: get(&format!("{p}.input_layernorm.weight"))?,
                ffn_norm: get(&format!("{p}.post_attention_layernorm.weight"))?,
                op,
                w1: q4(&format!("{p}.mlp.gate_proj.weight"))?,
                w2: q4(&format!("{p}.mlp.down_proj.weight"))?,
                w3: q4(&format!("{p}.mlp.up_proj.weight"))?,
                post_op_norm: None,
                post_ffn_norm: None,
                moe: None,
                edge: None,
            });
        }
        let mut cfg = cfg;
        cfg.n_layers = end - start;
        cfg.layer_is_attn = vec![true; end - start];
        cfg.layer_is_sliding = vec![false; end - start];
        cfg.layer_k_eq_v = vec![false; end - start];
        cfg.stage_first = stage_first;
        cfg.stage_last = stage_last;
        Ok(Self {
            cfg,
            embed,
            embed_q4,
            embedding_norm,
            layers,
            mtp: None,
            embed_f16: None,
            embed_split: None,
            ple: None,
        })
    }

    /// Load `config.json` + the checkpoint from `dir`, honouring the weight-precision policy in
    /// `OSFKB_DECODE_PRECISION` (default all-Q4 — see [`Precision`]).
    pub fn load(ctx: &GpuCtx, dir: impl AsRef<Path>) -> Result<Self> {
        Self::load_with_precision(ctx, dir, Precision::from_env()?)
    }

    /// Load with an explicit weight-precision policy. Quantization is a speed/fidelity trade the
    /// CALLER owns: 4-bit weights are the dominant decode-bandwidth lever, but they cost real
    /// accuracy (a 270M gemma-3 loses ~2× the per-step agreement a 1B does, measured against an
    /// f32 reference), so a caller serving prose or classification may want f16 where a caller
    /// serving grammar-constrained output does not.
    ///
    /// A precision this build cannot SERVE is REFUSED rather than silently served as Q4 —
    /// shipping a model the caller did not ask for is the one outcome a precision knob must never
    /// produce.
    ///
    /// Servable today: uniform `q4` (the default, byte-for-byte historical) and uniform `f16`
    /// (the reference arm the quantization ruler measures against — see `crate::kld`). MIXED
    /// policies and any Q8 group still refuse: the mixed case needs a per-SITE bind-group choice
    /// in the plan rather than one family switch, because the f16 kernels take four bindings
    /// where the Q4 kernels take five, so a mixed plan cannot share one bind-group shape.
    pub fn load_with_precision(
        ctx: &GpuCtx,
        dir: impl AsRef<Path>,
        precision: Precision,
    ) -> Result<Self> {
        // Fast path: the historical all-Q4 loader is preserved byte-for-byte.
        if precision.groups().iter().all(|d| *d == WDtype::Q4) && precision.layers.is_empty() {
            return Self::load_q4(ctx, dir);
        }
        Self::load_dtype(ctx, dir, &precision)
    }

    /// Load `config.json` + the checkpoint (single- or multi-file safetensors) from `dir` into
    /// GPU buffers, quantizing every projection to Q4_0. Tensors are read lazily per byte range —
    /// RAM stays flat even for 32B+.
    fn load_q4(ctx: &GpuCtx, dir: impl AsRef<Path>) -> Result<Self> {
        Self::load_dtype(ctx, dir, &Precision::uniform(WDtype::Q4))
    }

    /// [`Self::load_q4`] with the BODY and HEAD storage dtypes chosen by the caller. The body
    /// dtype rides `cfg.wdtype` (every body loader closure consults it via [`pack_weight`]); the
    /// head dtype rides `cfg.head_wdtype` (the head-packing sites and [`crate::Q4LmHead`] read
    /// it). The plan builder reads both to pick the matching kernel families.
    fn load_dtype(ctx: &GpuCtx, dir: impl AsRef<Path>, precision: &Precision) -> Result<Self> {
        let dir = dir.as_ref();
        let mut cfg = Lfm2Config::from_json(&std::fs::read(dir.join("config.json"))?)?;
        cfg.head_wdtype = precision.head;
        match precision.uniform_body() {
            Some(body) => {
                // Uniform body: the historical per-layer path (byte-identical when unmixed).
                cfg.wdtype = body;
                cfg.layer_wdtype = precision.layer_dtypes(cfg.n_layers)?;
            }
            None => {
                // BODY-MIXED policy (e.g. q4,gate_up=q8): resolve explicit per-SITE kinds; the
                // plan is per-site kind-aware and the pack cell is set per site below.
                cfg.wdtype = WDtype::Q4;
                cfg.layer_kinds = (0..cfg.n_layers)
                    .map(|li| precision.site_kinds(li, cfg.n_layers))
                    .collect::<Result<_>>()?;
            }
        }
        // Per-layer overrides are honored where the arch's layer loop sets the dtype cell below;
        // everywhere else they REFUSE (a silently uniform load would ship the wrong model).
        let kind_wired = matches!(
            cfg.arch,
            Arch::Gemma3 | Arch::Gemma4 | Arch::Qwen3 | Arch::Llama | Arch::Qwen2
        );
        anyhow::ensure!(
            (cfg.layer_wdtype.is_empty() && cfg.layer_kinds.is_empty()) || kind_wired,
            "per-layer (blk:) precision is not wired for {:?} yet — supported: Gemma3/Gemma4/\
             Qwen3/Llama/Qwen2",
            cfg.arch
        );
        // A Q8 BODY (uniform or per-layer) needs the kind-wired plans — their Attn sites pick
        // kernels per kind. A Q8 HEAD alone is arch-generic (Q4LmHead is kind-aware), but the
        // grammar-sparse constrained head has no Q8 twin (that path asserts separately).
        anyhow::ensure!(
            kind_wired
                || (cfg.wdtype != WDtype::Q8
                    && !cfg.layer_wdtype.iter().any(|d| *d == WDtype::Q8)
                    && cfg.layer_kinds.is_empty()),
            "q8 body precision is not wired for {:?} yet — supported: Gemma3/Gemma4/Qwen3/\
             Llama/Qwen2",
            cfg.arch
        );
        // The dtype the packing closures read; layer loops that honor per-layer overrides set it
        // to `cfg.wdtype_at(i)` before packing layer i. Everything outside a wired loop packs at
        // the uniform base.
        let layer_dtype = std::cell::Cell::new(cfg.wdtype);
        let st = LazySt::open(dir)?;
        let tens = |name: &str| -> Result<Vec<f32>> { st.tensor_f32(name) };
        let get = |name: &str| -> Result<wgpu::Buffer> { Ok(ctx.storage(&tens(name)?)) };
        // Gemma RMSNorm convention: `x/rms * (1+w)` — fold the +1 into the buffer so the kernels
        // stay convention-free.
        let get_p1 = |name: &str| -> Result<wgpu::Buffer> {
            let mut v = tens(name)?;
            for x in &mut v {
                *x += 1.0;
            }
            Ok(ctx.storage(&v))
        };
        // Q4_0 load (projection + MLP matrices — the bulk of decode bandwidth).
        let q4 = |name: &str| -> Result<Q4> {
            let shape = st.shape(name)?;
            let (rows, cols) = (shape[0], shape[1]);
            pack_weight(ctx, &tens(name)?, rows, cols, layer_dtype.get())
        };
        // The LM head packs at its OWN dtype (`head=f16` over a Q4 body is the recovery lever).
        // Even when the head is TIED to the embedding table, the head-as-matmul container is
        // independent of the f32 gather table, so packing it at head_wdtype is correct.
        let q4head = |name: &str| -> Result<Q4> {
            let shape = st.shape(name)?;
            let (rows, cols) = (shape[0], shape[1]);
            pack_weight(ctx, &tens(name)?, rows, cols, cfg.head_wdtype)
        };
        // Concatenated Q4 load: q|k|v rows stacked (each row independent in Q4) → one Q4 GEMV.
        let q4cat = |names: &[String]| -> Result<Q4> {
            let (mut scales, mut quants) = (Vec::new(), Vec::new());
            for name in names {
                let shape = st.shape(name)?;
                let (rows, cols) = (shape[0], shape[1]);
                let (s, q) = pack_weight_bytes(&tens(name)?, rows, cols, layer_dtype.get())?;
                scales.extend_from_slice(&s);
                quants.extend_from_slice(&q);
            }
            Ok(Q4 {
                scales: ctx.storage_bytes(if scales.is_empty() { &[0u8; 4] } else { &scales }),
                quants: ctx.storage_bytes(&quants),
            })
        };

        // Tied head unless the checkpoint carries a separate lm_head (Qwen3 ≥ 4B unties). When the
        // head is TIED it decodes through the embedding table — which on a multimodal Qwen3.5 lives
        // under `model.language_model.*`, not at `model.*`. (The MoE checkpoints untie, so only the
        // dense/tied ones — e.g. claim-extractor-4B — reach this arm.)
        let head_name = if st.has("lm_head.weight") {
            "lm_head.weight"
        } else if cfg.arch == Arch::Moshi {
            "text_linear.weight" // untied text head; kyutai naming
        } else if st.has("model.language_model.embed_tokens.weight") {
            "model.language_model.embed_tokens.weight"
        } else if cfg.arch == Arch::Falcon {
            "transformer.word_embeddings.weight" // Falcon's tied head + `transformer.*` naming
        } else if cfg.arch == Arch::Mamba {
            "backbone.embeddings.weight" // Mamba's tied head + `backbone.*` naming
        } else if cfg.arch == Arch::Rwkv {
            "head.weight" // RWKV has a SEPARATE (untied) head, and `rwkv.*` naming for the trunk
        } else {
            "model.embed_tokens.weight"
        };
        let embed_q4 = q4head(head_name)?;
        if cfg.arch == Arch::Qwen35 {
            // Multimodal checkpoint: text stack under `model.language_model.*`; `model.visual.*`
            // and `mtp.*` are ignored (vision input and the MTP draft head are out of scope).
            let lm = "model.language_model";
            eprintln!("qwen35 load: embed…");
            let embed_data = tens(&format!("{lm}.embed_tokens.weight"))?;
            let embed = ctx.storage(&embed_data);
            // NuExtract3's 2.37 GB f32 embedding exceeds the ~2 GB storage-buffer binding cap on
            // wgpu-Vulkan, so the on-GPU decode gather cannot bind it and `batch_step_chained`
            // panics. Build the f16 twin (1.18 GB) that the split `GATHER_K_F16` reads — exactly as
            // the GGUF/Q1 loader does (`gguf.rs`) — whenever the f32 table is over the cap. The f32
            // `embed` stays for the prefill/copy paths (single buffer: 2.37 GB < 4 GB max_buffer_size,
            // so no `embed_split` needed here, unlike the 5 GB Q1 table).
            let embed_f16 = {
                let cap = ctx.device.limits().max_storage_buffer_binding_size as u64;
                ((embed_data.len() as u64 * 4) > cap).then(|| {
                    eprintln!("qwen35 load: f16 embed twin for oversized on-GPU gather…");
                    ctx.storage_bytes(&f32_to_f16_bytes(&embed_data))
                })
            };
            ctx.queue.submit(std::iter::empty());
            let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
            let embedding_norm = get_p1_at(&st, ctx, &format!("{lm}.norm.weight"))?;
            let c = &cfg;
            let mut layers = Vec::with_capacity(c.n_layers);
            for i in 0..c.n_layers {
                layers.push(build_qwen35_layer(
                    &st,
                    ctx,
                    c,
                    &format!("{lm}.layers.{i}"),
                    c.layer_is_attn[i],
                )?);
                // Flush pending upload staging: `create_buffer_init` queues host→device copies
                // until a submit, so a 20 GB load would otherwise hold ~2× VRAM on Vulkan
                // (discrete memory; Metal's unified memory never showed it).
                ctx.queue.submit(std::iter::empty());
                let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
                eprintln!("qwen35 load: layer {i} done");
            }
            let mtp = load_mtp(&st, ctx, c)?;
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp,
                embed_f16,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Moshi {
            // Moshi temporal stack, kyutai-native names (see `Arch::Moshi`). The fused
            // `in_projs.0` is q|k|v row-stacked at full head width; q and k row ranges get the
            // GLM-OCR GPT-J→NeoX de-interleave IN PLACE (v untouched), then one Q4 buffer.
            let embed = get("text_emb.weight")?; // Step::Embed placeholder — see Arch::Moshi doc
            let embedding_norm = get("out_norm.alpha")?;
            let c = &cfg;
            let (hd, hidden) = (c.head_dim, c.hidden);
            let ones: Vec<f32> = vec![1.0f32; hd];
            let window = c.sliding_window as u32;
            let mut layers = Vec::with_capacity(c.n_layers);
            for i in 0..c.n_layers {
                let p = format!("transformer.layers.{i}");
                let mut fused = tens(&format!("{p}.self_attn.in_projs.0.weight"))?;
                anyhow::ensure!(fused.len() == 3 * hidden * hidden, "in_projs.0 shape");
                let half = hd / 2;
                for qk in 0..2 {
                    let base = qk * hidden * hidden;
                    let block = fused[base..base + hidden * hidden].to_vec();
                    for h in 0..c.n_heads {
                        for j in 0..hd {
                            let dst_lane = if j % 2 == 0 { j / 2 } else { half + j / 2 };
                            let src = (h * hd + j) * hidden;
                            let dst = base + (h * hd + dst_lane) * hidden;
                            fused[dst..dst + hidden].copy_from_slice(&block[src..src + hidden]);
                        }
                    }
                }
                let (sc, qz) = quantize_q4_0(&fused, 3 * hidden, hidden);
                // FFN at Q8_0 — 4-bit destroys the gating matrices (moshi_lm q4sim sweep:
                // FFN-q8 + attn-q4 = 78/81 tokens vs 50/81 all-q4); attention/head stay Q4.
                // (Attention-at-q8 was measured as nearly free — the engine is kernel-bound,
                // see tests/moshi_kernel_bw.rs — but the shared plan builder has FOUR qkv
                // dispatch sites and wiring only one drops the model to 9/81 tokens: the q4
                // kernel then reads q8 bytes. Do all four, or none.)
                let gu = tens(&format!("{p}.gating.linear_in.weight"))?;
                anyhow::ensure!(gu.len() == 2 * c.intermediate * hidden, "linear_in shape");
                let half_rows = c.intermediate * hidden;
                let (s1, z1) = quantize_q8_0(&gu[..half_rows], c.intermediate, hidden);
                let (s3, z3) = quantize_q8_0(&gu[half_rows..], c.intermediate, hidden);
                layers.push(Layer {
                    operator_norm: get(&format!("{p}.norm1.alpha"))?,
                    ffn_norm: get(&format!("{p}.norm2.alpha"))?,
                    op: Op::Attn {
                        qkv: Q4 {
                            scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                            quants: ctx.storage(bytemuck::cast_slice(&qz)),
                        },
                        o: q4(&format!("{p}.self_attn.out_projs.0.weight"))?,
                        q_norm: ctx.storage(&ones),
                        k_norm: ctx.storage(&ones),
                        window,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: None,
                    },
                    w1: Q4 {
                        scales: ctx.storage_bytes(&f32_to_f16_bytes(&s1)),
                        quants: ctx.storage(bytemuck::cast_slice(&z1)),
                    },
                    w2: {
                        let name = format!("{p}.gating.linear_out.weight");
                        let shape = st.shape(&name)?;
                        let (sc2, qz2) = quantize_q8_0(&tens(&name)?, shape[0], shape[1]);
                        Q4 {
                            scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc2)),
                            quants: ctx.storage(bytemuck::cast_slice(&qz2)),
                        }
                    },
                    w3: Q4 {
                        scales: ctx.storage_bytes(&f32_to_f16_bytes(&s3)),
                        quants: ctx.storage(bytemuck::cast_slice(&z3)),
                    },
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
                ctx.queue.submit(std::iter::empty());
                let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
            }
            return Ok(Self {
                cfg,
                embed,
                embed_f16: None,
                embed_split: None,
                ple: None,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
            });
        }
        if cfg.arch == Arch::GlmOcr {
            // GLM-OCR text stack: `model.language_model.*`, untied `lm_head` (head_name above
            // already picked it). `model.visual.*` (the tower — served by `vision_glm`) is never
            // touched; `model.language_model.layers.16.*` (the MTP draft head, which HF itself
            // ignores at inference) loads only when serving-MTP is enabled — see `load_glm_mtp`.
            let lm = "model.language_model";
            let embed = get(&format!("{lm}.embed_tokens.weight"))?;
            let embedding_norm = get(&format!("{lm}.norm.weight"))?; // plain `w`, not Gemma's (1+w)
            let c = &cfg;
            let mut layers = Vec::with_capacity(c.n_layers);
            for i in 0..c.n_layers {
                layers.push(build_glm_layer(&st, ctx, c, &format!("{lm}.layers.{i}"))?);
                ctx.queue.submit(std::iter::empty());
                let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
            }
            let mtp = load_glm_mtp(&st, ctx, c)?;
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Mamba {
            // Mamba-1: every layer is a selective-scan mixer (Op::Mamba). RMSNorm pre-norm rides
            // `operator_norm`; the post-mixer FFN slot is a tiny ZERO GLU (adds 0 — Mamba has no
            // FFN). `ssm_par` packs [A_log | D | dt_bias]. Mixer dims are re-read from config.json.
            let cj: serde_json::Value =
                serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?;
            let cg = |k: &str| cj.get(k).and_then(|x| x.as_u64()).map(|u| u as usize);
            let h = cfg.hidden;
            let d_state = cg("state_size").unwrap_or(16);
            let d_conv = cg("conv_kernel").unwrap_or(4);
            let expand = cg("expand").unwrap_or(2);
            let d_inner = expand * h;
            let dt_rank = cg("time_step_rank").unwrap_or(h.div_ceil(16));
            let ones32: Vec<f32> = vec![1.0f32; 32];
            let zero_gu = quantize_q4_0(&vec![0.0f32; 32 * h], 32, h);
            let zero_dn = quantize_q4_0(&vec![0.0f32; h * 32], h, 32);
            let zq = |sc: &[f32], qz: &[u32]| -> Q4 {
                Q4 {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(sc)),
                    quants: ctx.storage(bytemuck::cast_slice(qz)),
                }
            };
            let embed = get("backbone.embeddings.weight")?;
            let embedding_norm = get("backbone.norm_f.weight")?;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let mp = format!("backbone.layers.{i}.mixer");
                // conv1d.weight is [d_inner, 1, d_conv] → squeeze the singleton to [d_inner, d_conv].
                let conv_w = tens(&format!("{mp}.conv1d.weight"))?;
                anyhow::ensure!(conv_w.len() == d_inner * d_conv, "mamba conv1d shape");
                // ssm_par = [A_log(d_inner·d_state) | D(d_inner) | dt_bias(d_inner)].
                let mut ssm = tens(&format!("{mp}.A_log"))?;
                ssm.extend_from_slice(&tens(&format!("{mp}.D"))?);
                ssm.extend_from_slice(&tens(&format!("{mp}.dt_proj.bias"))?);
                // dt_proj is [d_inner, dt_rank]; Q4 needs cols % 32, so zero-pad the columns to the
                // next multiple of 32. The pad columns multiply the B/C part of the (dt|B|C) input in
                // the plan's gemv — but they are ZERO, so `dt_raw` is unchanged (the SSM reads B/C at
                // the REAL dt_rank offset).
                let dt_pad = dt_rank.div_ceil(32) * 32;
                let dt_q4 = {
                    let src = tens(&format!("{mp}.dt_proj.weight"))?; // [d_inner, dt_rank]
                    let mut padded = vec![0f32; d_inner * dt_pad];
                    for r in 0..d_inner {
                        padded[r * dt_pad..r * dt_pad + dt_rank]
                            .copy_from_slice(&src[r * dt_rank..(r + 1) * dt_rank]);
                    }
                    let (sc, qz) = quantize_q4_0(&padded, d_inner, dt_pad);
                    zq(&sc, &qz)
                };
                layers.push(Layer {
                    operator_norm: get(&format!("backbone.layers.{i}.norm.weight"))?,
                    ffn_norm: ctx.storage(&ones32), // unused (zero FFN)
                    op: Op::Mamba {
                        in_proj: q4(&format!("{mp}.in_proj.weight"))?,
                        conv_w: ctx.storage(&conv_w),
                        conv_b: get(&format!("{mp}.conv1d.bias"))?,
                        x_proj: q4(&format!("{mp}.x_proj.weight"))?,
                        dt_proj: dt_q4,
                        ssm_par: ctx.storage(&ssm),
                        out_proj: q4(&format!("{mp}.out_proj.weight"))?,
                        d_inner,
                        d_state,
                        d_conv,
                        dt_rank,
                    },
                    w1: zq(&zero_gu.0, &zero_gu.1),
                    w2: zq(&zero_dn.0, &zero_dn.1),
                    w3: zq(&zero_gu.0, &zero_gu.1),
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Rwkv {
            // RWKV-4: an RNN. Each block is time-mix (WKV recurrence) + channel-mix, both token-shift
            // mixed, both with an affine LayerNorm (ln1→operator_norm, ln2→ffn_norm packed [w|b]).
            // ln0 (extra affine LN) rides only block 0. Final norm is affine (embedding_norm=[w|b]).
            // The trailing GLU-FFN slot is a tiny ZERO GLU (adds 0 — the channel-mix is Op::Rwkv).
            let cj: serde_json::Value =
                serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?;
            let cg = |k: &str| cj.get(k).and_then(|x| x.as_u64()).map(|u| u as usize);
            let h = cfg.hidden;
            let ffn_dim = cg("intermediate_size").unwrap_or(4 * h);
            let zero_gu = quantize_q4_0(&vec![0.0f32; 32 * h], 32, h);
            let zero_dn = quantize_q4_0(&vec![0.0f32; h * 32], h, 32);
            let zq = |sc: &[f32], qz: &[u32]| -> Q4 {
                Q4 {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(sc)),
                    quants: ctx.storage(bytemuck::cast_slice(qz)),
                }
            };
            // Pack an affine LayerNorm as [weight | bias] for LAYERNORM_AFFINE.
            let ln_pack = |wn: &str, bn: &str| -> Result<wgpu::Buffer> {
                let mut p = tens(wn)?;
                p.extend_from_slice(&tens(bn)?);
                Ok(ctx.storage(&p))
            };
            let embed = get("rwkv.embeddings.weight")?;
            let embedding_norm = ln_pack("rwkv.ln_out.weight", "rwkv.ln_out.bias")?;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let bp = format!("rwkv.blocks.{i}");
                let ap = format!("{bp}.attention");
                let fp = format!("{bp}.feed_forward");
                // Time-mix token-shift interpolation: [time_mix_k | time_mix_v | time_mix_r].
                let mut att_mix = tens(&format!("{ap}.time_mix_key"))?;
                att_mix.extend_from_slice(&tens(&format!("{ap}.time_mix_value"))?);
                att_mix.extend_from_slice(&tens(&format!("{ap}.time_mix_receptance"))?);
                anyhow::ensure!(att_mix.len() == 3 * h, "rwkv att time_mix shape");
                // Channel-mix has NO value mix: [time_mix_k | 0(unused) | time_mix_r].
                let mut ffn_mix = tens(&format!("{fp}.time_mix_key"))?;
                ffn_mix.extend_from_slice(&vec![0f32; h]);
                ffn_mix.extend_from_slice(&tens(&format!("{fp}.time_mix_receptance"))?);
                // WKV uses `-exp(time_decay)`; fold the -exp at load (a static, exact transform).
                let time_decay: Vec<f32> = tens(&format!("{ap}.time_decay"))?
                    .iter()
                    .map(|x| -x.exp())
                    .collect();
                let ln0 = if i == 0 {
                    Some(ln_pack(
                        &format!("{bp}.pre_ln.weight"),
                        &format!("{bp}.pre_ln.bias"),
                    )?)
                } else {
                    None
                };
                layers.push(Layer {
                    operator_norm: ln_pack(&format!("{bp}.ln1.weight"), &format!("{bp}.ln1.bias"))?,
                    ffn_norm: ln_pack(&format!("{bp}.ln2.weight"), &format!("{bp}.ln2.bias"))?,
                    op: Op::Rwkv {
                        ln0,
                        att_mix: ctx.storage(&att_mix),
                        att_k: q4(&format!("{ap}.key.weight"))?,
                        att_v: q4(&format!("{ap}.value.weight"))?,
                        att_r: q4(&format!("{ap}.receptance.weight"))?,
                        att_o: q4(&format!("{ap}.output.weight"))?,
                        time_decay: ctx.storage(&time_decay),
                        time_first: get(&format!("{ap}.time_first"))?,
                        ffn_mix: ctx.storage(&ffn_mix),
                        ffn_k: q4(&format!("{fp}.key.weight"))?,
                        ffn_v: q4(&format!("{fp}.value.weight"))?,
                        ffn_r: q4(&format!("{fp}.receptance.weight"))?,
                        ffn_dim,
                    },
                    w1: zq(&zero_gu.0, &zero_gu.1),
                    w2: zq(&zero_dn.0, &zero_dn.1),
                    w3: zq(&zero_gu.0, &zero_gu.1),
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Olmo2 {
            // OLMo2: POST-norm RMSNorm. NO pre-norm (operator_norm/ffn_norm are unused ones — the
            // plan skips them), the sub-layer outputs are RMSNorm'd via the sandwich post_op_norm /
            // post_ffn_norm, and a FULL-VECTOR qk-norm (q_norm [nh·hd], k_norm [nkv·hd]) runs before
            // rope. SwiGLU, no biases. `model.norm` is the ordinary final RMSNorm.
            let ones_h: Vec<f32> = vec![1.0f32; cfg.hidden];
            let ones_hd: Vec<f32> = vec![1.0f32; cfg.head_dim];
            let embed = get("model.embed_tokens.weight")?;
            let embedding_norm = get("model.norm.weight")?;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                // `operator_norm` (otherwise unused on this post-norm arch) carries the FULL-vector
                // qk-norm weights packed [q_norm(nh·hd) | k_norm(nkv·hd)] for QK_FULLNORM (which
                // applies inv AND the weight); the Op's q_norm/k_norm are ONES so qkrc adds no norm.
                let qk_pack = {
                    let mut w = tens(&format!("{p}.self_attn.q_norm.weight"))?;
                    w.extend_from_slice(&tens(&format!("{p}.self_attn.k_norm.weight"))?);
                    ctx.storage(&w)
                };
                layers.push(Layer {
                    operator_norm: qk_pack,
                    ffn_norm: ctx.storage(&ones_h), // unused (skip_prenorm)
                    op: Op::Attn {
                        qkv: q4cat(&[
                            format!("{p}.self_attn.q_proj.weight"),
                            format!("{p}.self_attn.k_proj.weight"),
                            format!("{p}.self_attn.v_proj.weight"),
                        ])?,
                        o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                        q_norm: ctx.storage(&ones_hd), // qkrc identity (QK_FULLNORM did the norm)
                        k_norm: ctx.storage(&ones_hd),
                        window: 0,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: None,
                    },
                    w1: q4(&format!("{p}.mlp.gate_proj.weight"))?,
                    w2: q4(&format!("{p}.mlp.down_proj.weight"))?,
                    w3: q4(&format!("{p}.mlp.up_proj.weight"))?,
                    // Sandwich POST-norms — the attn / mlp outputs are RMSNorm'd before the residual.
                    post_op_norm: Some(get(&format!("{p}.post_attention_layernorm.weight"))?),
                    post_ffn_norm: Some(get(&format!("{p}.post_feedforward_layernorm.weight"))?),
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Phi2 {
            // Phi-2: parallel attn/MLP, non-gated tanh-GELU MLP, affine LayerNorm, biases everywhere.
            // Bias packing (so no per-layer struct field): each layer's `qkv_bias` holds
            //   [q_bias | k_bias | v_bias | fc1_bias(im) | (dense_bias + fc2_bias)(hidden)];
            // `embedding_norm` holds [final_ln.weight | final_ln.bias | lm_head.bias(vocab)].
            let (h, _im, vocab) = (cfg.hidden, cfg.intermediate, cfg.vocab);
            let ones: Vec<f32> = vec![1.0f32; cfg.head_dim];
            let packed_norm = |name: &str| -> Result<Vec<f32>> {
                let mut wb = tens(&format!("{name}.weight"))?;
                wb.extend_from_slice(&tens(&format!("{name}.bias"))?);
                Ok(wb)
            };
            let embed = get("model.embed_tokens.weight")?;
            // Final norm [weight | bias] + the untied lm_head bias appended.
            let embedding_norm = {
                let mut nb = packed_norm("model.final_layernorm")?;
                nb.extend_from_slice(&tens("lm_head.bias")?);
                anyhow::ensure!(nb.len() == 2 * h + vocab, "phi2 embedding_norm pack shape");
                ctx.storage(&nb)
            };
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                // qkv_bias packs the attention AND FFN biases (read by offset in the plan).
                let mut qkv_b = tens(&format!("{p}.self_attn.q_proj.bias"))?;
                qkv_b.extend_from_slice(&tens(&format!("{p}.self_attn.k_proj.bias"))?);
                qkv_b.extend_from_slice(&tens(&format!("{p}.self_attn.v_proj.bias"))?);
                qkv_b.extend_from_slice(&tens(&format!("{p}.mlp.fc1.bias"))?);
                // Parallel: o_proj(dense) bias and fc2 bias both add to the residual → combine.
                let dense_b = tens(&format!("{p}.self_attn.dense.bias"))?;
                let fc2_b = tens(&format!("{p}.mlp.fc2.bias"))?;
                qkv_b.extend((0..h).map(|j| dense_b[j] + fc2_b[j]));
                let ln = ctx.storage(&packed_norm(&format!("{p}.input_layernorm"))?);
                let fc1 = q4(&format!("{p}.mlp.fc1.weight"))?;
                layers.push(Layer {
                    operator_norm: ln.clone(),
                    ffn_norm: ln, // parallel: unused (hoisted norm reads operator_norm)
                    op: Op::Attn {
                        qkv: q4cat(&[
                            format!("{p}.self_attn.q_proj.weight"),
                            format!("{p}.self_attn.k_proj.weight"),
                            format!("{p}.self_attn.v_proj.weight"),
                        ])?,
                        o: q4(&format!("{p}.self_attn.dense.weight"))?,
                        q_norm: ctx.storage(&ones),
                        k_norm: ctx.storage(&ones),
                        window: 0,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: Some(ctx.storage(&qkv_b)),
                    },
                    w1: fc1.clone(), // non-gated: gate slot unused
                    w2: q4(&format!("{p}.mlp.fc2.weight"))?,
                    w3: fc1,
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Nemotron {
            // Nemotron / Minitron: LayerNorm1P (norm buffers packed [(weight+1) | bias]) + squared-
            // ReLU MLP (w1=w3=up_proj so the gated FFN computes relu(up·x)·(up·x) = relu(up·x)²) +
            // partial rotary. No biases; sequential (two norms per layer).
            let ones: Vec<f32> = vec![1.0f32; cfg.head_dim];
            // Pack [(weight + 1) | bias] — the "1P" LayerNorm applies (1 + weight).
            let packed_1p = |name: &str| -> Result<wgpu::Buffer> {
                let mut wb = tens(&format!("{name}.weight"))?;
                for x in &mut wb {
                    *x += 1.0;
                }
                wb.extend_from_slice(&tens(&format!("{name}.bias"))?);
                Ok(ctx.storage(&wb))
            };
            let embed = get("model.embed_tokens.weight")?;
            let embedding_norm = packed_1p("model.norm")?;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                let up = q4(&format!("{p}.mlp.up_proj.weight"))?; // fc1
                layers.push(Layer {
                    operator_norm: packed_1p(&format!("{p}.input_layernorm"))?,
                    ffn_norm: packed_1p(&format!("{p}.post_attention_layernorm"))?,
                    op: Op::Attn {
                        qkv: q4cat(&[
                            format!("{p}.self_attn.q_proj.weight"),
                            format!("{p}.self_attn.k_proj.weight"),
                            format!("{p}.self_attn.v_proj.weight"),
                        ])?,
                        o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                        q_norm: ctx.storage(&ones),
                        k_norm: ctx.storage(&ones),
                        window: 0,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: None,
                    },
                    w1: up.clone(), // squared-ReLU: gate = up
                    w2: q4(&format!("{p}.mlp.down_proj.weight"))?, // fc2
                    w3: up,
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Cohere {
            // Cohere / Command-R: parallel attn/MLP, gated SwiGLU, affine WEIGHT-ONLY LayerNorm
            // (packed [weight | zeros] for LAYERNORM_AFFINE), no biases, tied head. `logit_scale`
            // (logits ×= scale) is FOLDED into the LM head here — Q4-exact under positive scaling —
            // so there's no runtime step; the (un-scaled) gather table stays as `embed`.
            let (h, vocab) = (cfg.hidden, cfg.vocab);
            let zeros: Vec<f32> = vec![0.0f32; h];
            let ones: Vec<f32> = vec![1.0f32; cfg.head_dim];
            // Pack a norm's weight with a zero bias into one [2·hidden] buffer (weight-only LN).
            let packed_wz = |name: &str| -> Result<wgpu::Buffer> {
                let mut wz = tens(&format!("{name}.weight"))?;
                wz.extend_from_slice(&zeros);
                Ok(ctx.storage(&wz))
            };
            let cj: serde_json::Value =
                serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?;
            let logit_scale = cj
                .get("logit_scale")
                .and_then(|x| x.as_f64())
                .unwrap_or(1.0) as f32;
            let embed = get("model.embed_tokens.weight")?;
            // Fold logit_scale into the (tied) LM head only.
            let embed_q4 = {
                let mut head = tens("model.embed_tokens.weight")?;
                for x in &mut head {
                    *x *= logit_scale;
                }
                let (sc, qz) = quantize_q4_0(&head, vocab, h);
                Q4 {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                    quants: ctx.storage(bytemuck::cast_slice(&qz)),
                }
            };
            let embedding_norm = packed_wz("model.norm")?;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                let ln = packed_wz(&format!("{p}.input_layernorm"))?;
                layers.push(Layer {
                    operator_norm: ln.clone(),
                    ffn_norm: ln, // parallel: unused (hoisted norm reads operator_norm)
                    op: Op::Attn {
                        qkv: q4cat(&[
                            format!("{p}.self_attn.q_proj.weight"),
                            format!("{p}.self_attn.k_proj.weight"),
                            format!("{p}.self_attn.v_proj.weight"),
                        ])?,
                        o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                        q_norm: ctx.storage(&ones),
                        k_norm: ctx.storage(&ones),
                        window: 0,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: None,
                    },
                    w1: q4(&format!("{p}.mlp.gate_proj.weight"))?,
                    w2: q4(&format!("{p}.mlp.down_proj.weight"))?,
                    w3: q4(&format!("{p}.mlp.up_proj.weight"))?,
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Falcon {
            // Falcon-7B uses `transformer.*` naming, a FUSED query_key_value (loads straight into the
            // qkv slot: [q|k|v], n_kv=1), a single input_layernorm (parallel — the plan hoists it, so
            // ffn_norm is unused: point it at the same buffer), and a NON-GATED MLP (w3=dense_h_to_4h,
            // w2=dense_4h_to_h; w1 unused — filled from w3 so the slot is a valid Q4). No biases.
            let ones: Vec<f32> = vec![1.0f32; cfg.head_dim];
            let packed_norm = |name: &str| -> Result<wgpu::Buffer> {
                let mut wb = tens(&format!("{name}.weight"))?;
                wb.extend_from_slice(&tens(&format!("{name}.bias"))?);
                Ok(ctx.storage(&wb))
            };
            let embed = get("transformer.word_embeddings.weight")?;
            let embedding_norm = packed_norm("transformer.ln_f")?;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("transformer.h.{i}");
                let ln = packed_norm(&format!("{p}.input_layernorm"))?;
                let up = q4(&format!("{p}.mlp.dense_h_to_4h.weight"))?; // fc1 (w3 slot)
                layers.push(Layer {
                    operator_norm: ln.clone(),
                    ffn_norm: ln, // parallel: unused (hoisted norm reads operator_norm)
                    op: Op::Attn {
                        qkv: q4(&format!("{p}.self_attention.query_key_value.weight"))?,
                        o: q4(&format!("{p}.self_attention.dense.weight"))?,
                        q_norm: ctx.storage(&ones),
                        k_norm: ctx.storage(&ones),
                        window: 0,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: None,
                    },
                    w1: up.clone(), // gate slot unused (non-gated MLP)
                    w2: q4(&format!("{p}.mlp.dense_4h_to_h.weight"))?, // fc2
                    w3: up,
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::StableLm {
            // StableLM-2: affine LayerNorm (norm buffers packed [weight | bias] for LAYERNORM_AFFINE)
            // + Qwen2-style q/k/v bias + partial rotary. o_proj and the MLP are bias-free.
            let ones: Vec<f32> = vec![1.0f32; cfg.head_dim];
            // Pack a norm's weight and bias back-to-back into one [2·hidden] buffer.
            let packed_norm = |name: &str| -> Result<wgpu::Buffer> {
                let mut wb = tens(&format!("{name}.weight"))?;
                wb.extend_from_slice(&tens(&format!("{name}.bias"))?);
                Ok(ctx.storage(&wb))
            };
            let embed = get("model.embed_tokens.weight")?;
            let embedding_norm = packed_norm("model.norm")?;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                let mut qkv_b = tens(&format!("{p}.self_attn.q_proj.bias"))?;
                qkv_b.extend_from_slice(&tens(&format!("{p}.self_attn.k_proj.bias"))?);
                qkv_b.extend_from_slice(&tens(&format!("{p}.self_attn.v_proj.bias"))?);
                layers.push(Layer {
                    operator_norm: packed_norm(&format!("{p}.input_layernorm"))?,
                    ffn_norm: packed_norm(&format!("{p}.post_attention_layernorm"))?,
                    op: Op::Attn {
                        qkv: q4cat(&[
                            format!("{p}.self_attn.q_proj.weight"),
                            format!("{p}.self_attn.k_proj.weight"),
                            format!("{p}.self_attn.v_proj.weight"),
                        ])?,
                        o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                        q_norm: ctx.storage(&ones),
                        k_norm: ctx.storage(&ones),
                        window: 0,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: Some(ctx.storage(&qkv_b)),
                    },
                    w1: q4(&format!("{p}.mlp.gate_proj.weight"))?,
                    w2: q4(&format!("{p}.mlp.down_proj.weight"))?,
                    w3: q4(&format!("{p}.mlp.up_proj.weight"))?,
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Olmo {
            // OLMo 1: Llama attention (no bias, no qk-norm, NeoX rope, GQA) + SwiGLU MLP. The
            // NON-PARAMETRIC LayerNorm has no weights in the checkpoint, so every norm buffer is
            // ONES — the decode plan mean-centers the input, and RMSNorm(centered)·1 = LayerNorm.
            let ones_h: Vec<f32> = vec![1.0f32; cfg.hidden];
            let ones_hd: Vec<f32> = vec![1.0f32; cfg.head_dim];
            let embed = get("model.embed_tokens.weight")?;
            let embedding_norm = ctx.storage(&ones_h); // final norm weight = ones (centered in-plan)
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                layers.push(Layer {
                    operator_norm: ctx.storage(&ones_h),
                    ffn_norm: ctx.storage(&ones_h),
                    op: Op::Attn {
                        qkv: q4cat(&[
                            format!("{p}.self_attn.q_proj.weight"),
                            format!("{p}.self_attn.k_proj.weight"),
                            format!("{p}.self_attn.v_proj.weight"),
                        ])?,
                        o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                        q_norm: ctx.storage(&ones_hd),
                        k_norm: ctx.storage(&ones_hd),
                        window: 0,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: None,
                    },
                    w1: q4(&format!("{p}.mlp.gate_proj.weight"))?,
                    w2: q4(&format!("{p}.mlp.down_proj.weight"))?,
                    w3: q4(&format!("{p}.mlp.up_proj.weight"))?,
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if matches!(cfg.arch, Arch::Qwen3 | Arch::Llama | Arch::Qwen2) {
            // Llama/Mistral/Qwen2 have NO qk-norm tensors; the no-qk-norm kernel variant still
            // reads the q/k-norm buffers (skipping the norm), so feed it ones — like GLM-OCR.
            let no_qknorm = matches!(cfg.arch, Arch::Llama | Arch::Qwen2);
            let has_bias = cfg.arch == Arch::Qwen2; // Qwen2 alone carries q/k/v projection bias
            let ones: Vec<f32> = vec![1.0f32; cfg.head_dim];
            let embed = get("model.embed_tokens.weight")?;
            let embedding_norm = get("model.norm.weight")?;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let ks = cfg.kinds_at(i);
                let p = format!("model.layers.{i}");
                let (q_norm, k_norm) = if no_qknorm {
                    (ctx.storage(&ones), ctx.storage(&ones))
                } else {
                    (
                        get(&format!("{p}.self_attn.q_norm.weight"))?,
                        get(&format!("{p}.self_attn.k_norm.weight"))?,
                    )
                };
                // Qwen2 bias: concat q|k|v proj bias in the SAME order as the fused qkv weight.
                let qkv_bias = if has_bias {
                    let mut b = tens(&format!("{p}.self_attn.q_proj.bias"))?;
                    b.extend_from_slice(&tens(&format!("{p}.self_attn.k_proj.bias"))?);
                    b.extend_from_slice(&tens(&format!("{p}.self_attn.v_proj.bias"))?);
                    Some(ctx.storage(&b))
                } else {
                    None
                };
                layer_dtype.set(ks.qkv.pack_dtype());
                let qkv = q4cat(&[
                    format!("{p}.self_attn.q_proj.weight"),
                    format!("{p}.self_attn.k_proj.weight"),
                    format!("{p}.self_attn.v_proj.weight"),
                ])?;
                layer_dtype.set(ks.o.pack_dtype());
                let op = Op::Attn {
                    qkv,
                    o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                    q_norm,
                    k_norm,
                    // Mistral/Qwen2 sliding window (uniform, global rope base); 0 = full causal.
                    window: cfg.sliding_window as u32,
                    local_rope: false,
                    attn_gate: None,
                    qkv_bias,
                };
                layer_dtype.set(ks.gate_up.pack_dtype());
                let w1 = q4(&format!("{p}.mlp.gate_proj.weight"))?;
                let w3 = q4(&format!("{p}.mlp.up_proj.weight"))?;
                layer_dtype.set(ks.down.pack_dtype());
                let w2 = q4(&format!("{p}.mlp.down_proj.weight"))?;
                layers.push(Layer {
                    operator_norm: get(&format!("{p}.input_layernorm.weight"))?,
                    // HF/Llama naming: `post_attention_layernorm` IS the pre-FFN norm.
                    ffn_norm: get(&format!("{p}.post_attention_layernorm.weight"))?,
                    op,
                    w1,
                    w2,
                    w3,
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Qwen2Moe {
            // Qwen2 bias attention + Qwen3.5-style shared-gated MoE. Routed experts are SEPARATE
            // per-expert tensors (Mixtral layout), concatenated; the shared expert fills the dense
            // GLU slots and is scaled by sigmoid(shared_expert_gate·x) — the forward's shared-gate
            // path (keyed off `shared_gate.is_some()`) does the rest, unchanged from Qwen3.5.
            let (e, mi, h) = (cfg.num_experts, cfg.moe_intermediate, cfg.hidden);
            let ones: Vec<f32> = vec![1.0f32; cfg.head_dim];
            let embed = get("model.embed_tokens.weight")?;
            let embedding_norm = get("model.norm.weight")?;
            let q4raw = |w: &[f32], rows: usize, cols: usize| -> Q4 {
                let (sc, qn) = quantize_q4_0(w, rows, cols);
                Q4 {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                    quants: ctx.storage(bytemuck::cast_slice(&qn)),
                }
            };
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                // Qwen2 qkv bias: concat q|k|v in the fused-qkv order.
                let mut qkv_b = tens(&format!("{p}.self_attn.q_proj.bias"))?;
                qkv_b.extend_from_slice(&tens(&format!("{p}.self_attn.k_proj.bias"))?);
                qkv_b.extend_from_slice(&tens(&format!("{p}.self_attn.v_proj.bias"))?);
                // Routed experts, concatenated (gate=w1, up=w3, down=w2).
                let (mut w1f, mut w2f, mut w3f) = (Vec::new(), Vec::new(), Vec::new());
                for ex in 0..e {
                    let mp = format!("{p}.mlp.experts.{ex}");
                    w1f.extend_from_slice(&tens(&format!("{mp}.gate_proj.weight"))?);
                    w2f.extend_from_slice(&tens(&format!("{mp}.down_proj.weight"))?);
                    w3f.extend_from_slice(&tens(&format!("{mp}.up_proj.weight"))?);
                }
                let moe = Moe {
                    router: ctx.storage(&tens(&format!("{p}.mlp.gate.weight"))?),
                    shared_gate: Some(
                        ctx.storage(&tens(&format!("{p}.mlp.shared_expert_gate.weight"))?),
                    ),
                    per_expert_scale: ctx.storage(&vec![1.0f32; e]),
                    w1: q4raw(&w1f, e * mi, h),
                    w3: q4raw(&w3f, e * mi, h),
                    w2: q4raw(&w2f, e * h, mi),
                    router_v3: None,
                };
                layers.push(Layer {
                    operator_norm: get(&format!("{p}.input_layernorm.weight"))?,
                    ffn_norm: get(&format!("{p}.post_attention_layernorm.weight"))?,
                    op: Op::Attn {
                        qkv: q4cat(&[
                            format!("{p}.self_attn.q_proj.weight"),
                            format!("{p}.self_attn.k_proj.weight"),
                            format!("{p}.self_attn.v_proj.weight"),
                        ])?,
                        o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                        q_norm: ctx.storage(&ones),
                        k_norm: ctx.storage(&ones),
                        window: cfg.sliding_window as u32,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: Some(ctx.storage(&qkv_b)),
                    },
                    // Shared expert → dense GLU slots (Qwen3.5 layout).
                    w1: q4(&format!("{p}.mlp.shared_expert.gate_proj.weight"))?,
                    w2: q4(&format!("{p}.mlp.shared_expert.down_proj.weight"))?,
                    w3: q4(&format!("{p}.mlp.shared_expert.up_proj.weight"))?,
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: Some(moe),
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Phi3 {
            // Llama compute reached through Phi-3's FUSED projections: `qkv_proj`
            // [(nh+2nkv)·hd, hidden] is already q4cat's [q|k|v] row layout (plain rotate_half rope,
            // no de-interleave), and `gate_up_proj` [2·inter, hidden] splits gate/up like GLM-OCR.
            // No qk-norm (ones buffers), no bias; the untied head is handled by `embed_q4` above.
            let ones: Vec<f32> = vec![1.0f32; cfg.head_dim];
            let embed = get("model.embed_tokens.weight")?;
            let embedding_norm = get("model.norm.weight")?;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                let gu = tens(&format!("{p}.mlp.gate_up_proj.weight"))?;
                anyhow::ensure!(
                    gu.len() == 2 * cfg.intermediate * cfg.hidden,
                    "gate_up shape"
                );
                let half = cfg.intermediate * cfg.hidden;
                let (s1, z1) = quantize_q4_0(&gu[..half], cfg.intermediate, cfg.hidden);
                let (s3, z3) = quantize_q4_0(&gu[half..], cfg.intermediate, cfg.hidden);
                layers.push(Layer {
                    operator_norm: get(&format!("{p}.input_layernorm.weight"))?,
                    ffn_norm: get(&format!("{p}.post_attention_layernorm.weight"))?,
                    op: Op::Attn {
                        qkv: q4(&format!("{p}.self_attn.qkv_proj.weight"))?,
                        o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                        q_norm: ctx.storage(&ones),
                        k_norm: ctx.storage(&ones),
                        window: cfg.sliding_window as u32,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: None,
                    },
                    w1: Q4 {
                        scales: ctx.storage_bytes(&f32_to_f16_bytes(&s1)),
                        quants: ctx.storage(bytemuck::cast_slice(&z1)),
                    },
                    w2: q4(&format!("{p}.mlp.down_proj.weight"))?,
                    w3: Q4 {
                        scales: ctx.storage_bytes(&f32_to_f16_bytes(&s3)),
                        quants: ctx.storage(bytemuck::cast_slice(&z3)),
                    },
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Granite {
            // Granite = Llama compute + four scalar multipliers, all FOLDED into weights here so
            // NO kernel/plan changes are needed. `embed_q4` (the head, loaded above) stays RAW —
            // `embedding_multiplier` only scales the INPUT embedding; `logits_scaling` is folded
            // into the final norm below (`logits = head·(norm_out/logits_scaling)`).
            let cj: serde_json::Value =
                serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?;
            let mult = |k: &str| cj.get(k).and_then(|x| x.as_f64()).unwrap_or(1.0) as f32;
            let (emb_m, res_m, attn_m, logit_s) = (
                mult("embedding_multiplier"),
                mult("residual_multiplier"),
                mult("attention_multiplier"),
                mult("logits_scaling"),
            );
            let hidden = cfg.hidden;
            let ones: Vec<f32> = vec![1.0f32; cfg.head_dim];
            // Scaled Q4 load: multiply the f32 rows by `scale` before quantizing.
            let q4_scaled = |name: &str, scale: f32| -> Result<Q4> {
                let shape = st.shape(name)?;
                let (rows, cols) = (shape[0], shape[1]);
                let mut w = tens(name)?;
                for x in &mut w {
                    *x *= scale;
                }
                let (s, q) = quantize_q4_0(&w, rows, cols);
                Ok(Q4 {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(&s)),
                    quants: ctx.storage(bytemuck::cast_slice(&q)),
                })
            };
            // Input embedding × embedding_multiplier.
            let mut e = tens("model.embed_tokens.weight")?;
            for x in &mut e {
                *x *= emb_m;
            }
            let embed = ctx.storage(&e);
            drop(e);
            // Final norm ÷ logits_scaling (folds the logits scaling into the head input).
            let mut nrm = tens("model.norm.weight")?;
            for x in &mut nrm {
                *x /= logit_s;
            }
            let embedding_norm = ctx.storage(&nrm);
            // attention_multiplier folds into q_proj (·√head_dim, so the kernel's 1/√hd scale
            // produces `(q·k)·attention_multiplier`); residual_multiplier folds into o_proj/down.
            let q_scale = attn_m * (cfg.head_dim as f32).sqrt();
            let qkv_rows = (cfg.n_heads + 2 * cfg.n_kv_heads) * cfg.head_dim;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                // Fused qkv with q ROWS pre-scaled (k/v unscaled).
                let mut qkv_f = tens(&format!("{p}.self_attn.q_proj.weight"))?;
                for x in &mut qkv_f {
                    *x *= q_scale;
                }
                qkv_f.extend_from_slice(&tens(&format!("{p}.self_attn.k_proj.weight"))?);
                qkv_f.extend_from_slice(&tens(&format!("{p}.self_attn.v_proj.weight"))?);
                let (sc, qz) = quantize_q4_0(&qkv_f, qkv_rows, hidden);
                layers.push(Layer {
                    operator_norm: get(&format!("{p}.input_layernorm.weight"))?,
                    ffn_norm: get(&format!("{p}.post_attention_layernorm.weight"))?,
                    op: Op::Attn {
                        qkv: Q4 {
                            scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                            quants: ctx.storage(bytemuck::cast_slice(&qz)),
                        },
                        o: q4_scaled(&format!("{p}.self_attn.o_proj.weight"), res_m)?,
                        q_norm: ctx.storage(&ones),
                        k_norm: ctx.storage(&ones),
                        window: 0,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: None,
                    },
                    w1: q4(&format!("{p}.mlp.gate_proj.weight"))?, // gate unscaled
                    w2: q4_scaled(&format!("{p}.mlp.down_proj.weight"), res_m)?, // down × res_m
                    w3: q4(&format!("{p}.mlp.up_proj.weight"))?,   // up unscaled
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: None,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Mixtral {
            // Mistral attention + PURE sparse MoE. The dense GLU slots are ZEROED (tiny,
            // intermediate=32) so the always-computed dense path adds exactly 0 → block = routed.
            let (e, mi, h) = (cfg.num_experts, cfg.moe_intermediate, cfg.hidden);
            let ones: Vec<f32> = vec![1.0f32; cfg.head_dim];
            let embed = get("model.embed_tokens.weight")?;
            let embedding_norm = get("model.norm.weight")?;
            let q4raw = |w: &[f32], rows: usize, cols: usize| -> Q4 {
                let (sc, qn) = quantize_q4_0(w, rows, cols);
                Q4 {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                    quants: ctx.storage(bytemuck::cast_slice(&qn)),
                }
            };
            let zero_dense =
                |rows: usize, cols: usize| q4raw(&vec![0.0f32; rows * cols], rows, cols);
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                // Concatenate the SEPARATE per-expert matrices (w1=gate, w2=down, w3=up).
                let (mut w1f, mut w2f, mut w3f) = (Vec::new(), Vec::new(), Vec::new());
                for ex in 0..e {
                    let mp = format!("{p}.block_sparse_moe.experts.{ex}");
                    w1f.extend_from_slice(&tens(&format!("{mp}.w1.weight"))?);
                    w2f.extend_from_slice(&tens(&format!("{mp}.w2.weight"))?);
                    w3f.extend_from_slice(&tens(&format!("{mp}.w3.weight"))?);
                }
                let moe = Moe {
                    router: ctx.storage(&tens(&format!("{p}.block_sparse_moe.gate.weight"))?),
                    shared_gate: None,
                    per_expert_scale: ctx.storage(&vec![1.0f32; e]),
                    w1: q4raw(&w1f, e * mi, h),
                    w3: q4raw(&w3f, e * mi, h),
                    w2: q4raw(&w2f, e * h, mi),
                    router_v3: None,
                };
                layers.push(Layer {
                    operator_norm: get(&format!("{p}.input_layernorm.weight"))?,
                    ffn_norm: get(&format!("{p}.post_attention_layernorm.weight"))?,
                    op: Op::Attn {
                        qkv: q4cat(&[
                            format!("{p}.self_attn.q_proj.weight"),
                            format!("{p}.self_attn.k_proj.weight"),
                            format!("{p}.self_attn.v_proj.weight"),
                        ])?,
                        o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                        q_norm: ctx.storage(&ones),
                        k_norm: ctx.storage(&ones),
                        window: cfg.sliding_window as u32,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: None,
                    },
                    w1: zero_dense(cfg.intermediate, h),
                    w2: zero_dense(h, cfg.intermediate),
                    w3: zero_dense(cfg.intermediate, h),
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: Some(moe),
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::Qwen3Moe {
            // Qwen3 attention (REAL qk-norm) + the Mixtral pure-MoE. Experts are separate
            // `mlp.experts.N.{gate,up,down}_proj`; dense GLU zeroed tiny → pure routed.
            let (e, mi, h) = (cfg.num_experts, cfg.moe_intermediate, cfg.hidden);
            let embed = get("model.embed_tokens.weight")?;
            let embedding_norm = get("model.norm.weight")?;
            let q4raw = |w: &[f32], rows: usize, cols: usize| -> Q4 {
                let (sc, qn) = quantize_q4_0(w, rows, cols);
                Q4 {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                    quants: ctx.storage(bytemuck::cast_slice(&qn)),
                }
            };
            let zero_dense =
                |rows: usize, cols: usize| q4raw(&vec![0.0f32; rows * cols], rows, cols);
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                let (mut w1f, mut w2f, mut w3f) = (Vec::new(), Vec::new(), Vec::new());
                for ex in 0..e {
                    let mp = format!("{p}.mlp.experts.{ex}");
                    w1f.extend_from_slice(&tens(&format!("{mp}.gate_proj.weight"))?);
                    w2f.extend_from_slice(&tens(&format!("{mp}.down_proj.weight"))?);
                    w3f.extend_from_slice(&tens(&format!("{mp}.up_proj.weight"))?);
                }
                let moe = Moe {
                    router: ctx.storage(&tens(&format!("{p}.mlp.gate.weight"))?),
                    shared_gate: None,
                    per_expert_scale: ctx.storage(&vec![1.0f32; e]),
                    w1: q4raw(&w1f, e * mi, h),
                    w3: q4raw(&w3f, e * mi, h),
                    w2: q4raw(&w2f, e * h, mi),
                    router_v3: None,
                };
                layers.push(Layer {
                    operator_norm: get(&format!("{p}.input_layernorm.weight"))?,
                    ffn_norm: get(&format!("{p}.post_attention_layernorm.weight"))?,
                    op: Op::Attn {
                        qkv: q4cat(&[
                            format!("{p}.self_attn.q_proj.weight"),
                            format!("{p}.self_attn.k_proj.weight"),
                            format!("{p}.self_attn.v_proj.weight"),
                        ])?,
                        o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                        q_norm: get(&format!("{p}.self_attn.q_norm.weight"))?,
                        k_norm: get(&format!("{p}.self_attn.k_norm.weight"))?,
                        window: 0,
                        local_rope: false,
                        attn_gate: None,
                        qkv_bias: None,
                    },
                    w1: zero_dense(cfg.intermediate, h),
                    w2: zero_dense(h, cfg.intermediate),
                    w3: zero_dense(cfg.intermediate, h),
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe: Some(moe),
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if cfg.arch == Arch::DeepseekV2 {
            // MLA attention (projections + assembly in the forward) + dense SwiGLU. MLA dims come
            // from config (not Lfm2Config); o_proj is zero-column-padded to [hidden, nh·qk_head_dim].
            let cj: serde_json::Value =
                serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?;
            let cg = |k: &str| cj.get(k).and_then(|x| x.as_u64()).map(|u| u as usize);
            // Routed-expert scale: DeepSeek multiplies the (renormalized) top-k weights by
            // `routed_scaling_factor`; fold it into `per_expert_scale` (applied after renorm).
            let routed_scaling = cj
                .get("routed_scaling_factor")
                .and_then(|x| x.as_f64())
                .unwrap_or(1.0) as f32;
            let q4raw = |w: &[f32], rows: usize, cols: usize| -> Q4 {
                let (sc, qn) = quantize_q4_0(w, rows, cols);
                Q4 {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                    quants: ctx.storage(bytemuck::cast_slice(&qn)),
                }
            };
            let (nh, hdn) = (cfg.n_heads, cfg.hidden);
            let qk_nope = cg("qk_nope_head_dim").context("qk_nope_head_dim")?;
            let qk_rope = cg("qk_rope_head_dim").context("qk_rope_head_dim")?;
            let v_head_dim = cg("v_head_dim").context("v_head_dim")?;
            let kv_lora = cg("kv_lora_rank").context("kv_lora_rank")?;
            let q_lora = cg("q_lora_rank"); // None ⇒ direct q_proj
            let qk_head_dim = qk_nope + qk_rope;
            let embed = get("model.embed_tokens.weight")?;
            let embedding_norm = get("model.norm.weight")?;
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let p = format!("model.layers.{i}");
                // o_proj [hidden, nh·v_head_dim] → padded [hidden, nh·qk_head_dim] (zero cols).
                let o_src = tens(&format!("{p}.self_attn.o_proj.weight"))?;
                anyhow::ensure!(o_src.len() == hdn * nh * v_head_dim, "o_proj shape");
                let mut o_pad = vec![0f32; hdn * nh * qk_head_dim];
                for r in 0..hdn {
                    for h in 0..nh {
                        let src = &o_src[r * nh * v_head_dim + h * v_head_dim..][..v_head_dim];
                        o_pad[r * nh * qk_head_dim + h * qk_head_dim..][..v_head_dim]
                            .copy_from_slice(src);
                    }
                }
                let (sc, qz) = quantize_q4_0(&o_pad, hdn, nh * qk_head_dim);
                let (q_a, q_a_norm, q_or_qb) = match q_lora {
                    Some(ql) => (
                        Some(q4(&format!("{p}.self_attn.q_a_proj.weight"))?),
                        Some(get(&format!("{p}.self_attn.q_a_layernorm.weight"))?),
                        {
                            let w = tens(&format!("{p}.self_attn.q_b_proj.weight"))?;
                            let (s, z) = quantize_q4_0(&w, nh * qk_head_dim, ql);
                            Q4 {
                                scales: ctx.storage_bytes(&f32_to_f16_bytes(&s)),
                                quants: ctx.storage(bytemuck::cast_slice(&z)),
                            }
                        },
                    ),
                    None => (None, None, q4(&format!("{p}.self_attn.q_proj.weight"))?),
                };
                // FFN: dense SwiGLU, or (MoE) the always-on shared experts as the dense GLU +
                // routed experts on the `Moe` struct. DeepSeek's `shared_experts` is ONE
                // pre-concatenated MLP (intermediate = n_shared·moe_intermediate) → it lands in the
                // dense slots verbatim; the always-computed-and-added dense path IS the ungated
                // shared-expert sum (same reuse Mixtral/Gemma4 rely on). The routed experts are the
                // separate `mlp.experts.N.{gate,up,down}_proj`, concatenated like Mixtral.
                let (w1, w2, w3, moe) = if cfg.num_experts > 0 {
                    let (e, mi) = (cfg.num_experts, cfg.moe_intermediate);
                    let (mut w1f, mut w2f, mut w3f) = (Vec::new(), Vec::new(), Vec::new());
                    for ex in 0..e {
                        let mp = format!("{p}.mlp.experts.{ex}");
                        w1f.extend_from_slice(&tens(&format!("{mp}.gate_proj.weight"))?);
                        w2f.extend_from_slice(&tens(&format!("{mp}.down_proj.weight"))?);
                        w3f.extend_from_slice(&tens(&format!("{mp}.up_proj.weight"))?);
                    }
                    // DeepSeek-V3 aux-loss-free router (sigmoid scoring): read the group params +
                    // the additive selection bias; the V3 kernel bakes `routed_scaling` into its
                    // emitted weights (so per_expert_scale is inert on this path).
                    let router_v3 = if cj.get("scoring_func").and_then(|x| x.as_str())
                        == Some("sigmoid")
                    {
                        Some(V3Router {
                            bias: ctx
                                .storage(&tens(&format!("{p}.mlp.gate.e_score_correction_bias"))?),
                            n_group: cg("n_group").context("n_group")? as u32,
                            topk_group: cg("topk_group").context("topk_group")? as u32,
                            scaling: routed_scaling,
                        })
                    } else {
                        None
                    };
                    let moe = Moe {
                        router: ctx.storage(&tens(&format!("{p}.mlp.gate.weight"))?),
                        shared_gate: None,
                        per_expert_scale: ctx.storage(&vec![routed_scaling; e]),
                        w1: q4raw(&w1f, e * mi, hdn),
                        w3: q4raw(&w3f, e * mi, hdn),
                        w2: q4raw(&w2f, e * hdn, mi),
                        router_v3,
                    };
                    (
                        q4(&format!("{p}.mlp.shared_experts.gate_proj.weight"))?,
                        q4(&format!("{p}.mlp.shared_experts.down_proj.weight"))?,
                        q4(&format!("{p}.mlp.shared_experts.up_proj.weight"))?,
                        Some(moe),
                    )
                } else {
                    (
                        q4(&format!("{p}.mlp.gate_proj.weight"))?,
                        q4(&format!("{p}.mlp.down_proj.weight"))?,
                        q4(&format!("{p}.mlp.up_proj.weight"))?,
                        None,
                    )
                };
                layers.push(Layer {
                    operator_norm: get(&format!("{p}.input_layernorm.weight"))?,
                    ffn_norm: get(&format!("{p}.post_attention_layernorm.weight"))?,
                    op: Op::Mla {
                        q_a,
                        q_a_norm,
                        q_or_qb,
                        kv_a: q4(&format!("{p}.self_attn.kv_a_proj_with_mqa.weight"))?,
                        kv_a_norm: get(&format!("{p}.self_attn.kv_a_layernorm.weight"))?,
                        kv_b: q4(&format!("{p}.self_attn.kv_b_proj.weight"))?,
                        o: Q4 {
                            scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                            quants: ctx.storage(bytemuck::cast_slice(&qz)),
                        },
                        qk_nope,
                        qk_rope,
                        v_head_dim,
                        kv_lora,
                        q_lora: q_lora.unwrap_or(0),
                    },
                    w1,
                    w2,
                    w3,
                    post_op_norm: None,
                    post_ffn_norm: None,
                    moe,
                    edge: None,
                });
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple: None,
            });
        }
        if matches!(cfg.arch, Arch::Gemma3 | Arch::Gemma4) {
            // Norm convention diverges: Gemma-3 = `(1+w)` (folded at load); Gemma-4 = plain `*w`
            // (verbatim: Gemma3nRMSNorm.forward → `normed * weight`, which Gemma4RMSNorm inherits).
            let get_norm = |name: &str| -> Result<wgpu::Buffer> {
                if cfg.arch == Arch::Gemma4 {
                    get(name)
                } else {
                    get_p1(name)
                }
            };
            // Multimodal wrapper (Gemma4ForConditionalGeneration): text under
            // `model.language_model.*`; the audio/vision towers are never read. Bare text
            // checkpoints (and every synthetic fixture) keep `model.*`.
            let lm = if st.has("model.language_model.embed_tokens.weight") {
                "model.language_model"
            } else {
                "model"
            };
            // Gather rows pre-scaled by the embedding normalizer sqrt(hidden); the tied head
            // (embed_q4 above) stays unscaled. Gemma-4 rounds the scale through bf16 FIRST —
            // HF's ScaledWordEmbedding casts embed_scale to the weight dtype, so E2B multiplies
            // by 39.25, not sqrt(1536)=39.1918; missing that biases every residual stream by
            // ~0.15% on its embedding component.
            let mut e = tens(&format!("{lm}.embed_tokens.weight"))?;
            let scale = if cfg.arch == Arch::Gemma4 {
                bf16_round((cfg.hidden as f32).sqrt())
            } else {
                (cfg.hidden as f32).sqrt()
            };
            for x in &mut e {
                *x *= scale;
            }
            let embed = ctx.storage(&e);
            drop(e);
            let embedding_norm = get_norm(&format!("{lm}.norm.weight"))?;
            // Model-level PLE weights (Gemma-4 edge).
            let ple = if cfg.edge.ple_dim > 0 {
                let tname = format!("{lm}.embed_tokens_per_layer.weight");
                let shape = st.shape(&tname)?;
                anyhow::ensure!(
                    shape[0] == cfg.vocab && shape[1] == cfg.n_layers * cfg.edge.ple_dim,
                    "embed_tokens_per_layer shape mismatch"
                );
                // The table is pre-scaled by its embedding normalizer sqrt(ple_dim) (16.0 on
                // E2B — exact in bf16, no rounding subtlety) and quantized Q4_0: the bf16
                // original is 4.7 GB on E2B, above Vulkan's max_buffer_size, and its content is
                // embedding-class (already served Q4 as the tied head elsewhere). NOTE: this
                // reads the f32 table transiently (~9.4 GB on E2B) — chunked quantization is
                // the follow-up if load-RAM ever matters.
                let mut tbl = tens(&tname)?;
                let tscale = (cfg.edge.ple_dim as f32).sqrt();
                for x in &mut tbl {
                    *x *= tscale;
                }
                let (sc, qn) = quantize_q4_0(&tbl, cfg.vocab, cfg.n_layers * cfg.edge.ple_dim);
                drop(tbl);
                let table = Q4 {
                    scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                    quants: ctx.storage(bytemuck::cast_slice(&qn)),
                };
                // Flush the ~1.3 GB upload staging before the layer loop (see the per-layer
                // submit below).
                ctx.queue.submit(std::iter::empty());
                let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
                layer_dtype.set(cfg.wdtype);
                Some(PleWeights {
                    table,
                    model_proj: q4(&format!("{lm}.per_layer_model_projection.weight"))?,
                    proj_norm: get_norm(&format!("{lm}.per_layer_projection_norm.weight"))?,
                })
            } else {
                None
            };
            let mut layers = Vec::with_capacity(cfg.n_layers);
            for i in 0..cfg.n_layers {
                let ks = cfg.kinds_at(i);
                let p = format!("{lm}.layers.{i}");
                // K=V layers (Gemma-4 global): no v_proj tensor exists — concat [q|k] only; the
                // qkrc kernel synthesizes V from the K projection (weightless-normed, un-roped).
                // KV-shared CONSUMERS (Gemma-4 edge) run Q ONLY: they read the donor layer's
                // cache, and their on-disk k_proj/v_proj/k_norm are DEAD tensors that HF itself
                // drops at load (`_keys_to_ignore_on_load_unexpected`) — wiring them up would be
                // silently wrong. Gate on the config, never on tensor presence.
                let qkv_names: Vec<String> = if !cfg.owns_kv(i) {
                    vec![format!("{p}.self_attn.q_proj.weight")]
                } else if cfg.layer_k_eq_v[i] {
                    vec![
                        format!("{p}.self_attn.q_proj.weight"),
                        format!("{p}.self_attn.k_proj.weight"),
                    ]
                } else {
                    vec![
                        format!("{p}.self_attn.q_proj.weight"),
                        format!("{p}.self_attn.k_proj.weight"),
                        format!("{p}.self_attn.v_proj.weight"),
                    ]
                };
                layer_dtype.set(ks.qkv.pack_dtype());
                let qkv = q4cat(&qkv_names)?;
                layer_dtype.set(ks.o.pack_dtype());
                let op = Op::Attn {
                    qkv,
                    o: q4(&format!("{p}.self_attn.o_proj.weight"))?,
                    q_norm: get_norm(&format!("{p}.self_attn.q_norm.weight"))?,
                    k_norm: if cfg.owns_kv(i) {
                        get_norm(&format!("{p}.self_attn.k_norm.weight"))?
                    } else {
                        // Q-only consumer: the kernel never touches knw (its k_norm tensor is
                        // dead, see above) — bind the q_norm buffer as a layout placeholder.
                        get_norm(&format!("{p}.self_attn.q_norm.weight"))?
                    },
                    window: if cfg.layer_is_sliding[i] {
                        cfg.sliding_window as u32
                    } else {
                        0
                    },
                    local_rope: cfg.layer_is_sliding[i],
                    attn_gate: None,
                    qkv_bias: None,
                };
                // Gemma-4 MoE extras. Tensor names are PROVISIONAL (mirroring the transformers
                // module structure; no accessible checkpoint yet) — the synthetic e2e test writes
                // these names, and a real checkpoint may need a one-line rename here.
                let moe = if cfg.num_experts > 0 {
                    let (e, mi, h) = (cfg.num_experts, cfg.moe_intermediate, cfg.hidden);
                    // transformers fuses gate|up per expert ([E, 2·mi, h]); split into the engine's
                    // separate gate/up layout, concatenated across experts.
                    let gu = tens(&format!("{p}.moe.experts.gate_up_proj"))?;
                    anyhow::ensure!(
                        gu.len() == e * 2 * mi * h,
                        "gate_up_proj shape mismatch at layer {i}"
                    );
                    let (mut w1f, mut w3f) = (
                        Vec::with_capacity(e * mi * h),
                        Vec::with_capacity(e * mi * h),
                    );
                    for ex in 0..e {
                        let base = ex * 2 * mi * h;
                        w1f.extend_from_slice(&gu[base..base + mi * h]);
                        w3f.extend_from_slice(&gu[base + mi * h..base + 2 * mi * h]);
                    }
                    let q4raw = |w: &[f32], rows: usize, cols: usize| -> Q4 {
                        let (sc, qn) = quantize_q4_0(w, rows, cols);
                        Q4 {
                            scales: ctx.storage_bytes(&f32_to_f16_bytes(&sc)),
                            quants: ctx.storage(bytemuck::cast_slice(&qn)),
                        }
                    };
                    Some(Moe {
                        router: get(&format!("{p}.moe.router.weight"))?,
                        shared_gate: None,
                        per_expert_scale: get(&format!("{p}.moe.per_expert_scale"))?,
                        w1: q4raw(&w1f, e * mi, h),
                        w3: q4raw(&w3f, e * mi, h),
                        w2: q4raw(&tens(&format!("{p}.moe.experts.down_proj"))?, e * h, mi),
                        router_v3: None,
                    })
                } else {
                    None
                };
                layer_dtype.set(ks.gate_up.pack_dtype());
                let w1 = q4(&format!("{p}.mlp.gate_proj.weight"))?;
                let w3 = q4(&format!("{p}.mlp.up_proj.weight"))?;
                layer_dtype.set(ks.down.pack_dtype());
                let w2 = q4(&format!("{p}.mlp.down_proj.weight"))?;
                // Gemma-4 edge per-layer extras: the PLE injection weights, plus the residual
                // layer_scalar (the server checkpoints carry the scalar WITHOUT PLE).
                let ls_name = format!("{p}.layer_scalar");
                let edge = if cfg.edge.ple_dim > 0 || st.has(&ls_name) {
                    let layer_scalar = if st.has(&ls_name) {
                        let s = tens(&ls_name)?;
                        anyhow::ensure!(s.len() == 1, "layer_scalar must be a single element");
                        s[0]
                    } else {
                        1.0
                    };
                    let lple = if cfg.edge.ple_dim > 0 {
                        layer_dtype.set(cfg.wdtype);
                        Some(PleLayer {
                            gate: q4(&format!("{p}.per_layer_input_gate.weight"))?,
                            proj: q4(&format!("{p}.per_layer_projection.weight"))?,
                            post_norm: get_norm(&format!(
                                "{p}.post_per_layer_input_norm.weight"
                            ))?,
                        })
                    } else {
                        None
                    };
                    Some(LayerEdge {
                        ple: lple,
                        layer_scalar,
                    })
                } else {
                    None
                };
                layers.push(Layer {
                    operator_norm: get_norm(&format!("{p}.input_layernorm.weight"))?,
                    ffn_norm: get_norm(&format!("{p}.pre_feedforward_layernorm.weight"))?,
                    op,
                    w1,
                    w2,
                    w3,
                    post_op_norm: Some(get_norm(&format!("{p}.post_attention_layernorm.weight"))?),
                    post_ffn_norm: Some(get_norm(&format!(
                        "{p}.post_feedforward_layernorm.weight"
                    ))?),
                    moe,
                    edge,
                });
                // Flush pending upload staging per layer (see the qwen35 arm above for why).
                ctx.queue.submit(std::iter::empty());
                let _ = ctx.device.poll(wgpu::PollType::wait_indefinitely());
            }
            return Ok(Self {
                cfg,
                embed,
                embed_q4,
                embedding_norm,
                layers,
                mtp: None,
                embed_f16: None,
                embed_split: None,
                ple,
            });
        }
        let embed = get("model.embed_tokens.weight")?;
        let embedding_norm = get("model.embedding_norm.weight")?;
        let mut layers = Vec::with_capacity(cfg.n_layers);
        for i in 0..cfg.n_layers {
            let p = format!("model.layers.{i}");
            let op = if cfg.layer_is_attn[i] {
                Op::Attn {
                    qkv: q4cat(&[
                        format!("{p}.self_attn.q_proj.weight"),
                        format!("{p}.self_attn.k_proj.weight"),
                        format!("{p}.self_attn.v_proj.weight"),
                    ])?,
                    o: q4(&format!("{p}.self_attn.out_proj.weight"))?,
                    q_norm: get(&format!("{p}.self_attn.q_layernorm.weight"))?,
                    k_norm: get(&format!("{p}.self_attn.k_layernorm.weight"))?,
                    window: 0,
                    local_rope: false,
                    attn_gate: None,
                    qkv_bias: None,
                }
            } else {
                Op::Conv {
                    in_proj: q4(&format!("{p}.conv.in_proj.weight"))?,
                    conv_w: get(&format!("{p}.conv.conv.weight"))?,
                    out_proj: q4(&format!("{p}.conv.out_proj.weight"))?,
                }
            };
            layers.push(Layer {
                operator_norm: get(&format!("{p}.operator_norm.weight"))?,
                ffn_norm: get(&format!("{p}.ffn_norm.weight"))?,
                op,
                w1: q4(&format!("{p}.feed_forward.w1.weight"))?,
                w2: q4(&format!("{p}.feed_forward.w2.weight"))?,
                w3: q4(&format!("{p}.feed_forward.w3.weight"))?,
                post_op_norm: None,
                post_ffn_norm: None,
                moe: None,
                edge: None,
            });
        }
        Ok(Self {
            cfg,
            embed,
            embed_q4,
            embedding_norm,
            layers,
            mtp: None,
            embed_f16: None,
            embed_split: None,
            ple: None,
        })
    }
}

/// Decode a little-endian BF16 byte buffer to f32.
/// Round an f32 to the nearest bf16 (round-to-nearest-even) and back. Gemma-4's embedding scale
/// is applied at the WEIGHT dtype in HF (`embed_scale.to(weight.dtype)`), so parity requires
/// scaling by bf16(sqrt(hidden)) — 39.25 on E2B, not 39.1918.
pub fn bf16_round(x: f32) -> f32 {
    let bits = x.to_bits();
    let rounded = bits.wrapping_add(0x7FFF + ((bits >> 16) & 1)) & 0xFFFF_0000;
    f32::from_bits(rounded)
}

pub fn bf16_to_f32(bytes: &[u8]) -> Vec<f32> {
    bytes
        .chunks_exact(2)
        .map(|c| bf16::from_le_bytes([c[0], c[1]]).to_f32())
        .collect()
}

/// Q4_0-quantized weight matrix on the GPU: a per-32-block f32 scale + 4-bit quants (8 nibbles/u32),
/// dequantised inside the GEMV (`x ≈ d·(q-8)`). ~7× less weight bandwidth than f32 — the dominant
/// decode lever. `cols` must be a multiple of 32.
#[derive(Clone)]
pub struct Q4 {
    pub scales: wgpu::Buffer, // [rows * cols/32]
    pub quants: wgpu::Buffer, // [rows * cols/32 * 4] u32
}

/// Storage precision of ONE weight-matrix group. Quantization is a speed/fidelity trade the
/// CALLER owns, not a hard-coded engine choice: 4-bit weights are the dominant decode-bandwidth
/// lever, but they cost measurable accuracy — the smaller the model, the more (a 270M gemma-3
/// loses ~2× the per-step agreement a 1B does, MEASURED against an f32 reference; the loss is
/// the weights, not the arithmetic).
/// The three points of the measured speed/fidelity ladder (per-step top-1 agreement against an
/// f32 reference, gemma-3-270m — the most quantization-sensitive model in the bench matrix):
///
/// | dtype | agreement | weight bytes |
/// |-------|-----------|--------------|
/// | `Q4`  | 0.459     | 1.00×        |
/// | `Q8`  | 0.966     | 1.89×        |
/// | `F16` | 1.000     | 3.56×        |
///
/// Weight bytes are the decode-bandwidth (≈ tok/s) cost, so this is a real trade, not a
/// free upgrade — which is exactly why the caller picks rather than the engine.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum WDtype {
    /// llama.cpp Q4_0 blocks — maximum throughput, materially lossy on small models.
    Q4,
    /// llama.cpp Q8_0 blocks (8-bit symmetric, per-32-block f16 scale) — near-f16 fidelity at
    /// roughly half f16's bandwidth cost.
    Q8,
    /// Half precision — lossless in practice (checkpoints ship bf16; f16 round-trips them
    /// exactly at these magnitudes), at ~3.6× Q4's bytes.
    F16,
}

impl WDtype {
    fn parse(s: &str) -> Result<Self> {
        match s.trim() {
            "q4" | "q4_0" => Ok(Self::Q4),
            "q8" | "q8_0" => Ok(Self::Q8),
            "f16" | "fp16" | "half" => Ok(Self::F16),
            other => {
                anyhow::bail!("unknown weight dtype `{other}` (expected `q4`, `q8` or `f16`)")
            }
        }
    }
}

/// The weight-precision policy: a dtype per matrix GROUP, so a caller can spend bytes only where
/// they buy accuracy (e.g. an f16 LM head over Q4 layers) instead of paying for a whole f16 model.
///
/// Selected by the caller, or from `OSFKB_DECODE_PRECISION`:
/// - `q4` (the default — today's behaviour, byte-for-byte)
/// - `f16`
/// - a base plus per-group overrides: `q4,head=f16` · `f16,gate_up=q4,down=q4`
/// - per-LAYER whole-layer overrides: `q4,blk:0-1=f16,blk:25=f16` (the Dynamic-2.0 idea —
///   sensitivity is per-layer; first/last layers usually deserve more bits than the middle)
///
/// Groups: `qkv` (attention q/k/v), `o` (attention output), `gate_up` (MLP gate+up), `down`
/// (MLP down), `head` (the tied LM head; the embedding *gather* is always f32).
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Precision {
    /// Attention q/k/v projections.
    pub qkv: WDtype,
    /// Attention output projection.
    pub o: WDtype,
    /// MLP gate + up projections.
    pub gate_up: WDtype,
    /// MLP down projection.
    pub down: WDtype,
    /// Tied LM head.
    pub head: WDtype,
    /// Whole-layer dtype overrides: `(first, last, dtype)` inclusive layer ranges, applied in
    /// declaration order OVER the uniform body base. Whole-layer only — a per-group override
    /// *within* one layer would need a finer plan split than per-site bind groups give.
    pub layers: Vec<(usize, usize, WDtype)>,
}

impl Default for Precision {
    /// All-Q4 — the engine's historical behaviour, preserved byte-for-byte.
    fn default() -> Self {
        Self::uniform(WDtype::Q4)
    }
}

impl Precision {
    /// One dtype for every group.
    pub fn uniform(d: WDtype) -> Self {
        Self {
            qkv: d,
            o: d,
            gate_up: d,
            down: d,
            head: d,
            layers: Vec::new(),
        }
    }

    /// Parse the policy grammar (see the type docs). Unknown groups/dtypes are REFUSED — a
    /// silently-ignored precision request would ship the wrong model.
    pub fn parse(spec: &str) -> Result<Self> {
        let mut parts = spec.split(',').map(str::trim).filter(|p| !p.is_empty());
        let first = parts
            .next()
            .ok_or_else(|| anyhow::anyhow!("empty precision spec"))?;
        // A leading bare dtype sets the base; otherwise the base is the default (all-Q4).
        let (mut p, rest_first) = match WDtype::parse(first) {
            Ok(d) => (Self::uniform(d), None),
            Err(_) => (Self::default(), Some(first)),
        };
        for kv in rest_first.into_iter().chain(parts) {
            let (group, dtype) = kv.split_once('=').ok_or_else(|| {
                anyhow::anyhow!(
                    "bad precision term `{kv}` (expected `group=dtype`, e.g. `head=f16`)"
                )
            })?;
            let d = WDtype::parse(dtype)?;
            // Per-layer override: `blk:<i>` or `blk:<a>-<b>` (inclusive). Range sanity is checked
            // here; layer-count bounds only at resolve time (the model is not known yet).
            if let Some(spec) = group.trim().strip_prefix("blk:") {
                let (a, b) = match spec.split_once('-') {
                    Some((a, b)) => (
                        a.parse::<usize>()
                            .map_err(|_| anyhow::anyhow!("bad blk range `{spec}`"))?,
                        b.parse::<usize>()
                            .map_err(|_| anyhow::anyhow!("bad blk range `{spec}`"))?,
                    ),
                    None => {
                        let i = spec
                            .parse::<usize>()
                            .map_err(|_| anyhow::anyhow!("bad blk index `{spec}`"))?;
                        (i, i)
                    }
                };
                anyhow::ensure!(a <= b, "inverted blk range `{spec}`");
                p.layers.push((a, b, d));
                continue;
            }
            match group.trim() {
                "qkv" => p.qkv = d,
                "o" => p.o = d,
                "gate_up" => p.gate_up = d,
                "down" => p.down = d,
                "head" => p.head = d,
                other => anyhow::bail!(
                    "unknown weight group `{other}` (expected qkv|o|gate_up|down|head)"
                ),
            }
        }
        Ok(p)
    }

    /// The policy from `OSFKB_DECODE_PRECISION`, or all-Q4 when unset. An INVALID value is an
    /// error, never a silent fallback to the default.
    pub fn from_env() -> Result<Self> {
        match std::env::var("OSFKB_DECODE_PRECISION") {
            Ok(spec) => {
                Self::parse(&spec).with_context(|| format!("OSFKB_DECODE_PRECISION=`{spec}`"))
            }
            Err(_) => Ok(Self::default()),
        }
    }

    /// Every group's dtype, in a fixed order (qkv, o, gate_up, down, head).
    pub fn groups(&self) -> [WDtype; 5] {
        [self.qkv, self.o, self.gate_up, self.down, self.head]
    }

    /// Does any group ask for something other than the Q4 the kernels ship today?
    pub fn any_non_q4(&self) -> bool {
        self.groups().iter().any(|d| *d != WDtype::Q4)
    }

    /// The four BODY groups (everything but the head), in fixed order. The head decodes through a
    /// SEPARATE kernel path (`crate::Q4LmHead`), so it can carry a different dtype than the body
    /// without the body's plan needing a per-site bind-group choice — which is exactly why
    /// `q4,head=f16` is servable while a body-mixed policy is not yet.
    pub fn body_groups(&self) -> [WDtype; 4] {
        [self.qkv, self.o, self.gate_up, self.down]
    }

    /// `Some(dtype)` if every body group agrees (the body plan is uniform); `None` if they differ
    /// (a body-mixed policy, still unservable).
    pub fn uniform_body(&self) -> Option<WDtype> {
        let b = self.body_groups();
        b.iter().all(|d| *d == b[0]).then_some(b[0])
    }

    /// Resolve layer `li`'s per-SITE container kinds: each site takes its GROUP's dtype
    /// (qkv/o/gate_up/down may differ — the plan is per-site kind-aware), and a `blk:` override
    /// sets ALL FOUR sites of its layers (whole-layer semantics, matching `layer_dtypes`).
    pub fn site_kinds(&self, li: usize, n_layers: usize) -> Result<LayerKinds> {
        let mut blk = None;
        for &(a, b, d) in &self.layers {
            anyhow::ensure!(b < n_layers, "blk:{a}-{b} out of range for {n_layers} layers");
            if (a..=b).contains(&li) {
                blk = Some(d);
            }
        }
        let k = |d: WDtype| Lfm2Config::base_kind(d);
        Ok(match blk {
            Some(d) => LayerKinds::uniform(k(d)),
            None => LayerKinds {
                qkv: k(self.qkv),
                o: k(self.o),
                gate_up: k(self.gate_up),
                down: k(self.down),
            },
        })
    }

    /// Resolve the per-layer body dtypes for a model with `n_layers` layers: the uniform body
    /// base with the `blk:` overrides applied in declaration order. Returns an EMPTY vec when no
    /// override changes anything — the loader's uniform fast-path signal, which is also what
    /// keeps the no-override case byte-identical. An out-of-range override REFUSES rather than
    /// silently dropping (a dropped override would ship the wrong model).
    pub fn layer_dtypes(&self, n_layers: usize) -> Result<Vec<WDtype>> {
        let base = self.uniform_body().ok_or_else(|| {
            anyhow::anyhow!("per-layer overrides need a uniform body base (got {self:?})")
        })?;
        let mut per = vec![base; n_layers];
        for &(a, b, d) in &self.layers {
            anyhow::ensure!(
                b < n_layers,
                "blk:{a}-{b} out of range for a {n_layers}-layer model"
            );
            for slot in &mut per[a..=b] {
                *slot = d;
            }
        }
        if per.iter().all(|d| *d == base) {
            per.clear();
        }
        Ok(per)
    }
}

/// Quantize a row-major `[rows, cols]` f32 matrix to Q4_0 (llama.cpp scheme: per 32-block,
/// `d = max_mag / -8`, `q = clamp(round(x/d)+8, 0, 15)`). Packs byte `i` = `q[i] | q[i+16]<<4` so a
/// u32 word holds 4 low + 4 high nibbles, matching the reference's `unpack4xU8(word & 0x0F0F0F0F)-8`.
/// Pack one row-major `[rows, cols]` f32 weight into the container the plan binds, per `dtype`.
/// The single place the storage decision is realized — every architecture's loader closure routes
/// through here, so a new dtype is one arm rather than six edits.
///
/// f16 rides the SAME [`Q4`] container: the weights go in `quants` (plain row-major f16, the
/// layout `crate::gemv_f16_k_lcpp_src` expects) and `scales` becomes a 4-byte stub that the f16
/// kernels never bind — they have no scale table. Keeping one container type means the `Layer`/
/// `Op` structs, and every plan site that holds a `&Q4`, are untouched by the precision work.
/// Tool-facing safetensors access (`gguf-quantize` and friends): open a checkpoint directory
/// without any GPU involvement.
pub fn lazyst_open_for_tools(dir: &Path) -> Result<LazySt> {
    LazySt::open(dir)
}

pub fn pack_weight(ctx: &GpuCtx, w: &[f32], rows: usize, cols: usize, dtype: WDtype) -> Result<Q4> {
    let (scales, quants) = pack_weight_bytes(w, rows, cols, dtype)?;
    Ok(Q4 {
        // f16 carries no scale table; a 4-byte stub keeps the container shape uniform rather than
        // making `scales` an Option across the tree.
        scales: ctx.storage_bytes(if scales.is_empty() { &[0u8; 4] } else { &scales }),
        quants: ctx.storage_bytes(&quants),
    })
}

/// [`pack_weight`]'s byte-level primitive: `(scale_bytes, weight_bytes)` for one matrix.
///
/// Exists so the concatenated q|k|v loader can stream — it appends each projection's PACKED bytes
/// instead of holding all three in f32 at once. Concatenation is valid for both dtypes because
/// each is strictly per-row (Q4_0 blocks never straddle a row boundary, and f16 is elementwise),
/// so packing-then-concatenating equals concatenating-then-packing exactly.
pub fn pack_weight_bytes(
    w: &[f32],
    rows: usize,
    cols: usize,
    dtype: WDtype,
) -> Result<(Vec<u8>, Vec<u8>)> {
    match dtype {
        WDtype::Q4 => {
            let (s, q) = quantize_q4_0(w, rows, cols);
            Ok((
                f32_to_f16_bytes(&s),
                bytemuck::cast_slice::<u32, u8>(&q).to_vec(),
            ))
        }
        // Plain row-major f16 — the layout `crate::gemv_f16_k_lcpp_src` expects, no permutation.
        WDtype::F16 => Ok((Vec::new(), f32_to_f16_bytes(w))),
        // llama.cpp Q8_0 blocks in the engine's PADDED native layout (the Q8_0N containers the
        // native-GGUF path serves) — the same kernels, whether the bytes came from an artifact
        // or from quantizing a safetensors checkpoint here.
        WDtype::Q8 => Ok((Vec::new(), pack_q8_0n(w, rows, cols))),
    }
}

/// Quantize a row-major `[rows, cols]` f32 matrix straight into PADDED Q8_0N blocks (36 bytes /
/// 32 weights: f16 `d` + 2 pad + 32 `i8`), llama.cpp's Q8_0 scheme (`d = amax/127`,
/// `q = round(x/d)`). The layout [`crate::gemv_q8_0n_k_lcpp_src`] and its MLP/head twins read.
pub fn pack_q8_0n(w: &[f32], rows: usize, cols: usize) -> Vec<u8> {
    assert!(cols.is_multiple_of(32), "Q8_0 needs cols % 32 == 0");
    let nblk = cols / 32;
    let mut out = vec![0u8; rows * nblk * 36];
    for r in 0..rows {
        for b in 0..nblk {
            let blk = &w[r * cols + b * 32..r * cols + b * 32 + 32];
            let amax = blk.iter().fold(0f32, |a, v| a.max(v.abs()));
            let d = amax / 127.0;
            let id = if d != 0.0 { 1.0 / d } else { 0.0 };
            let o = (r * nblk + b) * 36;
            // store d exactly as the kernel will read it (f16), and quantize against THAT d so
            // encode and decode agree bit-for-bit on the scale.
            let d16 = half::f16::from_f32(d);
            out[o..o + 2].copy_from_slice(&d16.to_le_bytes());
            let df = d16.to_f32();
            let idf = if df != 0.0 { 1.0 / df } else { id };
            for (i, &v) in blk.iter().enumerate() {
                out[o + 4 + i] = (v * idf).round().clamp(-127.0, 127.0) as i8 as u8;
            }
        }
    }
    out
}

pub fn quantize_q4_0(w: &[f32], rows: usize, cols: usize) -> (Vec<f32>, Vec<u32>) {
    assert!(cols.is_multiple_of(32), "Q4_0 needs cols % 32 == 0");
    let nblk = cols / 32;
    let mut scales = vec![0f32; rows * nblk];
    let mut quants = vec![0u32; rows * nblk * 4];
    for r in 0..rows {
        for b in 0..nblk {
            let blk = &w[r * cols + b * 32..r * cols + b * 32 + 32];
            let (mut amax, mut max) = (0f32, 0f32);
            for &v in blk {
                if v.abs() > amax {
                    amax = v.abs();
                    max = v;
                }
            }
            let d = max / -8.0;
            let id = if d != 0.0 { 1.0 / d } else { 0.0 };
            scales[r * nblk + b] = d;
            let mut q = [0u8; 32];
            for i in 0..32 {
                q[i] = ((blk[i] * id).round() as i32 + 8).clamp(0, 15) as u8;
            }
            for wi in 0..4 {
                let mut word = 0u32;
                for byte in 0..4 {
                    let bi = wi * 4 + byte; // 0..16
                    let v = (q[bi] as u32) | ((q[bi + 16] as u32) << 4);
                    word |= v << (byte * 8);
                }
                quants[(r * nblk + b) * 4 + wi] = word;
            }
        }
    }
    (scales, quants)
}

/// Q8_0-quantized weight matrix on the GPU: a per-32-block scale + 8-bit symmetric quants
/// (4 per u32), dequantised inside the GEMV (`x ≈ d·q`). ~1.9× Q4_0's bytes for ~97% of f16's
/// fidelity (MEASURED on gemma-3-270m: per-step top-1 agreement vs an f32 reference — Q4 0.459,
/// Q8 0.966, f16 1.000) — the value rung of the precision ladder.
#[derive(Clone)]
pub struct Q8 {
    pub scales: wgpu::Buffer, // [rows * cols/32]
    pub quants: wgpu::Buffer, // [rows * cols/32 * 8] u32, 4 i8 each
}

/// Quantize a row-major `[rows, cols]` f32 matrix to Q8_0 (llama.cpp scheme: per 32-block,
/// `d = max|x| / 127`, `q = round(x/d)` clamped to i8). Packs 4 signed bytes per u32 so the GEMV
/// unpacks with `unpack4xI8` — the same block geometry as [`quantize_q4_0`], so the Q8 kernels are
/// a small delta from their Q4 twins rather than a new memory layout.
pub fn quantize_q8_0(w: &[f32], rows: usize, cols: usize) -> (Vec<f32>, Vec<u32>) {
    assert!(cols.is_multiple_of(32), "Q8_0 needs cols % 32 == 0");
    let nblk = cols / 32;
    let mut scales = vec![0f32; rows * nblk];
    let mut quants = vec![0u32; rows * nblk * 8];
    for r in 0..rows {
        for b in 0..nblk {
            let blk = &w[r * cols + b * 32..r * cols + b * 32 + 32];
            let amax = blk.iter().fold(0f32, |m, v| m.max(v.abs()));
            let d = amax / 127.0;
            let id = if d != 0.0 { 1.0 / d } else { 0.0 };
            scales[r * nblk + b] = d;
            for wi in 0..8 {
                let mut word = 0u32;
                for byte in 0..4 {
                    let q = ((blk[wi * 4 + byte] * id).round() as i32).clamp(-127, 127) as i8;
                    word |= (q as u8 as u32) << (byte * 8);
                }
                quants[(r * nblk + b) * 8 + wi] = word;
            }
        }
    }
    (scales, quants)
}

/// An f16 weight matrix on the GPU (`[rows*cols]` f16 little-endian). f16 packed math is ~2× the f32
/// ALU rate on the compute-bound decode GEMVs, with no Q4 dequant overhead.
pub struct F16 {
    pub buf: wgpu::Buffer,
}

/// Convert an f32 slice to little-endian f16 bytes.
pub fn f32_to_f16_bytes(v: &[f32]) -> Vec<u8> {
    let mut out = Vec::with_capacity(v.len() * 2);
    for &x in v {
        out.extend_from_slice(&half::f16::from_f32(x).to_le_bytes());
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A tiny mixed-dtype safetensors as raw bytes.
    fn toy_safetensors() -> Vec<u8> {
        use safetensors::tensor::TensorView;
        let f32v: Vec<u8> = bytemuck::cast_slice(&[1.0f32, 2.0, 3.0, 4.0]).to_vec();
        let f16v: Vec<u8> = [half::f16::from_f32(5.0), half::f16::from_f32(6.0)]
            .iter()
            .flat_map(|h| h.to_le_bytes())
            .collect();
        let views = vec![
            (
                "a".to_string(),
                TensorView::new(safetensors::Dtype::F32, vec![2, 2], &f32v).unwrap(),
            ),
            (
                "b".to_string(),
                TensorView::new(safetensors::Dtype::F16, vec![2], &f16v).unwrap(),
            ),
        ];
        safetensors::tensor::serialize(views, &None).unwrap()
    }

    #[test]
    fn should_read_tensors_identically_from_bytes_and_from_file() {
        let bytes = toy_safetensors();
        let dir = std::env::temp_dir().join(format!("lfm2-lazyst-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        std::fs::write(dir.join("model.safetensors"), &bytes).unwrap();

        let from_file = LazySt::open(&dir).expect("open");
        let from_bytes = LazySt::from_bytes(vec![bytes]).expect("from_bytes");

        for name in ["a", "b"] {
            assert_eq!(
                from_file.shape(name).unwrap(),
                from_bytes.shape(name).unwrap()
            );
            let f = from_file.tensor_f32(name).expect("file tensor");
            let b = from_bytes.tensor_f32(name).expect("bytes tensor");
            assert_eq!(f, b, "tensor {name} differs between file and bytes backing");
        }
        // The dtype conversions landed (F32 verbatim, F16 widened).
        assert_eq!(
            from_bytes.tensor_f32("a").unwrap(),
            vec![1.0, 2.0, 3.0, 4.0]
        );
        assert_eq!(from_bytes.tensor_f32("b").unwrap(), vec![5.0, 6.0]);
    }

    #[test]
    fn should_reject_a_truncated_blob() {
        let bytes = toy_safetensors();
        let err = LazySt::from_bytes(vec![bytes[..4].to_vec()])
            .err()
            .expect("truncated blob must error");
        assert!(err.to_string().contains("too small"), "{err}");
    }

    #[test]
    fn should_parse_the_gemma4_edge_e2b_config() {
        // The gemma-4-E2B-it shape, verbatim from the real checkpoint (wrapped text_config,
        // rope_parameters dict, 4:1 layer pattern, all four edge features on).
        let layer_types: Vec<&str> = (0..35)
            .map(|i| {
                if (i + 1) % 5 == 0 {
                    "full_attention"
                } else {
                    "sliding_attention"
                }
            })
            .collect();
        let cfg_json = serde_json::json!({
            "architectures": ["Gemma4ForConditionalGeneration"],
            "model_type": "gemma4",
            "text_config": {
                "hidden_size": 1536,
                "num_hidden_layers": 35,
                "num_attention_heads": 8,
                "num_key_value_heads": 1,
                "head_dim": 256,
                "global_head_dim": 512,
                "intermediate_size": 6144,
                "vocab_size": 262144,
                "rms_norm_eps": 1e-6,
                "hidden_activation": "gelu_pytorch_tanh",
                "final_logit_softcapping": 30.0,
                "sliding_window": 512,
                "attention_k_eq_v": false,
                "enable_moe_block": false,
                "use_double_wide_mlp": true,
                "hidden_size_per_layer_input": 256,
                "vocab_size_per_layer_input": 262144,
                "num_kv_shared_layers": 20,
                "layer_types": layer_types,
                "rope_parameters": {
                    "full_attention": {
                        "partial_rotary_factor": 0.25,
                        "rope_theta": 1000000.0,
                        "rope_type": "proportional"
                    },
                    "sliding_attention": {
                        "rope_theta": 10000.0,
                        "rope_type": "default"
                    }
                }
            }
        });
        let c = Lfm2Config::from_json(cfg_json.to_string().as_bytes()).expect("E2B must parse");
        assert_eq!(c.arch, Arch::Gemma4);
        assert_eq!((c.hidden, c.n_layers, c.vocab), (1536, 35, 262144));
        assert_eq!((c.n_heads, c.n_kv_heads, c.head_dim), (8, 1, 256));
        // Dual head geometry: full layers (4, 9, 14, …) run 512-wide heads.
        assert_eq!(c.head_dim_at(0), 256);
        assert_eq!(c.head_dim_at(4), 512);
        assert_eq!(c.head_dim_at(34), 512);
        assert_eq!(c.max_head_dim(), 512);
        // KV sharing: donors are the LAST layer of each type below 15 — sliding→13, full→14.
        assert!(c.owns_kv(14) && !c.owns_kv(15));
        assert_eq!(c.kv_src_at(15), 13, "shared sliding layers read layer 13");
        assert_eq!(c.kv_src_at(19), 14, "shared full layers read layer 14");
        assert_eq!(c.kv_src_at(33), 13);
        assert_eq!(c.kv_src_at(34), 14);
        // Double-wide MLP on exactly the shared layers.
        assert_eq!(c.intermediate_at(14), 6144);
        assert_eq!(c.intermediate_at(15), 12288);
        assert_eq!(c.max_intermediate(), 12288);
        // PLE + softcap + rope.
        assert_eq!(c.edge.ple_dim, 256);
        assert_eq!(c.edge.final_logit_softcapping, 30.0);
        assert_eq!(
            c.edge.rope_prop_pairs, 64,
            "0.25·512/2 rotated pairs, denominator = FULL head_dim"
        );
        assert_eq!((c.rope_theta, c.rope_local_theta), (1_000_000.0, 10_000.0));
        assert!(!c.layer_k_eq_v.iter().any(|&x| x), "E2B has no K=V layers");
        assert!(c.layer_is_sliding[0] && !c.layer_is_sliding[4]);

        // A NON-edge Gemma-4 config (the synthetic fixtures' shape) parses to empty edge
        // vectors — the uniform fast path, byte-identical plans.
        let plain = serde_json::json!({
            "architectures": ["Gemma4ForCausalLM"],
            "hidden_size": 64, "num_hidden_layers": 2, "num_attention_heads": 2,
            "num_key_value_heads": 1, "head_dim": 32, "intermediate_size": 96,
            "vocab_size": 128, "layer_types": ["sliding_attention", "full_attention"],
            "rope_local_base_freq": 10000.0, "sliding_window": 4
        });
        let p = Lfm2Config::from_json(plain.to_string().as_bytes()).expect("plain gemma4");
        assert!(p.edge.layer_head_dim.is_empty() && p.edge.layer_kv_src.is_empty());
        assert!(p.edge.layer_intermediate.is_empty());
        assert_eq!(p.edge.ple_dim, 0);
        assert_eq!(p.rope_local_theta, 10_000.0);
        assert!(p.owns_kv(1) && p.head_dim_at(1) == 32 && p.intermediate_at(1) == 96);
    }

    #[test]
    fn should_refuse_unknown_decoder_architectures_loudly() {
        // An UNPORTED checkpoint must NOT silently load as LFM2 — the old default arm decoded
        // garbage instead of refusing. (Llama/Mistral are now PORTED — see the arms above and
        // the `should_load_llama_and_mistral` gate below; the refusal set here is genuinely
        // unported families.)
        for arch in [
            "JambaForCausalLM",
            "Zamba2ForCausalLM",
            "Gemma3ForConditionalGeneration",
        ] {
            let cfg = format!(r#"{{"architectures":["{arch}"]}}"#);
            let err = detect_model_kind(cfg.as_bytes())
                .err()
                .unwrap_or_else(|| panic!("{arch} must be refused, not defaulted"));
            let msg = err.to_string();
            assert!(msg.contains(arch), "error names the offender: {msg}");
            assert!(
                msg.contains("Qwen3ForCausalLM"),
                "error lists supported archs: {msg}"
            );
        }
        let err = Lfm2Config::from_json(br#"{"architectures":["JambaForCausalLM"]}"#)
            .expect_err("from_json must refuse unknown archs too");
        assert!(err.to_string().contains("JambaForCausalLM"), "{err}");
    }

    #[test]
    fn should_load_qwen2_moe() {
        // Qwen2 attention + shared-gated MoE; intermediate = shared_expert_intermediate_size.
        let base = r#""architectures":["Qwen2MoeForCausalLM"],"hidden_size":64,
            "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":2,"head_dim":16,
            "vocab_size":256,"rms_norm_eps":1e-6,"rope_theta":1000000.0,
            "num_experts":60,"num_experts_per_tok":4,"moe_intermediate_size":128,
            "shared_expert_intermediate_size":512"#;
        let c = Lfm2Config::from_json(format!("{{{base}}}").as_bytes()).expect("qwen2-moe loads");
        assert_eq!(c.arch, Arch::Qwen2Moe);
        assert_eq!(c.num_experts, 60);
        assert_eq!(c.top_k_experts, 4);
        assert_eq!(c.moe_intermediate, 128);
        assert_eq!(c.intermediate, 512, "shared expert → dense GLU slots");
        // Mixed dense/MoE layouts and no-renorm routers are refused (single-slot-size assumption).
        for (extra, needle) in [
            (r#","decoder_sparse_step":2"#, "decoder_sparse_step"),
            (r#","mlp_only_layers":[0]"#, "mlp_only_layers"),
            (r#","norm_topk_prob":false"#, "norm_topk_prob"),
        ] {
            let err = Lfm2Config::from_json(format!("{{{base}{extra}}}").as_bytes())
                .expect_err("unsupported Qwen2-MoE variant must be refused");
            assert!(
                err.to_string().contains(needle),
                "error names {needle}: {err}"
            );
        }
    }

    #[test]
    fn should_load_olmo() {
        // OLMo 1: LayerNorm arch (non-parametric); head_dim = hidden/heads, clip_qkv refused.
        let base = r#""architectures":["OlmoForCausalLM"],"hidden_size":64,"intermediate_size":128,
            "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":2,
            "vocab_size":256,"layer_norm_eps":1e-5,"rope_theta":10000.0"#;
        let c = Lfm2Config::from_json(format!("{{{base}}}").as_bytes()).expect("olmo loads");
        assert_eq!(c.arch, Arch::Olmo);
        assert!(arch_is_layernorm(c.arch), "OLMo normalizes with LayerNorm");
        assert_eq!(c.head_dim, 16, "head_dim = hidden / heads");
        assert_eq!(c.num_experts, 0);
        let err = Lfm2Config::from_json(format!("{{{base},\"clip_qkv\":8.0}}").as_bytes())
            .expect_err("clip_qkv must be refused");
        assert!(err.to_string().contains("clip_qkv"), "{err}");
    }

    #[test]
    fn should_load_stablelm() {
        // StableLM-2: affine-LayerNorm arch, partial rotary from partial_rotary_factor, no qk-norm.
        let base = r#""architectures":["StableLmForCausalLM"],"hidden_size":64,"intermediate_size":128,
            "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":2,"head_dim":16,
            "vocab_size":256,"layer_norm_eps":1e-5,"rope_theta":10000.0,"partial_rotary_factor":0.25"#;
        let c = Lfm2Config::from_json(format!("{{{base}}}").as_bytes()).expect("stablelm loads");
        assert_eq!(c.arch, Arch::StableLm);
        assert!(
            arch_is_layernorm(c.arch) && arch_layernorm_affine(c.arch),
            "affine LayerNorm"
        );
        assert_eq!(c.rotary_dim, 4, "partial rotary = 0.25 · head_dim(16)");
        let err = Lfm2Config::from_json(format!("{{{base},\"qk_layernorm\":true}}").as_bytes())
            .expect_err("qk_layernorm must be refused");
        assert!(err.to_string().contains("qk_layernorm"), "{err}");
    }

    #[test]
    fn should_load_falcon() {
        // Falcon-7B: parallel + non-gated MLP + multiquery; MLP is 4·hidden; 40B variant refused.
        let base = r#""architectures":["FalconForCausalLM"],"hidden_size":64,
            "num_attention_heads":4,"num_hidden_layers":2,"vocab_size":256,
            "layer_norm_epsilon":1e-5,"multi_query":true,"parallel_attn":true,"bias":false"#;
        let c = Lfm2Config::from_json(format!("{{{base}}}").as_bytes()).expect("falcon loads");
        assert_eq!(c.arch, Arch::Falcon);
        assert!(arch_is_parallel(c.arch) && arch_mlp_nogate(c.arch) && arch_gelu_exact(c.arch));
        assert!(
            arch_layernorm_affine(c.arch),
            "Falcon uses affine LayerNorm"
        );
        assert_eq!(c.n_kv_heads, 1, "multiquery");
        assert_eq!(c.intermediate, 4 * 64, "MLP is 4·hidden");
        let err = Lfm2Config::from_json(
            format!("{{{base},\"new_decoder_architecture\":true}}").as_bytes(),
        )
        .expect_err("Falcon-40B must be refused");
        assert!(
            err.to_string().contains("new_decoder_architecture"),
            "{err}"
        );
    }

    #[test]
    fn should_load_cohere() {
        // Cohere: parallel + gated SwiGLU + weight-only affine LayerNorm; use_qk_norm refused.
        let base = r#""architectures":["CohereForCausalLM"],"hidden_size":64,"intermediate_size":128,
            "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":2,"head_dim":16,
            "vocab_size":256,"layer_norm_eps":1e-5,"rope_theta":10000.0,"logit_scale":0.0625"#;
        let c = Lfm2Config::from_json(format!("{{{base}}}").as_bytes()).expect("cohere loads");
        assert_eq!(c.arch, Arch::Cohere);
        assert!(arch_is_parallel(c.arch) && arch_layernorm_affine(c.arch));
        assert!(!arch_mlp_nogate(c.arch), "Cohere is a gated SwiGLU MLP");
        let err = Lfm2Config::from_json(format!("{{{base},\"use_qk_norm\":true}}").as_bytes())
            .expect_err("use_qk_norm must be refused");
        assert!(err.to_string().contains("use_qk_norm"), "{err}");
    }

    #[test]
    fn should_load_mamba() {
        // Mamba: state-space model, not a transformer. No attention (dummy head dims), SSM layers.
        let cfg = r#"{"architectures":["MambaForCausalLM"],"hidden_size":64,
            "num_hidden_layers":2,"vocab_size":256,"state_size":16,"conv_kernel":4,"expand":2,
            "layer_norm_epsilon":1e-5}"#;
        let c = Lfm2Config::from_json(cfg.as_bytes()).expect("mamba loads");
        assert_eq!(c.arch, Arch::Mamba);
        assert!(
            c.layer_is_attn.iter().all(|&a| !a),
            "Mamba layers are SSM, not attention"
        );
        assert_eq!(c.n_layers, 2);
    }

    #[test]
    fn should_load_rwkv() {
        // RWKV-4: an RNN (token-shift + WKV), not a transformer. No attention, affine LayerNorm.
        let cfg = r#"{"architectures":["RwkvForCausalLM"],"hidden_size":64,"intermediate_size":256,
            "num_hidden_layers":2,"vocab_size":256,"layer_norm_epsilon":1e-5}"#;
        let c = Lfm2Config::from_json(cfg.as_bytes()).expect("rwkv loads");
        assert_eq!(c.arch, Arch::Rwkv);
        assert!(arch_layernorm_affine(c.arch), "RWKV uses affine LayerNorm");
        assert!(
            c.layer_is_attn.iter().all(|&a| !a),
            "RWKV layers are recurrent, not attention"
        );
        assert_eq!(c.n_layers, 2);
    }

    #[test]
    fn should_load_olmo2() {
        // OLMo2: RMSNorm post-norm arch (skip_prenorm) — distinct from OLMo1 (LayerNorm pre-norm).
        let cfg = r#"{"architectures":["Olmo2ForCausalLM"],"hidden_size":64,"intermediate_size":128,
            "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":2,"head_dim":16,
            "vocab_size":256,"rms_norm_eps":1e-6,"rope_theta":500000.0}"#;
        let c = Lfm2Config::from_json(cfg.as_bytes()).expect("olmo2 loads");
        assert_eq!(c.arch, Arch::Olmo2);
        assert!(arch_skip_prenorm(c.arch), "OLMo2 is post-norm");
        assert!(
            !arch_is_layernorm(c.arch),
            "OLMo2 uses RMSNorm, not LayerNorm"
        );
    }

    #[test]
    fn should_load_phi2() {
        // Phi-2 (PhiForCausalLM): parallel + non-gated GELU + biases + partial rotary. Distinct
        // from Phi-3 (Arch::Phi3). act_gelu (tanh) set for the MLP_ACT_MUL mode-1 path.
        let cfg = r#"{"architectures":["PhiForCausalLM"],"hidden_size":64,"intermediate_size":128,
            "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":2,"head_dim":16,
            "vocab_size":256,"layer_norm_eps":1e-5,"rope_theta":10000.0,"partial_rotary_factor":0.5}"#;
        let c = Lfm2Config::from_json(cfg.as_bytes()).expect("phi2 loads");
        assert_eq!(c.arch, Arch::Phi2);
        assert!(arch_is_parallel(c.arch) && arch_has_biases(c.arch) && arch_mlp_nogate(c.arch));
        assert!(c.act_gelu, "gelu_new → tanh MLP_ACT_MUL mode");
        assert_eq!(c.rotary_dim, 8, "partial rotary = 0.5 · head_dim(16)");
    }

    #[test]
    fn should_load_nemotron() {
        // Nemotron: LayerNorm1P (affine) + squared-ReLU MLP + partial rotary.
        let cfg = r#"{"architectures":["NemotronForCausalLM"],"hidden_size":64,
            "intermediate_size":128,"num_hidden_layers":2,"num_attention_heads":4,
            "num_key_value_heads":2,"head_dim":16,"vocab_size":256,"norm_eps":1e-5,
            "rope_theta":10000.0,"partial_rotary_factor":0.5}"#;
        let c = Lfm2Config::from_json(cfg.as_bytes()).expect("nemotron loads");
        assert_eq!(c.arch, Arch::Nemotron);
        assert!(
            arch_is_layernorm(c.arch) && arch_layernorm_affine(c.arch) && arch_act_relu(c.arch)
        );
        assert_eq!(c.rotary_dim, 8, "partial rotary = 0.5 · head_dim(16)");
    }

    #[test]
    fn should_load_llama_and_mistral() {
        // Both map to Arch::Llama; Mistral's sliding_window is read, Llama's stays 0.
        for (name, want_window) in [("LlamaForCausalLM", 0usize), ("MistralForCausalLM", 4096)] {
            let cfg = format!(
                r#"{{"architectures":["{name}"],"hidden_size":64,"intermediate_size":128,
                    "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":2,
                    "vocab_size":256,"rms_norm_eps":1e-5,"rope_theta":500000.0{}}}"#,
                if want_window > 0 {
                    format!(r#","sliding_window":{want_window}"#)
                } else {
                    String::new()
                }
            );
            let c = Lfm2Config::from_json(cfg.as_bytes()).expect("llama/mistral config loads");
            assert_eq!(c.arch, Arch::Llama, "{name}");
            assert_eq!(c.sliding_window, want_window, "{name} sliding window");
            assert!(c.layer_is_attn.iter().all(|&a| a), "all attention layers");
            assert_eq!(c.head_dim, 16, "{name} head_dim = hidden/heads");
        }
    }

    #[test]
    fn should_load_deepseek_mla() {
        // MLA: head_dim = qk_nope + qk_rope (MHA-expanded); the MLA-specific dims live in Op::Mla.
        let cfg = r#"{"architectures":["DeepseekV2ForCausalLM"],"hidden_size":64,
            "intermediate_size":128,"num_hidden_layers":2,"num_attention_heads":4,
            "vocab_size":256,"rms_norm_eps":1e-6,"rope_theta":10000.0,
            "q_lora_rank":96,"kv_lora_rank":64,"qk_nope_head_dim":16,"qk_rope_head_dim":16,
            "v_head_dim":32}"#;
        let c = Lfm2Config::from_json(cfg.as_bytes()).expect("deepseek-v2 config loads");
        assert_eq!(c.arch, Arch::DeepseekV2);
        assert_eq!(c.head_dim, 32, "head_dim = qk_nope + qk_rope");
        assert_eq!(c.n_kv_heads, c.n_heads, "MLA expands to full MHA");
        assert_eq!(c.num_experts, 0, "no n_routed_experts ⇒ dense FFN");
    }

    #[test]
    fn should_load_deepseek_moe() {
        // MLA attention + fine-grained MoE: shared experts fold into the dense GLU
        // (intermediate = n_shared·moe_intermediate), routed experts on the Moe struct.
        let base = r#""architectures":["DeepseekV2ForCausalLM"],"hidden_size":64,
            "intermediate_size":999,"num_hidden_layers":2,"num_attention_heads":4,
            "vocab_size":256,"rms_norm_eps":1e-6,"rope_theta":10000.0,
            "q_lora_rank":96,"kv_lora_rank":64,"qk_nope_head_dim":16,"qk_rope_head_dim":16,
            "v_head_dim":32,"n_routed_experts":64,"num_experts_per_tok":6,
            "moe_intermediate_size":128,"n_shared_experts":2,"routed_scaling_factor":16.0"#;
        let c =
            Lfm2Config::from_json(format!("{{{base}}}").as_bytes()).expect("deepseek moe loads");
        assert_eq!(c.arch, Arch::DeepseekV2);
        assert_eq!(c.num_experts, 64);
        assert_eq!(c.top_k_experts, 6);
        assert_eq!(c.moe_intermediate, 128);
        assert_eq!(
            c.intermediate, 256,
            "shared experts fold into dense GLU (2·128), NOT intermediate_size"
        );
        // Unsupported router/layout variants are refused loudly, not silently mis-computed.
        for (extra, needle) in [
            (r#","first_k_dense_replace":1"#, "first_k_dense_replace"),
            (r#","norm_topk_prob":false"#, "norm_topk_prob"),
            // sigmoid scoring (V3) WITHOUT the group params can't build the V3 router.
            (r#","scoring_func":"sigmoid""#, "n_group"),
            (r#","topk_method":"group_limited_greedy""#, "topk_method"),
        ] {
            let err = Lfm2Config::from_json(format!("{{{base}{extra}}}").as_bytes())
                .expect_err("unsupported MoE variant must be refused");
            assert!(
                err.to_string().contains(needle),
                "error names {needle}: {err}"
            );
        }
    }

    #[test]
    fn should_load_deepseek_v3() {
        // V3 = MLA (same arch) + the sigmoid/group-limited router. The group params make the
        // sigmoid config load; the config layer only needs them present (the loader reads them).
        let cfg = r#"{"architectures":["DeepseekV3ForCausalLM"],"hidden_size":64,
            "num_hidden_layers":2,"num_attention_heads":4,"vocab_size":256,"rms_norm_eps":1e-6,
            "rope_theta":10000.0,"q_lora_rank":96,"kv_lora_rank":64,"qk_nope_head_dim":16,
            "qk_rope_head_dim":16,"v_head_dim":32,"n_routed_experts":256,"num_experts_per_tok":8,
            "moe_intermediate_size":128,"n_shared_experts":1,"routed_scaling_factor":2.5,
            "n_group":8,"topk_group":4,"scoring_func":"sigmoid","topk_method":"noaux_tc"}"#;
        let c = Lfm2Config::from_json(cfg.as_bytes()).expect("deepseek-v3 config loads");
        assert_eq!(c.arch, Arch::DeepseekV2, "V3 reuses the MLA arch");
        assert_eq!(c.num_experts, 256);
        assert_eq!(c.top_k_experts, 8);
        assert_eq!(c.intermediate, 128, "1 shared expert · moe_intermediate");
    }

    #[test]
    fn should_load_qwen3_moe() {
        let cfg = r#"{"architectures":["Qwen3MoeForCausalLM"],"hidden_size":64,
            "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":2,"head_dim":16,
            "vocab_size":256,"rms_norm_eps":1e-6,"rope_theta":1000000.0,
            "num_experts":128,"num_experts_per_tok":8,"moe_intermediate_size":768}"#;
        let c = Lfm2Config::from_json(cfg.as_bytes()).expect("qwen3-moe config loads");
        assert_eq!(c.arch, Arch::Qwen3Moe);
        assert_eq!(c.num_experts, 128);
        assert_eq!(c.top_k_experts, 8);
        assert_eq!(c.moe_intermediate, 768);
        assert_eq!(c.intermediate, 32, "tiny zero dense → pure routed");
    }

    #[test]
    fn should_load_mixtral_as_pure_moe() {
        let cfg = r#"{"architectures":["MixtralForCausalLM"],"hidden_size":64,
            "intermediate_size":128,"num_hidden_layers":2,"num_attention_heads":4,
            "num_key_value_heads":2,"vocab_size":256,"rms_norm_eps":1e-5,"rope_theta":1000000.0,
            "num_local_experts":8,"num_experts_per_tok":2,"sliding_window":4096}"#;
        let c = Lfm2Config::from_json(cfg.as_bytes()).expect("mixtral config loads");
        assert_eq!(c.arch, Arch::Mixtral);
        assert_eq!(c.num_experts, 8);
        assert_eq!(c.top_k_experts, 2);
        assert_eq!(c.moe_intermediate, 128, "experts use intermediate_size");
        assert_eq!(
            c.intermediate, 32,
            "dense GLU is tiny (zeroed → pure routed)"
        );
        assert_eq!(c.sliding_window, 4096);
    }

    #[test]
    fn should_load_granite() {
        // Granite maps to Arch::Granite (its multipliers are folded into weights at load, so they
        // don't surface in the config struct — this just checks detection + the Llama-family dims).
        let cfg = r#"{"architectures":["GraniteForCausalLM"],"hidden_size":64,
            "intermediate_size":128,"num_hidden_layers":2,"num_attention_heads":4,
            "num_key_value_heads":2,"vocab_size":256,"rms_norm_eps":1e-5,"rope_theta":10000.0,
            "embedding_multiplier":12.0,"residual_multiplier":0.22,"attention_multiplier":0.015625,
            "logits_scaling":8.0}"#;
        let c = Lfm2Config::from_json(cfg.as_bytes()).expect("granite config loads");
        assert_eq!(c.arch, Arch::Granite);
        assert!(c.layer_is_attn.iter().all(|&a| a));
        assert_eq!(c.head_dim, 16);
    }

    #[test]
    fn should_load_phi3_with_partial_rotary() {
        // Full rotary (factor 1.0 / absent) → rotary_dim 0; a partial factor → head_dim·factor.
        for (factor, want_rd) in [(None, 0usize), (Some(1.0), 0), (Some(0.5), 32)] {
            let extra = factor
                .map(|f| format!(r#","partial_rotary_factor":{f}"#))
                .unwrap_or_default();
            let cfg = format!(
                r#"{{"architectures":["Phi3ForCausalLM"],"hidden_size":256,"intermediate_size":512,
                    "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":4,
                    "vocab_size":256,"rms_norm_eps":1e-5,"rope_theta":10000.0{extra}}}"#,
            );
            let c = Lfm2Config::from_json(cfg.as_bytes()).expect("phi3 config loads");
            assert_eq!(c.arch, Arch::Phi3);
            assert_eq!(c.head_dim, 64, "head_dim = hidden/heads");
            assert_eq!(c.rotary_dim, want_rd, "partial_rotary_factor {factor:?}");
        }
    }

    #[test]
    fn should_load_qwen2_and_gate_its_sliding_window() {
        // Qwen2 = Arch::Qwen2 (bias plumbing is a loader/plan detail); `use_sliding_window`
        // gates whether `sliding_window` takes effect (Qwen2.5 ships it false).
        for (use_sliding, want) in [(false, 0usize), (true, 4096)] {
            let cfg = format!(
                r#"{{"architectures":["Qwen2ForCausalLM"],"hidden_size":64,"intermediate_size":128,
                    "num_hidden_layers":2,"num_attention_heads":4,"num_key_value_heads":2,
                    "vocab_size":256,"rms_norm_eps":1e-6,"rope_theta":1000000.0,
                    "sliding_window":4096,"use_sliding_window":{use_sliding}}}"#,
            );
            let c = Lfm2Config::from_json(cfg.as_bytes()).expect("qwen2 config loads");
            assert_eq!(c.arch, Arch::Qwen2);
            assert_eq!(c.sliding_window, want, "use_sliding_window={use_sliding}");
        }
    }

    #[test]
    fn body_mixed_policies_resolve_to_per_site_kinds() {
        // Body-mixed policies now SERVE (the plan is per-site kind-aware); the loader resolves
        // each group to its container kind, and a blk: override takes its whole layers.
        let p = Precision::parse("q4,gate_up=q8,head=f16,blk:2=f16").unwrap();
        assert_eq!(p.uniform_body(), None, "gate_up differs — body is mixed");
        let k0 = p.site_kinds(0, 4).unwrap();
        assert_eq!(
            (k0.qkv, k0.o, k0.gate_up, k0.down),
            (WKind::Q4_0, WKind::Q4_0, WKind::Q8_0N, WKind::Q4_0)
        );
        let k2 = p.site_kinds(2, 4).unwrap();
        assert_eq!(k2, LayerKinds::uniform(WKind::F16), "blk override takes all four sites");
        // an out-of-range OVERRIDE refuses (a silently dropped override ships the wrong model)
        let oor = Precision::parse("q4,gate_up=q8,blk:9=f16").unwrap();
        assert!(oor.site_kinds(0, 4).is_err(), "out-of-range blk refuses");
    }

    /// Dequantize Q4_0 exactly as the GEMV does (nibble `i` / `i+16` share a byte; the block
    /// scale round-trips through f16 as it does in the GPU buffer).
    fn dequant_q4(scales: &[f32], quants: &[u32], rows: usize, cols: usize) -> Vec<f32> {
        let nblk = cols / 32;
        let mut out = vec![0f32; rows * cols];
        for r in 0..rows {
            for b in 0..nblk {
                let d = half::f16::from_f32(scales[r * nblk + b]).to_f32();
                for wi in 0..4 {
                    let word = quants[(r * nblk + b) * 4 + wi];
                    for byte in 0..4 {
                        let v = (word >> (byte * 8)) & 0xff;
                        let i = wi * 4 + byte;
                        out[r * cols + b * 32 + i] = d * ((v & 0x0f) as i32 - 8) as f32;
                        out[r * cols + b * 32 + i + 16] = d * (((v >> 4) & 0x0f) as i32 - 8) as f32;
                    }
                }
            }
        }
        out
    }

    /// Dequantize Q8_0 exactly as the GEMV does (`unpack4xI8` per word, f16 block scale).
    fn dequant_q8(scales: &[f32], quants: &[u32], rows: usize, cols: usize) -> Vec<f32> {
        let nblk = cols / 32;
        let mut out = vec![0f32; rows * cols];
        for r in 0..rows {
            for b in 0..nblk {
                let d = half::f16::from_f32(scales[r * nblk + b]).to_f32();
                for wi in 0..8 {
                    let word = quants[(r * nblk + b) * 8 + wi];
                    for byte in 0..4 {
                        let q = ((word >> (byte * 8)) & 0xff) as u8 as i8;
                        out[r * cols + b * 32 + wi * 4 + byte] = d * q as f32;
                    }
                }
            }
        }
        out
    }

    #[test]
    fn should_round_trip_q8_within_its_quantization_step() {
        // Q8_0's worst-case error is half a step: d/2 = max|x|/254 per block.
        let (rows, cols) = (3usize, 64usize);
        let w: Vec<f32> = (0..rows * cols)
            .map(|i| ((i * 37 % 101) as f32 - 50.0) / 50.0)
            .collect();
        let (scales, quants) = quantize_q8_0(&w, rows, cols);
        let back = dequant_q8(&scales, &quants, rows, cols);
        for b in 0..(rows * cols) / 32 {
            let blk = &w[b * 32..b * 32 + 32];
            let amax = blk.iter().fold(0f32, |m, v| m.max(v.abs()));
            // The f16 block scale itself rounds, so allow a small multiple of the step.
            let tol = amax / 254.0 * 1.05 + 1e-6;
            for i in 0..32 {
                let e = (back[b * 32 + i] - blk[i]).abs();
                assert!(e <= tol, "block {b} elem {i}: err {e} > tol {tol}");
            }
        }
    }

    #[test]
    fn should_quantize_more_accurately_in_q8_than_q4() {
        // The ORDERING the whole precision ladder rests on (Q4 0.459 → Q8 0.966 → f16 1.000
        // agreement in the end-to-end bench). Pin it here so a kernel/format change can never
        // silently invert it.
        let (rows, cols) = (4usize, 128usize);
        let w: Vec<f32> = (0..rows * cols)
            .map(|i| (((i * 61 % 197) as f32 - 98.0) / 98.0) * 0.3)
            .collect();
        let (s4, q4) = quantize_q4_0(&w, rows, cols);
        let (s8, q8) = quantize_q8_0(&w, rows, cols);
        let rmse = |back: &[f32]| -> f32 {
            (back
                .iter()
                .zip(&w)
                .map(|(a, b)| (a - b).powi(2))
                .sum::<f32>()
                / w.len() as f32)
                .sqrt()
        };
        let e4 = rmse(&dequant_q4(&s4, &q4, rows, cols));
        let e8 = rmse(&dequant_q8(&s8, &q8, rows, cols));
        assert!(
            e8 * 4.0 < e4,
            "Q8 must be far more accurate than Q4 (rmse q8 {e8:.2e} vs q4 {e4:.2e})"
        );
    }

    #[test]
    fn pack_q8_0n_round_trips_within_the_quantization_step() {
        // encode → decode must agree with the SOURCE within half a step of the f16-stored scale
        // (the same bytes the native Q8_0N kernels read; decode mirrors decode_q8_0n in the
        // legacy_native_kernel gate).
        let (rows, cols) = (3usize, 96usize);
        let w: Vec<f32> = (0..rows * cols)
            .map(|i| ((i * 37 % 113) as f32 - 56.0) * 0.013)
            .collect();
        let bytes = pack_q8_0n(&w, rows, cols);
        assert_eq!(bytes.len(), rows * (cols / 32) * 36);
        for r in 0..rows {
            for b in 0..cols / 32 {
                let o = (r * (cols / 32) + b) * 36;
                let d = half::f16::from_le_bytes([bytes[o], bytes[o + 1]]).to_f32();
                for i in 0..32 {
                    let got = d * (bytes[o + 4 + i] as i8 as f32);
                    let want = w[r * cols + b * 32 + i];
                    assert!(
                        (got - want).abs() <= d * 0.5 + 1e-7,
                        "[{r},{b},{i}] {got} vs {want} (d={d})"
                    );
                }
            }
        }
    }

    #[test]
    fn should_default_precision_to_all_q4() {
        // The historical behaviour must survive as the default: no caller opting in, no change.
        let p = Precision::default();
        assert_eq!(p, Precision::uniform(WDtype::Q4));
        assert!(!p.any_non_q4());
    }

    #[test]
    fn should_split_body_from_head_for_the_head_override_lever() {
        // `head=f16` over a Q4 body is Unsloth's canonical recovery lever, and it is servable
        // BECAUSE the head is a separate kernel path: the body stays uniform (one plan family),
        // only the head differs. `uniform_body` is the predicate the loader gates on.
        let hf16 = Precision::parse("q4,head=f16").unwrap();
        assert_eq!(hf16.uniform_body(), Some(WDtype::Q4), "body is uniform Q4");
        assert_eq!(hf16.head, WDtype::F16, "only the head is upgraded");

        // f16 body with a Q4 head is equally uniform-bodied (the symmetric case).
        let f16_q4head = Precision::parse("f16,head=q4").unwrap();
        assert_eq!(f16_q4head.uniform_body(), Some(WDtype::F16));
        assert_eq!(f16_q4head.head, WDtype::Q4);

        // A body-mixed policy has NO uniform body → still refused by the loader.
        let body_mixed = Precision::parse("q4,gate_up=f16").unwrap();
        assert_eq!(body_mixed.uniform_body(), None, "gate_up differs from qkv/o/down");

        // Uniform everything: body agrees AND equals the head.
        assert_eq!(Precision::uniform(WDtype::F16).uniform_body(), Some(WDtype::F16));
        assert_eq!(Precision::default().uniform_body(), Some(WDtype::Q4));
    }

    #[test]
    fn should_parse_and_resolve_per_layer_overrides() {
        // `blk:<i>=<dtype>` / `blk:<a>-<b>=<dtype>` is the per-LAYER grammar — the Dynamic-2.0
        // core idea (sensitivity is per-layer, not just per-group). Overrides are whole-layer and
        // apply over the uniform body base, in order.
        let p = Precision::parse("q4,blk:0-1=f16,blk:5=f16,head=f16").unwrap();
        assert_eq!(p.uniform_body(), Some(WDtype::Q4), "base body stays uniform");
        assert_eq!(p.head, WDtype::F16);
        let per = p.layer_dtypes(8).unwrap();
        assert_eq!(
            per,
            vec![
                WDtype::F16, // blk 0
                WDtype::F16, // blk 1
                WDtype::Q4,
                WDtype::Q4,
                WDtype::Q4,
                WDtype::F16, // blk 5
                WDtype::Q4,
                WDtype::Q4,
            ]
        );
        // A later override wins over an earlier one (declaration order = application order).
        let p2 = Precision::parse("q4,blk:0-3=f16,blk:2=q4").unwrap();
        assert_eq!(
            p2.layer_dtypes(4).unwrap(),
            vec![WDtype::F16, WDtype::F16, WDtype::Q4, WDtype::F16]
        );
        // No overrides ⇒ an EMPTY vec, the loader's uniform fast-path signal (byte-identity).
        assert!(Precision::parse("q4").unwrap().layer_dtypes(4).unwrap().is_empty());
        assert!(
            Precision::parse("q4,blk:0-0=q4").unwrap().layer_dtypes(4).unwrap().is_empty(),
            "an override that changes nothing must not defeat the uniform fast path"
        );
        // Out-of-range and inverted ranges REFUSE (a silently dropped override would ship the
        // wrong model); so does a bad term.
        assert!(Precision::parse("q4,blk:9=f16").unwrap().layer_dtypes(4).is_err());
        assert!(Precision::parse("q4,blk:3-1=f16").is_err());
        assert!(Precision::parse("q4,blk:x=f16").is_err());
    }

    #[test]
    fn should_parse_the_whole_speed_fidelity_ladder() {
        // q4 → q8 → f16 is a MEASURED ladder (0.459 → 0.966 → 1.000 agreement at 1.0× → 1.9× →
        // 3.6× the weight bytes); every rung must be selectable.
        for (spec, want) in [
            ("q4", WDtype::Q4),
            ("q8", WDtype::Q8),
            ("q8_0", WDtype::Q8),
            ("f16", WDtype::F16),
        ] {
            assert_eq!(Precision::parse(spec).unwrap(), Precision::uniform(want));
        }
        assert!(!Precision::parse("q4").unwrap().any_non_q4());
        assert!(Precision::parse("q8").unwrap().any_non_q4());
        assert!(Precision::parse("f16").unwrap().any_non_q4());
    }

    #[test]
    fn should_parse_a_base_with_per_group_overrides() {
        // Mixed policies are the point of the knob: the ablation showed 4-bit damage is
        // sub-additive across groups, so a caller may want e.g. an 8-bit MLP under 4-bit
        // attention — the engine must not second-guess the mix.
        let p = Precision::parse("q4,head=f16").unwrap();
        assert_eq!(p.head, WDtype::F16);
        assert_eq!(p.qkv, WDtype::Q4);
        assert_eq!(p.gate_up, WDtype::Q4);
        assert!(p.any_non_q4());
        let p = Precision::parse("q4,gate_up=q8,down=q8").unwrap();
        assert_eq!(p.gate_up, WDtype::Q8);
        assert_eq!(p.down, WDtype::Q8);
        assert_eq!(p.qkv, WDtype::Q4);
        // …and the inverse: an f16 model that keeps the (bandwidth-heavy) MLP in 4-bit.
        let p = Precision::parse("f16,gate_up=q4,down=q4").unwrap();
        assert_eq!(p.qkv, WDtype::F16);
        assert_eq!(p.head, WDtype::F16);
        assert_eq!(p.gate_up, WDtype::Q4);
        assert_eq!(p.down, WDtype::Q4);
    }

    #[test]
    fn should_parse_overrides_without_an_explicit_base() {
        let p = Precision::parse("head=f16,o=f16").unwrap();
        assert_eq!(p.head, WDtype::F16);
        assert_eq!(p.o, WDtype::F16);
        assert_eq!(p.down, WDtype::Q4); // base = default
    }

    #[test]
    fn should_refuse_unknown_precision_groups_and_dtypes() {
        // A silently-ignored precision request would ship a different model than the user asked
        // for — refuse, and name the valid options.
        let err = Precision::parse("q4,heads=f16").expect_err("unknown group must not parse");
        assert!(err.to_string().contains("heads"), "{err}");
        let err = Precision::parse("q4,head=int4").expect_err("unknown dtype must not parse");
        assert!(err.to_string().contains("int4"), "{err}");
        let err = Precision::parse("q4,head").expect_err("malformed term must not parse");
        assert!(err.to_string().contains("group=dtype"), "{err}");
        assert!(Precision::parse("").is_err(), "empty spec must not parse");
    }

    #[test]
    fn should_derive_gemma3_sliding_layout_from_pattern_when_layer_types_absent() {
        // HF-vintage gemma-3 checkpoints (e.g. google/gemma-3-1b-it) ship
        // `sliding_window_pattern` and NO `layer_types`; transformers derives
        // `is_sliding = (layer_idx + 1) % pattern != 0`. Defaulting to all-full here
        // mis-ropes 5/6 of the layers with the GLOBAL theta (caught as a quality gap by
        // the generation-bench q4-actuals pins: HF 1b 0.766 vs layer_types-twin 0.821).
        let cfg = Lfm2Config::from_json(
            br#"{"architectures":["Gemma3ForCausalLM"],"hidden_size":1152,
                 "num_attention_heads":4,"num_hidden_layers":12,"vocab_size":262144,
                 "intermediate_size":6912,"sliding_window_pattern":6}"#,
        )
        .expect("pattern-only gemma3 config parses");
        let want: Vec<bool> = (0..12).map(|i| (i + 1) % 6 != 0).collect();
        assert_eq!(cfg.layer_is_sliding, want);
    }

    #[test]
    fn should_prefer_explicit_layer_types_over_pattern_for_gemma3() {
        // When both fields are present (e.g. data/gemma-osfql-1b), `layer_types` is the
        // explicit source of truth.
        let cfg = Lfm2Config::from_json(
            br#"{"architectures":["Gemma3ForCausalLM"],"hidden_size":64,
                 "num_attention_heads":2,"num_hidden_layers":3,"vocab_size":128,
                 "intermediate_size":128,"sliding_window_pattern":2,
                 "layer_types":["sliding_attention","sliding_attention","full_attention"]}"#,
        )
        .expect("both-fields gemma3 config parses");
        assert_eq!(cfg.layer_is_sliding, vec![true, true, false]);
    }

    #[test]
    fn should_refuse_gemma3_config_with_neither_layer_types_nor_pattern() {
        // A silent all-full default assigns the wrong per-layer RoPE theta — refuse loudly
        // instead (no known checkpoint omits both fields).
        let err = Lfm2Config::from_json(
            br#"{"architectures":["Gemma3ForCausalLM"],"hidden_size":64,
                 "num_attention_heads":2,"num_hidden_layers":4,"vocab_size":128,
                 "intermediate_size":128}"#,
        )
        .expect_err("neither layer_types nor sliding_window_pattern must not parse silently");
        assert!(err.to_string().contains("layer_types"), "{err}");
    }

    #[test]
    fn should_dispatch_known_decoder_architectures_and_legacy_default() {
        let kind = |j: &[u8]| detect_model_kind(j).expect("known arch");
        for (cfg, want) in [
            (
                br#"{"architectures":["Lfm2ForCausalLM"]}"# as &[u8],
                Arch::Lfm2,
            ),
            (br#"{"architectures":["Gemma3ForCausalLM"]}"#, Arch::Gemma3),
            (br#"{"architectures":["Gemma4ForCausalLM"]}"#, Arch::Gemma4),
            (br#"{"architectures":["Qwen3ForCausalLM"]}"#, Arch::Qwen3),
            (
                br#"{"architectures":["Qwen3_5MoeForConditionalGeneration"]}"#,
                Arch::Qwen35,
            ),
        ] {
            assert!(
                matches!(kind(cfg), ModelKind::Decoder(a) if a == want),
                "{} should dispatch to {want:?}",
                String::from_utf8_lossy(cfg)
            );
        }
        // No `architectures` key at all = the documented legacy default (synthetic fixtures).
        assert!(matches!(kind(b"{}"), ModelKind::Decoder(Arch::Lfm2)));
    }
}