metrale_model_layers/layers/dense_ffn_m16_tc_oracle.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The numerics contract for the tensor-core decode tiers, graded against the scalar GEMV.
4//!
5//! Owner: model-layers (dense FFN).
6//! Invariants: none beyond the types.
7//!
8//! [`within_m16_tc_budget`] is the one predicate the GPU microtests (for example
9//! `examples/native_fp8_ffn_m16_tc_microtest.rs` and
10//! `examples/native_bf16_lm_head_m16_microtest.rs`) and the host tests
11//! (`dense_ffn_m16_tc_m32_tests.rs`, `ops/dense_gemm_m16_bf16_tests.rs`,
12//! `ops/dense_gemm_m16_bf16_floor_tests.rs`) import. An m16n8k16 MMA sums K in
13//! a different order from the scalar GEMV, so the contract is a tolerance: an
14//! ordinal-ULP budget plus an absolute floor scaled to the reduction
15//! ([`m16_tc_acc_floor`]).
16
17/// 2026-09-25: Ordinal BF16 ULP budget against the scalar GEMV reference.
18pub const M16_TC_MAX_ULP: i32 = 2;
19
20/// 2026-09-25: FP32 unit roundoff, `2^-24`.
21pub const F32_UNIT_ROUNDOFF: f64 = 5.960_464_477_539_063e-8;
22
23/// 2026-09-25: Margin over the `u * sqrt(k) * row_rms` accumulation scale in
24/// [`m16_tc_acc_floor`].
25///
26/// Calibrated, not derived: the MMA's internal summation order is unspecified
27/// hardware. Tests pin both ends of the calibration.
28/// `the_round9_lm_head_outlier_is_a_cancelled_logit`
29/// (`ops/dense_gemm_m16_bf16_floor_tests.rs`) admits an H100 LM-head element
30/// 100 ordinal ULP from a reference of -1.173e-4 at an RMS of 23.80, where the
31/// floor is 8.1e-4. `the_k_aware_floor_still_rejects_every_structural_error`
32/// (same file) and `the_accumulation_floor_still_rejects_a_structural_error`
33/// (`dense_ffn_m16_tc_m32_tests.rs`) still reject a misplaced output and a
34/// three-ordinal error on a value near 1.
35pub const M16_TC_ACC_FLOOR_MARGIN: f64 = 8.0;
36
37/// 2026-09-25: Absolute error below which an ordinal-ULP budget says nothing,
38/// for one output of a length-`k` FP32 reduction whose row has RMS `row_rms`.
39///
40/// An ordinal BF16 ULP is a relative unit, and an output that cancelled to
41/// near zero has no relative accuracy left to grade. The floor is
42/// `M16_TC_ACC_FLOOR_MARGIN * u * sqrt(k) * row_rms`: two FP32 summation orders
43/// over `k` terms drift apart like a random walk of `k` roundings, each about
44/// `u * row_rms`. `the_fixed_floor_does_not_track_the_reduction_depth_and_the_new_one_does`
45/// (`ops/dense_gemm_m16_bf16_floor_tests.rs`) measures that drift on the host
46/// and checks the floor stays more than 50x above it.
47///
48/// `row_rms` is the RMS of the reference row, not of the whole block: an
49/// element's accumulation noise scales with the norm of the activation row
50/// that produced it, and `the_floor_follows_the_row_and_not_the_block` pins
51/// that a block-wide RMS over-forgives a quiet row.
52pub fn m16_tc_acc_floor(k: usize, row_rms: f64) -> f64 {
53 M16_TC_ACC_FLOOR_MARGIN * F32_UNIT_ROUNDOFF * (k as f64).sqrt() * row_rms
54}
55
56/// 2026-09-25: BF16 bits -> a monotone integer, so `|ord(a) - ord(b)|` is the
57/// ULP distance and +0/-0 are the same point.
58pub fn bf16_ord(bits: u16) -> i32 {
59 if bits & 0x8000 != 0 {
60 -((bits & 0x7FFF) as i32)
61 } else {
62 bits as i32
63 }
64}
65
66/// 2026-09-25: The tier's numerics contract as one predicate.
67///
68/// An element passes if it is within [`M16_TC_MAX_ULP`] ordinal BF16 ULP of the
69/// reference, or if its absolute error is under [`m16_tc_acc_floor`] for the
70/// reduction depth `k` and the reference row's RMS.
71///
72/// `k` is the reduction depth the kernel ran, not the output width.
73pub fn within_m16_tc_budget(actual_bits: u16, reference_bits: u16, k: usize, row_rms: f64) -> bool {
74 if (bf16_ord(actual_bits) - bf16_ord(reference_bits)).abs() <= M16_TC_MAX_ULP {
75 return true;
76 }
77 let a = f64::from(half::bf16::from_bits(actual_bits).to_f32());
78 let b = f64::from(half::bf16::from_bits(reference_bits).to_f32());
79 (a - b).abs() <= m16_tc_acc_floor(k, row_rms)
80}
81
82/// 2026-09-25: One element the comparison rejected: where it is, the two
83/// values, the ordinal distance and the row RMS the floor was evaluated at.
84#[derive(Debug, Clone, Copy)]
85pub struct M16TcOutlier {
86 pub row: usize,
87 pub col: usize,
88 pub reference: f32,
89 pub actual: f32,
90 pub ulp: i32,
91 /// 2026-09-25: RMS of the reference row this element sits in, the scale
92 /// [`m16_tc_acc_floor`] was evaluated at.
93 pub row_rms: f64,
94}
95
96/// 2026-09-25: The result of comparing an `m x n` BF16 block against the
97/// scalar reference.
98#[derive(Debug, Clone, Default)]
99pub struct M16TcDiff {
100 /// 2026-09-25: Largest ordinal BF16 ULP distance; sign flips are counted in
101 /// `sign_flips` and excluded here.
102 pub max_ulp: i32,
103 /// 2026-09-25: Elements the full criterion rejected (ordinal budget and
104 /// accumulation floor both exceeded).
105 pub over_budget: Vec<M16TcOutlier>,
106 /// 2026-09-25: Elements over the ordinal budget alone, whether or not the
107 /// floor admitted them.
108 pub over_ulp_only: usize,
109 pub sign_flips: usize,
110 pub max_abs: f64,
111 pub rel_rms: f64,
112 /// 2026-09-25: RMS of the whole reference block. The floor uses the per-row
113 /// values in `row_rms`, and this one only for a row that has none.
114 pub rms: f64,
115 /// 2026-09-25: Per-row reference RMS, in row order: the scales the floor was
116 /// evaluated at.
117 pub row_rms: Vec<f64>,
118}
119
120/// 2026-09-25: A sign change on a reference smaller than this in magnitude is
121/// counted in `sign_flips` and not graded.
122pub const M16_TC_SIGN_FLIP_BAND: f64 = 0.05;
123
124/// 2026-09-25: RMS of a little-endian BF16 slice, in f64; 0 for an empty slice.
125fn bf16_rms(block: &[u8]) -> f64 {
126 let count = block.len() / 2;
127 if count == 0 {
128 return 0.0;
129 }
130 let sum: f64 = block
131 .chunks_exact(2)
132 .map(|b| {
133 let v = f64::from(half::bf16::from_bits(u16::from_le_bytes([b[0], b[1]])).to_f32());
134 v * v
135 })
136 .sum();
137 (sum / count as f64).sqrt()
138}
139
140/// 2026-09-25: Compare an `m x n` BF16 block against the scalar reference
141/// under the tier's contract. Both slices are `m * n` little-endian BF16
142/// elements; `k` is the reduction depth the kernel ran.
143///
144/// Two passes: the per-row reference RMS is computed first, then each element
145/// is graded against its row's floor.
146pub fn compare_m16_tc_block(actual: &[u8], reference: &[u8], n: usize, k: usize) -> M16TcDiff {
147 let bits = |b: &[u8]| u16::from_le_bytes([b[0], b[1]]);
148 let val = |b: u16| f64::from(half::bf16::from_bits(b).to_f32());
149 let row_bytes = n * 2;
150 let row_rms: Vec<f64> = if row_bytes == 0 {
151 Vec::new()
152 } else {
153 reference.chunks(row_bytes).map(bf16_rms).collect()
154 };
155 let mut d = M16TcDiff {
156 rms: bf16_rms(reference),
157 row_rms,
158 ..Default::default()
159 };
160 let (mut err_sq, mut ref_sq) = (0.0_f64, 0.0_f64);
161 for (i, (a, b)) in actual
162 .chunks_exact(2)
163 .zip(reference.chunks_exact(2))
164 .enumerate()
165 {
166 let (ab, bb) = (bits(a), bits(b));
167 let (av, bv) = (val(ab), val(bb));
168 err_sq += (av - bv) * (av - bv);
169 ref_sq += bv * bv;
170 d.max_abs = d.max_abs.max((av - bv).abs());
171 let ulp = (bf16_ord(ab) - bf16_ord(bb)).abs();
172 if av.signum() != bv.signum() && bv.abs() < M16_TC_SIGN_FLIP_BAND {
173 d.sign_flips += 1;
174 continue;
175 }
176 d.max_ulp = d.max_ulp.max(ulp);
177 if ulp > M16_TC_MAX_ULP {
178 d.over_ulp_only += 1;
179 }
180 let row = if n == 0 { 0 } else { i / n };
181 let scale = d.row_rms.get(row).copied().unwrap_or(d.rms);
182 if !within_m16_tc_budget(ab, bb, k, scale) {
183 d.over_budget.push(M16TcOutlier {
184 row,
185 col: if n == 0 { 0 } else { i % n },
186 reference: bv as f32,
187 actual: av as f32,
188 ulp,
189 row_rms: scale,
190 });
191 }
192 }
193 d.rel_rms = if ref_sq > 0.0 {
194 (err_sq / ref_sq).sqrt()
195 } else {
196 err_sq.sqrt()
197 };
198 d
199}