metrale_model_engine/model/
lm_head_nvfp4_rows.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The NVFP4 LM head at its declared W4A16 on the tensor-core row tiles
4//! (`ops::w4a16_tc_rows`), at every row count in 64-row calls. Installed under
5//! `--weight-quantization declared` for a checkpoint that declares an NVFP4 head with 16-bit
6//! activations (nvidia/Qwen3.6-35B-A3B-NVFP4: `W4A16_NVFP4`). It stands in for every other NVFP4
7//! head path, among them the transposed tile GEMM, which casts activations to E4M3 (below the
8//! declared precision), and the by-rows GEMV/GEMM switch, whose summation order depends on the
9//! row count. A row's logits here do not depend on the row count.
10//!
11//! Owner: model-engine (LM head).
12//! Invariants: `lm_head_nvfp4_rows_run` writes `logits[rows, vocab]` BF16 and returns `true`
13//! exactly when the row tiles are installed.
14
15use anyhow::{Context, Result};
16use metrale_gpu_runtime::gpu::DevicePtr;
17use metrale_model_layers::layers::ops;
18
19use super::types::TransformerModel;
20
21impl TransformerModel {
22    /// 2026-10-02: Install the W4A16 row-tile head when the declared-precision policy declares an
23    /// NVFP4 head with 16-bit activations, the model loaded the NVFP4 head, the target ships the
24    /// kernels, the head fits their shape contract, and the vocab is not sharded.
25    pub fn install_declared_lm_head_w4a16_rows(&mut self) -> Result<()> {
26        use metrale_config::weight_quantization::{LmHeadChoice, LmHeadFormat};
27        let policy = metrale_config::WeightQuantPolicy::for_checkpoint(
28            metrale_model_layers::layers::weight_quantization(),
29            self.config.quantization_config.as_ref(),
30            metrale_model_layers::layers::kernel_caps(),
31        );
32        let declared_a16 = policy.lm_head() == LmHeadChoice::Declared(LmHeadFormat::Nvfp4)
33            && policy.declared("lm_head").activation.is_none();
34        let (h, v) = (
35            self.config.hidden_size as u32,
36            self.config.vocab_size as u32,
37        );
38        if !declared_a16
39            || self.lm_head_nvfp4.is_none()
40            || self.config.tp_world_size > 1
41            || !self.gpu.has_module(ops::W4A16_TC_ROWS_MODULE)
42            || !ops::w4a16_tc_rows_shape_ok(1, v, h, h, v)
43        {
44            return Ok(());
45        }
46        self.lm_head_rows.nvfp4_rows = true;
47        tracing::info!(
48            "lm_head: declared NVFP4 W4A16 on the tensor-core row tiles at every row count"
49        );
50        Ok(())
51    }
52
53    /// 2026-10-02: `logits[rows, vocab]` from `hidden[rows, hidden]` through the W4A16 row tiles;
54    /// `Ok(false)`, launching nothing, unless they are installed.
55    pub(super) fn lm_head_nvfp4_rows_run(
56        &self,
57        hidden: DevicePtr,
58        rows: usize,
59        logits: DevicePtr,
60        stream: u64,
61    ) -> Result<bool> {
62        if !self.lm_head_rows.nvfp4_rows {
63            return Ok(false);
64        }
65        let head = self
66            .lm_head_nvfp4
67            .as_ref()
68            .context("W4A16 row-tile lm_head needs the NVFP4 head")?;
69        let (h, v) = (self.config.hidden_size, self.config.vocab_size);
70        let mut done = 0;
71        while done < rows {
72            let m = (rows - done).min(ops::W4A16_TC_ROWS_MAX_M as usize);
73            ops::w4a16_tc_rows(
74                self.gpu.as_ref(),
75                hidden.offset(done * h * 2),
76                head,
77                logits.offset(done * v * 2),
78                m as u32,
79                v as u32,
80                h as u32,
81                h as u32,
82                v as u32,
83                stream,
84            )?;
85            done += m;
86        }
87        Ok(true)
88    }
89}
metrale_model_engine/model/lm_head_nvfp4_rows.rs

metrale_model_engine/model/
lm_head_nvfp4_rows.rs
