metrale_model_layers/layers/mtp_head/
new.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `MtpHead::new`: quantizes the head's projections, builds its
4//! FFN (dense, NVFP4 MoE, native-FP8 MoE or per-expert MoE), and allocates
5//! its KV pool, kernel handles and scratch.
6//!
7//! Owner: model-layers (MTP head).
8//! Invariants:
9//! - A head has one FFN form: the MoE parts are built only when the weights
10//!   carry no dense FFN.
11//! - The drafter KV pool and the `propose_meta` stride use the same block
12//!   size.
13
14use anyhow::Result;
15use metrale_cache::kv_cache::{KvCacheConfig, KvCacheDtype, PagedKvCache};
16use metrale_gpu_runtime::gpu::GpuBackend;
17use parking_lot::Mutex;
18
19use super::{MtpHead, MtpQuantization, ProjectionWeight};
20use crate::layers::MoeLayer;
21use crate::weight_map::{DenseWeight, MoeWeights, MtpWeights, QuantizedWeight, quantize_to_nvfp4};
22
23impl MtpHead {
24    pub fn new(
25        weights: MtpWeights,
26        embed_tokens: DenseWeight,
27        target_final_norm: DenseWeight,
28        lm_head_nvfp4: QuantizedWeight,
29        // 2026-09-25: Transposed twin of the main LM head; the caller passes
30        // `None` when the drafter has a dedicated draft head.
31        lm_head_nvfp4_t: Option<(QuantizedWeight, u32)>,
32        config: &metrale_config::ModelConfig,
33        gpu: &dyn GpuBackend,
34        quant: MtpQuantization,
35        mtp_vocab_size: u32,
36        max_seq_len: usize,
37        main_kv_blocks: usize,
38        levers: &crate::layers::ops::ModelLevers,
39    ) -> Result<Self> {
40        let stream = gpu.default_stream();
41        // 2026-09-25: A dense-FFN head under `--mtp-quantization nvfp4` runs
42        // the BF16 forward (`effective_for_head`): q/o and the dense
43        // gate/up/down are stored NVFP4 and run W4A16, while fc/k/v stay BF16,
44        // which the drafter prefill requires.
45        let wquant = quant;
46        let quant = wquant.effective_for_head(weights.dense_ffn.is_some());
47        let dense_nvfp4 = quant != wquant;
48        let absmax_k = gpu.kernel("quantize_nvfp4", "nvfp4_global_absmax")?;
49        let nvfp4_k = gpu.kernel("quantize_nvfp4", "quantize_bf16_to_nvfp4")?;
50        let fp8_k = gpu.kernel("gemv_fp8w", "quantize_bf16_to_fp8")?;
51
52        let h = config.hidden_size;
53        let nq = config.num_attention_heads;
54        let nkv = config.num_key_value_heads;
55        let hd = config.head_dim;
56        let inter = if config.moe_intermediate_size > 0 {
57            config.moe_intermediate_size
58        } else {
59            config.intermediate_size
60        };
61
62        let q = |bf16: &DenseWeight, n: usize, k: usize| -> Result<ProjectionWeight> {
63            Self::quantize_proj(bf16, n, k, quant, gpu, absmax_k, nvfp4_k, fp8_k, stream)
64        };
65        // 2026-09-25: The requested precision, for the weight-only NVFP4
66        // projections of a dense head; identical to `q` for every other head.
67        let qw = |bf16: &DenseWeight, n: usize, k: usize| -> Result<ProjectionWeight> {
68            Self::quantize_proj(bf16, n, k, wquant, gpu, absmax_k, nvfp4_k, fp8_k, stream)
69        };
70
71        let fc = q(&weights.fc, h, h * 2)?;
72        let q_proj = qw(&weights.q_proj, nq * hd * 2, h)?;
73        let k_proj = q(&weights.k_proj, nkv * hd, h)?;
74        let v_proj = q(&weights.v_proj, nkv * hd, h)?;
75        let o_proj = qw(&weights.o_proj, h, nq * hd)?;
76
77        let dense_ffn_generic = if let Some(dense_ffn) = weights.dense_ffn.as_ref() {
78            Some((
79                qw(&dense_ffn.gate_proj, inter, h)?,
80                qw(&dense_ffn.up_proj, inter, h)?,
81                qw(&dense_ffn.down_proj, h, inter)?,
82            ))
83        } else {
84            None
85        };
86
87        // 2026-09-25: MoE: an NVFP4 head gets a fused `MoeLayer`; an FP8 or
88        // BF16 head gets a `MoeLayer` on the checkpoint's FP8 tables when the
89        // weights carry them (`moe_grouped`), else per-expert weights.
90        let mut weights = weights;
91        // 2026-10-02: `--mtp-experts-nvfp4` applies to one head form only: a BF16 MoE head on
92        // BF16 experts. Asked for on any other head, the serve refuses rather than ignore it.
93        let draft_nvfp4 = super::mtp_experts_nvfp4();
94        anyhow::ensure!(
95            !draft_nvfp4
96                || (dense_ffn_generic.is_none()
97                    && quant == MtpQuantization::Bf16
98                    && weights.fp8_experts.is_none()),
99            "--mtp-experts-nvfp4 needs a BF16 MoE MTP head on BF16 experts (this head: \
100             {quant:?}, dense FFN {}, FP8 experts {})",
101            dense_ffn_generic.is_some(),
102            weights.fp8_experts.is_some()
103        );
104        let moe_parts = if dense_ffn_generic.is_some() {
105            (None, None, None, None)
106        } else {
107            match quant {
108                MtpQuantization::Nvfp4 => {
109                    let gate_nvfp4 = quantize_to_nvfp4(
110                        &weights.moe_gate,
111                        config.num_experts,
112                        h,
113                        gpu,
114                        absmax_k,
115                        nvfp4_k,
116                        stream,
117                    )?;
118                    let mut experts = Vec::with_capacity(weights.experts.len());
119                    for (i, de) in weights.experts.iter().enumerate() {
120                        let gate_proj = quantize_to_nvfp4(
121                            &de.gate_proj,
122                            inter,
123                            h,
124                            gpu,
125                            absmax_k,
126                            nvfp4_k,
127                            stream,
128                        )?;
129                        let up_proj = quantize_to_nvfp4(
130                            &de.up_proj,
131                            inter,
132                            h,
133                            gpu,
134                            absmax_k,
135                            nvfp4_k,
136                            stream,
137                        )?;
138                        let down_proj = quantize_to_nvfp4(
139                            &de.down_proj,
140                            h,
141                            inter,
142                            gpu,
143                            absmax_k,
144                            nvfp4_k,
145                            stream,
146                        )?;
147                        experts.push(crate::weight_map::ExpertWeight {
148                            gate_proj,
149                            up_proj,
150                            down_proj,
151                        });
152                        if (i + 1) % 128 == 0 {
153                            tracing::info!(
154                                "  MTP experts quantized: {}/{}",
155                                i + 1,
156                                weights.experts.len()
157                            );
158                        }
159                    }
160                    let shared_gate = quantize_to_nvfp4(
161                        &weights.shared_expert.gate_proj,
162                        inter,
163                        h,
164                        gpu,
165                        absmax_k,
166                        nvfp4_k,
167                        stream,
168                    )?;
169                    let shared_up = quantize_to_nvfp4(
170                        &weights.shared_expert.up_proj,
171                        inter,
172                        h,
173                        gpu,
174                        absmax_k,
175                        nvfp4_k,
176                        stream,
177                    )?;
178                    let shared_down = quantize_to_nvfp4(
179                        &weights.shared_expert.down_proj,
180                        h,
181                        inter,
182                        gpu,
183                        absmax_k,
184                        nvfp4_k,
185                        stream,
186                    )?;
187                    let moe_weights = MoeWeights {
188                        gate: weights.moe_gate,
189                        shared_expert: crate::weight_map::ExpertWeight {
190                            gate_proj: shared_gate,
191                            up_proj: shared_up,
192                            down_proj: shared_down,
193                        },
194                        shared_expert_gate: weights.shared_expert_gate,
195                        experts,
196                        router_pre_norm: None,
197                        correction_bias: None,
198                    };
199                    let moe = MoeLayer::new(
200                        moe_weights,
201                        config.num_experts,
202                        Some(gate_nvfp4),
203                        gpu,
204                        config,
205                    )?;
206                    (Some(moe), None, None, None)
207                }
208                MtpQuantization::Fp8 | MtpQuantization::Bf16 if weights.fp8_experts.is_some() => {
209                    let moe = Self::new_native_fp8_moe(&mut weights, config, gpu)?;
210                    (None, None, None, Some(moe))
211                }
212                // 2026-10-02: A BF16 head on BF16 experts: the grouped BF16 MoE layer.
213                MtpQuantization::Bf16 => {
214                    let moe = if draft_nvfp4 {
215                        Self::new_nvfp4_draft_moe(&weights, config, gpu, absmax_k, nvfp4_k, stream)?
216                    } else {
217                        Self::new_bf16_moe(&weights, config, gpu)?
218                    };
219                    (None, None, None, Some(moe))
220                }
221                // 2026-10-02: An FP8 head on BF16 experts: per-expert FP8 projections.
222                MtpQuantization::Fp8 => {
223                    let mut experts_g = Vec::with_capacity(weights.experts.len());
224                    for (i, de) in weights.experts.iter().enumerate() {
225                        let gate_proj = q(&de.gate_proj, inter, h)?;
226                        let up_proj = q(&de.up_proj, inter, h)?;
227                        let down_proj = q(&de.down_proj, h, inter)?;
228                        experts_g.push((gate_proj, up_proj, down_proj));
229                        if (i + 1) % 128 == 0 {
230                            tracing::info!(
231                                "  MTP experts quantized: {}/{}",
232                                i + 1,
233                                weights.experts.len()
234                            );
235                        }
236                    }
237                    let shared = (
238                        q(&weights.shared_expert.gate_proj, inter, h)?,
239                        q(&weights.shared_expert.up_proj, inter, h)?,
240                        q(&weights.shared_expert.down_proj, h, inter)?,
241                    );
242                    (None, Some(experts_g), Some(shared), None)
243                }
244            }
245        };
246        let (moe_nvfp4, moe_experts_generic, moe_shared_generic, moe_grouped) = moe_parts;
247
248        // 2026-09-25: One attention layer. BF16 KV for BF16 and FP8 heads, FP8
249        // KV for NVFP4 heads; the FP8 KV forward passes unit K/V scales
250        // (`forward.rs`).
251        let kv_bf16 = matches!(quant, MtpQuantization::Bf16 | MtpQuantization::Fp8);
252        let kv_config = KvCacheConfig {
253            block_size: 16,
254            num_kv_heads: nkv,
255            head_dim: hd,
256            num_layers: 1,
257            dtype: if kv_bf16 {
258                KvCacheDtype::Bf16
259            } else {
260                KvCacheDtype::Fp8
261            },
262            layer_dtypes: vec![],
263            layer_dims: vec![],
264            cache_blocks_per_seq: None,
265        };
266        // 2026-09-26: `max_seq_len / block + 1` blocks for each of
267        // `mtp_max_seqs()` sequences, capped at the main pool's block count
268        // (and never below one sequence). model-engine `factory/build/kv_budget.rs`
269        // (`mtp_pool_reserve_bytes`) reserves the same arithmetic from the KV
270        // budget; change both together.
271        let per_seq_blocks = max_seq_len / kv_config.block_size + 1;
272        let mtp_num_blocks = per_seq_blocks
273            .saturating_mul(crate::speculative::mtp_max_seqs())
274            .min(main_kv_blocks.max(per_seq_blocks));
275        // 2026-09-25: The per-sequence `propose_meta` stride, from the drafter
276        // pool's block size (`batch_caps::propose_meta_stride_env`: computed
277        // from `max_seq_len`, floor 2048, override
278        // `METRALE_PROPOSE_META_STRIDE`).
279        let propose_meta_stride =
280            super::batch_caps::propose_meta_stride_env(max_seq_len, kv_config.block_size);
281        let kv_cache = PagedKvCache::new(kv_config, mtp_num_blocks, gpu)?;
282
283        let (
284            dense_gemv_k,
285            dense_gemv_fp8w_k,
286            deinterleave_qg_k,
287            moe_topk_k,
288            moe_silu_mul_k,
289            moe_weighted_sum_blend_k,
290        ) = match quant {
291            MtpQuantization::Nvfp4 => (None, None, None, None, None, None),
292            MtpQuantization::Fp8 => (
293                // 2026-09-25: BF16 GEMV for the router (`moe_gate` is BF16).
294                Some(gpu.kernel("gemv", "dense_gemv_bf16")?),
295                Some(gpu.kernel("gemv_fp8w", "dense_gemv_fp8w")?),
296                Some(gpu.kernel("ssm_preprocess", "deinterleave_qg")?),
297                Some(gpu.kernel("moe_topk", "moe_topk_softmax")?),
298                Some(gpu.kernel("moe_silu_mul", "moe_silu_mul")?),
299                Some(gpu.kernel("moe_expert_gemv", "moe_weighted_sum_blend")?),
300            ),
301            MtpQuantization::Bf16 => (
302                Some(gpu.kernel("gemv", "dense_gemv_bf16")?),
303                None,
304                Some(gpu.kernel("ssm_preprocess", "deinterleave_qg")?),
305                Some(gpu.kernel("moe_topk", "moe_topk_softmax")?),
306                Some(gpu.kernel("moe_silu_mul", "moe_silu_mul")?),
307                Some(gpu.kernel("moe_expert_gemv", "moe_weighted_sum_blend")?),
308            ),
309        };
310
311        let effective_vocab = if mtp_vocab_size > 0 {
312            (mtp_vocab_size as usize).min(config.vocab_size)
313        } else {
314            config.vocab_size
315        };
316        let ffn_kind: &str = if dense_ffn_generic.is_some() {
317            "dense FFN"
318        } else if moe_nvfp4.is_some() {
319            "MoE (NVFP4 fused)"
320        } else if moe_grouped.is_some() {
321            "MoE (native FP8 tables; batched propose grouped)"
322        } else {
323            "MoE (per-expert)"
324        };
325        tracing::info!(
326            "MTP head: quant={:?}{wo}, fc=[{h},{h2}], attn Q=[{qd},{h}], ffn={ffn}, \
327             {ne} experts, vocab={ev}/{fv} (LM head {lm:.1} MB)",
328            quant,
329            wo = if dense_nvfp4 {
330                " (weight-only NVFP4 q/o + dense FFN)"
331            } else {
332                ""
333            },
334            h2 = h * 2,
335            qd = nq * hd * 2,
336            ffn = ffn_kind,
337            ne = if dense_ffn_generic.is_some() {
338                0
339            } else {
340                config.num_experts
341            },
342            ev = effective_vocab,
343            fv = config.vocab_size,
344            lm = (effective_vocab * h / 2) as f64 / (1024.0 * 1024.0),
345        );
346
347        // 2026-09-25: Dedicated `PREFILL_CHUNK`-row scratch for the batched
348        // row writer (`drafter_rows_impl`), allocated only when the drafter
349        // prefill or the catch-up feed (`METRALE_MTP_CATCHUP`) is on.
350        let prefill_scratch = if super::mtp_drafter_prefill_enabled(levers)
351            || crate::speculative::mtp_catchup_enabled()
352        {
353            let c = super::prefill::PREFILL_CHUNK;
354            let bf16 = 2usize;
355            Some(super::MtpPrefillScratch {
356                embed: gpu.alloc(c * h * bf16)?,
357                normed_embed: gpu.alloc(c * h * bf16)?,
358                normed_hidden: gpu.alloc(c * h * bf16)?,
359                concat: gpu.alloc(c * 2 * h * bf16)?,
360                fc_out: gpu.alloc(c * h * bf16)?,
361                normed2: gpu.alloc(c * h * bf16)?,
362                k_out: gpu.alloc(c * nkv * hd * bf16)?,
363                v_out: gpu.alloc(c * nkv * hd * bf16)?,
364                q_scratch: gpu.alloc(c * nq * hd * bf16)?,
365                pos_dev: gpu.alloc(c * 4)?,
366                slot_dev: gpu.alloc(c * 8)?,
367            })
368        } else {
369            None
370        };
371
372        Ok(Self {
373            pre_fc_norm_embedding: weights.pre_fc_norm_embedding,
374            pre_fc_norm_hidden: weights.pre_fc_norm_hidden,
375            input_layernorm: weights.input_layernorm,
376            post_attn_layernorm: weights.post_attn_layernorm,
377            norm: weights.norm,
378            fc,
379            q_proj,
380            k_proj,
381            v_proj,
382            o_proj,
383            q_norm: weights.q_norm,
384            k_norm: weights.k_norm,
385            moe_nvfp4,
386            moe_experts_generic,
387            moe_shared_generic,
388            moe_grouped,
389            moe_gate: weights.moe_gate,
390            shared_expert_gate: weights.shared_expert_gate,
391            dense_ffn_generic,
392            quant,
393            mtp_vocab_size,
394            embed_tokens,
395            target_final_norm,
396            lm_head_nvfp4,
397            kv_cache: Mutex::new(kv_cache),
398            attn_layer_idx: 0,
399            rms_norm_k: gpu.kernel("norm", "rms_norm")?,
400            rms_norm_residual_k: gpu.kernel("norm", "rms_norm_residual")?,
401            w4a16_gemv_k: gpu.kernel("w4a16_gemv", "w4a16_gemv")?,
402            w4a16_gemv_sw_k: crate::layers::try_kernel(gpu, "w4a16_gemv", "w4a16_gemv_sw"),
403            gemv_sw: crate::layers::ops::gemv_sw_from(
404                std::env::var("METRALE_NO_GEMV_SW").ok().as_deref(),
405            ),
406            w4a16_gemv_qg_k: gpu.kernel("w4a16_gemv", "w4a16_gemv_qg")?,
407            w4a16_gemv_dual_k: gpu.kernel("w4a16_gemv_fused", "w4a16_gemv_dual")?,
408            rope_k: gpu.kernel("rope", "rope_forward")?,
409            rope_strided_k: crate::layers::try_kernel(gpu, "rope", "rope_forward_strided"),
410            rms_norm_strided_k: crate::layers::try_kernel(gpu, "norm", "rms_norm_strided"),
411            sigmoid_gate_mul_batched_k: crate::layers::try_kernel(
412                gpu,
413                "residual_add",
414                "sigmoid_gate_mul_batched",
415            ),
416            reshape_cache_k: if kv_bf16 {
417                gpu.kernel("reshape_and_cache", "reshape_and_cache_flash")?
418            } else {
419                gpu.kernel("reshape_and_cache", "reshape_and_cache_flash_fp8")?
420            },
421            paged_decode_k: if kv_bf16 {
422                gpu.kernel("paged_decode", "paged_decode_attn")?
423            } else {
424                gpu.kernel("paged_decode_fp8", "paged_decode_attn_fp8")?
425            },
426            kv_bf16,
427            kv_exact: levers.mtp_kv_exact,
428            residual_add_k: gpu.kernel("residual_add", "bf16_residual_add")?,
429            residual_add_rms_norm_k: gpu.kernel("norm", "residual_add_rms_norm")?,
430            sigmoid_gate_mul_k: gpu.kernel("residual_add", "sigmoid_gate_mul")?,
431            bf16_concat_k: gpu.kernel("residual_add", "bf16_concat")?,
432            argmax_k: gpu.kernel("argmax", "argmax_bf16")?,
433            embed_from_argmax_k: gpu.kernel("embed_from_argmax", "embed_from_argmax")?,
434            draft_token_id_dev: gpu.alloc(4)?,
435            last_conf_bits: std::sync::atomic::AtomicU32::new(1.0f32.to_bits()),
436            dense_gemv_k,
437            dense_gemv_fp8w_k,
438            w8a16_gemv_k: gpu.kernel("w8a16_gemv", "w8a16_gemv").ok(),
439            deinterleave_qg_k,
440            moe_topk_k,
441            moe_silu_mul_k,
442            moe_weighted_sum_blend_k,
443            // 2026-09-25: 0 when absent; the drafter prefill then writes no
444            // rows.
445            dense_gemm_k: crate::layers::try_kernel(gpu, "gemm", "dense_gemm_bf16"),
446            batched_embed_k: crate::layers::try_kernel(gpu, "embed_from_argmax", "batched_embed"),
447            dense_gemm_pipelined_k: crate::layers::try_kernel(
448                gpu,
449                "gemm",
450                "dense_gemm_bf16_pipelined",
451            ),
452            // 2026-09-25: 0 when absent; `row_dispatch::drafter_row_kernel`
453            // then picks the pipelined GEMM or the per-row loop.
454            dense_gemv_batchm_k: crate::layers::try_kernel(
455                gpu,
456                "dense_gemv_bf16_batchm",
457                "dense_gemv_bf16_batchm",
458            ),
459            w4a16_batchm: crate::layers::w4a16_gemv_tiers::W4a16BatchmTiers::resolve(gpu),
460            w4a16_gemv_batch16_k: crate::layers::try_kernel(
461                gpu,
462                "w4a16_gemv",
463                "w4a16_gemv_batch16",
464            ),
465            w4a16_gemv_batch32_k: crate::layers::try_kernel(
466                gpu,
467                "w4a16_gemv",
468                "w4a16_gemv_batch32",
469            ),
470            // 2026-09-25: Dropped when `METRALE_NO_MTP_LMHEAD_TGEMM` is set to
471            // any value, `0` included.
472            lm_head_nvfp4_t: lm_head_nvfp4_t
473                .filter(|_| std::env::var_os("METRALE_NO_MTP_LMHEAD_TGEMM").is_none()),
474            w4a16_gemm_t_k: crate::layers::tgemm_kernel(gpu),
475            argmax_batch_k: crate::layers::try_kernel(gpu, "argmax", "argmax_bf16_batch"),
476            argmax_batch_lp_k: crate::layers::try_kernel(gpu, "argmax", "argmax_bf16_batch_lp"),
477            // 2026-09-25: `PROPOSE_META_SEQS` slabs of `propose_meta_stride`
478            // bytes, a dedicated allocation rather than an offset into
479            // `scratch`.
480            propose_meta: gpu.alloc(super::batch_caps::PROPOSE_META_SEQS * propose_meta_stride)?,
481            propose_meta_stride,
482            prefill_scratch,
483            circuit_draft: parking_lot::RwLock::new(None),
484        })
485    }
486}
metrale_model_layers/layers/mtp_head/new.rs

metrale_model_layers/layers/mtp_head/
new.rs
