spark_model/weight_loader/

qwen35_dense.rs

// SPDX-License-Identifier: AGPL-3.0-only

use anyhow::Result;
use atlas_core::config::{LayerType, ModelConfig};
use spark_runtime::gpu::GpuBackend;
use spark_runtime::kv_cache::KvCacheDtype;
use spark_runtime::weights::{WeightDtype, WeightStore};

use super::{ModelWeightLoader, WeightFormat};
use crate::layer::TransformerLayer;
use crate::layers::{DenseFfnLayer, FfnComponent, Qwen3AttentionLayer, Qwen3SsmLayer};
use crate::tp_shard::{
    TpGdnDims, TpShardKind, load_qkvo_tp, shard_dense_bf16, shard_gdn_ba_rows, shard_gdn_conv_rows,
    shard_gdn_out_proj_row_parallel, shard_gdn_qkvz_rows, shard_gdn_value_vector,
    shard_quantized_nvfp4,
};
use crate::weight_map::{
    AttentionWeights, DenseWeight, Fp8Weight, MtpWeights, Nvfp4Variant, PackedQ2Weight, SsmWeights,
    dense, dense_auto, dense_f32_safe, dense_keep_f32, dequant_nvfp4_to_bf16, detect_nvfp4_variant,
    gpu_concat_rows, interleave_ba, load_dense_ffn, load_fp8_block_scaled_as_fp8weight,
    load_kv_scales, load_mtp, quantize_to_nvfp4, quantized_auto,
};

/// True when `{prefix}.weight` is FP8 E4M3 on disk with a 2D block scale
/// (`weight_scale_inv` or 2D `weight_scale`) — i.e. a native FP8 checkpoint
/// projection that should load as `Fp8Weight` rather than be requantized to
/// NVFP4. Mirrors `qwen35::load_layers::proj_is_native_fp8`.
/// The Q2_0 group size if `{prefix}.weight` is a keep-packed ternary tensor
/// (`WeightDtype::PackedQ2_0`, produced by the GGUF loader under
/// `ATLAS_GGUF_NATIVE_Q2=1`), else `None`. When `None` for every projection the
/// FFN takes the unchanged BF16→NVFP4 path, so the default (flag-off) behavior
/// is byte-identical.
fn proj_q2_group(store: &WeightStore, prefix: &str) -> Option<u16> {
    store
        .get(&format!("{prefix}.weight"))
        .ok()
        .and_then(|w| w.q2_group())
}

/// Build a [`PackedQ2Weight`] borrowing the store's packed `block_q2_0` buffer.
/// The buffer is owned by the `WeightStore` (freed with it), so this only wraps
/// the pointer + `[n, k]` + group; no dequant, no allocation.
fn packed_q2_from_store(store: &WeightStore, prefix: &str) -> Result<PackedQ2Weight> {
    let w = store.get(&format!("{prefix}.weight"))?;
    let group = w
        .q2_group()
        .ok_or_else(|| anyhow::anyhow!("{prefix}.weight is not keep-packed Q2_0"))?;
    anyhow::ensure!(
        w.shape.len() == 2,
        "packed Q2_0 {prefix}.weight must be 2D, got {:?}",
        w.shape
    );
    Ok(PackedQ2Weight {
        weight: w.ptr,
        n: w.shape[0] as u32,
        k: w.shape[1] as u32,
        group,
    })
}

fn proj_is_native_fp8(store: &WeightStore, prefix: &str) -> bool {
    let is_fp8_weight = store
        .get(&format!("{prefix}.weight"))
        .map(|w| w.dtype == WeightDtype::FP8E4M3)
        .unwrap_or(false);
    let has_block_scale = store.contains(&format!("{prefix}.weight_scale_inv"))
        || store
            .get(&format!("{prefix}.weight_scale"))
            .map(|s| s.shape.len() == 2)
            .unwrap_or(false);
    is_fp8_weight && has_block_scale
}

/// True when `{prefix}.weight` is on-disk FP8 with a scale the native-FP8
/// `w8a16` path can actually consume — i.e. one that is (or broadcasts to) the
/// `[ceil(N/128), ceil(K/128)]` FP32 block grid the kernel indexes as
/// `block_scale[n/128, k/128]`:
///
///   * `weight_scale_inv` / 2-D `weight_scale` shaped as that block grid
///     (DeepSeek-V3 / Qwen-native convention), or
///   * a per-tensor SCALAR `weight_scale` (ModelOpt; e.g. the nvidia
///     Qwen3.6-27B-NVFP4 GDN projections) — `load_fp8_block_scaled_as_fp8weight`
///     broadcasts it into a uniform grid, which is exact.
///
/// A **per-row** `weight_scale` (`[N,1]`, e.g. unsloth's re-quantized
/// Qwen3.6-*-NVFP4, 2026-07-10) is deliberately REJECTED. It is not a block
/// grid: the kernel would read row `n`'s multiplier from grid cell `n/128`, so
/// 127 of every 128 rows get some other row's scale. That is in-bounds — the
/// widened `[N]` buffer is *larger* than the `[N/128, K/128]` grid — so it does
/// not fault; it silently produces garbage logits. Returning false here drops
/// the projection to the default `dequant_fp8_blockscaled_to_bf16` →
/// `quantize_to_nvfp4` path, which reads a `[N,1]` scale correctly
/// (`block_n = N/N = 1`, i.e. one multiplier per row).
fn proj_is_fp8_any_scale(store: &WeightStore, prefix: &str) -> bool {
    let Ok(w) = store.get(&format!("{prefix}.weight")) else {
        return false;
    };
    if w.dtype != WeightDtype::FP8E4M3 || w.shape.len() != 2 {
        return false;
    }
    let (n, k) = (w.shape[0], w.shape[1]);

    for key in [
        format!("{prefix}.weight_scale_inv"),
        format!("{prefix}.weight_scale"),
    ] {
        let Ok(s) = store.get(&key) else { continue };
        // Per-tensor scalar → broadcast to a uniform grid: exact.
        if s.num_elements() == 1 {
            return true;
        }
        // 2-D scale: only a genuine 128×128 block grid is consumable.
        if s.shape.len() == 2 && s.shape[0] == n.div_ceil(128) && s.shape[1] == k.div_ceil(128) {
            return true;
        }
    }
    false
}

/// Concatenate two block-scaled FP8 weights along rows (dim 0):
/// `[n_a, k] ++ [n_b, k] -> [n_a+n_b, k]`. FP8 bytes (1 B/elem) and the
/// `[n/128, k/128]` FP32 block-scale grids are contiguous row-major, so each is
/// a straight device-to-device append. Requires `n_a % 128 == 0` so the
/// scale-grid rows meet at a block boundary (GDN qkv rows are a multiple of
/// 128). Produces the `[Q|K|V|Z]` sequential order the SSM layer expects.
fn concat_fp8_block_scaled(
    a: &Fp8Weight,
    b: &Fp8Weight,
    k: usize,
    gpu: &dyn GpuBackend,
) -> Result<Fp8Weight> {
    let kb = k.div_ceil(128);
    let a_w = a.n as usize * k;
    let b_w = b.n as usize * k;
    let weight = gpu.alloc(a_w + b_w)?;
    gpu.copy_d2d(a.weight, weight, a_w)?;
    gpu.copy_d2d(b.weight, weight.offset(a_w), b_w)?;
    let a_s = (a.n as usize).div_ceil(128) * kb * 4;
    let b_s = (b.n as usize).div_ceil(128) * kb * 4;
    let row_scale = gpu.alloc(a_s + b_s)?;
    gpu.copy_d2d(a.row_scale, row_scale, a_s)?;
    gpu.copy_d2d(b.row_scale, row_scale.offset(a_s), b_s)?;
    Ok(Fp8Weight {
        weight,
        row_scale,
        n: a.n + b.n,
        k: k as u32,
        scale_format: crate::weight_map::WeightQuantFormat::Fp8BlockScaled,
    })
}

/// Opt-in gate for native dense-FP8 attention + FFN dispatch (Qwythos / dense
/// Ornith-FP8). Default OFF.
///
/// VERIFIED 2026-06-29 on Qwythos-9B-FP8 (gb10/ornith-1.0-9b): with the flag
/// on, the FP8 arms fire for all 32 FFN + 8 full-attn layers and text is
/// correct (coherence/fib/tools 3/3). BUT it is NOT a perf win — ~30 tok/s vs
/// ~40 for the NVFP4 fallback — because this target's NVFP4 W4A16 kernels
/// (fused dual-GEMV decode, transposed m128 prefill) are more optimized than
/// its FP8 W8A16 kernels (unfused per-projection GEMV, non-transposed
/// `w8a16_gemm` prefill; the attention FP8 prefill transpose also does not
/// engage). Vision prefill additionally hits a CUDA-700. Making FP8 pay off
/// here needs dedicated dense-FP8 kernels (fused FP8 dual-GEMV + fast
/// transposed FP8 prefill GEMM), not loader wiring. Until then NVFP4 autoquant
/// is the better dense runtime. `ATLAS_DENSE_FP8=1` opts in for that kernel work.
fn dense_fp8_enabled() -> bool {
    std::env::var("ATLAS_DENSE_FP8").as_deref() == Ok("1")
}

/// Whether the native block-scaled FP8 GDN arm runs for this SSM layer.
///
/// ONE predicate, called by `load_layers` and by `prune_after_load`: the
/// second frees the store tensors the first copied into the fused `[QKV|Z]`
/// weight, and a drift between the two is a use-after-free with no
/// diagnostic. The keep-packed Q2 arm `continue`s ahead of this one, so it is
/// part of the condition (#915).
fn gdn_fp8_arm_selected(store: &WeightStore, la: &str, tp_size: usize) -> bool {
    let q2 = tp_size.max(1) == 1
        && std::env::var_os("ATLAS_NO_Q2_GDN").is_none()
        && proj_q2_group(store, &format!("{la}.in_proj_qkv")).is_some()
        && proj_q2_group(store, &format!("{la}.in_proj_z")).is_some();
    !q2 && std::env::var_os("ATLAS_NO_GDN_FP8").is_none()
        && proj_is_fp8_any_scale(store, &format!("{la}.in_proj_qkv"))
        && proj_is_fp8_any_scale(store, &format!("{la}.in_proj_z"))
        && proj_is_fp8_any_scale(store, &format!("{la}.out_proj"))
}

/// The dense-FFN width. `moe_intermediate_size` is the PER-EXPERT width and is
/// unset (=0) on dense Qwen3.6/3.8-*-FP8, so reading it first would size a
/// 0-byte allocation; `intermediate_size` is unset on the older MoE-style
/// configs. Same fallback order as `weight_map/fp8_lut.rs::load_dense_ffn`.
fn ffn_inter(config: &ModelConfig) -> usize {
    if config.intermediate_size > 0 {
        config.intermediate_size
    } else {
        config.moe_intermediate_size
    }
}

/// Whether the native-FP8 dense-FFN overlay runs for THIS layer.
///
/// The condition `load_layers` applied inline before `prune_after_load` needed
/// the same answer. Pure — a function of the store, the config and the
/// detected variant — so the two sites cannot drift, which is the failure a
/// prune predicate must not have: freeing a store tensor a layer still aliases
/// is a use-after-free with no diagnostic.
fn ffn_fp8_arm_selected(
    store: &WeightStore,
    config: &ModelConfig,
    variant: Nvfp4Variant,
    lp: &str,
) -> bool {
    dense_fp8_enabled()
        && config.tp_world_size.max(1) == 1
        && matches!(variant, Nvfp4Variant::Fp8Dequanted)
        && proj_is_native_fp8(store, &format!("{lp}.mlp.gate_proj"))
}

/// Whether THIS layer's gate and up are fused into one `[2*inter, hidden]`
/// block-scaled FP8 weight (#927).
///
/// Beyond the FP8 overlay itself: the target must arm the arm (`[defaults]
/// ffn_gateup_fused`), the model must be DENSE (a MoE config has no
/// `ffn_gate_up_fused` arena buffer and never reaches the dense-FFN path), and
/// both extents must be whole 128-blocks — `concat_fp8_block_scaled` appends
/// the `[N/128, K/128]` scale grids, which is only the fused grid when the
/// seam falls on a block boundary (`ceil` of a sum is not the sum of `ceil`s).
///
/// The SECOND caller is `prune_after_load`, which releases exactly the store
/// tensors this returned true for. Rule and the receipt:
/// `layers/dense_ffn_gateup_fused.rs`.
fn ffn_gateup_fused_selected(
    store: &WeightStore,
    config: &ModelConfig,
    variant: Nvfp4Variant,
    lp: &str,
) -> bool {
    let inter = ffn_inter(config);
    let hidden = config.hidden_size;
    // The CONFIG's widths are what the ledger terms, the prediction and the
    // view offsets are all computed from, so a checkpoint whose tensors
    // disagree with them must DECLINE rather than be concatenated at the wrong
    // stride. Checked here and not at the call site because `prune_after_load`
    // asks this same question and frees on the answer: a predicate the loader
    // narrows locally is a store tensor freed out from under a live view.
    let on_disk = |name: &str| {
        store
            .get(&format!("{lp}.mlp.{name}.weight"))
            .is_ok_and(|w| w.shape == [inter, hidden])
    };
    ffn_fp8_arm_selected(store, config, variant, lp)
        && crate::layers::dense_ffn::gateup_fused::ffn_gateup_fused()
        && config.num_experts == 0
        && inter > 0
        && inter.is_multiple_of(128)
        && hidden.is_multiple_of(128)
        && on_disk("gate_proj")
        && on_disk("up_proj")
}

// `pub` (re-exported from `weight_loader/mod.rs`): the pre-load residency
// PREDICTION in `predicted_residency` is read by spark-server's preflight,
// and it prices its terms with this module's shape helpers so the prediction
// and the loader's own tally cannot be two different arithmetics (#915).
pub mod fp8_residency;
mod loaders_b;
pub mod predicted_residency;
mod rowwise_fp8;

use crate::layers::qwen3_attention::Fp8TwinSet;
use fp8_residency::{DerivedResidency, RouteEnv};

pub struct Qwen35DenseWeightLoader;

impl ModelWeightLoader for Qwen35DenseWeightLoader {
    fn supports_tp(&self) -> bool {
        // FullAttention layers are TP-sharded (NVFP4-from-disk and BF16
        // → NVFP4 paths). LinearAttention (GDN SSM) layers run
        // full-replica per rank — see qwen35.rs for the rationale.
        true
    }

    fn load_layers(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
        layer_kv_dtypes: &[KvCacheDtype],
    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
        let layer_types = if config.layer_types.is_empty() {
            (0..config.num_hidden_layers)
                .map(|i| config.layer_type(i))
                .collect::<Vec<_>>()
        } else {
            config.layer_types.clone()
        };

        let mut layers: Vec<Box<dyn TransformerLayer>> =
            Vec::with_capacity(config.num_hidden_layers);
        let mut attn_idx = 0usize;

        let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
        let quantize_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
        let stream = gpu.default_stream();
        let h = config.hidden_size;

        let variant = detect_nvfp4_variant(store, config);
        let weight_format = WeightFormat::detect(store, config);
        tracing::info!(
            "Weight format: {:?}, NVFP4 variant: {:?}",
            weight_format,
            variant
        );

        // Native FP8 SSM prefill GEMM (Qwen3.6-27B-FP8 root-cause fix,
        // commit 3ebc08a). Atlas's prior SSM in_proj_qkv path was
        // FP8 → BF16 → NVFP4 → BF16 (in `w4a16_gemm` dequant) → MMA — a
        // double-quant chain whose NVFP4 hop's ~4-bit per-group precision
        // is dominated by signal at q/v but attenuated into a k-direction
        // error (HF conv-k ‖6.3‖ vs conv-v ‖117.2‖, ~18× smaller). For
        // every FP8-on-disk checkpoint we install a single-scale FP8 copy
        // of the stacked `[QKV|Z]` and `out_proj` weights for prefill,
        // bypassing the NVFP4 intermediate. Prefill dispatches via the
        // existing `fp8_gemm_n128` (BF16 act × FP8 weight) — same path
        // the MoE shared-expert FP8 prefill uses. Decode/GEMV unchanged.
        // Originally env-gated `ATLAS_FP8_SSM_PREFILL=1`; promoted to
        // unconditional 2026-05-20 after live verification (commit
        // dfb4e8a era, tokens_to_first_degeneration 1,196 → 16,968).
        // 2026-07-03 precision policy: GDN projections stay ≥FP8 — the
        // NVFP4 requant of qkvz/out_proj flattens tensors that both
        // checkpoint toolchains deliberately keep high-precision (modelopt
        // sensitivity analysis ships them FP8; unsloth ships BF16). Applies
        // to ALL NVFP4 variants now, not just Fp8Dequanted. Kill-switch
        // ATLAS_NO_GDN_FP8_PREFILL restores the pre-policy behavior for A/B.
        let fp8_ssm_prefill = std::env::var_os("ATLAS_NO_GDN_FP8_PREFILL").is_none();
        let bf16_to_fp8_k = if fp8_ssm_prefill {
            tracing::info!(
                // The old wording — "NVFP4 kept as structural fallback for
                // decode batch paths" — was DISPROVEN by nsys at C=8
                // (2026-08-17): the batched verify QKVZ was taking this FP8
                // copy at every M>8, 26.3% of the step. It is true again only
                // because that arm now reads NVFP4; say the actual rule so the
                // log cannot re-drift from the dispatch.
                "SSM in_proj_qkv + out_proj via native FP8 prefill GEMM \
                 (BF16 act × FP8 weight via fp8_gemm_n128). PREFILL ONLY: \
                 decode + batched verify read the NVFP4 copy — weight-streaming \
                 GEMV at M<=8, tile GEMM above it (trait_decode_batched.rs). \
                 The FP8 copy reaches decode only at M<=8 on a build whose \
                 batched NVFP4 GEMVs are absent"
            );
            Some(gpu.kernel("w4a16", "bf16_to_fp8")?)
        } else {
            None
        };

        // ATLAS_MEM_PROFILE: per-phase GPU-free trace to pin the strix/APU
        // load-time footprint (FP8-source persistence vs NVFP4 steady-state vs
        // BF16 requant transients). Gated env so it's a no-op in production.
        let mem_profile = std::env::var("ATLAS_MEM_PROFILE").is_ok();
        let log_free = |tag: &str| {
            if mem_profile && let Ok(free) = gpu.free_memory() {
                tracing::info!("MEM_PROFILE[{tag}]: {:.2} GB GPU-free", free as f64 / 1e9);
            }
        };
        log_free("dense-load-start");

        // #915: resolved ONCE, from the same resolvers the forward pass uses —
        // the loader runs before `ForwardContext` exists. Contract and WHY:
        // `fp8_residency.rs`.
        let route_env = RouteEnv::from_env();
        let mut residency = DerivedResidency::default();

        for (i, lt) in layer_types.iter().enumerate() {
            if i % 8 == 0 {
                log_free(&format!("layer-{i}"));
            }
            let lp = config.layer_prefix(i);
            let input_norm = dense(store, &format!("{lp}.input_layernorm.weight"))?;
            let post_attn_norm = dense(store, &format!("{lp}.post_attention_layernorm.weight"))?;

            // Dense FFN instead of MoE. Native FP8 checkpoints (single-GPU)
            // load gate/up/down directly as block-scaled `Fp8Weight` and
            // dispatch w8a16 — no NVFP4 requant. TP>1 still uses the NVFP4
            // path (FP8 FFN sharding is a follow-up).
            let ffn_fp8 = ffn_fp8_arm_selected(store, config, variant, &lp);
            // 2026-09-11 (#915), replacing "always load the NVFP4 weights so
            // every dispatch path has a valid weight to fall back to": there is
            // no such path any more. `forward_k2`/`k3`/`km` redirect to
            // `forward_prefill` whenever an FP8 overlay is installed
            // (`native_small_batch_uses_prefill`), and both `forward` and
            // `forward_prefill_inner` return from inside their `fp8_weights`
            // arm. Building the fallback cost 18.4 GiB on Qwen3.8-27B-FP8 that
            // nothing could read. The CUDA-700-at-concurrency hazard the old
            // rule guarded is real and is why this is decided by
            // `fp8_residency.rs` rather than inline: a NULL NVFP4 weight is
            // only sound when the plan proves no w4a16 dispatch is reachable.
            // Native-BF16 dense-FFN prefill (Holo/Ornith Bf16Raw): the fast
            // tensor-core `dense_gemm_tc` path (dense_ffn::forward_prefill) needs
            // LIVE BF16 gate/up/down. `load_dense_ffn`'s Bf16Raw arm runtime-
            // quantizes each proj via `quantized_any`, whose Bf16Raw branch
            // FREES the store's cached BF16 buffer (`gpu.free(w.ptr)` in
            // nvfp4_detect.rs:288-309, the Bf16Raw arm). `dense_auto` returns that
            // SAME (now-freed) store ptr for a BF16 tensor — quant_helpers.rs:290
            // hands back `w.ptr` uncopied — so overlaying it AFTER load_dense_ffn hands
            // `set_bf16_weights` freed GPU memory -> dense_gemm_tc CUDA-700 at
            // grid=[div_ceil(intermediate_size,64),..] on the first prefill.
            // Snapshot fresh D2D copies BEFORE the free (the clone is an
            // independent allocation the layer owns for its lifetime).
            // Native keep-packed ternary Q2_0 (ATLAS_GGUF_NATIVE_Q2=1): when the
            // GGUF loader tagged gate/up/down as `PackedQ2_0`, DON'T requant to
            // NVFP4 — install the raw 2-bit blocks and dispatch `q2_0_gemv` at
            // decode. Requires tp_size=1 (packed-block sharding is unimplemented).
            // Flag off → tensors are BF16, `ffn_q2` is false, path unchanged.
            let ffn_q2 = config.tp_world_size.max(1) == 1
                && proj_q2_group(store, &format!("{lp}.mlp.gate_proj")).is_some()
                && proj_q2_group(store, &format!("{lp}.mlp.up_proj")).is_some()
                && proj_q2_group(store, &format!("{lp}.mlp.down_proj")).is_some();
            // Keep-packed projections are 2-bit blocks, not BF16 — there is
            // nothing to snapshot (dense_auto has no PackedQ2_0 arm and would
            // abort the load), and the ffn_q2 arm below owns their compute.
            let ffn_bf16_snapshot = if !ffn_q2 && matches!(variant, Nvfp4Variant::Bf16Raw) {
                let inter = if config.intermediate_size > 0 {
                    config.intermediate_size
                } else {
                    config.moe_intermediate_size
                };
                let clone_bf16 = |name: &str, rows: usize, cols: usize| -> Result<DenseWeight> {
                    let src = dense_auto(store, &format!("{lp}.mlp.{name}.weight"), gpu)?;
                    let bytes = rows * cols * 2; // BF16 = 2 bytes/elem
                    let dst = gpu.alloc(bytes)?;
                    // The whole point of this snapshot is an allocation
                    // INDEPENDENT of the store's cached ptr (which load_dense_ffn
                    // frees below); an aliased dst would silently re-create the
                    // freed-memory CUDA-700 this helper exists to prevent.
                    debug_assert_ne!(
                        dst, src.weight,
                        "ffn_bf16_snapshot must be a fresh allocation, not an alias of the store ptr"
                    );
                    gpu.copy_d2d(src.weight, dst, bytes)?;
                    Ok(DenseWeight { weight: dst })
                };
                Some((
                    clone_bf16("gate_proj", inter, h)?,
                    clone_bf16("up_proj", inter, h)?,
                    clone_bf16("down_proj", h, inter)?,
                ))
            } else {
                None
            };

            // #915: the NVFP4 gate/up/down fallback and its transposed twins
            // are 18.4 GiB on Qwen3.8-27B-FP8 that NO native-FP8 dispatch can
            // reach — every `DenseFfnLayer` entry point returns from inside its
            // `self.fp8_weights` arm and `w8_gemm!` binds the transposed
            // operand to a literal `None`. The comment above ("always load the
            // NVFP4 weights ... the rare batched paths fall back to real
            // NVFP4") described the pre-#927 dispatch; `forward_k2`/`k3`/`km`
            // have redirected to `forward_prefill` since
            // `native_small_batch_uses_prefill` landed. Decision table and the
            // loader-runs-before-dispatch contract: `fp8_residency.rs`.
            let plan = route_env.plan(ffn_fp8, false, false);
            let ffn_nvfp4 = !ffn_q2 && plan.ffn_nvfp4;
            let ffn_weights = if ffn_nvfp4 {
                load_dense_ffn(
                    store, &lp, gpu, variant, absmax_k, quantize_k, stream, config,
                )?
            } else {
                // NULL NVFP4 fallback. Packed-Q2 (Tier-1c): decode uses the
                // packed weights; prefill / batched paths bail (Tier-2).
                // Native FP8: every arm reads the E4M3 bytes installed by
                // `set_fp8_weights` below. No NVFP4 allocation → memory win.
                if ffn_fp8 {
                    residency.skip(fp8_residency::dense_ffn_nvfp4_bytes(h, ffn_inter(config)));
                }
                use crate::weight_map::QuantizedWeight;
                crate::layers::dense_ffn::DenseFfnWeights {
                    gate_proj: QuantizedWeight::null(),
                    up_proj: QuantizedWeight::null(),
                    down_proj: QuantizedWeight::null(),
                    gate_proj_t: None,
                    up_proj_t: None,
                    down_proj_t: None,
                }
            };
            residency.twins.ffn_nvfp4 |= ffn_nvfp4;
            let mut dffn = DenseFfnLayer::new(ffn_weights, gpu)?;
            if ffn_q2 {
                dffn.set_q2_weights(
                    packed_q2_from_store(store, &format!("{lp}.mlp.gate_proj"))?,
                    packed_q2_from_store(store, &format!("{lp}.mlp.up_proj"))?,
                    packed_q2_from_store(store, &format!("{lp}.mlp.down_proj"))?,
                    gpu,
                );
            }
            if ffn_fp8 {
                // The FP8 `.weight` bytes are store-owned (zero-copy from the
                // checkpoint); only the widened FP32 block-scale grid is a
                // fresh allocation (`loaders_fp8.rs:109`). Adopt it — #736
                // lists the loader sites that allocate without an owner.
                let load_ffn_fp8 = |name: &str| -> Result<Fp8Weight> {
                    let w = load_fp8_block_scaled_as_fp8weight(
                        store,
                        &format!("{lp}.mlp.{name}"),
                        gpu,
                    )?;
                    let bytes = (w.n as usize).div_ceil(128) * (w.k as usize).div_ceil(128) * 4;
                    store
                        .derived()
                        .adopt("ffn fp8 block scale (widened)", w.row_scale, bytes);
                    Ok(w)
                };
                let mut gate = load_ffn_fp8("gate_proj")?;
                let mut up = load_ffn_fp8("up_proj")?;
                let down = load_ffn_fp8("down_proj")?;
                // FUSED gate+up (#927). ONE `[2*inter, hidden]` E4M3 buffer
                // and ONE `[2*inter/128, hidden/128]` FP32 scale grid, with
                // `gate` and `up` re-pointed at VIEWS inside them — so the
                // fused decode GEMM and every un-fused rung of `w8_gemm!` read
                // the SAME bytes and neither needs a second copy.
                //
                // RESIDENCY IS NET ZERO, and that is the whole design
                // constraint: a second copy is 178.3 MB x 64 layers = 11.4 GB,
                // straight out of the bs32 KV budget. `prune_after_load`
                // releases the two source store tensors this copied
                // (`{lp}.mlp.{gate,up}_proj.weight` + their scale tensors),
                // exactly as the SSM `[QKV|Z]` concat above has done since
                // #915, and `predicted_residency` prices the pair at zero for
                // the preflight ring fit.
                let inter = ffn_inter(config);
                let fused = if ffn_gateup_fused_selected(store, config, variant, &lp) {
                    let fused = concat_fp8_block_scaled(&gate, &up, h, gpu)?;
                    let (w_bytes, s_bytes) = fp8_residency::ffn_gateup_fused_parts(h, inter);
                    // The concat COPIED both widened grids, so the two
                    // per-projection allocations are dead. They were adopted a
                    // few lines up; disown before freeing, or teardown frees
                    // them a second time.
                    let d = store.derived();
                    let grid = inter.div_ceil(128) * h.div_ceil(128) * 4;
                    for ptr in [gate.row_scale, up.row_scale] {
                        d.disown(ptr);
                        gpu.free(ptr)?;
                    }
                    residency.free(2 * grid);
                    // The views. `gate` is the fused buffer's head and `up`
                    // starts one `[inter, hidden]` block in; the scale grids
                    // meet at the same boundary because `inter % 128 == 0` is
                    // a clause of the selector above.
                    gate.weight = fused.weight;
                    gate.row_scale = fused.row_scale;
                    up.weight = fused.weight.offset(inter * h);
                    up.row_scale = fused.row_scale.offset(grid);
                    d.adopt("ffn gate+up fp8 concat", fused.weight, w_bytes);
                    d.adopt("ffn gate+up fp8 block scale", fused.row_scale, s_bytes);
                    residency.keep(w_bytes + s_bytes);
                    residency.twins.ffn_gateup_fused = true;
                    Some(fused)
                } else {
                    None
                };
                dffn.set_fp8_weights(gate, up, down);
                if let Some(fused) = fused {
                    dffn.set_fp8_gate_up_fused(fused);
                }
            }
            // ATLAS_FFN_MMQ: eagerly materialize Q4_K + free the dead `_t` copies at load,
            // BEFORE KV cache sizing, so net FFN footprint == NVFP4 baseline (no decode OOM-throttle).
            dffn.finalize_q4k_load(gpu, h as u32, config.intermediate_size as u32, stream)?;
            // ATLAS_FFN_NVFP4_MMQ: same discipline for the W4A4 FP4-MMQ arm — repack
            // gate/up to block_nvfp4 + free their `_t` copies (net ~0 footprint).
            //
            // SKIPPED when a LoRA adapter is pending. The forward-time FP4-MMQ
            // arm is disabled while an adapter is installed (it leaves gate/up
            // UNSCALED and folds weight_scale_2 inside the SiLU-mul, which
            // would silently scale a true-valued delta). But this finalize
            // FREES the transposed `_t` copies on the assumption that the MMQ
            // arm will serve prefill — so running it and then disabling the arm
            // left prefill on the slowest non-transposed GEMM with nothing to
            // fall back to: 176 tok/s against 841 on a 2K prompt, and it had
            // NOTHING to do with the cost of applying the deltas (measured with
            // the deltas skipped entirely).
            //
            // `adapter_max_rank` is the load-time signal that `--lora-adapter`
            // was given; the adapter itself is installed later (build step 8),
            // so this is the only point where the decision can be made before
            // the twins are freed.
            if config.adapter_max_rank == 0 {
                dffn.finalize_nvfp4_mmq_load(
                    gpu,
                    h as u32,
                    config.intermediate_size as u32,
                    stream,
                )?;
            }
            // Native-BF16 dense-FFN overlay (Bf16Raw, no-metadata Holo dense): install the
            // live BF16 gate/up/down snapshot so forward/forward_prefill's bf16 branch
            // (preferred over the NVFP4 fallback) reads valid memory. The NVFP4 weights built
            // by load_dense_ffn stay as the spec-decode/batched fallback (never null -> no
            // CUDA-700 at concurrency). Snapshot was taken before load_dense_ffn freed the
            // store's BF16 buffer (see ffn_bf16_snapshot above).
            if let Some((g, u, d)) = ffn_bf16_snapshot {
                dffn.set_bf16_weights(g, u, d);
            }
            let ffn = FfnComponent::Dense(dffn);

            match lt {
                LayerType::FullAttention => {
                    let p = format!("{lp}.self_attn");
                    let tp_rank = config.tp_rank;
                    let tp_size = config.tp_world_size.max(1);
                    // Bf16Raw installs a BF16 dense O-proj after the layer is
                    // built; other variants leave this None (NVFP4/FP8 dispatch).
                    let mut o_dense_bf16: Option<DenseWeight> = None;

                    // Native keep-packed ternary Q2_0 (Tier-1c): when the GGUF
                    // loader tagged q/k/v/o as `PackedQ2_0` (transform-free
                    // full-attention projections), install the raw 2-bit blocks
                    // and dispatch `q2_0_gemv_vec` at decode / transient-dequant
                    // at prefill — no NVFP4 requant, no `_t` copies. Requires
                    // tp_size=1 (packed-block sharding is unimplemented). Bonsai
                    // has no kill-switch here; the whole path is gated upstream
                    // by ATLAS_GGUF_NATIVE_Q2 (else these tensors are BF16 and
                    // `attn_q2` is false). ATLAS_NO_Q2_ATTN forces the BF16 path
                    // for A/B bisection.
                    let attn_q2 = tp_size == 1
                        && std::env::var_os("ATLAS_NO_Q2_ATTN").is_none()
                        && proj_q2_group(store, &format!("{p}.q_proj")).is_some()
                        && proj_q2_group(store, &format!("{p}.k_proj")).is_some()
                        && proj_q2_group(store, &format!("{p}.v_proj")).is_some()
                        && proj_q2_group(store, &format!("{p}.o_proj")).is_some();
                    if attn_q2 {
                        let (k_scale, v_scale) = load_kv_scales(store, &p, gpu);
                        let attn = AttentionWeights {
                            q_proj: DenseWeight {
                                weight: spark_runtime::gpu::DevicePtr::NULL,
                            },
                            k_proj: DenseWeight {
                                weight: spark_runtime::gpu::DevicePtr::NULL,
                            },
                            v_proj: DenseWeight {
                                weight: spark_runtime::gpu::DevicePtr::NULL,
                            },
                            o_proj: crate::weight_map::QuantizedWeight::null(),
                            q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
                            k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
                            q_norm_full: None,
                            k_norm_full: None,
                            k_scale,
                            v_scale,
                        };
                        let mut attn_layer = Qwen3AttentionLayer::new(
                            input_norm,
                            attn,
                            post_attn_norm,
                            ffn,
                            attn_idx,
                            None,
                            None,
                            None,
                            gpu,
                            layer_kv_dtypes[attn_idx],
                            config.fp8_kv_calibration_tokens,
                            config,
                        )?;
                        attn_layer.set_packed_q2_weights(
                            packed_q2_from_store(store, &format!("{p}.q_proj"))?,
                            packed_q2_from_store(store, &format!("{p}.k_proj"))?,
                            packed_q2_from_store(store, &format!("{p}.v_proj"))?,
                            packed_q2_from_store(store, &format!("{p}.o_proj"))?,
                            gpu,
                        );
                        tracing::info!(
                            "ATTN[{lp}] native keep-packed Q2_0: q/k/v/o 2-bit \
                             (q2_0_gemv_vec decode; transient-dequant prefill)"
                        );
                        layers.push(Box::new(attn_layer));
                        attn_idx += 1;
                        if (i + 1) % 10 == 0 {
                            tracing::info!("Loaded layers 0..{}", i + 1);
                        }
                        continue;
                    }
                    // #915: does the native FP8 attention overlay replace
                    // these weights below? Hoisted out of the overlay block so
                    // the NVFP4 build can be skipped rather than built and
                    // immediately orphaned — `set_fp8_weights` OVERWRITES
                    // `q/k/v/o_weight`, so on this route the NVFP4 q/k/v/o were
                    // unreachable the moment they were installed.
                    let attn_fp8 = dense_fp8_enabled()
                        && config.tp_world_size.max(1) == 1
                        && matches!(variant, Nvfp4Variant::Fp8Dequanted)
                        && proj_is_native_fp8(store, &format!("{p}.q_proj"));
                    let attn_nvfp4 = route_env.attn_nvfp4(attn_fp8);
                    let (attn, q_nvfp4, k_nvfp4, v_nvfp4) = match variant {
                        Nvfp4Variant::CompressedTensors => {
                            // NVFP4-from-disk path: column-parallel Q/K/V, row-parallel O.
                            let group_size = 16usize;
                            let load_nvfp4 = |name: &str,
                                              full_n: usize,
                                              full_k: usize,
                                              kind: TpShardKind|
                             -> Result<crate::weight_map::QuantizedWeight> {
                                let prefix = format!("{p}.{name}");
                                // Mixed-precision compressed-tensors checkpoints
                                // (unsloth Qwen3.6-*-NVFP4, re-quantized 2026-07-10)
                                // NVFP4-pack most of the net but keep attention
                                // q/k/v/o as FP8 (`.weight` FP8E4M3 + a per-row
                                // `.weight_scale`, no `.weight_packed`). Without the
                                // pack metadata, dequant and runtime-quantize to NVFP4
                                // instead of failing on the absent
                                // `weight_global_scale`. Mirrors the MoE loader's
                                // attention arm (weight_loader/qwen35/load_layers/
                                // attention_arms.rs).
                                let src = if store.contains(&format!("{prefix}.weight_packed")) {
                                    quantized_auto(store, &prefix, gpu, variant)?
                                } else {
                                    let dense_bf16 =
                                        dense_auto(store, &format!("{prefix}.weight"), gpu)?;
                                    quantize_to_nvfp4(
                                        &dense_bf16,
                                        full_n,
                                        full_k,
                                        gpu,
                                        absmax_k,
                                        quantize_k,
                                        stream,
                                    )?
                                };
                                if tp_size == 1 {
                                    return Ok(src);
                                }
                                let sharded = shard_quantized_nvfp4(
                                    &src, full_n, full_k, kind, tp_rank, tp_size, group_size, gpu,
                                )?;
                                gpu.free(src.weight)?;
                                gpu.free(src.weight_scale)?;
                                Ok(sharded)
                            };
                            let [q, k, v, o] = load_qkvo_tp(config, load_nvfp4)?;
                            let dummy = DenseWeight {
                                weight: spark_runtime::gpu::DevicePtr::NULL,
                            };
                            let (k_scale, v_scale) = load_kv_scales(store, &p, gpu);
                            let attn = AttentionWeights {
                                q_proj: dummy,
                                k_proj: dummy,
                                v_proj: dummy,
                                o_proj: o,
                                q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
                                k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
                                q_norm_full: None,
                                k_norm_full: None,
                                k_scale,
                                v_scale,
                            };
                            (attn, Some(q), Some(k), Some(v))
                        }
                        Nvfp4Variant::Standard | Nvfp4Variant::Fp8Dequanted if !attn_nvfp4 => {
                            // #915: the FP8 overlay below replaces q/k/v/o, and
                            // nothing reads the NVFP4 copies afterwards — the
                            // transposed twins only under `ATLAS_CUTLASS_NVFP4_*`
                            // and the base `o_proj` only under `ATLAS_ATTN_W4A4`,
                            // both of which `RouteEnv::attn_nvfp4` accounts for.
                            // Skipping the build also skips the BF16 dequant of
                            // every FP8 projection that fed `quantize_to_nvfp4`.
                            // Same NULL-o_proj shape the Bf16Raw arm below uses.
                            let (k_scale, v_scale) = load_kv_scales(store, &p, gpu);
                            let (nh, hd) = (config.num_attention_heads, config.head_dim);
                            let (nkv, hh) = (config.num_key_value_heads, config.hidden_size);
                            let q_n = nh * hd * if config.attn_gated { 2 } else { 1 };
                            residency.skip(fp8_residency::attn_nvfp4_bytes(
                                q_n,
                                nkv * hd,
                                nh * hd,
                                hh,
                            ));
                            let null = || DenseWeight {
                                weight: spark_runtime::gpu::DevicePtr::NULL,
                            };
                            let attn = AttentionWeights {
                                q_proj: null(),
                                k_proj: null(),
                                v_proj: null(),
                                o_proj: crate::weight_map::QuantizedWeight::null(),
                                q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
                                k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
                                q_norm_full: None,
                                k_norm_full: None,
                                k_scale,
                                v_scale,
                            };
                            // The NULL o_proj is only sound while the FP8
                            // overlay is installed below — a null NVFP4 weight
                            // reached by a w4a16 dispatch is the
                            // CUDA-700-at-concurrency failure mode.
                            debug_assert!(
                                attn_fp8,
                                "attn_nvfp4=false must imply a native FP8 overlay"
                            );
                            (attn, None, None, None)
                        }
                        Nvfp4Variant::Standard | Nvfp4Variant::Fp8Dequanted => {
                            // BF16 → NVFP4 path: shard BF16 then quantize per-rank.
                            let load_bf16_then_nvfp4 = |name: &str,
                                                        full_n: usize,
                                                        full_k: usize,
                                                        kind: TpShardKind|
                             -> Result<(
                                DenseWeight,
                                crate::weight_map::QuantizedWeight,
                            )> {
                                // Pre-quantized Standard NVFP4 (e.g. sakamakismile): weight is U8
                                // on disk. Load directly as QuantizedWeight without BF16 roundtrip.
                                // TP sharding of pre-quantized NVFP4 is not yet supported: enforce
                                // tp_size=1 explicitly, or every rank silently loads the full
                                // unsharded weight (duplicated, not sharded — wrong results with
                                // no error).
                                let weight_key = format!("{p}.{name}.weight");
                                if matches!(
                                    store.get(&weight_key).map(|w| w.dtype),
                                    Ok(WeightDtype::UInt8)
                                ) {
                                    anyhow::ensure!(
                                        tp_size == 1,
                                        "pre-quantized NVFP4 weight '{weight_key}' (U8 on disk) \
                                         cannot be loaded under tensor parallelism (tp_size={tp_size}): \
                                         TP sharding of pre-quantized NVFP4 checkpoints is not yet \
                                         implemented. Use tp_size=1, or dequantize this checkpoint to \
                                         BF16 first so it goes through the shard-then-requantize path."
                                    );
                                    let null_dense = DenseWeight {
                                        weight: spark_runtime::gpu::DevicePtr::NULL,
                                    };
                                    let qw = quantized_auto(
                                        store,
                                        &format!("{p}.{name}"),
                                        gpu,
                                        Nvfp4Variant::Standard,
                                    )?;
                                    return Ok((null_dense, qw));
                                }
                                let src = dense_auto(store, &weight_key, gpu)?;
                                let (sharded_ptr, local_n, local_k) = shard_dense_bf16(
                                    src.weight, full_n, full_k, kind, tp_rank, tp_size, gpu,
                                )?;
                                let sharded = DenseWeight {
                                    weight: sharded_ptr,
                                };
                                let q = quantize_to_nvfp4(
                                    &sharded, local_n, local_k, gpu, absmax_k, quantize_k, stream,
                                )?;
                                if sharded_ptr != src.weight {
                                    gpu.free(sharded_ptr)?;
                                }
                                Ok((src, q))
                            };
                            let [
                                (q_dense, q_nvfp4),
                                (k_dense, k_nvfp4),
                                (v_dense, v_nvfp4),
                                (o_dense, o_nvfp4),
                            ] = load_qkvo_tp(config, load_bf16_then_nvfp4)?;

                            let (k_scale, v_scale) = load_kv_scales(store, &p, gpu);

                            // The BF16 q/k/v/o dense tensors are only the intermediate
                            // fed to the GPU quantize_to_nvfp4 above. Prefill AND decode
                            // always dispatch the NVFP4 weights, so the BF16 copies are
                            // dead once quantized. Free them instead of retaining a full
                            // second copy of every projection (Atlas issue #A1).
                            gpu.free(q_dense.weight)?;
                            gpu.free(k_dense.weight)?;
                            gpu.free(v_dense.weight)?;
                            gpu.free(o_dense.weight)?;

                            let attn = AttentionWeights {
                                q_proj: DenseWeight {
                                    weight: spark_runtime::gpu::DevicePtr::NULL,
                                },
                                k_proj: DenseWeight {
                                    weight: spark_runtime::gpu::DevicePtr::NULL,
                                },
                                v_proj: DenseWeight {
                                    weight: spark_runtime::gpu::DevicePtr::NULL,
                                },
                                o_proj: o_nvfp4,
                                q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
                                k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
                                q_norm_full: None,
                                k_norm_full: None,
                                k_scale,
                                v_scale,
                            };
                            (attn, Some(q_nvfp4), Some(k_nvfp4), Some(v_nvfp4))
                        }
                        Nvfp4Variant::Bf16Raw => {
                            // Native BF16 dense attention: keep Q/K/V/O in BF16
                            // and dispatch the dense_gemv/dense_gemm kernels that
                            // ship in the nvfp4 bundle (common/ dense_*_bf16). No
                            // runtime BF16 -> NVFP4 quant — that lossily quantized
                            // these no-metadata Holo dense checkpoints. Mirrors
                            // qwen35/load_layers.rs:552 (BF16-dequant attention)
                            // and gemma4/loader_a.rs:300/418.
                            let load_bf16_dense =
                                |name: &str,
                                 full_n: usize,
                                 full_k: usize,
                                 kind: TpShardKind|
                                 -> Result<DenseWeight> {
                                    let src =
                                        dense_auto(store, &format!("{p}.{name}.weight"), gpu)?;
                                    if tp_size == 1 {
                                        return Ok(src);
                                    }
                                    let (sharded_ptr, _local_n, _local_k) = shard_dense_bf16(
                                        src.weight, full_n, full_k, kind, tp_rank, tp_size, gpu,
                                    )?;
                                    if sharded_ptr != src.weight {
                                        gpu.free(src.weight)?;
                                    }
                                    Ok(DenseWeight {
                                        weight: sharded_ptr,
                                    })
                                };
                            let [q_dense, k_dense, v_dense, o_dense] =
                                load_qkvo_tp(config, load_bf16_dense)?;

                            let (k_scale, v_scale) = load_kv_scales(store, &p, gpu);

                            // Keep q/k/v BF16 dense ALIVE (do NOT free): the dense
                            // forward reads them directly. o_proj stays NULL — the
                            // set_o_dense_bf16(o_dense) call after layer build
                            // installs the BF16 O-proj that decode/prefill prefer.
                            let attn = AttentionWeights {
                                q_proj: q_dense,
                                k_proj: k_dense,
                                v_proj: v_dense,
                                o_proj: crate::weight_map::QuantizedWeight::null(),
                                q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
                                k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
                                q_norm_full: None,
                                k_norm_full: None,
                                k_scale,
                                v_scale,
                            };
                            o_dense_bf16 = Some(o_dense);
                            // The NULL o_proj above is only sound while the BF16
                            // dense O-proj is installed after layer build — a null
                            // NVFP4 weight reached by a w4a16 dispatch is the
                            // CUDA-700-at-concurrency failure mode.
                            debug_assert!(
                                o_dense_bf16.is_some(),
                                "Bf16Raw attention must install a dense O-proj to cover the null o_proj"
                            );
                            // BF16 native: no NVFP4 q/k/v weights → dense fallback.
                            (attn, None, None, None)
                        }
                    };

                    let mut attn_layer = Qwen3AttentionLayer::new(
                        input_norm,
                        attn,
                        post_attn_norm,
                        ffn,
                        attn_idx,
                        q_nvfp4,
                        k_nvfp4,
                        v_nvfp4,
                        gpu,
                        layer_kv_dtypes[attn_idx],
                        config.fp8_kv_calibration_tokens,
                        config,
                    )?;
                    // Fast-prefill: transposed NVFP4 copies route the 16 full-attn
                    // layers' q/k/v/o prefill GEMMs onto w4a16_gemm_t_m128 (28.8%
                    // of prefill GPU time on the base w4a16_gemm path; ~1.3x e2e).
                    // predequant_for_prefill() is deliberately NOT called: the FP8
                    // predequant route is slower for these bandwidth-bound GEMMs.
                    if let (Some(qw), Some(kw), Some(vw)) = (q_nvfp4, k_nvfp4, v_nvfp4) {
                        let (nh, hd) = (config.num_attention_heads, config.head_dim);
                        let (nkv, hh) = (config.num_key_value_heads, config.hidden_size);
                        let q_n = nh * hd * if config.attn_gated { 2 } else { 1 };
                        let qt = qw.transpose_for_gemm(gpu, q_n, hh)?;
                        let kt = kw.transpose_for_gemm(gpu, nkv * hd, hh)?;
                        let vt = vw.transpose_for_gemm(gpu, nkv * hd, hh)?;
                        let op = &attn_layer.attn.o_proj;
                        let ot = op.transpose_for_gemm(gpu, hh, nh * hd)?;
                        attn_layer.set_prefill_weights(Some(qt), Some(kt), Some(vt), Some(ot));
                        // Fused [q|k|v] twin: k/v are N=1024, which against the
                        // 128-wide N tile is 8 CTAs on 48 SMs (40 idle, 23.6 GB/s,
                        // 9.75x off floor). Concatenated N=14336 runs 112 CTAs in
                        // ONE launch. Bit-identical ONLY when the three share a
                        // single `weight_scale_2` — the GEMM applies one scale2 per
                        // launch — so verify the device values rather than assume.
                        // Bit-exact float comparison, not an epsilon: the GEMM
                        // applies ONE scale2, so anything but exact equality
                        // changes results.
                        let scales_equal = qw.weight_scale_2.to_bits()
                            == kw.weight_scale_2.to_bits()
                            && kw.weight_scale_2.to_bits() == vw.weight_scale_2.to_bits();
                        if scales_equal {
                            let fused =
                                crate::weight_map::QuantizedWeight::transpose_concat_for_gemm(
                                    gpu,
                                    &[(&qw, q_n), (&kw, nkv * hd), (&vw, nkv * hd)],
                                    hh,
                                )?;
                            attn_layer.set_fused_qkv_prefill_weight(Some(fused));
                        } else if attn_idx == 0 {
                            tracing::warn!(
                                "attention q/k/v have differing weight_scale_2 — fused QKV GEMM disabled (3 separate launches per layer)"
                            );
                        }
                    }
                    // Native-BF16 (Bf16Raw): install the dense O-proj so decode +
                    // prefill prefer it over the (NULL) NVFP4 o_proj. Mutually
                    // exclusive with the transposed-NVFP4 block above (q_nvfp4 is
                    // None on the Bf16Raw path).
                    if let Some(o_dense) = o_dense_bf16 {
                        attn_layer.set_o_dense_bf16(o_dense);
                    }
                    // Install native FP8 q/k/v/o (single-GPU FP8 checkpoint).
                    // `set_fp8_weights` REPLACES `q/k/v/o_weight`, so on this
                    // route the NVFP4 weights above are not a fallback — they
                    // are unreachable, which is why `attn_nvfp4` skipped
                    // building them (#915). The paths that could still read
                    // NVFP4 (`ATLAS_CUTLASS_NVFP4_*`, `ATLAS_ATTN_W4A4`) are
                    // exactly what `RouteEnv::attn_nvfp4` checks.
                    if attn_fp8 {
                        let load_fp8_proj = |name: &str,
                                             _n: usize,
                                             _k: usize,
                                             _kind: TpShardKind|
                         -> Result<Fp8Weight> {
                            // Only `row_scale` is allocated here; the E4M3
                            // bytes are store-owned. Adopt it (#736).
                            let w = load_fp8_block_scaled_as_fp8weight(
                                store,
                                &format!("{p}.{name}"),
                                gpu,
                            )?;
                            let bytes =
                                (w.n as usize).div_ceil(128) * (w.k as usize).div_ceil(128) * 4;
                            store.derived().adopt(
                                "attn fp8 block scale (widened)",
                                w.row_scale,
                                bytes,
                            );
                            Ok(w)
                        };
                        let [q_fp8, k_fp8, v_fp8, o_fp8] = load_qkvo_tp(config, load_fp8_proj)?;
                        attn_layer.set_fp8_weights(
                            Some(q_fp8),
                            Some(k_fp8),
                            Some(v_fp8),
                            Some(o_fp8),
                        );
                        // #915: build only the twins a selected kernel reads,
                        // and hand each to the store so teardown RELEASES it.
                        // K/V are NOT optional — `prefill/cache_skip_qkv.rs` has
                        // no W8A8 arm and is the first-chunk path of every
                        // request. Table: `Fp8TwinSet` + `fp8_residency.rs`.
                        let want =
                            route_env.attn_fp8_twins(true, attn_layer.has_w8a8_prefill_kernels());
                        let (nh, hd) = (config.num_attention_heads, config.head_dim);
                        let (nkv, hh) = (config.num_key_value_heads, config.hidden_size);
                        let q_n = nh * hd * if config.attn_gated { 2 } else { 1 };
                        let twin_bytes = |set| {
                            fp8_residency::attn_fp8_twin_bytes(set, q_n, nkv * hd, nh * hd, hh)
                        };
                        residency.twins.attn_fp8 |= want.any();
                        residency.skip(twin_bytes(Fp8TwinSet::ALL) - twin_bytes(want));
                        residency.keep(twin_bytes(want));
                        if let Err(e) = attn_layer.transpose_fp8_for_prefill_selected(
                            gpu,
                            stream,
                            want,
                            Some(store.derived()),
                        ) {
                            tracing::warn!("Layer {i}: dense FP8 transpose failed: {e}");
                        }
                    }
                    layers.push(Box::new(attn_layer));
                    attn_idx += 1;
                }
                LayerType::LinearAttention => {
                    let nv = config.linear_num_value_heads;
                    let nk = config.linear_num_key_heads;
                    // GDN HeadParallel: config holds per-rank-LOCAL linear head counts
                    // (topology.rs divided them by tp_size). TpGdnDims rebuilds the FULL
                    // pre-shard sizes so load/concat/interleave run at FULL, then the
                    // shard_gdn_* slicers cut this rank's contiguous head range. value_dim
                    // stays LOCAL (sizes the downstream quantize of the sharded buffers).
                    let tp_size = config.tp_world_size.max(1);
                    let dims = TpGdnDims::from_config(config);
                    let qkv_rows = dims.full_conv_dim();
                    let z_rows = dims.full_value_dim();
                    let value_dim = nv * config.linear_value_head_dim;
                    let la = format!("{lp}.linear_attn");

                    // Native keep-packed ternary Q2_0 GDN (Tier-1c): the GGUF
                    // loader kept `in_proj_qkv` (V-region row-permuted) and
                    // `in_proj_z` (row-permuted) 2-bit. Byte-concat them into the
                    // fused [Q|K|V|Z] `qkvz` and dispatch `q2_0_gemv_vec` at decode
                    // / transient-dequant at prefill. `out_proj` (a within-row
                    // COLUMN reorder) is NOT packed here — it stays NVFP4. a/b/
                    // conv1d/norm/A_log stay BF16/F32. Requires tp_size=1.
                    // ATLAS_NO_Q2_GDN forces the BF16/NVFP4 path for A/B bisection.
                    let gdn_q2 = config.tp_world_size.max(1) == 1
                        && std::env::var_os("ATLAS_NO_Q2_GDN").is_none()
                        && proj_q2_group(store, &format!("{la}.in_proj_qkv")).is_some()
                        && proj_q2_group(store, &format!("{la}.in_proj_z")).is_some();
                    if gdn_q2 {
                        let qkv_q2 = packed_q2_from_store(store, &format!("{la}.in_proj_qkv"))?;
                        let z_q2 = packed_q2_from_store(store, &format!("{la}.in_proj_z"))?;
                        anyhow::ensure!(
                            qkv_q2.group == z_q2.group && qkv_q2.k == z_q2.k,
                            "GDN packed qkv/z group|k mismatch ({},{} vs {},{})",
                            qkv_q2.group,
                            qkv_q2.k,
                            z_q2.group,
                            z_q2.k
                        );
                        // Byte-concat packed rows: [Q|K|V] ++ [Z]. Each row is
                        // (k/group)*block_bytes; whole-row copy never splits a block.
                        let group = qkv_q2.group as usize;
                        let block_bytes = 2 + group / 4;
                        let row_bytes = (qkv_q2.k as usize / group) * block_bytes;
                        let qkv_bytes = qkv_q2.n as usize * row_bytes;
                        let z_bytes = z_q2.n as usize * row_bytes;
                        let qkvz_buf = gpu.alloc(qkv_bytes + z_bytes)?;
                        gpu.copy_d2d(qkv_q2.weight, qkvz_buf, qkv_bytes)?;
                        gpu.copy_d2d(z_q2.weight, qkvz_buf.offset(qkv_bytes), z_bytes)?;
                        let qkvz_q2 = PackedQ2Weight {
                            weight: qkvz_buf,
                            n: qkv_q2.n + z_q2.n,
                            k: qkv_q2.k,
                            group: qkv_q2.group,
                        };
                        // out_proj + a/b/conv1d/norm are BF16/F32 in the store
                        // (sidecar dequanted the reorder tensors). out_proj → NVFP4.
                        let in_proj_a = dense_auto(store, &format!("{la}.in_proj_a.weight"), gpu)?;
                        let in_proj_b = dense_auto(store, &format!("{la}.in_proj_b.weight"), gpu)?;
                        let conv1d = dense(store, &format!("{la}.conv1d.weight"))?;
                        let a_log = dense_keep_f32(store, &format!("{la}.A_log"), gpu)?;
                        let dt_bias = dense_keep_f32(store, &format!("{la}.dt_bias"), gpu)?;
                        let norm = dense_f32_safe(store, &format!("{la}.norm.weight"), gpu)?;
                        let ba_dense = interleave_ba(&in_proj_a, &in_proj_b, nv, nk, h, gpu)?;
                        let out_proj_dense =
                            dense_auto(store, &format!("{la}.out_proj.weight"), gpu)?;
                        let out_proj_nvfp4 = quantize_to_nvfp4(
                            &out_proj_dense,
                            h,
                            value_dim,
                            gpu,
                            absmax_k,
                            quantize_k,
                            stream,
                        )?;
                        let out_proj_nvfp4_t =
                            out_proj_nvfp4.transpose_for_gemm(gpu, h, value_dim)?;
                        gpu.free(out_proj_dense.weight)?;
                        let ssm = SsmWeights {
                            in_proj_qkvz: DenseWeight {
                                weight: spark_runtime::gpu::DevicePtr::NULL,
                            },
                            in_proj_ba: ba_dense,
                            conv1d,
                            a_log,
                            dt_bias,
                            norm,
                            out_proj: out_proj_nvfp4,
                        };
                        let mut layer = Qwen3SsmLayer::new_sequential(
                            input_norm,
                            ssm,
                            post_attn_norm,
                            ffn,
                            None,
                            None,
                            Some(out_proj_nvfp4_t),
                            config,
                            gpu,
                        )?;
                        layer.set_packed_q2_qkvz(qkvz_q2, gpu);
                        layer.predequant_for_prefill(gpu, config, stream)?;
                        tracing::info!(
                            "SSM[{lp}] native keep-packed Q2_0 GDN: qkvz 2-bit \
                             (concat qkv+z row-permuted), out_proj NVFP4"
                        );
                        layers.push(Box::new(layer));
                        continue;
                    }

                    // SSM projections are loaded per-projection by on-disk dtype:
                    // each of in_proj_qkv / in_proj_z / out_proj may independently
                    // be NVFP4-packed (`weight_packed`) or plain (`weight`, routed
                    // by `dense_auto` → BF16/FP32/FP8). The unsloth NVFP4 re-quant
                    // of Qwen3.6-27B quantizes ONLY out_proj while keeping the
                    // in_proj_* in BF16; the old all-or-nothing gate (keyed on
                    // in_proj_qkv.weight_packed) then looked for a non-existent
                    // out_proj.weight and failed to build. `dense_auto` is dequant-
                    // to-BF16 for the concat pipeline regardless of source dtype.
                    let load_ssm_proj =
                        |name: &str, rows: usize, cols: usize| -> Result<DenseWeight> {
                            if store.contains(&format!("{name}.weight_packed")) {
                                dequant_nvfp4_to_bf16(store, name, rows, cols, gpu)
                            } else if matches!(
                                store.get(&format!("{name}.weight")).map(|w| w.dtype),
                                Ok(WeightDtype::UInt8)
                            ) {
                                // Standard-convention NVFP4 (packed bytes at
                                // `.weight`, not `.weight_packed`) — same dequant,
                                // different on-disk key.
                                dequant_nvfp4_to_bf16(store, name, rows, cols, gpu)
                            } else {
                                dense_auto(store, &format!("{name}.weight"), gpu)
                            }
                        };
                    // Native FP8 GDN (nvidia mixed-precision checkpoint): the
                    // in_proj_qkv / in_proj_z / out_proj projections ship as
                    // F8_E4M3 + per-tensor scale — modelopt's sensitivity
                    // analysis keeps the SSM projections high-precision. The
                    // default path (`load_ssm_proj` → `dense_auto`) dequants to
                    // BF16 then RE-quantizes to NVFP4 (4-bit), a lossy
                    // double-quant of these 48/64 layers that regressed BFCL-ST
                    // ~7pt (non_live 85.4→76.6). Load the on-disk FP8 directly
                    // (concat qkv+z on-device into [Q|K|V|Z] order) and route
                    // BOTH prefill (w8a16_gemm_pipelined) and decode
                    // (w8a16_gemv) through the fp8w fields — no requant, decode
                    // stays fast (FP8 = half BF16's weight bytes). MUST run
                    // BEFORE `load_ssm_proj` consumes the store tensors.
                    // Internal opt-out for the FP8-vs-NVFP4 GDN A/B + KL-drift
                    // gate (not a user choice; mirrors the `ATLAS_NO_*` debug
                    // levers). Default engages native FP8.
                    if gdn_fp8_arm_selected(store, &la, config.tp_world_size) {
                        let in_proj_a = dense(store, &format!("{la}.in_proj_a.weight"))?;
                        let in_proj_b = dense(store, &format!("{la}.in_proj_b.weight"))?;
                        let conv1d = dense(store, &format!("{la}.conv1d.weight"))?;
                        let a_log = dense_keep_f32(store, &format!("{la}.A_log"), gpu)?;
                        let dt_bias = dense_keep_f32(store, &format!("{la}.dt_bias"), gpu)?;
                        let norm = dense_f32_safe(store, &format!("{la}.norm.weight"), gpu)?;
                        let ba_dense = interleave_ba(&in_proj_a, &in_proj_b, nv, nk, h, gpu)?;
                        let qkv_f = load_fp8_block_scaled_as_fp8weight(
                            store,
                            &format!("{la}.in_proj_qkv"),
                            gpu,
                        )?;
                        let z_f = load_fp8_block_scaled_as_fp8weight(
                            store,
                            &format!("{la}.in_proj_z"),
                            gpu,
                        )?;
                        let out_f = load_fp8_block_scaled_as_fp8weight(
                            store,
                            &format!("{la}.out_proj"),
                            gpu,
                        )?;
                        let qkvz_f = concat_fp8_block_scaled(&qkv_f, &z_f, h, gpu)?;
                        // The concat copied both grids; free the per-projection
                        // scale allocs (weight bytes are store-owned, not freed).
                        let scale_bytes =
                            |n: usize, k: usize| n.div_ceil(128) * k.div_ceil(128) * 4;
                        gpu.free(qkv_f.row_scale)?;
                        gpu.free(z_f.row_scale)?;
                        residency.free(
                            scale_bytes(qkv_f.n as usize, h) + scale_bytes(z_f.n as usize, h),
                        );
                        // #736/#915: the fused `[QKV|Z]` weight is the SSM's
                        // hottest tensor and there is no un-fused dispatch to
                        // fall back to, so it STAYS — but it is a derived copy
                        // that no `ModelResource` owned, which is the
                        // 3,840 MB x48 row of the H100 teardown sweep. Adopt it.
                        let qkvz_bytes = (qkvz_f.n as usize) * h;
                        let qkvz_scale_bytes =
                            scale_bytes(qkv_f.n as usize, h) + scale_bytes(z_f.n as usize, h);
                        let ba_bytes = nv * 2 * h * 2; // [2*nv, h] BF16
                        let out_scale_bytes = scale_bytes(out_f.n as usize, out_f.k as usize);
                        let d = store.derived();
                        d.adopt("ssm qkvz fp8 concat", qkvz_f.weight, qkvz_bytes);
                        d.adopt(
                            "ssm qkvz fp8 block scale",
                            qkvz_f.row_scale,
                            qkvz_scale_bytes,
                        );
                        d.adopt(
                            "ssm out_proj fp8 block scale",
                            out_f.row_scale,
                            out_scale_bytes,
                        );
                        d.adopt("ssm in_proj_ba interleaved", ba_dense.weight, ba_bytes);
                        residency.keep(qkvz_bytes + qkvz_scale_bytes + out_scale_bytes + ba_bytes);
                        residency.twins.ssm_fp8_concat = true;
                        let ssm = SsmWeights {
                            in_proj_qkvz: DenseWeight {
                                weight: spark_runtime::gpu::DevicePtr::NULL,
                            },
                            in_proj_ba: ba_dense,
                            conv1d,
                            a_log,
                            dt_bias,
                            norm,
                            out_proj: crate::weight_map::QuantizedWeight::null(),
                        };
                        let mut layer = Qwen3SsmLayer::new_sequential(
                            input_norm,
                            ssm,
                            post_attn_norm,
                            ffn,
                            None,
                            None,
                            None,
                            config,
                            gpu,
                        )?;
                        layer.set_fp8_decode_weights(Some(qkvz_f), Some(out_f));
                        tracing::info!(
                            "SSM[{lp}] native FP8 GDN: qkvz+out_proj block-scaled FP8 \
                             (no NVFP4 requant; prefill+decode via w8a16)"
                        );
                        layers.push(Box::new(layer));
                        continue;
                    }

                    // A, B, conv1d, A_log, dt_bias, norm are independent of the
                    // qkv/z/out_proj on-disk format below — load them once up
                    // front so both the native-NVFP4 fast path and the legacy
                    // dequant/requant path can share them.
                    //
                    // A_log and dt_bias MUST be FP32 — consumer kernels in
                    // `ssm_preprocess.cu` and `mamba2_ssm_decode.cu` declare
                    // them `const float*`. Loading via `dense()` kept BF16
                    // storage, reinterpreting 48-elt BF16 (96B) as 48-elt
                    // FP32 → per-head scrambled decay gates and exponential
                    // error amplification through GDR recurrence at long
                    // context. The MoE sister loader (`ssm_qwen35.rs`)
                    // already promotes these; dense was missing the mirror.
                    //
                    // in_proj_a/b: route through `load_ssm_proj` (not the raw
                    // `dense()` byte-reinterpret) so a Standard-NVFP4 A/B
                    // (U8-packed) checkpoint dequants correctly instead of
                    // being read as BF16 garbage.
                    let in_proj_a = load_ssm_proj(&format!("{la}.in_proj_a"), nv, h)?;
                    let in_proj_b = load_ssm_proj(&format!("{la}.in_proj_b"), nv, h)?;
                    let conv1d = dense(store, &format!("{la}.conv1d.weight"))?;
                    let a_log = dense_keep_f32(store, &format!("{la}.A_log"), gpu)?;
                    let dt_bias = dense_keep_f32(store, &format!("{la}.dt_bias"), gpu)?;
                    // norm.weight: use `dense_f32_safe` (FP32-aware: detects
                    // a fp32 checkpoint and truncates to BF16 with logging;
                    // bf16 passes through). Mirrors `weight_map/ssm_qwen35.rs`
                    // MoE sister loader (backported here 2026-05-20).
                    let norm = dense_f32_safe(store, &format!("{la}.norm.weight"), gpu)?;
                    let ba_dense = interleave_ba(&in_proj_a, &in_proj_b, nv, nk, h, gpu)?;
                    let qkvz_size = config.ssm_qkvz_size();

                    // Native Standard-NVFP4 GDN (pre-quantized checkpoint, e.g.
                    // sakamakismile): in_proj_qkv / in_proj_z / out_proj ship
                    // U8-packed NVFP4 directly on disk (`.weight` dtype UInt8,
                    // not `.weight_packed` — that's the compressed-tensors
                    // convention `load_ssm_proj` already dequants above). Load
                    // them straight into `QuantizedWeight` and concat on GPU,
                    // skipping the BF16-dequant→re-quantize roundtrip entirely
                    // (that roundtrip is what the FP8/BF16-opt-in paths above
                    // exist to avoid for FP8-native and BF16-preferring
                    // checkpoints; here there's no lossy step to avoid in the
                    // first place — the data is already NVFP4). Requires all
                    // three projections to be U8; a partial-U8 checkpoint
                    // falls through to the legacy path below, where the
                    // `load_ssm_proj` UInt8 branch added above still dequants
                    // each U8 tensor correctly on its own.
                    let native_nvfp4 = matches!(
                        store
                            .get(&format!("{la}.in_proj_qkv.weight"))
                            .map(|w| w.dtype),
                        Ok(WeightDtype::UInt8)
                    ) && matches!(
                        store
                            .get(&format!("{la}.in_proj_z.weight"))
                            .map(|w| w.dtype),
                        Ok(WeightDtype::UInt8)
                    ) && matches!(
                        store.get(&format!("{la}.out_proj.weight")).map(|w| w.dtype),
                        Ok(WeightDtype::UInt8)
                    );
                    if native_nvfp4 {
                        let qkv_qw = quantized_auto(
                            store,
                            &format!("{la}.in_proj_qkv"),
                            gpu,
                            Nvfp4Variant::Standard,
                        )?;
                        let z_qw = quantized_auto(
                            store,
                            &format!("{la}.in_proj_z"),
                            gpu,
                            Nvfp4Variant::Standard,
                        )?;
                        let qkvz_nvfp4 = qkv_qw.concat_rows(&z_qw, qkv_rows, z_rows, h, gpu)?;
                        let qkvz_nvfp4_t = qkvz_nvfp4.transpose_for_gemm(gpu, qkvz_size, h)?;

                        let out_proj_nvfp4 = quantized_auto(
                            store,
                            &format!("{la}.out_proj"),
                            gpu,
                            Nvfp4Variant::Standard,
                        )?;
                        let out_proj_nvfp4_t =
                            out_proj_nvfp4.transpose_for_gemm(gpu, h, value_dim)?;

                        let ssm = SsmWeights {
                            in_proj_qkvz: DenseWeight {
                                weight: spark_runtime::gpu::DevicePtr::NULL,
                            },
                            in_proj_ba: ba_dense,
                            conv1d,
                            a_log,
                            dt_bias,
                            norm,
                            out_proj: out_proj_nvfp4,
                        };
                        let mut layer = Qwen3SsmLayer::new_sequential(
                            input_norm,
                            ssm,
                            post_attn_norm,
                            ffn,
                            Some(qkvz_nvfp4),
                            Some(qkvz_nvfp4_t),
                            Some(out_proj_nvfp4_t),
                            config,
                            gpu,
                        )?;
                        layer.predequant_for_prefill(gpu, config, stream)?;
                        tracing::info!(
                            "SSM[{lp}] native NVFP4 GDN: qkvz+out_proj loaded pre-quantized \
                             (U8-packed on disk; no BF16 dequant/requant roundtrip)"
                        );
                        layers.push(Box::new(layer));
                        continue;
                    }

                    // PER-ROW FP8 for the row-wise cuBLASLt PREFILL arm
                    // (`ATLAS_FP8_ROWWISE=1`). Read from the store BEFORE the
                    // dequant below, though the bytes are store-owned either
                    // way; the NVFP4 build continues underneath because decode
                    // still needs it — `w8a16_gemv` cannot index a per-row
                    // scale. See weight_loader/qwen35_dense/rowwise_fp8.rs.
                    let rowwise_gdn = rowwise_fp8::rowwise_fp8_enabled()
                        && rowwise_fp8::proj_is_fp8_per_row(store, &format!("{la}.in_proj_qkv"))
                        && rowwise_fp8::proj_is_fp8_per_row(store, &format!("{la}.in_proj_z"))
                        && rowwise_fp8::proj_is_fp8_per_row(store, &format!("{la}.out_proj"));
                    let (qkvz_rowwise, out_proj_rowwise) = if rowwise_gdn && tp_size == 1 {
                        let qkv_r = rowwise_fp8::load_fp8_per_row(
                            store,
                            &format!("{la}.in_proj_qkv"),
                            gpu,
                        )?;
                        let z_r =
                            rowwise_fp8::load_fp8_per_row(store, &format!("{la}.in_proj_z"), gpu)?;
                        let out_r =
                            rowwise_fp8::load_fp8_per_row(store, &format!("{la}.out_proj"), gpu)?;
                        let qkvz_r = rowwise_fp8::concat_fp8_per_row(&qkv_r, &z_r, h, gpu)?;
                        // The concat copied both scale vectors; free the
                        // per-projection allocs. Weight bytes are store-owned.
                        gpu.free(qkv_r.row_scale)?;
                        gpu.free(z_r.row_scale)?;
                        (Some(qkvz_r), Some(out_r))
                    } else {
                        (None, None)
                    };

                    let qkv_dense = load_ssm_proj(&format!("{la}.in_proj_qkv"), qkv_rows, h)?;
                    let z_dense = load_ssm_proj(&format!("{la}.in_proj_z"), z_rows, h)?;
                    let out_proj_dense =
                        load_ssm_proj(&format!("{la}.out_proj"), h, dims.full_value_dim())?;

                    let qkvz_dense =
                        gpu_concat_rows(&qkv_dense, qkv_rows, &z_dense, z_rows, h, gpu)?;
                    // qkv/z BF16 are only inputs to the concat above; free them now
                    // rather than leaking them for the layer's lifetime (Atlas issue #A1).
                    gpu.free(qkv_dense.weight)?;
                    gpu.free(z_dense.weight)?;

                    let ba_dense =
                        interleave_ba(&in_proj_a, &in_proj_b, dims.full_nv, dims.full_nk, h, gpu)?;

                    // GDN HeadParallel shard: cut the FULL concat/interleave/on-disk buffers
                    // to this rank's contiguous head range (mirrors the MoE loader,
                    // linear_attn_arms.rs). Runs BEFORE both consumers below (the Bf16Raw
                    // branch AND the NVFP4/FP8 path), so both get sharded weights. qkvz
                    // (segmented [Q|K|V|Z]) + conv (segmented [Q|K|V]) slice per-block; ba
                    // slices contiguously; a_log/dt_bias are per-value-head FP32 scalars;
                    // out_proj is row-parallel on value_dim (partials summed by the
                    // post-out_proj all-reduce already in Qwen3SsmLayer::forward). norm is
                    // [vd] shared across value heads -> REPLICATE (never sliced). Free only
                    // the fresh qkvz/ba concat buffers; conv/a_log/dt_bias/out_proj sources
                    // are store aliases or dequant bufs freed downstream (the NVFP4 path
                    // frees the sharded qkvz/out_proj later). At tp==1 the else arm is a
                    // pure pass-through and every slicer no-ops -> byte-identical.
                    let (qkvz_dense, ba_dense, conv1d, a_log, dt_bias, out_proj_dense) = if tp_size
                        > 1
                    {
                        let d_conv = config.linear_conv_kernel_dim;
                        let (qkvz_ptr, _, _) = shard_gdn_qkvz_rows(qkvz_dense.weight, &dims, gpu)?;
                        gpu.free(qkvz_dense.weight)?;
                        let (ba_ptr, _, _) = shard_gdn_ba_rows(ba_dense.weight, &dims, gpu)?;
                        gpu.free(ba_dense.weight)?;
                        let (conv_ptr, _, _) =
                            shard_gdn_conv_rows(conv1d.weight, &dims, d_conv, gpu)?;
                        let (a_log_ptr, _) =
                            shard_gdn_value_vector(a_log.weight, &dims, 1, 4, gpu)?;
                        let (dt_bias_ptr, _) =
                            shard_gdn_value_vector(dt_bias.weight, &dims, 1, 4, gpu)?;
                        let (out_ptr, _, _) =
                            shard_gdn_out_proj_row_parallel(out_proj_dense.weight, &dims, gpu)?;
                        (
                            DenseWeight { weight: qkvz_ptr },
                            DenseWeight { weight: ba_ptr },
                            DenseWeight { weight: conv_ptr },
                            DenseWeight { weight: a_log_ptr },
                            DenseWeight {
                                weight: dt_bias_ptr,
                            },
                            DenseWeight { weight: out_ptr },
                        )
                    } else {
                        (qkvz_dense, ba_dense, conv1d, a_log, dt_bias, out_proj_dense)
                    };

                    // Native-BF16 SSM arm (no-metadata dense Holo checkpoints:
                    // Nvfp4Variant::Bf16Raw). The standard nvfp4 bundle ships the
                    // common/ BF16 dense kernels (dense_gemv_bf16 / dense_gemm_bf16
                    // / dense_gemm_bf16_pipelined) that every SSM forward arm's
                    // dense fallback already dispatches, so keep the concatenated
                    // qkvz_dense [Q|K|V|Z] and out_proj_dense ALIVE and route
                    // through those instead of the lossy BF16->NVFP4 runtime
                    // requant. No NVFP4/FP8 copy is built at all: in_proj_qkvz +
                    // out_proj_dense feed dense_gemv (per-seq decode,
                    // ssm_forward.rs:107/412), dense_gemm (batched decode,
                    // trait_decode_batched.rs:113/350 + ssm_batched.rs:180/279)
                    // and dense_gemm_bf16_pipelined (prefill,
                    // trait_prefill_proj.rs:298 + trait_prefill_helper.rs:89).
                    // ssm.out_proj stays null — every out_proj arm prefers
                    // out_proj_dense when Some; predequant_for_prefill /
                    // set_fp8_prefill_only_weights are skipped (NVFP4/FP8 only).
                    if matches!(variant, Nvfp4Variant::Bf16Raw) {
                        let ssm = SsmWeights {
                            in_proj_qkvz: qkvz_dense,
                            in_proj_ba: ba_dense,
                            conv1d,
                            a_log,
                            dt_bias,
                            norm,
                            out_proj: crate::weight_map::QuantizedWeight::null(),
                        };
                        let mut layer = Qwen3SsmLayer::new_sequential(
                            input_norm,
                            ssm,
                            post_attn_norm,
                            ffn,
                            None, // qkvz_nvfp4  — BF16 dense fallback used instead
                            None, // qkvz_nvfp4_t
                            None, // out_proj_nvfp4_t
                            config,
                            gpu,
                        )?;
                        // pub field (qwen3_ssm/mod.rs:46); selected by every
                        // out_proj arm (ssm_forward.rs:412,
                        // trait_decode_batched.rs:350, ssm_batched.rs:279,
                        // trait_prefill_helper.rs:89).
                        layer.out_proj_dense = Some(out_proj_dense);
                        layers.push(Box::new(layer));
                        continue;
                    }

                    let qkvz_size = config.ssm_qkvz_size();

                    // GDN ≥FP8 precision policy (2026-07-04). The nvidia
                    // Qwen3.6-27B-NVFP4 checkpoint ships GDN in_proj_qkv /
                    // out_proj as native F8_E4M3 (modelopt sensitivity
                    // analysis deliberately keeps the SSM projections
                    // high-precision); `load_ssm_proj` dequants them to BF16,
                    // and the code below then RE-quantizes to NVFP4 (4-bit) —
                    // a lossy double-quant of the exact tensors the toolchain
                    // protected. That regressed BFCL-ST ~7pt (non_live 85.4→
                    // 76.6) vs the 06-15 reference. When enabled, keep the
                    // BF16 dequant (≥FP8) and route qkvz + out_proj through the
                    // dense_gemv / dense_gemm dispatch (in_proj_qkvz +
                    // out_proj_dense fields), mirroring the MoE sister loader's
                    // arm (qwen35/load_layers/linear_attn_arms.rs). Gated for a
                    // clean A/B + KL-drift gate before flipping the default.
                    let gdn_bf16 = matches!(
                        std::env::var("ATLAS_GDN_BF16_WEIGHTS").ok().as_deref(),
                        Some("1")
                    );
                    if gdn_bf16 {
                        let ssm = SsmWeights {
                            in_proj_qkvz: DenseWeight {
                                weight: qkvz_dense.weight,
                            },
                            in_proj_ba: ba_dense,
                            conv1d,
                            a_log,
                            dt_bias,
                            norm,
                            // Unused: out_proj_dense (set below) has higher
                            // dispatch priority in both prefill and decode.
                            out_proj: crate::weight_map::QuantizedWeight::null(),
                        };
                        let mut layer = Qwen3SsmLayer::new_sequential(
                            input_norm,
                            ssm,
                            post_attn_norm,
                            ffn,
                            None,
                            None,
                            None,
                            config,
                            gpu,
                        )?;
                        layer.out_proj_dense = Some(out_proj_dense);
                        tracing::info!(
                            "SSM[{lp}] ATLAS_GDN_BF16_WEIGHTS: qkvz + out_proj kept BF16 \
                             (≥FP8; NVFP4 requant skipped)"
                        );
                        layers.push(Box::new(layer));
                        continue;
                    }

                    let qkvz_nvfp4 = quantize_to_nvfp4(
                        &qkvz_dense,
                        qkvz_size,
                        h,
                        gpu,
                        absmax_k,
                        quantize_k,
                        stream,
                    )?;

                    let qkvz_nvfp4_t = qkvz_nvfp4.transpose_for_gemm(gpu, qkvz_size, h)?;

                    let out_proj_nvfp4 = quantize_to_nvfp4(
                        &out_proj_dense,
                        h,
                        value_dim,
                        gpu,
                        absmax_k,
                        quantize_k,
                        stream,
                    )?;

                    let out_proj_nvfp4_t = out_proj_nvfp4.transpose_for_gemm(gpu, h, value_dim)?;

                    // Native FP8 SSM prefill GEMM: build a single-scale FP8
                    // copy of `qkvz_dense` [qkvz_size, h] and `out_proj_dense`
                    // [h, value_dim] by direct BF16→FP8 truncation. SSM weight
                    // magnitudes fit in FP8 E4M3 range (|w| ≤ 448), so no
                    // separate scalar dequant is needed at GEMM time — the
                    // `fp8_gemm_n128` kernel interprets the FP8 bytes as
                    // values directly (mirrors how `predequant_nvfp4_to_fp8`
                    // bakes `scale2` into the FP8 stream). PCND: gated.
                    let (qkvz_fp8_prefill, out_proj_fp8_prefill) =
                        if let Some(b2f_k) = bf16_to_fp8_k {
                            let qkvz_total = (qkvz_size * h) as u32;
                            let qkvz_fp8 = gpu.alloc(qkvz_size * h)?;
                            crate::layers::ops::bf16_to_fp8(
                                gpu,
                                b2f_k,
                                qkvz_dense.weight,
                                qkvz_fp8,
                                qkvz_total,
                                stream,
                            )?;
                            let out_total = (h * value_dim) as u32;
                            let out_fp8 = gpu.alloc(h * value_dim)?;
                            crate::layers::ops::bf16_to_fp8(
                                gpu,
                                b2f_k,
                                out_proj_dense.weight,
                                out_fp8,
                                out_total,
                                stream,
                            )?;
                            gpu.synchronize(stream)?;
                            (Some(qkvz_fp8), Some(out_fp8))
                        } else {
                            (None, None)
                        };

                    // SSM prefill/decode always dispatch qkvz_nvfp4/_t and the NVFP4
                    // out_proj; the BF16 qkvz_dense / out_proj_dense were only quantize
                    // inputs. Free them rather than keep a third full-precision copy of
                    // the largest SSM tensor across every layer (Atlas issue #A1).
                    gpu.free(qkvz_dense.weight)?;
                    gpu.free(out_proj_dense.weight)?;

                    let ssm = SsmWeights {
                        in_proj_qkvz: DenseWeight {
                            weight: spark_runtime::gpu::DevicePtr::NULL,
                        },
                        in_proj_ba: ba_dense,
                        conv1d,
                        a_log,
                        dt_bias,
                        norm,
                        out_proj: out_proj_nvfp4,
                    };

                    let mut layer = Qwen3SsmLayer::new_sequential(
                        input_norm,
                        ssm,
                        post_attn_norm,
                        ffn,
                        Some(qkvz_nvfp4),
                        Some(qkvz_nvfp4_t),
                        Some(out_proj_nvfp4_t),
                        config,
                        gpu,
                    )?;
                    layer.predequant_for_prefill(gpu, config, stream)?;
                    // Install the FP8 prefill weights AFTER `predequant_for_prefill`
                    // (which sets `out_proj_fp8` from NVFP4 + scale2). The
                    // native-FP8 path overrides both pointers when active,
                    // routing prefill through `fp8_gemm_n128` instead of
                    // `w4a16_gemm_t`. Decode batch paths keep their NVFP4
                    // fallback (the `qkvz_nvfp4*` fields above).
                    if qkvz_fp8_prefill.is_some() || out_proj_fp8_prefill.is_some() {
                        layer.set_fp8_prefill_only_weights(qkvz_fp8_prefill, out_proj_fp8_prefill);
                    }
                    // …and LAST, so it wins over both of the prefill installs
                    // above: the checkpoint's own per-row FP8, which reaches
                    // the GEMM with no conversion at all. Decode is untouched.
                    if qkvz_rowwise.is_some() {
                        layer.set_fp8_rowwise_prefill_weights(qkvz_rowwise, out_proj_rowwise);
                        if i == 0 {
                            tracing::info!(
                                "SSM[{lp}] ATLAS_FP8_ROWWISE: qkvz + out_proj prefill via \
                                 native per-row FP8 (no BF16 dequant, no NVFP4 requant); \
                                 decode keeps NVFP4"
                            );
                        }
                    }
                    layers.push(Box::new(layer));
                }
                LayerType::SlidingAttention => {
                    unreachable!("unexpected SlidingAttention in this loader")
                }
                LayerType::Moe => unreachable!("Qwen3.5 dense has no standalone MoE layers"),
                // GLM-5.3's `deepseek_sparse_attention`: a full-rank mixer whose visible key set
                // is chosen at runtime by an indexer. Hard error, not a silent fallthrough into
                // the dense-attention arm -- that would attend over the WHOLE cache and look right.
                LayerType::SparseAttention => anyhow::bail!(
                    "layer {i}: SparseAttention needs a DSA indexer and per-query top-k; Qwen3.5 dense has neither"
                ),
            }

            if (i + 1) % 10 == 0 {
                tracing::info!("Loaded layers 0..{}", i + 1);
                spark_runtime::progress::layer(i + 1, config.num_hidden_layers);
            }
        }

        tracing::info!(
            "Qwen3.5 dense weight loader: {} layers ({} attention, {} SSM, dense FFN)",
            layers.len(),
            attn_idx,
            layers.len() - attn_idx,
        );
        // #915: the line the next H100 serve is read against. `weights` is the
        // checkpoint as the ledger sees it; `derived` is everything this loader
        // built on top of it and now OWNS (released at teardown, not swept);
        // `not built` is what the pre-#915 loader allocated here and this one
        // proved unreachable. Emitted unconditionally — a residency regression
        // that only shows under a debug flag is one nobody sees.
        tracing::info!("{}", residency.summary(store.resident_bytes()));
        for (label, bytes, count) in store.derived().by_label() {
            tracing::debug!(
                "  derived-weight owner: {:>9.1} MB x{:<5} {label}",
                bytes as f64 / (1024.0 * 1024.0),
                count,
            );
        }

        Ok(layers)
    }

    /// Drop the SSM source tensors the native-FP8 GDN arm copied and no longer
    /// reads (#915).
    ///
    /// That arm builds a fused `[QKV|Z]` FP8 weight by device-to-device
    /// appending `in_proj_qkv` and `in_proj_z` (`concat_fp8_block_scaled`), and
    /// builds `in_proj_ba` by a D2H -> host interleave -> H2D round trip
    /// (`interleave_ba`). All four store originals are dead the moment
    /// `load_layers` returns — 80 MiB + ~1 MB per SSM layer, ~3.9 GB across the
    /// 48 GDN layers of Qwen3.8-27B-FP8, which is what the fused copy costs.
    /// Keeping both is the duplicate that came straight out of the KV budget.
    ///
    /// 🪤 Narrow on purpose. `out_proj.weight` IS `out_proj_fp8w.weight`
    /// (zero-copy from the store), `conv1d`, `A_log`, `dt_bias` and
    /// `norm.weight` are aliased or conditionally aliased depending on their
    /// on-disk dtype, and every attention tensor — and `mlp.down_proj` — is
    /// bound zero-copy. Only the names below are freed, and only for layers
    /// where the module-private predicate that SELECTED the corresponding
    /// concat (`gdn_fp8_arm_selected` for the SSM `[QKV|Z]` weight,
    /// `ffn_gateup_fused_selected` for the dense-FFN gate+up weight, #927)
    /// says that arm actually ran.
    fn prune_after_load(
        &self,
        store: &mut WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<()> {
        // Same resolution `load_layers` uses: an explicit `layer_types` list
        // wins over the computed pattern, and pruning must be decided from the
        // list that actually selected the arms.
        let layer_types = if config.layer_types.is_empty() {
            (0..config.num_hidden_layers)
                .map(|i| config.layer_type(i))
                .collect::<Vec<_>>()
        } else {
            config.layer_types.clone()
        };
        let mut doomed: std::collections::HashSet<String> = std::collections::HashSet::new();
        // The dense-FFN gate+up fusion (#927) consumed these the same way the
        // SSM `[QKV|Z]` concat below consumes its two: the fused buffer holds
        // the bytes, and `gate_proj`/`up_proj` are VIEWS inside it, so nothing
        // aliases the store tensors any more. Releasing them is what makes the
        // fusion residency-neutral; keeping them would be the 11.4 GB the
        // fused arm exists to not spend. `down_proj` is NOT pruned — it is
        // still bound zero-copy.
        //
        // Same `variant` the load used (`detect_nvfp4_variant` is a pure
        // function of the store and the config), and the SAME predicate, so a
        // layer that did not fuse cannot be pruned here.
        let variant = detect_nvfp4_variant(store, config);
        for i in 0..layer_types.len() {
            let lp = config.layer_prefix(i);
            if !ffn_gateup_fused_selected(store, config, variant, &lp) {
                continue;
            }
            for proj in ["gate_proj", "up_proj"] {
                for leaf in ["weight", "weight_scale_inv", "weight_scale"] {
                    doomed.insert(format!("{lp}.mlp.{proj}.{leaf}"));
                }
            }
        }
        for (i, lt) in layer_types.iter().enumerate() {
            if *lt != LayerType::LinearAttention {
                continue;
            }
            let la = format!("{}.linear_attn", config.layer_prefix(i));
            if !gdn_fp8_arm_selected(store, &la, config.tp_world_size) {
                continue;
            }
            for leaf in [
                "in_proj_qkv.weight",
                "in_proj_qkv.weight_scale_inv",
                "in_proj_qkv.weight_scale",
                "in_proj_z.weight",
                "in_proj_z.weight_scale_inv",
                "in_proj_z.weight_scale",
                "in_proj_a.weight",
                "in_proj_b.weight",
            ] {
                doomed.insert(format!("{la}.{leaf}"));
            }
        }
        if doomed.is_empty() {
            return Ok(());
        }
        let (count, bytes) = store.free_matching(gpu, |name| doomed.contains(name))?;
        tracing::info!(
            "native FP8: released {count} store tensors ({:.2} GB) consumed by the fused              [QKV|Z] SSM concat, the BA interleave and the dense-FFN gate+up fusion;              out_proj/conv1d/A_log/dt_bias/norm/down_proj kept (still aliased)",
            bytes as f64 / 1e9,
        );
        Ok(())
    }

    fn load_embedding(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        _gpu: &dyn GpuBackend,
    ) -> Result<DenseWeight> {
        loaders_b::load_embedding(store, config)
    }

    fn load_final_norm(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        _gpu: &dyn GpuBackend,
    ) -> Result<DenseWeight> {
        loaders_b::load_final_norm(store, config)
    }

    fn load_lm_head(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<DenseWeight> {
        loaders_b::load_lm_head(store, config, gpu)
    }

    fn load_mtp_weights(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<Option<MtpWeights>> {
        if !store.contains("mtp.fc.weight") {
            return Ok(None);
        }
        let variant = detect_nvfp4_variant(store, config);
        tracing::info!(
            "Loading dense MTP weights (variant={:?}, hidden={}, inter={})",
            variant,
            config.hidden_size,
            config.intermediate_size,
        );
        // `load_mtp` auto-detects MoE vs dense FFN by inspecting the weight
        // names. For dense Qwen3.6-27B-FP8 it returns a MtpWeights with
        // `dense_ffn = Some(...)` and NULL placeholders for the MoE fields.
        let mtp = load_mtp(store, config.num_experts, gpu, variant)?;
        if mtp.dense_ffn.is_some() {
            tracing::info!("Dense MTP head ready (FP8 e4m3 projections + dense gate/up/down MLP)");
        } else {
            tracing::info!(
                "MoE MTP head ready ({} experts) — dense loader sees MoE bundle",
                mtp.experts.len(),
            );
        }
        Ok(Some(mtp))
    }

    fn load_vision_encoder(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<Option<crate::layers::VisionEncoder>> {
        // Dense Qwen3.5 / Holo VL checkpoints (e.g. Holo-3.1-0.8B, Ornith-1.0-9B)
        // ship the SAME Qwen3-VL ViT tower as their MoE siblings. The MoE
        // loader's `load_vision_encoder` reads only `store` + `config.vision`
        // (no MoE-specific state), so reuse it verbatim. The shared model
        // forward (`model/trait_impl/*`, gated on `vision_encoder.is_some()`)
        // then merges image embeddings — no dense-specific forward changes.
        super::qwen35::Qwen35WeightLoader.load_vision_encoder(store, config, gpu)
    }
}