metrale_model_engine/model/
lm_head_fp8_rows.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The FP8 E4M3 LM head (per-row F32 scales, `lm_head_fp8`) at any row count, in
4//! as few weight passes as the kernels allow:
5//! - declared W8A8 (`set_lm_head_w8a8`, installed when the checkpoint declares the head FP8
6//!   W8A8): `ops::w8a8_proj` at every row count, in `ops::W8A8_MAX_ROWS`-row calls (one weight
7//!   pass per 128 rows), never per row; a row's logits do not depend on the row count;
8//! - otherwise W8A16: `dense_gemv_fp8w` for 1 row, `dense_gemv_fp8w_batch2` for 2, and the
9//!   register-tiled `fp8_gemv_rowscale_batch{8,16}_rt2` for 3..=16 and in 16-row chunks above.
10//!   Before this, every row count but 2 ran one GEMV per row, reading the 1.27 GB head once
11//!   per row (4x at the decode-floor recipe's 4 verify rows).
12//!
13//! Owner: model-engine (LM head).
14//! Invariants: `lm_head_fp8_run` writes `logits[rows, vocab]` BF16 and returns `true` exactly
15//! when `lm_head_fp8` is set.
16
17use anyhow::{Context, Result, ensure};
18use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
19use metrale_model_layers::layers::W8a8Ctx;
20use metrale_model_layers::layers::ops::{self, W8a8Weight};
21use metrale_model_layers::weight_map::{Fp8Weight, WeightQuantFormat};
22
23use super::types::TransformerModel;
24
25/// 2026-09-28: The multi-row FP8 head kernels, and the W8A8 head when installed. 2026-10-02:
26/// `nvfp4_rows`, set when the NVFP4 head runs the declared W4A16 row tiles
27/// (`lm_head_nvfp4_rows.rs`).
28pub(crate) struct LmHeadRows {
29    rt8: KernelHandle,
30    rt16: KernelHandle,
31    w8a8: Option<(W8a8Ctx, W8a8Weight)>,
32    pub(super) nvfp4_rows: bool,
33}
34
35impl LmHeadRows {
36    pub(crate) fn resolve(gpu: &dyn GpuBackend) -> Self {
37        let rt = |f| metrale_model_layers::layers::try_kernel(gpu, "fp8_gemv_rt", f);
38        Self {
39            rt8: rt("fp8_gemv_rowscale_batch8_rt2"),
40            rt16: rt("fp8_gemv_rowscale_batch16_rt2"),
41            w8a8: None,
42            nvfp4_rows: false,
43        }
44    }
45}
46
47impl TransformerModel {
48    /// 2026-09-28: Install the W8A8 head when the declared-precision policy wants the head run
49    /// W8A8 (`WeightQuantPolicy::fp8_decode_act("lm_head") == Some(Fp8)`) and the model loaded
50    /// the FP8 head (`lm_head_fp8`, the declared head format under `declared`).
51    pub fn install_declared_lm_head_w8a8(&mut self) -> Result<()> {
52        let policy = metrale_config::WeightQuantPolicy::for_checkpoint(
53            metrale_model_layers::layers::weight_quantization(),
54            self.config.quantization_config.as_ref(),
55            metrale_model_layers::layers::kernel_caps(),
56        );
57        let fp8 = policy.fp8_decode_act("lm_head")
58            == Some(metrale_config::weight_quantization::ActFormat::Fp8);
59        if !fp8 || self.lm_head_fp8.is_none() {
60            return Ok(());
61        }
62        let ctx = W8a8Ctx::new(self.gpu.as_ref(), self.config.hidden_size as u32)?;
63        self.set_lm_head_w8a8(ctx)?;
64        tracing::info!("W8A8 lm_head (declared FP8 W8A8, per-row E4M3 read in place)");
65        Ok(())
66    }
67
68    /// 2026-09-28: Run the FP8 head W8A8 (the checkpoint's declared precision) from now on.
69    /// Needs `lm_head_fp8`, and kernels and scratch that serve one row.
70    pub fn set_lm_head_w8a8(&mut self, ctx: W8a8Ctx) -> Result<()> {
71        let fp8 = self
72            .lm_head_fp8
73            .context("W8A8 lm_head needs the FP8 head (lm_head_fp8)")?;
74        let w = W8a8Weight::new(&[Fp8Weight {
75            weight: fp8.weight,
76            row_scale: fp8.row_scale,
77            n: self.config.vocab_size as u32,
78            k: self.config.hidden_size as u32,
79            scale_format: WeightQuantFormat::Fp8PerRow,
80        }])?;
81        ensure!(
82            ctx.available(&w, 1),
83            "W8A8 lm_head kernels or scratch unavailable"
84        );
85        self.lm_head_rows.w8a8 = Some((ctx, w));
86        Ok(())
87    }
88
89    /// 2026-09-28: `logits[rows, vocab]` from `hidden[rows, hidden]` through the FP8 head;
90    /// `Ok(false)`, launching nothing, when the model has no FP8 head.
91    pub(super) fn lm_head_fp8_run(
92        &self,
93        hidden: DevicePtr,
94        rows: usize,
95        logits: DevicePtr,
96        stream: u64,
97    ) -> Result<bool> {
98        let Some(ref fp8) = self.lm_head_fp8 else {
99            return Ok(false);
100        };
101        let gpu = self.gpu.as_ref();
102        let (h, v) = (self.config.hidden_size, self.config.vocab_size);
103        if let Some((ref ctx, ref w)) = self.lm_head_rows.w8a8 {
104            // 2026-09-28: Any row count: `W8A8_MAX_ROWS`-row calls (the scratch's capacity); the
105            // quantization is per row, so chunking does not change any row's bits.
106            let mut done = 0;
107            while done < rows {
108                let m = (rows - done).min(ops::W8A8_MAX_ROWS);
109                let (x, out) = (hidden.offset(done * h * 2), logits.offset(done * v * 2));
110                ensure!(
111                    ctx.proj(gpu, w, x, h as u32, m, out, v as u32, stream)?,
112                    "W8A8 lm_head: {m} rows not servable"
113                );
114                done += m;
115            }
116            return Ok(true);
117        }
118        let k = &self.lm_head_rows;
119        let mut done = 0;
120        while done < rows {
121            let (x, out, left) = (
122                hidden.offset(done * h * 2),
123                logits.offset(done * v * 2),
124                rows - done,
125            );
126            let step = match left {
127                1 => {
128                    ops::dense_gemv_fp8w(
129                        gpu,
130                        self.dense_gemv_fp8w_kernel,
131                        x,
132                        fp8,
133                        out,
134                        v as u32,
135                        h as u32,
136                        stream,
137                    )?;
138                    1
139                }
140                2 if self.dense_gemv_fp8w_batch2_kernel.0 != 0 => {
141                    let kh = self.dense_gemv_fp8w_batch2_kernel;
142                    ops::dense_gemv_fp8w_batch2(gpu, kh, x, fp8, out, v as u32, h as u32, stream)?;
143                    2
144                }
145                3..=8 if k.rt8.0 != 0 => {
146                    let m = left as u32;
147                    ops::fp8_gemv_rowscale_batch8_rt2(
148                        gpu, k.rt8, x, fp8, out, m, v as u32, h as u32, stream,
149                    )?;
150                    left
151                }
152                _ if left > 2 && k.rt16.0 != 0 => {
153                    let m = left.min(16);
154                    let (mu, (vu, hu)) = (m as u32, (v as u32, h as u32));
155                    ops::fp8_gemv_rowscale_batch16_rt2(
156                        gpu, k.rt16, x, fp8, out, mu, vu, hu, stream,
157                    )?;
158                    m
159                }
160                _ => {
161                    ops::dense_gemv_fp8w(
162                        gpu,
163                        self.dense_gemv_fp8w_kernel,
164                        x,
165                        fp8,
166                        out,
167                        v as u32,
168                        h as u32,
169                        stream,
170                    )?;
171                    1
172                }
173            };
174            done += step;
175        }
176        Ok(true)
177    }
178}
metrale_model_engine/model/lm_head_fp8_rows.rs

metrale_model_engine/model/
lm_head_fp8_rows.rs
