Module feed_argmax 

2026-09-25: Host reference for one row of argmax_bf16_batch_feed (kernels/gb10/common/argmax_feed.cu), computed from the same BF16 bytes.

Owner: metrale-sampling. Invariants: none beyond the types. NaN input is outside the contract.

The kernel runs two reductions over a row and combines them:

• p: the plain argmax, the same scan and tree reduction as argmax_bf16_batch. Each of the BLOCK threads keeps the first strict maximum of its strided slice above the -1e30 floor. At each merge of threads t and t + s, t keeps its value unless t + s holds a strictly larger one. Ties therefore follow the tree order, not the lowest index. plain_kernel_argmax simulates this, floor included.
• q: the row with the two masked ids set to -inf (as PostCloseThinkMask does) and the last index holding the maximum, the tie rule of the temperature-0 host sampler (masked_last_wins).
• The answer is q when p is a masked id, else p (feed_argmax). The synchronous decode step does the same: it takes the device argmax unless a think_ended row’s argmax landed on a masked id, and then samples on the host.

feed_argmax_tests.rs checks q against the host sampler, p against argmax_bf16 for rows with a unique maximum, and pins the tie cases.

Constants

BLOCK

2026-09-25: Threads per row block of the kernel; the tie order depends on it.

NO_ID

2026-09-25: “No id” in a row mask.

PLAIN_FLOOR

2026-09-25: The floor the plain kernel starts its running maximum at.

Functions

bf16_to_f32

2026-09-25: BF16 bits to f32, exact (what __bfloat162float computes).

feed_argmax

2026-09-25: One row of argmax_bf16_batch_feed, from the BF16 bits.

feed_argmax_f32

2026-09-25: feed_argmax over an f32 row.

masked_last_wins

2026-09-25: q: the last index holding the maximum of masked_row, with the maximum starting at -inf; an empty row gives 0.

masked_row

2026-09-25: The row with the masked ids at -inf; out-of-range ids are ignored.

plain_kernel_argmax

2026-09-25: p: the plain kernel argmax of row, tie order and floor included.