metrale_sampling/
feed_argmax.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host reference for one row of `argmax_bf16_batch_feed`
4//! (`kernels/gb10/common/argmax_feed.cu`), computed from the same BF16 bytes.
5//!
6//! Owner: metrale-sampling.
7//! Invariants: none beyond the types. NaN input is outside the contract.
8//!
9//! The kernel runs two reductions over a row and combines them:
10//!
11//! * `p`: the plain argmax, the same scan and tree reduction as
12//!   `argmax_bf16_batch`. Each of the `BLOCK` threads keeps the first strict
13//!   maximum of its strided slice above the `-1e30` floor. At each merge of
14//!   threads `t` and `t + s`, `t` keeps its value unless `t + s` holds a
15//!   strictly larger one. Ties therefore follow the tree order, not the
16//!   lowest index. [`plain_kernel_argmax`] simulates this, floor included.
17//! * `q`: the row with the two masked ids set to `-inf` (as
18//!   `PostCloseThinkMask` does) and the last index holding the maximum, the
19//!   tie rule of the temperature-0 host sampler ([`masked_last_wins`]).
20//! * The answer is `q` when `p` is a masked id, else `p` ([`feed_argmax`]).
21//!   The synchronous decode step does the same: it takes the device argmax
22//!   unless a `think_ended` row's argmax landed on a masked id, and then
23//!   samples on the host.
24//!
25//! `feed_argmax_tests.rs` checks `q` against the host sampler, `p` against
26//! `argmax_bf16` for rows with a unique maximum, and pins the tie cases.
27
28/// 2026-09-25: "No id" in a row mask.
29pub const NO_ID: u32 = u32::MAX;
30/// 2026-09-25: The floor the plain kernel starts its running maximum at.
31pub const PLAIN_FLOOR: f32 = -1e30;
32/// 2026-09-25: Threads per row block of the kernel; the tie order depends on it.
33pub const BLOCK: usize = 1024;
34
35/// 2026-09-25: BF16 bits to f32, exact (what `__bfloat162float` computes).
36#[inline]
37pub fn bf16_to_f32(bits: u16) -> f32 {
38    f32::from_bits((bits as u32) << 16)
39}
40
41/// 2026-09-25: `p`: the plain kernel argmax of `row`, tie order and floor included.
42pub fn plain_kernel_argmax(row: &[f32]) -> u32 {
43    let mut val = [PLAIN_FLOOR; BLOCK];
44    let mut idx = [0u32; BLOCK];
45    for (i, &v) in row.iter().enumerate() {
46        let t = i % BLOCK;
47        if v > val[t] {
48            val[t] = v;
49            idx[t] = i as u32;
50        }
51    }
52    let mut s = BLOCK / 2;
53    while s > 0 {
54        for t in 0..s {
55            if val[t + s] > val[t] {
56                val[t] = val[t + s];
57                idx[t] = idx[t + s];
58            }
59        }
60        s /= 2;
61    }
62    idx[0]
63}
64
65/// 2026-09-25: The row with the masked ids at `-inf`; out-of-range ids are ignored.
66pub fn masked_row(row: &[f32], mask: [u32; 2]) -> Vec<f32> {
67    let mut out = row.to_vec();
68    for m in mask {
69        if let Some(v) = out.get_mut(m as usize) {
70            *v = f32::NEG_INFINITY;
71        }
72    }
73    out
74}
75
76/// 2026-09-25: `q`: the last index holding the maximum of [`masked_row`], with
77/// the maximum starting at `-inf`; an empty row gives 0.
78pub fn masked_last_wins(row: &[f32], mask: [u32; 2]) -> u32 {
79    let mut best = f32::NEG_INFINITY;
80    let mut idx = 0u32;
81    for (i, &v) in masked_row(row, mask).iter().enumerate() {
82        if v > best || (v == best && i as u32 > idx) {
83            best = v;
84            idx = i as u32;
85        }
86    }
87    idx
88}
89
90/// 2026-09-25: One row of `argmax_bf16_batch_feed`, from the BF16 bits.
91pub fn feed_argmax(row_bf16: &[u16], mask: [u32; 2]) -> u32 {
92    let row: Vec<f32> = row_bf16.iter().map(|&b| bf16_to_f32(b)).collect();
93    feed_argmax_f32(&row, mask)
94}
95
96/// 2026-09-25: [`feed_argmax`] over an f32 row.
97pub fn feed_argmax_f32(row: &[f32], mask: [u32; 2]) -> u32 {
98    let p = plain_kernel_argmax(row);
99    if p == mask[0] || p == mask[1] {
100        masked_last_wins(row, mask)
101    } else {
102        p
103    }
104}
metrale_sampling/feed_argmax.rs

metrale_sampling/
feed_argmax.rs
