metrale_model_engine/model/
impl_b3_accessors.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Post-construction wiring for [`TransformerModel`]: borrow the GPU and
4//! config, install a proposer or an n-gram embedding, and the MLA prefix-skip check.
5//!
6//! Owner: model-engine.
7//! Invariants:
8//! - `set_dflash_proposer` never grows the MTP prefill capture buffer. If shrinking it
9//!   fails, the buffer is left NULL with capacity 0.
10
11use metrale_config::ModelConfig;
12use metrale_gpu_runtime::gpu::GpuBackend;
13
14use super::types::TransformerModel;
15use metrale_model_layers::speculative::DraftProposer;
16
17impl TransformerModel {
18    /// 2026-09-25: Borrow the GPU backend. The factory uses it to build the
19    /// proposers installed after construction and to read allocation reports.
20    pub fn gpu_backend(&self) -> &dyn GpuBackend {
21        self.gpu.as_ref()
22    }
23
24    /// 2026-09-25: Borrow the model config, for building the proposers that
25    /// are installed after construction.
26    pub fn config_ref(&self) -> &ModelConfig {
27        &self.config
28    }
29
30    /// 2026-09-25: Install a proposer built after construction (the
31    /// DeepSeek-V4 or GLM-5.3 MTP head, or a DFlash drafter), replacing any
32    /// MTP proposer `TransformerModel::new` built. It also shrinks the MTP
33    /// prefill capture buffer to the rows this proposer can use.
34    ///
35    /// It does not check flag combinations; clap rejects `--dflash` with
36    /// `--speculative`.
37    pub fn set_dflash_proposer(&mut self, proposer: std::sync::Arc<dyn DraftProposer>) {
38        if self.proposer.is_some() {
39            tracing::info!("DFlash: replacing existing MTP proposer with BlockDiffusionDraftHead");
40        }
41        // 2026-09-25: `new()` sized `mtp_prefill_hidden` before this proposer
42        // existed. Ask the proposer how many rows it can be handed
43        // (`DraftProposer::prefill_hidden_rows`; the default keeps the size) and
44        // shrink to that; never grow. The old buffer is freed before the smaller
45        // one is allocated, so the two are never held at once.
46        let rows = proposer.prefill_hidden_rows(self.mtp_prefill_capacity);
47        if !self.mtp_prefill_hidden.is_null() && rows < self.mtp_prefill_capacity {
48            let was = self.mtp_prefill_capacity;
49            let bytes = rows * self.config.hidden_size * 2;
50            let old = std::mem::replace(
51                &mut self.mtp_prefill_hidden,
52                metrale_gpu_runtime::gpu::DevicePtr::NULL,
53            );
54            self.mtp_prefill_capacity = 0;
55            match self.gpu.free(old).and_then(|_| self.gpu.alloc(bytes)) {
56                Ok(smaller) => {
57                    self.mtp_prefill_hidden = smaller;
58                    self.mtp_prefill_capacity = rows;
59                    tracing::info!(
60                        "MTP drafter context: capture buffer rightsized {was} -> {rows} rows \
61                         ({:.0} -> {:.0} MB) — the proposer cannot be handed a position past \
62                         {rows} (A59)",
63                        (was * self.config.hidden_size * 2) as f64 / 1e6,
64                        bytes as f64 / 1e6,
65                    );
66                }
67                // 2026-09-25: NULL with capacity 0 is the feature's off state: the
68                // capture epilogue and the propose site both check it, so drafter
69                // prefill is disabled and the serve keeps running.
70                Err(e) => tracing::warn!(
71                    "MTP drafter context: rightsizing the capture buffer failed ({e:#}) — \
72                     drafter prefill and carry are DISABLED for this serve"
73                ),
74            }
75        }
76        self.proposer = Some(proposer);
77    }
78
79    /// 2026-09-25: Install the fused n-gram input embedding (only the LongCat
80    /// loader builds one). Once set, `embed_ctx` embeds through it and the
81    /// bare-token `embed` refuses to run.
82    pub fn set_ngram_embedding(
83        &mut self,
84        ngram: metrale_model_layers::layers::ngram_embed::NgramEmbedding,
85    ) {
86        tracing::info!("set_ngram_embedding: installed on the served model");
87        self.ngram_embed = Some(std::sync::Mutex::new(ngram));
88    }
89
90    /// 2026-09-25: True when this model fuses n-gram lookups into its input embedding.
91    pub fn has_ngram_embedding(&self) -> bool {
92        self.ngram_embed.is_some()
93    }
94
95    /// 2026-09-25: True when prefill must not skip a cached prefix because a
96    /// layer's local MLA prefill (`uses_local_mla_prefill`) attends only over
97    /// the tokens it processes: the skipped prefix would be absent from
98    /// attention. The prefix lookup then uses an empty match and prefill
99    /// recomputes the whole prompt. `METRALE_MLA_PREFIX_SKIP=1` makes it
100    /// return false.
101    pub(crate) fn mla_prefill_needs_full_recompute(&self) -> bool {
102        if std::env::var("METRALE_MLA_PREFIX_SKIP").as_deref() == Ok("1") {
103            return false;
104        }
105        self.layers.iter().any(|l| l.uses_local_mla_prefill())
106    }
107}
metrale_model_engine/model/impl_b3_accessors.rs

metrale_model_engine/model/
impl_b3_accessors.rs
