spark_model/factory/

build.rs

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

//! `build_model` — entry point that wires up the configured loader,
//! buffers, KV cache, and (optional) DFlash drafter into a `TransformerModel`.

use anyhow::Result;
use atlas_core::config::ModelConfig;
use spark_runtime::buffers::BufferArena;
use spark_runtime::gpu::GpuBackend;
use spark_runtime::kv_cache::{KvCacheConfig, KvCacheDtype, PagedKvCache};
use spark_runtime::prefix_cache::PrefixCache;
use spark_runtime::weights::WeightStore;

use super::loader_for_config;
use super::m2_setup::maybe_run_minimax_m2_moe_transpose;
use super::{DflashBuildArgs, LoraBuildArgs};
use crate::layers::MtpQuantization;
use crate::model::TransformerModel;
use crate::traits::Model;
use crate::weight_loader::load_dflash_weights;

mod kv_summary;

pub fn build_model(
    mut config: ModelConfig,
    // BY VALUE. Every reader below still takes `&WeightStore` — the change is
    // only in who OWNS it. The store is the one structure that knows every
    // weight pointer, and it used to be a local in `startup()` that was dropped
    // once the layers had copied pointers out of it: the memory stayed live
    // with nothing able to free it. The model owns it now, so `teardown` can.
    // `mut` for `prune_after_load` (Step 3c), which lets a loader drop the
    // originals of tensors it re-uploaded before the KV budget is computed.
    mut store: WeightStore,
    gpu: Box<dyn GpuBackend>,
    max_batch_tokens: usize,
    kv_block_size: usize,
    max_seq_len: usize,
    max_batch_size: usize,
    mtp_quant: MtpQuantization,
    use_speculative: bool,
    prefix_cache: Box<dyn PrefixCache>,
    mtp_vocab_size: u32,
    comm: Option<std::sync::Arc<dyn spark_comm::CommBackend>>,
    self_speculative: bool,
    num_drafts: usize,
    kv_dtype: KvCacheDtype,
    inference_reserve: usize,
    gpu_memory_utilization: f64,
    ssm_cache_slots: usize,
    layer_dtypes: Vec<KvCacheDtype>,
    ssm_checkpoint_interval: usize,
    // Phase 6.1.f: per-sequence HBM cache cap. `Some(N)` enables
    // `--high-speed-swap` HBM-shrink behavior. `None` preserves the
    // pre-Phase-6 unbounded behavior.
    hss_cache_blocks_per_seq: Option<u32>,
    // DFlash speculative-decoding pairing. `None` = no DFlash; existing
    // MTP / no-spec paths unchanged.
    dflash_args: Option<DflashBuildArgs<'_>>,
    // Startup-static LoRA adapter (`--lora-adapter`). `None` = base-only.
    lora_args: Option<LoraBuildArgs<'_>>,
    // NLLB / M2M-100 translation language pair (tokenizer-resolved
    // `(src_lang_id, tgt_lang_id)`), resolved server-side. `None` for all other
    // model types.
    nllb_lang: Option<(u32, u32)>,
    // NLLB / M2M-100 PEFT LoRA adapter directory (`--lora-adapter` for an
    // encoder-decoder checkpoint). `None` = base model.
    nllb_lora_dir: Option<std::path::PathBuf>,
) -> Result<Box<dyn Model>> {
    // NLLB / M2M-100 is an encoder-decoder model that cannot be represented by
    // the decoder-only TransformerModel stack. Serve it with the dedicated
    // `NllbGpuModel`, which reads its weights from the standard `store` — this
    // returns BEFORE `loader_for_config`, so the decoder-only weight loader
    // (and its fail-fast) never runs on this path.
    #[cfg(feature = "cuda")]
    if matches!(config.model_type.as_str(), "m2m_100" | "nllb") {
        let (src, tgt) = nllb_lang.ok_or_else(|| {
            anyhow::anyhow!(
                "NLLB serving requires --src-lang and --tgt-lang (translation language pair)"
            )
        })?;
        let lang = crate::model::nllb::NllbLang {
            src_lang_id: src,
            tgt_lang_id: tgt,
            decoder_start_id: config.eos_token_id,
            eos_id: config.eos_token_id,
            pad_id: 1,
        };
        let model = crate::model::nllb::NllbGpuModel::new(
            &config,
            &store,
            gpu,
            lang,
            max_seq_len,
            max_batch_size,
            nllb_lora_dir.as_deref(),
        )?;
        return Ok(Box::new(model));
    }
    #[cfg(not(feature = "cuda"))]
    let _ = (nllb_lang, nllb_lora_dir);

    // ── Step 1: Select weight loader (only model-specific dispatch) ──
    let loader = loader_for_config(&config)?;

    // Entry free-memory sample for the KV-budget transient correction below.
    // Taken BEFORE the LoRA load and buffer arena — the delta from here to
    // the budget-time sample then isolates exactly what build_model itself
    // allocated, with the weight loader's transient footprint present in
    // BOTH samples so it cancels out. The LoRA-ordering invariant in the
    // next comment is about the BUDGET sample and is unaffected: LoRA still
    // lands between these two samples and is therefore charged.
    let build_entry_free = gpu.free_memory().ok();

    // ── LoRA adapter load (pre-arena, pre-KV-sizing) ──
    // MUST run before `BufferArena::new` and the `gpu.free_memory()`
    // snapshot below: the pool allocation then lands in `used_so_far`, so
    // the KV-cache budget shrinks automatically (positional budgeting —
    // no arithmetic edit needed). Do NOT move this later. Setting
    // `config.adapter_max_rank` here also lets `BufferSizes` size the
    // lora_xa/lora_delta/lora_hact scratch.
    let lora_weights: Option<crate::lora::LoraWeights> = if let Some(ref la) = lora_args {
        config.adapter_max_rank = la.max_lora_rank;
        loader.load_lora_adapters(
            &la.adapters,
            &config,
            gpu.as_ref(),
            la.max_loras,
            la.max_lora_rank,
        )?
    } else {
        None
    };

    // Pre-construction: when DFlash is active, populate the target's
    // capture-layer indices from the drafter's `dflash_config.target_layer_ids`
    // so `TransformerModel::new` allocates the 5×hidden_size capture buffer.
    //
    // The drafter's `target_layer_ids` are used DIRECTLY as Atlas capture
    // indices. An earlier implementation subtracted 1 from each id on HF
    // `output_hidden_states` reasoning; measurement on the z-lab drafters
    // shows that adjustment feeds the drafter hidden states one layer early
    // on every capture layer and costs ~1 full accepted row per step (27B
    // dense MinHeap: mean accept 6.61 direct vs 5.56 shifted; the 2026-07-19
    // 35B accept hunt found the same).
    if let Some(ref args) = dflash_args
        && let Some(ref sub) = args.drafter_config.dflash_config
    {
        config.dflash_capture_layers = sub.target_layer_ids.clone();
        // gamma for the pool sizing below, resolved from the drafter itself:
        // a DFlash2 checkpoint states its trained block size inside
        // `dflash_config`, and the top-level field's default of 16 must not
        // shadow it. --dflash-gamma still wins over both.
        //
        // MUST go through `default_dflash_gamma`, this value sizes the SSM
        // pools (uniform K = γ+1), and the HEAD resolves its own γ through
        // the same helper. When this said `block_size` while the head said
        // block+2, the pools came up one verify row short and the first
        // DFlash step died with "SSM MTP intermediate buffers not allocated
        // (h=8, conv=9, num_tokens=10)" mid-graph-capture.
        config.dflash_gamma = Some(args.gamma.unwrap_or_else(|| {
            crate::layers::qwen3_ssm::default_dflash_gamma(
                args.drafter_config.effective_block_size(),
            )
        }));
        tracing::info!(
            "DFlash: target layer capture indices = {:?} (drafter target_layer_ids, \
             used directly), γ = {:?}",
            config.dflash_capture_layers,
            config.dflash_gamma,
        );
    }

    // ── Step 2: Load weights (model-agnostic from here) ──
    let attn_layer_dtypes: Vec<KvCacheDtype> = if layer_dtypes.is_empty() {
        vec![kv_dtype; config.num_attention_layers()]
    } else {
        layer_dtypes.clone()
    };

    // Populate per-layer KV dims for heterogeneous-attention models (Gemma-4).
    // Homogeneous models return an empty Vec which the KV cache treats as
    // "use global num_kv_heads/head_dim for all layers" (backward compatible).
    config.kv_layer_dims = loader.kv_layer_dims(&config);

    // Attribute the memory the BUILD spends, not just the shards.
    //
    // Weight upload reports itself per shard, and the buffer arena reports its
    // own total, but everything between — per-layer construction, runtime
    // requantization, derived weights — was invisible. On qwen4_exp that gap
    // is ~8.5 GB: shards end at 85.2 GB and the KV budget sees 94.7 GB
    // pre-KV, of which the arena (872 MB) and the GDN prefill scratch (88 MB)
    // explain under a gigabyte. Without these three lines the only way to
    // find the rest is to guess.
    //
    // `MemTrace` (campaign) marks EVERY build step; upstream's free_before/after
    // pair below is kept because it is the one step with a per-layer average.
    // Reconciling the two into a single reporter belongs to the M0 telemetry
    // commit, not to this rebase.
    let mut mem = MemTrace::new(gpu.as_ref());
    let free_before_layers = gpu.free_memory().unwrap_or(0);
    let mut layers = loader.load_layers(&store, &config, gpu.as_ref(), &attn_layer_dtypes)?;
    let free_after_layers = gpu.free_memory().unwrap_or(0);
    tracing::info!(
        "Layer construction: {:.2} GB consumed ({:.2} GB free -> {:.2} GB free) \
         across {} layers, {:.1} MB/layer average",
        (free_before_layers.saturating_sub(free_after_layers)) as f64 / 1e9,
        free_before_layers as f64 / 1e9,
        free_after_layers as f64 / 1e9,
        config.num_hidden_layers,
        (free_before_layers.saturating_sub(free_after_layers)) as f64
            / 1e6
            / config.num_hidden_layers.max(1) as f64,
    );
    mem.mark("load_layers");
    let embed = loader.load_embedding(&store, &config, gpu.as_ref())?;
    mem.mark("load_embedding");
    // n-gram fused embedding (LongCat family; None everywhere else). Built
    // before `config` is moved into the model. Staged for `max_batch_tokens`
    // because that is exactly the widest embed the arena can be handed.
    let ngram_embed =
        loader.load_ngram_embedding(&store, &config, gpu.as_ref(), max_batch_tokens)?;
    let final_norm = loader.load_final_norm(&store, &config, gpu.as_ref())?;
    mem.mark("load_final_norm");
    let lm_head = loader.load_lm_head(&store, &config, gpu.as_ref())?;
    mem.mark("load_lm_head");
    let mtp_weights = loader.load_mtp_weights_multi(&store, &config, gpu.as_ref())?;
    mem.mark("load_mtp_weights_multi");

    // DeepSeek-V4 ships an architecturally distinct MTP module (MLA + mHC), not
    // the Qwen-shaped `MtpWeights`. Load it via the V4-specific path and keep it
    // — the `DeepseekV4MtpHead` proposer is built from it after the model is
    // constructed (it needs the resolved draft NVFP4 LM head + the model's
    // owned GPU backend) and installed via `set_dflash_proposer`. Only built
    // when `--speculative` is set; otherwise the module is loaded for
    // verification then dropped.
    // Only rank 0 runs the MTP draft (no-EP, all experts local). Skip loading it
    // on the worker ranks — they never call propose(), so it would be dead weight.
    // GLM-5.3's MTP block is architecturally distinct in a THIRD way: neither the Qwen-shaped
    // `MtpWeights` nor DeepSeek's `mtp.0.*` module, but `layers.45` — a DSA mixer + the same
    // 288-expert routed MoE + `shared_head.norm`, with NO hyper-connection.
    //
    // 🔴 Loaded on EVERY rank, unlike the V4 module below. GLM's MTP MoE is EP-sharded exactly
    // like the text stack, so both ranks hold a half and the block's own all-reduce assembles
    // it; a rank-0-only drafter would silently drop half the routed sum and draft from a
    // half-computed hidden.
    let glm_mtp_module = if config.model_type == "glm5_next" && use_speculative {
        match crate::weight_loader::load_glm5next_mtp_module(&store, &config, gpu.as_ref()) {
            Ok(Some(m)) => {
                tracing::info!(
                    "GLM-5.3 MTP draft module loaded (layers.{})",
                    config.num_hidden_layers
                );
                Some(m)
            }
            Ok(None) => {
                tracing::info!("GLM-5.3: no MTP block in checkpoint (MTP off)");
                None
            }
            Err(e) => {
                tracing::error!("GLM-5.3 MTP module load FAILED: {e:#}");
                None
            }
        }
    } else {
        None
    };
    let glm_mtp_embed = embed;
    let glm_mtp_lm_head = lm_head;

    let v4_mtp_module =
        if config.model_type == "deepseek_v4" && use_speculative && config.ep_rank == 0 {
            match crate::weight_loader::deepseek_v4::load_v4_mtp_module(
                &store,
                &config,
                gpu.as_ref(),
                &attn_layer_dtypes,
            ) {
                Ok(Some(m)) => {
                    tracing::info!(
                        "DeepSeek-V4 MTP draft module loaded OK (num_mtp_modules={})",
                        config.num_mtp_modules
                    );
                    Some(m)
                }
                Ok(None) => {
                    tracing::info!("DeepSeek-V4: no MTP module in checkpoint (MTP off)");
                    None
                }
                Err(e) => {
                    tracing::error!("DeepSeek-V4 MTP module load FAILED: {e:#}");
                    None
                }
            }
        } else {
            None
        };

    // Capability warning: user asked for `--speculative` but nothing bound an
    // MTP head, so speculative decoding will silently no-op.
    //
    // 🪤 `mtp_weights` is only the GENERIC (`load_mtp_weights_multi`) path.
    // GLM-5.3 and DeepSeek-V4 bind architecturally distinct modules above and
    // leave that vec empty, so testing it alone printed "no MTP weights were
    // loaded" two lines under "GLM-5.3 MTP draft module loaded (layers.45)".
    // Every binding path has to be consulted, and when none of them bound
    // anything the checkpoint still has to be asked whether it SHIPS an MTP
    // head — "Atlas can't read this layout" and "there is no head here" are
    // different faults and want different messages.
    if use_speculative
        && mtp_weights.is_empty()
        && glm_mtp_module.is_none()
        && v4_mtp_module.is_none()
    {
        match crate::mtp_layout::detect_in_store(&store, &config) {
            None => tracing::warn!(
                "`--speculative` was requested but this checkpoint ships no MTP head — \
                 speculative decoding will be disabled. Either drop `--speculative` or \
                 use a checkpoint that ships one (e.g. `mtp.safetensors`)."
            ),
            Some(layout) => tracing::error!(
                "`--speculative` was requested and this checkpoint DOES ship MTP weights \
                 ({layout:?}), but no loader bound them for model_type '{}' — speculative \
                 decoding will be disabled. This is an Atlas capability gap, not a \
                 checkpoint problem.",
                config.model_type,
            ),
        }
    }
    mem.mark("mtp modules (glm/v4)");
    let vision_encoder = loader.load_vision_encoder(&store, &config, gpu.as_ref())?;
    mem.mark("load_vision_encoder");

    // A multimodal checkpoint's vision tower is read by the weight loader like
    // everything else, but only a loader that implements `load_vision_encoder`
    // ever binds it. GLM-5.3's port is text-only by design
    // (`weight_loader/glm5_next.rs`: "Vision tower — present in the checkpoint,
    // out of scope for the text port"), so its 1.05 GiB of `model.visual.*`
    // sat resident on BOTH ranks for the life of the process, bound to nothing,
    // subtracted from the KV budget computed below.
    //
    // Freeing is keyed off the bind result, not off a model list: if the encoder
    // was built, `vision_encoder` is `Some` and nothing is touched — including
    // the loaders that bind zero-copy from these very pointers. The day a GLM
    // vision encoder lands, this stops firing on its own.
    if vision_encoder.is_none() {
        let (n, bytes) = store.free_matching(gpu.as_ref(), |name| {
            name.starts_with("model.visual.")
                || name.starts_with("model.vision")
                || name.starts_with("visual.")
        })?;
        if n > 0 {
            tracing::info!(
                "Vision tower: {n} tensors ({:.2} GiB) released — this build binds no vision \
                 encoder for model_type '{}', so the tower was resident and unreachable. \
                 Text capability is unchanged; image input was already unsupported here.",
                bytes as f64 / (1024.0 * 1024.0 * 1024.0),
                config.model_type,
            );
        }
    }
    mem.mark("vision reclaim");

    // If the checkpoint's `quantization_config.ignore_modules` lists MTP
    // (e.g. Sehyo/Qwen3.5-35B-A3B-NVFP4 ignores `mtp.*`), the MTP weights
    // were stored as BF16 on disk. Runtime-quantizing them to NVFP4
    // anyway — which is what `mtp_quant` would otherwise do — produces
    // garbage drafts (vllm PR #38832). Force BF16 in that case.
    let effective_mtp_quant = if !mtp_weights.is_empty() {
        let quant_fmt = crate::quant_format::detect_quant_format(&config, &store);
        if quant_fmt.is_ignored("mtp.fc.weight")
            || quant_fmt.is_ignored("mtp.layers.0.self_attn.q_proj.weight")
        {
            if mtp_quant != MtpQuantization::Bf16 {
                tracing::info!(
                    "MTP head listed in checkpoint ignore_modules — overriding \
                     --mtp-quantization {:?} → Bf16 to preserve precision",
                    mtp_quant,
                );
            }
            MtpQuantization::Bf16
        } else {
            mtp_quant
        }
    } else {
        mtp_quant
    };

    // ── Step 3: LM-head quantization (NVFP4 / FP8 / BF16-skip) + the
    // draft-only NVFP4 head for MTP — extracted to lm_head_setup.rs
    // (file-size cap; pure code move).
    let (lm_head_nvfp4, lm_head_fp8, mtp_lm_head_nvfp4) = super::lm_head_setup::setup_lm_heads(
        &store,
        &lm_head,
        &config,
        gpu.as_ref(),
        use_speculative,
        !mtp_weights.is_empty(),
    )?;

    // Capture the shared embed + resolved draft NVFP4 head for the DeepSeek-V4
    // MTP proposer BEFORE `embed` / `lm_head_nvfp4` / `mtp_lm_head_nvfp4` are
    // moved into `TransformerModel::new`. All are `Copy` (DenseWeight /
    // QuantizedWeight). The draft head resolves to the separate draft-only
    // NVFP4 head (main head kept BF16) or the main NVFP4 head. `None` ⇒ no
    // NVFP4 head available ⇒ the V4 proposer can't draft and is skipped.
    let v4_mtp_embed = embed;
    // DeepSeek-V4-Flash keeps the LM head in BF16; the proposer drafts with the
    // same BF16 head via dense_gemv (drafts are re-verified by the target, so the
    // draft head only affects acceptance). DenseWeight is Copy.
    let v4_mtp_lm_head = lm_head;

    // ── Step 3b: Post-load MoE prefill transpose (MiniMax EP=2 TTFT fix) ──
    //
    // MiniMax M2.7-NVFP4 EP=2 has ~46 GB free at layer-0 load time but
    // ~65 GB free here (the BF16 lm_head just freed ~22 GB during NVFP4
    // quantization). The transpose costs ~59 GB — fits in the post-load
    // window but not the pre-load one. Other loaders (qwen35, qwen3,
    // gemma4) still call `transpose_for_prefill` inline during layer
    // construction; this default-no-op hook doesn't perturb them.
    maybe_run_minimax_m2_moe_transpose(&config, gpu.as_ref(), &mut layers)?;

    // ── Step 3c: Let the loader drop store tensors it has finished with ──
    //
    // Default is a no-op. Loaders that upload their OWN copies (a TP shard, a
    // host round-trip) leave the store's originals resident for nothing; on
    // unified-memory GB10 that duplicate is subtracted from the KV budget
    // computed a few lines below, so it has to happen HERE — after every
    // `load_*` reader above, before `BufferArena::new` and `gpu.free_memory()`.
    loader.prune_after_load(&mut store, &config, gpu.as_ref())?;
    mem.mark("prune_after_load");
    tracing::info!(
        "WeightStore after prune: {} tensors, {:.3} GiB still resident",
        store.len(),
        store.resident_bytes() as f64 / (1024.0 * 1024.0 * 1024.0),
    );

    // ── Step 4: Create buffer arena ──
    let buffers = BufferArena::new(
        &config,
        max_batch_tokens,
        max_seq_len,
        kv_block_size,
        max_batch_size,
        gpu.as_ref(),
    )?;

    // ── Step 5: Size KV cache from actual free memory ──
    // MLA absorbed: cache compressed latent [kv_lora + rope] instead of expanded [nkv * hd]
    // This gives 12.8x smaller KV cache AND better precision (no expand→cache→read roundtrip)
    let (kv_num_heads, kv_head_dim) = if config.kv_lora_rank > 0 {
        let mla_cache_dim = config.kv_lora_rank + config.qk_rope_head_dim;
        tracing::info!(
            "MLA absorbed KV cache: 1 head × {} dims ({}+{}) per token (vs {} heads × {})",
            mla_cache_dim,
            config.kv_lora_rank,
            config.qk_rope_head_dim,
            config.num_key_value_heads,
            config.head_dim,
        );
        (1, mla_cache_dim)
    } else {
        (config.num_key_value_heads, config.head_dim)
    };
    let kv_config = KvCacheConfig {
        block_size: kv_block_size,
        num_kv_heads: kv_num_heads,
        head_dim: kv_head_dim,
        num_layers: config.num_attention_layers(),
        dtype: kv_dtype,
        layer_dtypes: layer_dtypes.clone(),
        layer_dims: config.kv_layer_dims.clone(),
        cache_blocks_per_seq: hss_cache_blocks_per_seq,
    };

    if hss_cache_blocks_per_seq.is_some() {
        kv_summary::log_hss_kv_summary(&kv_config);
    }
    // ── gpu_memory_utilization as fraction of TOTAL GPU memory ──
    //
    // User-facing contract (matches vLLM / sparkrun convention):
    //   total_memory × gpu_memory_utilization = hard ceiling on everything
    //   this process consumes (weights + buffers + KV cache + reserves).
    //
    // KV cache gets whatever remains inside that ceiling after deducting
    // prior allocations (model weights, buffer arena, CUDA context/driver)
    // and the inference reserve (SSM state pools, CUDA headroom).  A safety
    // clamp ensures we never exceed what the device can physically provide
    // right now (handles external memory pressure on shared-memory /
    // unified-memory systems like GB10).
    let total_mem = gpu.total_memory()?;
    let actual_free = gpu.free_memory()?;
    let gib = |b: usize| b as f64 / (1024.0 * 1024.0 * 1024.0);
    let mut used_so_far = total_mem.saturating_sub(actual_free);
    // GB10 is shared (ComfyUI/voxel/etc.). Raw `used_so_far` counts those
    // co-tenants against our --gpu-memory-utilization budget, so a low util
    // needlessly starves the KV pool (vs vLLM, whose util is self-relative).
    //
    // We want the KV pool sized against Atlas's OWN footprint (weights +
    // buffers), excluding co-tenants. Two ways to find that footprint:
    //
    //   1. AUTO via LEDGER (default, preferred): the alloc ledger's live
    //      bytes — every allocation this backend made and hasn't freed
    //      (issue #740). The former free-memory delta (baseline-at-init
    //      minus free-now) counted OS page cache against us on unified
    //      memory: streaming ~20 GB of safetensors depresses MemFree
    //      without being an allocation Atlas owns, inflating "Atlas-own"
    //      by tens of GB on a cold-cache boot and refusing serves with
    //      >100 GB actually available. The ledger is immune to page-cache
    //      noise, co-tenant churn, and mid-load sampling by construction.
    //      Slight undercount (driver context, cuBLAS workspaces are not
    //      ledgered) is absorbed by the inference reserve and the physical
    //      `.min(actual_free - reserve)` clamp below.
    //
    //   2. AUTO via FREE-DELTA (fallback when the backend has no ledger):
    //      free-at-context-init minus free-now. Requires
    //      `set_baseline_free_bytes` to have run (it does under the real
    //      server; absent under the mock backend → we skip it).
    //
    //   3. MANUAL override: ATLAS_KV_EXTERNAL_RESERVE_GB=<co-tenant GB> still
    //      wins when explicitly set (>0), for operators who want to RESERVE
    //      headroom for co-tenants that will arrive LATER (the auto measures
    //      only see current state).
    //
    // The `.min(actual_free - reserve)` clamp below still guarantees a physical
    // fit regardless of which path set `used_so_far`.
    let manual_reserve_gb = std::env::var("ATLAS_KV_EXTERNAL_RESERVE_GB")
        .ok()
        .and_then(|v| v.parse::<f64>().ok())
        .filter(|&gb| gb > 0.0);
    if let Some(gb) = manual_reserve_gb {
        let ext = (gb * 1024.0 * 1024.0 * 1024.0) as usize;
        let discounted = used_so_far.saturating_sub(ext);
        tracing::info!(
            "ATLAS_KV_EXTERNAL_RESERVE_GB={gb} (manual override): discounting \
             external/co-tenant memory from KV budget — used_so_far {:.1} GB → \
             Atlas-own {:.1} GB",
            gib(used_so_far),
            gib(discounted),
        );
        used_so_far = discounted;
    } else if let Some(ledger_live) = gpu.live_bytes() {
        // AUTO via LEDGER: what this backend actually allocated and still
        // holds. Sanity-gate mirrors the free-delta path: the ledger can
        // only be a subset of total used (it can't see co-tenants), so a
        // value above `used_so_far` means the ledger and the device
        // disagree — fall through to raw rather than oversize the pool.
        if ledger_live > 0 && ledger_live <= used_so_far {
            tracing::info!(
                "KV budget self-relative (ledger): Atlas-own {:.1} GB live in \
                 the alloc ledger; {:.1} GB of co-tenant/page-cache use \
                 excluded (set ATLAS_KV_EXTERNAL_RESERVE_GB to override)",
                gib(ledger_live),
                gib(used_so_far - ledger_live),
            );
            used_so_far = ledger_live;
        } else if ledger_live > used_so_far {
            // Not "the ledger is implausible" — it is the DEVICE figure that
            // cannot be true. The ledger counts only allocations this backend
            // made and still holds, so real device usage is always at least
            // `ledger_live`; `used_so_far` (total − free_memory()) coming out
            // SMALLER means `free_memory()` over-reported free memory. That is
            // exactly what a discrete GPU did when `free_memory()` substituted
            // host MemAvailable: ~990 GB "free" on a 95 GB card → used_so_far
            // 0 → a KV pool sized as if nothing were allocated → OOM at
            // cuMemAlloc. Charging the ledger is strictly more conservative
            // than charging the smaller (impossible) device-derived figure:
            // it can only shrink the KV budget, never grow it.
            tracing::warn!(
                "KV budget: free_memory() looks wrong — the alloc ledger holds \
                 {:.1} GB but the device implies only {:.1} GB used, and real \
                 device usage can never be below the ledger. Charging the \
                 ledger's {:.1} GB (the larger, safer figure) instead.",
                gib(ledger_live),
                gib(used_so_far),
                gib(ledger_live),
            );
            used_so_far = ledger_live;
        } else {
            // ledger_live == 0: nothing ledgered yet (or a backend that does
            // not ledger). Nothing better to charge than the raw figure.
            tracing::warn!(
                "KV budget: alloc ledger reports 0 GB live against {:.1} GB \
                 used on the device — using raw used_so_far",
                gib(used_so_far),
            );
        }
    } else if let Some(baseline) = spark_runtime::gpu::baseline_free_bytes() {
        // AUTO: bytes this process consumed since context init.
        let atlas_own = baseline.saturating_sub(actual_free);
        // The free-delta above charges the weight loader's TRANSIENT
        // footprint — checkpoint mapping/staging still resident at this
        // instant — as if it were permanent. On a 27B NVFP4 load the delta
        // reads ~61 GB while the process's steady state is ~27 GB; the
        // difference is the size of the safetensors file, released once the
        // loader settles, and it is charged here regardless of page-cache
        // warmth. The footprint that actually PERSISTS is knowable without
        // heuristics:
        //   * `store.total_bytes()` — every weight tensor on the GPU;
        //   * entry-free − free-now — what build_model itself allocated
        //     (LoRA pool + buffer arena), transient-cancelling because the
        //     transient is present in both samples.
        // Charge the smaller of measured and known: the delta can only
        // OVER-count (transients), the known sum can only UNDER-count (CUDA
        // context overhead), and the `.min(actual_free − reserve)` clamp
        // below still guarantees a physical fit at allocation time either
        // way.
        let build_own = build_entry_free
            .map(|e| e.saturating_sub(actual_free))
            .unwrap_or(0);
        let known_own = store.total_bytes().saturating_add(build_own);
        let settled = atlas_own.min(known_own);
        // Sanity-gate: baseline must be ≥ free-now, the charge positive and
        // no larger than total used (co-tenants can't be negative). If a
        // co-tenant *freed* memory during our load, baseline > free-now
        // still holds and the charge just slightly overcounts (conservative
        // — fine). If the numbers are implausible, fall back to raw
        // used_so_far.
        if settled > 0 && settled <= used_so_far {
            tracing::info!(
                "KV budget self-relative (auto): baseline-free {:.1} GB − free-now \
                 {:.1} GB = {:.1} GB measured; charging settled Atlas-own {:.1} GB \
                 (weights {:.1} GB + build allocs {:.1} GB, loader transient \
                 {:.1} GB released from the charge); co-tenants {:.1} GB excluded \
                 (set ATLAS_KV_EXTERNAL_RESERVE_GB to override)",
                gib(baseline),
                gib(actual_free),
                gib(atlas_own),
                gib(settled),
                gib(store.total_bytes()),
                gib(build_own),
                gib(atlas_own.saturating_sub(settled)),
                gib(used_so_far - atlas_own.min(used_so_far)),
            );
            used_so_far = settled;
        } else {
            tracing::warn!(
                "KV budget auto-measure implausible (baseline {:.1} GB, free-now \
                 {:.1} GB, used {:.1} GB) — using raw used_so_far",
                gib(baseline),
                gib(actual_free),
                gib(used_so_far),
            );
        }
    }
    // DFlash drafter head allocations happen at Step 7 — AFTER this sizing —
    // so without a reserve they land OUTSIDE the util pledge (the documented
    // dflash-oom hazard; 2026-08-19 256K/C8 boot ledger measured ~10.5 GB of
    // post-sizing drafter allocs on a boot whose planner believed it had
    // honored a 79 GB budget, leaving 14 GB on the whole box before the first
    // request). Estimate mirrors serve's load_dflash_drafter pre-flight:
    // drafter KV (max_seq_len·L·2·kv_dim·bf16) + fused_kv + prompt-hidden
    // capture + FP8 MLP mirrors (~store/2; the lm_head mirror is shared) +
    // scratch.
    let dflash_reserve: usize = dflash_args
        .as_ref()
        .map(|a| {
            let c = &a.drafter_config;
            let kv_dim = c.num_key_value_heads * c.head_dim;
            let drafter_kv = max_seq_len * c.num_hidden_layers * 2 * kv_dim * 2;
            let fused_kv = c.num_hidden_layers * 2 * kv_dim * c.hidden_size * 2;
            let capture = max_seq_len * config.hidden_size * 2;
            // Same predicate as the allocating gate (`!= Some("0")`).
            // FP8 drafter weights are default-ON, so `.is_some()` made the
            // KV budget under-reserve by the mirror size on the default path.
            let fp8_mirrors =
                if std::env::var("ATLAS_DFLASH_DRAFTER_FP8").ok().as_deref() != Some("0") {
                    a.drafter_store.total_bytes() / 2
                } else {
                    0
                };
            drafter_kv + fused_kv + capture + fp8_mirrors + (300 << 20)
        })
        .unwrap_or(0);
    if dflash_reserve > 0 {
        tracing::info!(
            "KV budget: reserving {:.1} GB for post-sizing DFlash drafter allocations",
            gib(dflash_reserve),
        );
    }
    let total_budget = (total_mem as f64 * gpu_memory_utilization) as usize;
    let kv_budget = total_budget
        .saturating_sub(used_so_far)
        .saturating_sub(inference_reserve)
        .saturating_sub(dflash_reserve)
        .min(
            actual_free
                .saturating_sub(inference_reserve)
                .saturating_sub(dflash_reserve),
        );
    // ── MTP propose-pool pre-charge ──
    // `MtpHead::new` allocates its own paged KV pool AFTER this sizing
    // (per-seq blocks × the MTP concurrency cap, bounded by the MAIN pool's
    // block count) and its comment's "well inside the serve reserve" named a
    // reserve that never existed — ~0.97 GB at 128K/bs8 landed OUTSIDE the
    // util pledge (the last tracked allocation that did, 2026-08-22 ledger).
    // Mirror the pool arithmetic here and charge it. Two-pass on the cap:
    // the bound uses the PRE-charge block count, which is >= the final one,
    // so the miss direction is a slightly larger reserve, never a smaller
    // pool than reserved. Gate matches `build_mtp_proposer` minus the
    // LM-head-dtype refusal — if that refusal fires the head is skipped and
    // this over-reserves one pool, which is the safe direction.
    let mtp_pool_reserve: usize = if use_speculative && !mtp_weights.is_empty() {
        // Mirrors the head's kv_config: block 16, target attention dims,
        // K+V, BF16 KV for Bf16/Fp8 heads and FP8 KV for NVFP4
        // (`kv_bf16` in mtp_head/new.rs).
        let block = 16usize;
        let per_seq_blocks = max_seq_len / block + 1;
        let elem = match mtp_quant {
            MtpQuantization::Nvfp4 => 1usize,
            MtpQuantization::Fp8 | MtpQuantization::Bf16 => 2,
        };
        let mtp_block_bytes = block * config.num_key_value_heads * config.head_dim * elem * 2;
        let blocks0 = PagedKvCache::compute_num_blocks(&kv_config, kv_budget).unwrap_or(0);
        let pool_blocks = per_seq_blocks
            .saturating_mul(crate::speculative::mtp_max_seqs())
            .min(blocks0.max(per_seq_blocks));
        pool_blocks * mtp_block_bytes
    } else {
        0
    };
    let kv_budget = kv_budget.saturating_sub(mtp_pool_reserve);
    if mtp_pool_reserve > 0 {
        tracing::info!(
            "KV budget: reserving {:.1} GB for the MTP propose pool (post-sizing alloc in MtpHead::new)",
            gib(mtp_pool_reserve),
        );
    }
    // Phase 6.1.f: when HBM-shrink is active, size the production cache to
    // `max_batch_size × cache_blocks_per_seq` rather than the unbounded
    // budget-driven sum. This is the *whole point* of the HBM-shrink
    // feature — the production cache becomes write staging only; older
    // blocks live on disk under the orchestrator's control.
    let num_kv_blocks = match hss_cache_blocks_per_seq {
        Some(cap) => {
            // Phase 6.3 (original): pool = max_batch × cap + 1 dummy + 1 spare per seq.
            // Issue #31 (2026-05-08): the cap×bs sizing assumed prefill would
            // fit in cap blocks AND the slide-during-prefill path would handle
            // any overflow. Live-tested: slides during prefill produce silently
            // wrong attention output (the orchestrator-fed disk-read path is
            // wired up for DECODE attention only — Phase 6.2.a — not for
            // prefill — Phase 6.2.b deferred). The companion change in
            // `block_mgmt::ensure_blocks_through_prefill` removes the broken
            // slide; this change resizes the pool so prefill can grow up to
            // `max_seq_len` blocks without hitting "no free blocks". HBM-shrink
            // remains in effect post-prefill: the FIRST decode step finds
            // bt_len > cap and slides down via the orchestrator-aware path
            // (which IS correct).
            //
            // Sizing rationale:
            //   * Per-seq blocks: `max(cap + 1, ceil(max_seq_len / block_size))`
            //     so prefill of any prompt up to max_seq_len fits in HBM.
            //   * +1 dummy slot for OOB-safe paged-kernel reads.
            //
            // For multi-seq HSS where the user wanted strict HBM-shrink, this
            // increases pool size by `(max_seq_len_blocks - cap) × max_batch`
            // bytes per block. The existing post-load OOM check (line 304+)
            // catches infeasible configs at startup with a clear message.
            let max_seq_blocks = max_seq_len.div_ceil(kv_block_size);
            let per_seq = (cap as usize + 1).max(max_seq_blocks);
            let n = max_batch_size * per_seq + 1;
            tracing::info!(
                "--high-speed-swap: HBM cache sized to {n} blocks ({} batch × max(cap={cap}+1, max_seq_len_blocks={max_seq_blocks}) + 1 dummy); \
                 prefill grows monotonically, decode shrinks to cap × bs and streams older blocks from disk via the orchestrator",
                max_batch_size
            );
            n
        }
        None => {
            if kv_budget == 0 {
                anyhow::bail!(
                    "No memory left for KV cache: total GPU = {:.1} GB, \
                     --gpu-memory-utilization {:.0}% → budget {:.1} GB, \
                     but {:.1} GB already consumed + {:.1} GB inference reserve \
                     = {:.1} GB committed.  Raise --gpu-memory-utilization or \
                     use a smaller model.",
                    total_mem as f64 / (1024.0 * 1024.0 * 1024.0),
                    gpu_memory_utilization * 100.0,
                    total_budget as f64 / (1024.0 * 1024.0 * 1024.0),
                    used_so_far as f64 / (1024.0 * 1024.0 * 1024.0),
                    inference_reserve as f64 / (1024.0 * 1024.0 * 1024.0),
                    (used_so_far + inference_reserve) as f64 / (1024.0 * 1024.0 * 1024.0),
                );
            }
            let budget_blocks = PagedKvCache::compute_num_blocks(&kv_config, kv_budget)?;
            // ── Clamp the pool to blocks the engine can actually reach ──
            //
            // `compute_num_blocks` spends the ENTIRE residual budget, and
            // nothing downstream caps it: the `max_concurrent` check below is a
            // warn/bail only, never a cap. So the pool is sized by "what is
            // left over", not by "what can be addressed".
            //
            // Measured on GLM-5.3, 2xGB10, `--max-seq-len 2048
            // --max-batch-size 1`, prefix caching off: 45,386 blocks = 7.6 GiB
            // = 726,176 KV tokens, against a reachable ceiling of
            // `1 x ceil(2048/16) = 128` blocks. ~99.7 % of the pool could never
            // be addressed by any request.
            //
            // On a discrete GPU that waste is merely idle VRAM. On unified
            // memory it is host RAM taken from the kernel, the page cache and
            // every co-tenant — and it is the reason correcting an
            // over-reservation elsewhere frees nothing: `kv_budget` is a
            // residual, so every byte released by a smaller `inference_reserve`
            // is immediately re-absorbed here. This clamp is what turns a
            // reserve correction into recovered headroom.
            //
            // Only applied when the prefix cache is INACTIVE. An active cache
            // makes surplus blocks genuinely reachable (they hold shared
            // prefixes), which is exactly the case the unbounded sizing was
            // written for. `+ max_batch_size + 1` mirrors the HBM-shrink arm
            // above: one spare block per sequence plus the dummy slot the
            // OOB-safe paged kernels read.
            //
            // Kill switch: `ATLAS_KV_POOL_UNCLAMPED` (presence — `=0` is NOT
            // "off") restores the budget-driven pool.
            let n = if prefix_cache.is_active() || std::env::var("ATLAS_KV_POOL_UNCLAMPED").is_ok()
            {
                budget_blocks
            } else {
                let per_seq = max_seq_len.div_ceil(kv_block_size);
                let reachable = max_batch_size
                    .saturating_mul(per_seq)
                    .saturating_add(max_batch_size)
                    .saturating_add(1);
                let clamped = budget_blocks.min(reachable);
                if clamped < budget_blocks {
                    let freed = (budget_blocks - clamped) * kv_config.block_bytes_kv_all_layers();
                    tracing::info!(
                        "KV pool clamped to reachable demand: {} -> {} blocks \
                         ({} seq x {} blocks/seq + {} spare + 1 dummy); \
                         {:.2} GB not allocated (prefix caching inactive, so surplus \
                         blocks are unreachable). Restore with --enable-prefix-caching \
                         or ATLAS_KV_POOL_UNCLAMPED.",
                        budget_blocks,
                        clamped,
                        max_batch_size,
                        per_seq,
                        max_batch_size,
                        freed as f64 / (1024.0 * 1024.0 * 1024.0),
                    );
                }
                clamped
            };
            let max_kv_tokens = n * kv_block_size;
            tracing::info!(
                "KV cache: {:.1} GB total × {:.0}% util = {:.1} GB budget; \
                 {:.1} GB pre-KV + {:.1} GB reserve → {:.1} GB for KV \
                 → {} blocks × {} tok/block = {} max KV tokens",
                total_mem as f64 / (1024.0 * 1024.0 * 1024.0),
                gpu_memory_utilization * 100.0,
                total_budget as f64 / (1024.0 * 1024.0 * 1024.0),
                used_so_far as f64 / (1024.0 * 1024.0 * 1024.0),
                inference_reserve as f64 / (1024.0 * 1024.0 * 1024.0),
                kv_budget as f64 / (1024.0 * 1024.0 * 1024.0),
                n,
                kv_block_size,
                max_kv_tokens,
            );
            n
        }
    };
    let _max_kv_tokens = num_kv_blocks * kv_block_size;
    // Phase 6.1.f / 6.2.c — when --high-speed-swap is on with HBM-shrink, the
    // production KV cache only has to fit the per-seq HBM window, not the full
    // sequence (older blocks live on disk). Compare against `cache_blocks_per_seq`
    // in that mode; the legacy "blocks per max_seq_len" check is invalid for
    // HBM-shrunk pools by design.
    let blocks_per_seq = match hss_cache_blocks_per_seq {
        Some(cap) => cap as usize,
        None => max_seq_len.div_ceil(kv_block_size),
    };
    let max_concurrent = num_kv_blocks / blocks_per_seq.max(1);
    if max_concurrent < max_batch_size {
        // Suggest a max_seq_len that lets the requested batch size fit.
        let suggested_max_seq_len = (num_kv_blocks / max_batch_size.max(1)) * kv_block_size;
        // The check is WORST-CASE: it assumes every concurrent sequence reaches
        // --max-seq-len. With paged KV (blocks allocated on demand) that almost
        // never holds for real agent traffic (mixed/shorter sequences), so a high
        // --max-seq-len (e.g. 64K for long agent contexts) needlessly caps
        // --max-batch-size. Overcommit (DEFAULT ON since wave 10: config of
        // record for the native bs=32 rung) downgrades the hard error to a
        // warning: the scheduler admits up to max_batch_size and the pool fills on
        // demand (a genuinely over-long burst gets back-pressured by the block
        // allocator, not a boot-time refusal). Kill switch: ATLAS_KV_OVERCOMMIT=0
        // (or =false) restores the boot-time hard refusal. Value is parsed, not
        // presence-checked.
        let overcommit = !matches!(
            std::env::var("ATLAS_KV_OVERCOMMIT").as_deref(),
            Ok("0") | Ok("false")
        );
        if overcommit {
            tracing::warn!(
                "KV OVERCOMMIT: pool fits {} seq(s) at full --max-seq-len={} but \
                 --max-batch-size={} requested ({} block(s)/seq, {} block(s) total). \
                 Paged KV allocates on demand; long-context bursts are back-pressured \
                 at the block allocator, not refused at boot.",
                max_concurrent,
                max_seq_len,
                max_batch_size,
                blocks_per_seq,
                num_kv_blocks,
            );
        } else {
            anyhow::bail!(
                "KV cache can hold at most {} concurrent sequence(s) at --max-seq-len={}, \
                 but --max-batch-size={} was requested. \
                 KV pool has {} block(s) of {} tokens each; each sequence needs {} block(s). \
                 Try --max-seq-len {} (keeps max_batch_size={}), reduce --max-batch-size, \
                 or unset ATLAS_KV_OVERCOMMIT=0 to allow on-demand paged allocation (default).",
                max_concurrent,
                max_seq_len,
                max_batch_size,
                num_kv_blocks,
                kv_block_size,
                blocks_per_seq,
                suggested_max_seq_len.max(kv_block_size),
                max_batch_size,
            );
        }
    }
    let kv_cache = PagedKvCache::new(kv_config, num_kv_blocks, gpu.as_ref())?;

    // ── Step 6: Assemble model ──
    // Capture pointers for any post-construction sharing (DFlash drafter
    // shares embed_tokens + lm_head with the target). DenseWeight is Copy
    // so this clones the device pointer cheaply.
    let target_embed_for_dflash = embed.weight;
    let target_lm_head_for_dflash = lm_head.weight;
    // NVFP4 lm_head (Copy) shared with the DFlash drafter so its final logits
    // GEMM uses w4a16 instead of a BF16 dense_gemm on NVFP4-packed bytes.
    let target_lm_head_nvfp4_for_dflash = lm_head_nvfp4;
    let target_hidden_for_dflash = config.hidden_size;
    // Native FP8 lm_head share for the DFlash drafter tail: when the
    // checkpoint ships lm_head as FP8 E4M3 + per-row scale, the drafter's
    // Phase-G tail reads THOSE bytes instead of building a 1.27 GB
    // runtime-requantized mirror (see lm_head_setup::native_fp8_lm_head_share).
    // Built here because `store` is dropped into the model right after.
    let target_lm_head_native_fp8_for_dflash = if dflash_args.is_some() {
        super::lm_head_setup::native_fp8_lm_head_share(&store, &config, gpu.as_ref())?
    } else {
        None
    };

    let mut model = TransformerModel::new(
        config,
        embed,
        final_norm,
        lm_head,
        lm_head_nvfp4,
        lm_head_fp8,
        mtp_lm_head_nvfp4,
        layers,
        buffers,
        kv_cache,
        mtp_weights,
        gpu,
        max_seq_len,
        max_batch_size,
        effective_mtp_quant,
        use_speculative,
        prefix_cache,
        mtp_vocab_size,
        comm,
        self_speculative,
        num_drafts,
        vision_encoder,
        ssm_cache_slots,
        ssm_checkpoint_interval,
    )?;

    // ── Step 6b: DeepSeek-V4 MTP proposer (optional, post-construction) ──
    //
    // Built here (not inside `new()`, which only knows the Qwen-shaped
    // `MtpWeights`) because it needs the model's owned GPU backend, the
    // resolved draft NVFP4 head, and the shared embedding. Installed via the
    // existing proposer setter. DFlash (below) is CLI-exclusive with
    // `--speculative`, so the two never both install.
    if let Some(v4_module) = v4_mtp_module {
        match crate::layers::DeepseekV4MtpHead::new(
            v4_module,
            v4_mtp_embed,
            v4_mtp_lm_head,
            model.config_ref(),
            model.gpu_backend(),
            mtp_vocab_size,
            max_seq_len,
        ) {
            Ok(head) => {
                model.set_dflash_proposer(std::sync::Arc::new(head));
                tracing::info!("DeepSeek-V4 MTP speculative decoding: ENABLED (single-module)");
            }
            Err(e) => tracing::warn!(
                "Failed to build DeepSeek-V4 MTP proposer: {e:#}. Speculative decoding disabled."
            ),
        }
    }

    // ── Step 6c: GLM-5.3 MTP proposer (optional, post-construction) ──
    //
    // Built here for the same reason as the V4 head: it needs the model's owned GPU backend and
    // the shared embedding + LM head, neither of which `new()` hands out.
    if let Some(m) = glm_mtp_module {
        match crate::layers::Glm5NextMtpHead::new(
            m,
            glm_mtp_embed,
            glm_mtp_lm_head,
            model.config_ref(),
            model.gpu_backend(),
            max_seq_len,
        ) {
            Ok(head) => {
                model.set_dflash_proposer(std::sync::Arc::new(head));
                tracing::info!("GLM-5.3 MTP speculative decoding: ENABLED");
            }
            Err(e) => tracing::warn!(
                "Failed to build GLM-5.3 MTP proposer: {e:#}. Speculative decoding disabled."
            ),
        }
    }

    // ── Step 7: DFlash drafter (optional, post-construction) ──
    //
    // Loaded last because it depends on the target's `embed_tokens` and
    // `lm_head` pointers (the drafter checkpoint omits these — they're
    // shared at runtime, mirroring vLLM PR #40898's `skip_substrs` flow).
    if let Some(args) = dflash_args {
        let weights = load_dflash_weights(
            args.drafter_store,
            &args.drafter_config,
            model.gpu_backend(),
            1, // tp_size for the drafter side: replicated, so always 1
        )?;
        if let Some(weights) = weights {
            let head = crate::layers::BlockDiffusionDraftHead::from_weights(
                weights,
                target_embed_for_dflash,
                target_lm_head_for_dflash,
                target_lm_head_nvfp4_for_dflash,
                target_lm_head_native_fp8_for_dflash,
                target_hidden_for_dflash,
                args.gamma,
                args.window_size,
                model.gpu_backend(),
                max_seq_len,
                max_batch_size,
            )?;
            model.set_dflash_proposer(std::sync::Arc::new(head));
            tracing::info!("DFlash drafter installed as the active proposer");
        } else {
            tracing::warn!(
                "DFlash drafter store had no fc.weight — proposer not installed; \
                 falling back to whatever proposer (if any) the target's MTP path built"
            );
        }
    }

    // ── Step 8: LoRA adapter install (optional, post-construction) ──
    // The pool/tables were loaded up top (pre-KV-sizing); this walk copies
    // the per-layer pairs into the layer structs. M0: layers only STORE the
    // adapter — base output is unchanged until the M1 compute insertions.
    if let Some(ngram) = ngram_embed {
        model.set_ngram_embedding(ngram);
    }
    model.set_lora_weights(lora_weights)?;

    // Every layer has taken the pointers it needs; hand the ledger to the model
    // so `teardown` can free the weights. Dropping it here — which is what used
    // to happen — orphaned the memory: live, referenced by the layers, with
    // nothing owning the ability to release it.
    model.adopt_weight_store(store);

    // ── Allocation attribution, once, at the end of load ──
    // Everything the serve will hold is allocated by now: weights, KV, the SSM
    // pools, the Marconi snapshots, the arenas, the vision encoder. This is
    // the only point where the ledger describes the STEADY STATE rather than
    // some midpoint of the build.
    //
    // It exists because the KV-budget line above reports `pre-KV` as one
    // opaque number, and on a 27B that number is ~59 GB against 22 GB of
    // weights — the rest was unattributable until the ledger learned sizes.
    // Logged at INFO, not behind a flag: every OOM and every mis-sized-pool
    // investigation so far has begun by wanting exactly this table, and a
    // once-per-load table is not a cost worth flagging off.
    if let Some(report) = model.gpu_backend().alloc_report(12, 64) {
        for line in report.lines() {
            tracing::info!("{line}");
        }
    }
    // ── Pledge reconciliation, same place, same reason ──
    // The table above attributes the spend; this line judges it against the
    // promise. Tracked-live > budget means some allocation family was never
    // modeled by the preflight/KV sizing (the 2026-08-22 case: the DFlash
    // verify pools, 13.7 GB against a 1.3 GB reserve) — capacity the
    // scheduler will happily promise to requests it cannot actually fund.
    // WARN, not error: the serve is already up, and the operator's fix is a
    // sizing/reserve change, not a restart loop.
    if let Some(live) = model.gpu_backend().live_bytes() {
        if live > total_budget {
            tracing::warn!(
                "util pledge exceeded: {:.1} GB tracked live vs {:.1} GB pledged \
                 (--gpu-memory-utilization {:.0}% of {:.1} GB) — an allocation \
                 family above is missing from the preflight reserve",
                gib(live),
                gib(total_budget),
                gpu_memory_utilization * 100.0,
                gib(total_mem),
            );
        } else {
            tracing::info!(
                "util pledge honored: {:.1} GB tracked live within the {:.1} GB \
                 budget ({:.1} GB pledge headroom)",
                gib(live),
                gib(total_budget),
                gib(total_budget - live),
            );
        }
    }
    Ok(Box::new(model))
}

/// Per-step GPU residency ledger for model construction.
///
/// Every `load_*` step below allocates into the same GB10 unified pool the KV
/// cache is later sized from, but until now the only numbers in the log were
/// the loader's on-disk estimate ("Weights: 99.64 GB" — which is
/// `estimate_load_bytes`, an ON-DISK byte sum of the tensors this rank reads,
/// NOT residency) and one aggregate free-memory reading. Anything between them
/// — binder re-uploads, dtype conversions, the store originals a loader forgot
/// to drop — was unattributable, and a 4-5 GB residual is the difference
/// between K=3 fitting and not.
///
/// This walks `gpu.free_memory()` across the build and prints a signed delta
/// per step. Log-only: it allocates nothing and changes no semantics.
struct MemTrace<'a> {
    gpu: &'a dyn GpuBackend,
    last: usize,
    start: usize,
}

impl<'a> MemTrace<'a> {
    fn new(gpu: &'a dyn GpuBackend) -> Self {
        let f = gpu.free_memory().unwrap_or(0);
        Self {
            gpu,
            last: f,
            start: f,
        }
    }

    /// Log the free-memory delta since the previous mark. Negative = allocated.
    fn mark(&mut self, step: &str) {
        let Ok(now) = self.gpu.free_memory() else {
            return;
        };
        let gib = |b: usize| b as f64 / (1024.0 * 1024.0 * 1024.0);
        let delta = now as i128 - self.last as i128;
        tracing::info!(
            "build residency: {step:<26} {:+9.3} GiB   (cumulative {:8.3} GiB, free {:7.3} GiB)",
            delta as f64 / (1024.0 * 1024.0 * 1024.0),
            gib(self.start) - gib(now),
            gib(now),
        );
        self.last = now;
    }
}