metrale_model_layers/layers/ops/
sampling.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the device-side argmax kernels (single row,
4//! batched, with top-1 log-probability, and the feed-cell variant), the embedding
5//! gathers, and the device token feed resolver.
6//!
7//! Owner: model-layers ops.
8//! Invariants: none beyond the types.
9
10#![allow(unused_imports)]
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16use crate::layers::moe;
17use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
18
19use super::*;
20
21/// 2026-09-25: Index of the largest of `vocab_size` BF16 logits, written as one
22/// u32 to `out`, by a single block.
23pub fn argmax_bf16(
24    gpu: &dyn GpuBackend,
25    kernel: KernelHandle,
26    logits: DevicePtr,
27    out: DevicePtr,
28    vocab_size: u32,
29    stream: u64,
30) -> Result<()> {
31    KernelLaunch::new(gpu, kernel)
32        .grid([1, 1, 1])
33        .block([1024, 1, 1])
34        .arg_ptr(logits)
35        .arg_ptr(out)
36        .arg_u32(vocab_size)
37        .launch(stream)
38}
39
40/// 2026-09-25: [`argmax_bf16`] for `n_rows` rows `row_stride` apart, one block
41/// per row. Each block runs the per-row body of `argmax_bf16`
42/// (`argmax_bf16.cu`), so a row's index equals what `argmax_bf16` returns for it.
43#[allow(clippy::too_many_arguments)]
44pub fn argmax_bf16_batch(
45    gpu: &dyn GpuBackend,
46    kernel: KernelHandle,
47    logits: DevicePtr,
48    out: DevicePtr,
49    vocab_size: u32,
50    n_rows: u32,
51    row_stride: u32,
52    stream: u64,
53) -> Result<()> {
54    KernelLaunch::new(gpu, kernel)
55        .grid([n_rows, 1, 1])
56        .block([1024, 1, 1])
57        .arg_ptr(logits)
58        .arg_ptr(out)
59        .arg_u32(vocab_size)
60        .arg_u32(row_stride)
61        .launch(stream)
62}
63
64/// 2026-09-25: [`argmax_bf16_batch`] that also writes each row's top-1
65/// log-probability, `out_logprob[row] = log softmax(row)[argmax]` in FP32,
66/// from an online softmax in the same pass over the row. D-Cut ranks
67/// verification depths by prefix sums of these values.
68#[allow(clippy::too_many_arguments)]
69pub fn argmax_bf16_batch_lp(
70    gpu: &dyn GpuBackend,
71    kernel: KernelHandle,
72    logits: DevicePtr,
73    out: DevicePtr,
74    out_logprob: DevicePtr,
75    vocab_size: u32,
76    n_rows: u32,
77    row_stride: u32,
78    stream: u64,
79) -> Result<()> {
80    KernelLaunch::new(gpu, kernel)
81        .grid([n_rows, 1, 1])
82        .block([1024, 1, 1])
83        .arg_ptr(logits)
84        .arg_ptr(out)
85        .arg_ptr(out_logprob)
86        .arg_u32(vocab_size)
87        .arg_u32(row_stride)
88        .launch(stream)
89}
90
91/// 2026-09-25: Read the token id at `argmax_out`, copy its BF16 row of
92/// `embed_table` to `embed_out`, and copy the id to `token_id_out`, all on the
93/// device.
94pub fn embed_from_argmax(
95    gpu: &dyn GpuBackend,
96    kernel: KernelHandle,
97    argmax_out: DevicePtr,
98    embed_table: DevicePtr,
99    embed_out: DevicePtr,
100    token_id_out: DevicePtr,
101    hidden_size: u32,
102    stream: u64,
103) -> Result<()> {
104    let grid_x = hidden_size.div_ceil(256);
105    KernelLaunch::new(gpu, kernel)
106        .grid([grid_x, 1, 1])
107        .block([256, 1, 1])
108        .arg_ptr(argmax_out)
109        .arg_ptr(embed_table)
110        .arg_ptr(embed_out)
111        .arg_ptr(token_id_out)
112        .arg_u32(hidden_size)
113        .launch(stream)
114}
115
116/// 2026-09-25: Gather the BF16 embedding rows of `num_tokens` token ids
117/// (`token_ids_dev`, device `[num_tokens]` u32) into `output`, one block per
118/// token.
119pub fn batched_embed(
120    gpu: &dyn GpuBackend,
121    kernel: KernelHandle,
122    token_ids_dev: DevicePtr,
123    embed_table: DevicePtr,
124    output: DevicePtr,
125    num_tokens: u32,
126    hidden_size: u32,
127    stream: u64,
128) -> Result<()> {
129    KernelLaunch::new(gpu, kernel)
130        .grid([num_tokens, 1, 1])
131        .block([256, 1, 1])
132        .arg_ptr(token_ids_dev)
133        .arg_ptr(embed_table)
134        .arg_ptr(output)
135        .arg_u32(hidden_size)
136        .launch(stream)
137}
138
139/// 2026-09-25: [`batched_embed`] over an FP8 E4M3 table with one f32 scale per
140/// row (the `quantize_bf16_to_fp8` layout); rows are dequantized to BF16 on
141/// read.
142pub fn batched_embed_fp8(
143    gpu: &dyn GpuBackend,
144    kernel: KernelHandle,
145    token_ids_dev: DevicePtr,
146    embed_table: DevicePtr,
147    row_scale: DevicePtr,
148    output: DevicePtr,
149    num_tokens: u32,
150    hidden_size: u32,
151    stream: u64,
152) -> Result<()> {
153    KernelLaunch::new(gpu, kernel)
154        .grid([num_tokens, 1, 1])
155        .block([256, 1, 1])
156        .arg_ptr(token_ids_dev)
157        .arg_ptr(embed_table)
158        .arg_ptr(row_scale)
159        .arg_ptr(output)
160        .arg_u32(hidden_size)
161        .launch(stream)
162}
163
164/// 2026-09-25: Bit of a [`feed_resolve`] source word that marks an inline host
165/// token id in the low 31 bits; without it the word indexes the previous
166/// step's feed cell (`argmax_feed.cu`).
167pub const FEED_HOST_BIT: u32 = 0x8000_0000;
168
169/// 2026-09-25: Batched argmax into the feed `cells` (device u32 per row) with a
170/// per-row pair of masked ids `masks` (`[m0, m1]` per row, `u32::MAX` = none):
171/// when the plain argmax is a masked id, the row gets the argmax with both ids
172/// excluded. `argmax_feed.cu` states the tie rules. One block per row.
173pub fn argmax_bf16_batch_feed(
174    gpu: &dyn GpuBackend,
175    kernel: KernelHandle,
176    logits: DevicePtr,
177    masks: DevicePtr,
178    cells: DevicePtr,
179    vocab_size: u32,
180    n_rows: u32,
181    row_stride: u32,
182    stream: u64,
183) -> Result<()> {
184    KernelLaunch::new(gpu, kernel)
185        .grid([n_rows, 1, 1])
186        .block([1024, 1, 1])
187        .arg_ptr(logits)
188        .arg_ptr(masks)
189        .arg_ptr(cells)
190        .arg_u32(vocab_size)
191        .arg_u32(row_stride)
192        .launch(stream)
193}
194
195/// 2026-09-25: `ids_out[i]` is the inline id when `sources[i]` has
196/// [`FEED_HOST_BIT`] set, else `cells[sources[i]]`.
197pub fn feed_resolve(
198    gpu: &dyn GpuBackend,
199    kernel: KernelHandle,
200    sources: DevicePtr,
201    cells: DevicePtr,
202    ids_out: DevicePtr,
203    n_rows: u32,
204    stream: u64,
205) -> Result<()> {
206    KernelLaunch::new(gpu, kernel)
207        .grid([1, 1, 1])
208        .block([128, 1, 1])
209        .arg_ptr(sources)
210        .arg_ptr(cells)
211        .arg_ptr(ids_out)
212        .arg_u32(n_rows)
213        .launch(stream)
214}
metrale_model_layers/layers/ops/sampling.rs

metrale_model_layers/layers/ops/
sampling.rs
