spark_model/layers/ops/
sampling.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `ops.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
9use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
10
11use crate::layers::moe;
12use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
13
14use super::*;
15
16/// GPU-side argmax over BF16 logits.
17///
18/// Finds the index of the maximum value, writes a single u32 to `out`.
19///
20/// Kernel: `argmax_bf16(logits, out, n)`
21/// Grid: (1, 1, 1)  Block: (1024, 1, 1)
22pub fn argmax_bf16(
23    gpu: &dyn GpuBackend,
24    kernel: KernelHandle,
25    logits: DevicePtr,
26    out: DevicePtr,
27    vocab_size: u32,
28    stream: u64,
29) -> Result<()> {
30    KernelLaunch::new(gpu, kernel)
31        .grid([1, 1, 1])
32        .block([1024, 1, 1])
33        .arg_ptr(logits)
34        .arg_ptr(out)
35        .arg_u32(vocab_size)
36        .launch(stream)
37}
38
39/// Batched argmax: ONE launch, one block per row, instead of n serial launches of
40/// the single-row `argmax_bf16` (which is a one-CTA reduction and so uses 1 of 48
41/// SMs). Byte-identical — each block runs the identical per-row body.
42#[allow(clippy::too_many_arguments)]
43pub fn argmax_bf16_batch(
44    gpu: &dyn GpuBackend,
45    kernel: KernelHandle,
46    logits: DevicePtr,
47    out: DevicePtr,
48    vocab_size: u32,
49    n_rows: u32,
50    row_stride: u32,
51    stream: u64,
52) -> Result<()> {
53    KernelLaunch::new(gpu, kernel)
54        .grid([n_rows, 1, 1])
55        .block([1024, 1, 1])
56        .arg_ptr(logits)
57        .arg_ptr(out)
58        .arg_u32(vocab_size)
59        .arg_u32(row_stride)
60        .launch(stream)
61}
62
63/// Batched argmax that ALSO writes each row's top-1 log-probability
64/// (`out_logprob[row] = log softmax(row)[argmax]`, FP32), computed by online
65/// softmax in the same pass — same bandwidth as `argmax_bf16_batch`, same
66/// index semantics.
67///
68/// Consumer: D-Cut verification-depth pruning, whose ranking key is the prefix
69/// SUM of these log-probabilities (= the log of the prefix product of survival
70/// probabilities). Separate kernel so every existing `argmax_bf16_batch` caller
71/// stays byte-identical and an unresolved handle is a silent 0 the caller gates
72/// on.
73#[allow(clippy::too_many_arguments)]
74pub fn argmax_bf16_batch_lp(
75    gpu: &dyn GpuBackend,
76    kernel: KernelHandle,
77    logits: DevicePtr,
78    out: DevicePtr,
79    out_logprob: DevicePtr,
80    vocab_size: u32,
81    n_rows: u32,
82    row_stride: u32,
83    stream: u64,
84) -> Result<()> {
85    KernelLaunch::new(gpu, kernel)
86        .grid([n_rows, 1, 1])
87        .block([1024, 1, 1])
88        .arg_ptr(logits)
89        .arg_ptr(out)
90        .arg_ptr(out_logprob)
91        .arg_u32(vocab_size)
92        .arg_u32(row_stride)
93        .launch(stream)
94}
95
96/// GPU-side argmax + embedding lookup — eliminates D2H sync in MTP propose.
97///
98/// Reads the argmax result from `argmax_out`, looks up the embedding row
99/// from `embed_table`, and writes it to `embed_out`. Also copies the token
100/// ID to `token_id_out` for deferred CPU readback.
101pub fn embed_from_argmax(
102    gpu: &dyn GpuBackend,
103    kernel: KernelHandle,
104    argmax_out: DevicePtr,
105    embed_table: DevicePtr,
106    embed_out: DevicePtr,
107    token_id_out: DevicePtr,
108    hidden_size: u32,
109    stream: u64,
110) -> Result<()> {
111    let grid_x = hidden_size.div_ceil(256);
112    KernelLaunch::new(gpu, kernel)
113        .grid([grid_x, 1, 1])
114        .block([256, 1, 1])
115        .arg_ptr(argmax_out)
116        .arg_ptr(embed_table)
117        .arg_ptr(embed_out)
118        .arg_ptr(token_id_out)
119        .arg_u32(hidden_size)
120        .launch(stream)
121}
122
123/// Batched embedding: gather N rows from embedding table in one launch.
124///
125/// Replaces N individual D2D copies with a single kernel.
126/// `token_ids_dev` must point to `[num_tokens]` u32 on device.
127///
128/// Kernel: `batched_embed(token_ids, embed_table, output, hidden_size)`
129/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
130pub fn batched_embed(
131    gpu: &dyn GpuBackend,
132    kernel: KernelHandle,
133    token_ids_dev: DevicePtr,
134    embed_table: DevicePtr,
135    output: DevicePtr,
136    num_tokens: u32,
137    hidden_size: u32,
138    stream: u64,
139) -> Result<()> {
140    KernelLaunch::new(gpu, kernel)
141        .grid([num_tokens, 1, 1])
142        .block([256, 1, 1])
143        .arg_ptr(token_ids_dev)
144        .arg_ptr(embed_table)
145        .arg_ptr(output)
146        .arg_u32(hidden_size)
147        .launch(stream)
148}
149
150/// FP8-table variant of [`batched_embed`]: rows are FP8 E4M3 bytes with a
151/// per-row f32 dequant scale (the `quantize_bf16_to_fp8` layout); the
152/// kernel dequantizes on read and writes BF16 rows.
153///
154/// Kernel: `batched_embed_fp8(token_ids, table, row_scale, output, hidden)`
155/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
156pub fn batched_embed_fp8(
157    gpu: &dyn GpuBackend,
158    kernel: KernelHandle,
159    token_ids_dev: DevicePtr,
160    embed_table: DevicePtr,
161    row_scale: DevicePtr,
162    output: DevicePtr,
163    num_tokens: u32,
164    hidden_size: u32,
165    stream: u64,
166) -> Result<()> {
167    KernelLaunch::new(gpu, kernel)
168        .grid([num_tokens, 1, 1])
169        .block([256, 1, 1])
170        .arg_ptr(token_ids_dev)
171        .arg_ptr(embed_table)
172        .arg_ptr(row_scale)
173        .arg_ptr(output)
174        .arg_u32(hidden_size)
175        .launch(stream)
176}
177
178// ── MoE routing ──────────────────────────────────────────────────
spark_model/layers/ops/sampling.rs

spark_model/layers/ops/
sampling.rs
