mlx_native/ops/
gpu_sample.rs

1//! ADR-040 §26 iter-M — GPU-side first-max argmax + threshold candidate collect.
2//!
3//! Replaces the host full-vocab argmax + candidate-threshold scans (~0.92ms/step
4//! on the autoregressive critical path) with one GPU dispatch that reads back
5//! only the per-slot top1 + the few threshold candidates. The host keeps the
6//! cheap F64 rerank over those candidates (Metal has no f64). Byte-matches the
7//! host `argmax_f32_first_max` (first-max, lower-index tie-break) + the finalize
8//! threshold scan (logits >= top1_val - 0.5f).
9
10use metal::MTLSize;
11
12use crate::buffer::MlxBuffer;
13use crate::encoder::CommandEncoder;
14use crate::error::{MlxError, Result};
15use crate::kernel_registry::KernelRegistry;
16
17pub static GPU_SAMPLE_SHADER_SOURCE: &str =
18    include_str!("../shaders/gpu_sample_argmax_candidates.metal");
19
20pub fn register(registry: &mut KernelRegistry) {
21    registry.register_source("gpu_sample_argmax_candidates", GPU_SAMPLE_SHADER_SOURCE);
22}
23
24/// Dispatch GPU argmax+candidate-collect over `[n_slots, vocab]` logits.
25///
26/// Outputs (per slot): `out_top1_idx[n]`, `out_top1_val[n]`,
27/// `out_cand_count[n]` (atomic u32 — total count, may exceed `cap`),
28/// `out_overflow[n]` (1 if count>cap), `out_cand_ids[n*cap]`.
29#[allow(clippy::too_many_arguments)]
30pub fn dispatch_gpu_sample_argmax_candidates(
31    encoder: &mut CommandEncoder,
32    registry: &mut KernelRegistry,
33    device: &metal::DeviceRef,
34    logits: &MlxBuffer,
35    out_top1_idx: &MlxBuffer,
36    out_top1_val: &MlxBuffer,
37    out_cand_count: &MlxBuffer,
38    out_overflow: &MlxBuffer,
39    out_cand_ids: &MlxBuffer,
40    params_buf: &MlxBuffer,
41    n_slots: u32,
42    vocab: u32,
43    cap: u32,
44) -> Result<()> {
45    if n_slots == 0 || vocab == 0 || cap == 0 {
46        return Err(MlxError::InvalidArgument(
47            "gpu_sample: n_slots, vocab, cap must all be > 0".into(),
48        ));
49    }
50    if logits.element_count() < (n_slots * vocab) as usize {
51        return Err(MlxError::InvalidArgument(format!(
52            "gpu_sample: logits {} < n_slots*vocab {}",
53            logits.element_count(),
54            n_slots * vocab
55        )));
56    }
57    if out_cand_ids.element_count() < (n_slots * cap) as usize {
58        return Err(MlxError::InvalidArgument(
59            "gpu_sample: out_cand_ids too small".into(),
60        ));
61    }
62
63    let pipeline = registry.get_pipeline("gpu_sample_argmax_candidates", device)?;
64
65    // Power-of-two threadgroup for the tree reduction; 1024 (each thread scans
66    // ~256 cols at vocab=262144).
67    let tg_size: u64 = std::cmp::min(1024, vocab.next_power_of_two() as u64).max(1);
68    let float_shared = tg_size * 4;
69    let uint_shared = tg_size * 4;
70
71    encoder.encode_threadgroups_with_shared(
72        pipeline,
73        &[
74            (0, logits),
75            (1, out_top1_idx),
76            (2, out_top1_val),
77            (3, out_cand_count),
78            (4, out_overflow),
79            (5, out_cand_ids),
80            (6, params_buf),
81        ],
82        &[(0, float_shared), (1, uint_shared)],
83        MTLSize::new(n_slots as u64, 1, 1), // one threadgroup per slot
84        MTLSize::new(tg_size, 1, 1),
85    );
86
87    Ok(())
88}
mlx_native/ops/gpu_sample.rs

mlx_native/ops/
gpu_sample.rs
