metrale_model_layers/layers/ops/
e4m3_range.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-29: Debug-build range check on the activations an unscaled E4M3 cast reads.
4//!
5//! Owner: model-layers ops.
6//! Invariants: release builds compile the check out; it never changes a launch.
7//!
8//! Several W4A16 and FP8 prefill kernels cast their BF16 A operand to E4M3 with no scale
9//! (`bf16x4_to_e4m3x4`, `cvt.rn.satfinite`), and the `bf16_to_fp8` op casts a buffer the
10//! same way (the FP8 activation path, and some weights at load). E4M3's largest finite value is 448, so a larger
11//! activation saturates without an error. A kernel declares the cast with
12//! `<entry>_a_e4m3` (`GpuBackend::kernel_casts_a_to_e4m3`). Measured on the certified
13//! dense 27B models (Qwen3.8-27B-NVFP4 and Qwen3.6-27B-NVFP4; short, 5k, 27.5k-token,
14//! tool-call and code prompts), the largest activation at these sites was 134, a margin of
15//! 3.3x. In debug builds every such launch, and every `bf16_to_fp8`, checks its source and
16//! panics past 448.
17//!
18//! 2026-09-30: Two more things.
19//! - The grouped MoE launchers check the A rows their kernel reads (the rows the expert
20//!   offsets span, through `sorted_token_ids` when one is passed), the same way.
21//! - A saturation that is disclosed or known does not panic. The call site opens an
22//!   [`allow_e4m3_saturation`] scope naming one of the [`E4m3Saturation`] constants below,
23//!   which are the complete list. Inside it, a violation up to that entry's bound is logged
24//!   once per entry and counted ([`e4m3_saturations`]); past the bound, or outside any
25//!   scope, it still panics.
26
27use std::cell::RefCell;
28use std::collections::BTreeMap;
29
30use anyhow::Result;
31use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
32use parking_lot::Mutex;
33
34/// 2026-09-29: E4M3's largest finite magnitude; an unscaled cast saturates above it.
35pub const E4M3_MAX: f32 = 448.0;
36
37/// 2026-09-30: An unscaled E4M3 cast whose saturation is disclosed or known, with the
38/// largest `|x|` the disclosure covers.
39#[derive(Debug, Clone, Copy, PartialEq)]
40pub struct E4m3Saturation {
41    /// 2026-09-30: Where it happens; the key its count is kept under.
42    pub site: &'static str,
43    /// 2026-09-30: Why it is not a defect to panic on.
44    pub why: &'static str,
45    /// 2026-09-30: The largest `|x|` covered. Past it the check panics as for an undeclared
46    /// site.
47    pub bound: f32,
48}
49
50/// 2026-09-30: `--nemotron-shared-expert-e4m3`: the Nemotron-H shared expert's prefill tile
51/// GEMM casts its activations to E4M3 unscaled. The flag's help discloses the saturation
52/// (1-30% local error measured on Nemotron-3-Nano), so any magnitude is covered.
53pub const NEMOTRON_SHARED_EXPERT_E4M3: E4m3Saturation = E4m3Saturation {
54    site: "Nemotron-H shared expert, --nemotron-shared-expert-e4m3",
55    why: "opted in on the command line; the flag discloses the saturating cast",
56    bound: f32::INFINITY,
57};
58
59/// 2026-09-30: The routed experts' down projection reads the SiLU product. On
60/// Qwen3.6-35B-A3B-NVFP4 it reaches 510 at position 0 of the last layer (measured
61/// 2026-09-29, no effect on the output); every other certified model stays under 448. The
62/// scaled cast that removes it is follow-up work, so this covers exactly the known value.
63pub const MOE_ROUTED_DOWN_KNOWN: E4m3Saturation = E4m3Saturation {
64    site: "MoE routed-down input",
65    why: "known: Qwen3.6-35B-A3B-NVFP4 reaches 510 at position 0 of its last layer",
66    bound: 512.0,
67};
68
69thread_local! {
70    static SCOPES: RefCell<Vec<E4m3Saturation>> = const { RefCell::new(Vec::new()) };
71}
72
73/// 2026-09-30: Saturating launches seen per site, and whether the site has been logged.
74static SEEN: Mutex<BTreeMap<&'static str, u64>> = Mutex::new(BTreeMap::new());
75
76/// 2026-09-30: While the returned scope lives, a saturation on this thread up to
77/// `saturation.bound` is logged once and counted instead of panicking. Scopes nest; the
78/// innermost one applies.
79#[must_use = "the allowance ends when the scope is dropped"]
80pub fn allow_e4m3_saturation(saturation: E4m3Saturation) -> E4m3SaturationScope {
81    SCOPES.with(|s| s.borrow_mut().push(saturation));
82    E4m3SaturationScope { _private: () }
83}
84
85/// 2026-09-30: See [`allow_e4m3_saturation`].
86pub struct E4m3SaturationScope {
87    _private: (),
88}
89
90impl Drop for E4m3SaturationScope {
91    fn drop(&mut self) {
92        SCOPES.with(|s| s.borrow_mut().pop());
93    }
94}
95
96/// 2026-09-30: Saturating launches counted under `site` in this process (debug builds only).
97pub fn e4m3_saturations(site: &str) -> u64 {
98    SEEN.lock().get(site).copied().unwrap_or(0)
99}
100
101/// 2026-09-30: The verdict on one launch's largest activation magnitude.
102fn judge(amax: f32) {
103    if amax <= E4M3_MAX {
104        return;
105    }
106    match SCOPES.with(|s| s.borrow().last().copied()) {
107        Some(allowed) if amax <= allowed.bound => {
108            let mut seen = SEEN.lock();
109            let count = seen.entry(allowed.site).or_insert(0);
110            if *count == 0 {
111                tracing::warn!(
112                    site = allowed.site,
113                    why = allowed.why,
114                    amax,
115                    "unscaled E4M3 cast saturates (|x| = {amax} > {E4M3_MAX}) at a disclosed \
116                     site; counted, logged once"
117                );
118            }
119            *count += 1;
120        }
121        Some(allowed) => panic!(
122            "unscaled E4M3 cast would saturate: activation |x| = {amax} > {E4M3_MAX}, past the \
123             {} that `{}` discloses",
124            allowed.bound, allowed.site
125        ),
126        None => panic!("unscaled E4M3 cast would saturate: activation |x| = {amax} > {E4M3_MAX}"),
127    }
128}
129
130/// 2026-09-30: The largest finite-or-infinite magnitude among BF16 `bytes`; NaN is skipped.
131fn bf16_amax<'a>(rows: impl Iterator<Item = &'a [u8]>) -> f32 {
132    rows.flat_map(|r| r.chunks_exact(2))
133        .map(|b| f32::from_bits(u32::from(u16::from_le_bytes([b[0], b[1]])) << 16).abs())
134        .filter(|v| !v.is_nan())
135        .fold(0.0f32, f32::max)
136}
137
138/// 2026-09-29: When `kernel` casts its A operand to E4M3 with no scale, check the
139/// `[m, k]` BF16 `input` it reads (see [`check_e4m3_range`]). A no-op otherwise, and in
140/// release builds.
141pub fn check_e4m3_activation_range(
142    gpu: &dyn GpuBackend,
143    kernel: KernelHandle,
144    input: DevicePtr,
145    m: u32,
146    k: u32,
147    stream: u64,
148) -> Result<()> {
149    if cfg!(debug_assertions) && gpu.kernel_casts_a_to_e4m3(kernel) {
150        check_e4m3_range(gpu, input, m as usize * k as usize, stream)?;
151    }
152    Ok(())
153}
154
155/// 2026-09-29: Debug builds: copy the `elements` BF16 values at `input` to the host on
156/// `stream` and panic when one exceeds [`E4M3_MAX`] in magnitude, unless an
157/// [`allow_e4m3_saturation`] scope covers it. NaN is not checked here. Skipped under graph
158/// capture, where a copy to the host cannot run. Release builds: a no-op.
159pub fn check_e4m3_range(
160    gpu: &dyn GpuBackend,
161    input: DevicePtr,
162    elements: usize,
163    stream: u64,
164) -> Result<()> {
165    if !cfg!(debug_assertions) || elements == 0 || gpu.stream_is_capturing(stream) {
166        return Ok(());
167    }
168    let mut buf = vec![0u8; elements * 2];
169    gpu.copy_d2h_on_stream(input, &mut buf, stream)?;
170    judge(bf16_amax(std::iter::once(buf.as_slice())));
171    Ok(())
172}
173
174/// 2026-09-30: The grouped (pointer-table) MoE form of [`check_e4m3_activation_range`]. The
175/// kernel reads A row `sorted_token_ids[p]`, or row `p` when `sorted_token_ids` is null, for
176/// every position `p` in `expert_offsets[e]..expert_offsets[e + 1]`, `e < num_experts`; only
177/// those rows are checked, each `k` BF16 values wide. A no-op for a kernel that keeps A in
178/// BF16, in release builds, and under graph capture.
179#[allow(clippy::too_many_arguments)]
180pub fn check_e4m3_grouped(
181    gpu: &dyn GpuBackend,
182    kernel: KernelHandle,
183    a: DevicePtr,
184    expert_offsets: DevicePtr,
185    sorted_token_ids: DevicePtr,
186    num_experts: u32,
187    k: u32,
188    stream: u64,
189) -> Result<()> {
190    if !cfg!(debug_assertions)
191        || !gpu.kernel_casts_a_to_e4m3(kernel)
192        || gpu.stream_is_capturing(stream)
193    {
194        return Ok(());
195    }
196    let offsets = read_i32s(gpu, expert_offsets, num_experts as usize + 1, stream)?;
197    let positions: Vec<usize> = offsets
198        .windows(2)
199        .flat_map(|w| w[0].max(0) as usize..w[1].max(0) as usize)
200        .collect();
201    let Some(&last) = positions.iter().max() else {
202        return Ok(());
203    };
204    let rows: Vec<usize> = if sorted_token_ids.is_null() {
205        positions
206    } else {
207        let ids = read_i32s(gpu, sorted_token_ids, last + 1, stream)?;
208        positions
209            .iter()
210            .map(|&p| usize::try_from(ids[p]))
211            .collect::<Result<_, _>>()
212            .map_err(|_| {
213                anyhow::anyhow!("a negative sorted token id at a position the kernel reads")
214            })?
215    };
216    let Some(&top) = rows.iter().max() else {
217        return Ok(());
218    };
219    let width = k as usize * 2;
220    let mut buf = vec![0u8; (top + 1) * width];
221    gpu.copy_d2h_on_stream(a, &mut buf, stream)?;
222    judge(bf16_amax(
223        rows.iter().map(|&r| &buf[r * width..(r + 1) * width]),
224    ));
225    Ok(())
226}
227
228fn read_i32s(gpu: &dyn GpuBackend, ptr: DevicePtr, n: usize, stream: u64) -> Result<Vec<i32>> {
229    let mut buf = vec![0u8; n * 4];
230    gpu.copy_d2h_on_stream(ptr, &mut buf, stream)?;
231    Ok(buf
232        .chunks_exact(4)
233        .map(|b| i32::from_le_bytes([b[0], b[1], b[2], b[3]]))
234        .collect())
235}
236
237#[cfg(test)]
238#[path = "e4m3_range_tests.rs"]
239mod tests;
metrale_model_layers/layers/ops/e4m3_range.rs

metrale_model_layers/layers/ops/
e4m3_range.rs
