mlx_native/ops/
softmax.rs

1//! Numerically stable softmax GPU dispatch.
2//!
3//! Computes softmax along the last dimension of a 2D tensor using the
4//! subtract-max trick for numerical stability.  All accumulations use f32
5//! even when inputs are f16 to prevent overflow.
6
7use metal::MTLSize;
8
9use crate::buffer::MlxBuffer;
10use crate::dtypes::DType;
11use crate::encoder::{CapturedOpKind, CommandEncoder};
12use crate::error::{MlxError, Result};
13use crate::kernel_registry::KernelRegistry;
14
15/// MSL source for the softmax kernels (embedded at compile time).
16pub static SOFTMAX_SHADER_SOURCE: &str = include_str!("../shaders/softmax.metal");
17
18/// Register softmax shader sources with the given kernel registry.
19pub fn register(registry: &mut KernelRegistry) {
20    registry.register_source("softmax_f32", SOFTMAX_SHADER_SOURCE);
21    registry.register_source("softmax_f16", SOFTMAX_SHADER_SOURCE);
22    registry.register_source("softmax_bf16", SOFTMAX_SHADER_SOURCE);
23}
24
25/// Dispatch a softmax operation on the GPU.
26///
27/// # Arguments
28///
29/// * `encoder`    - Command encoder to record the dispatch into.
30/// * `registry`   - Kernel registry (must have softmax sources registered).
31/// * `device`     - Metal device for pipeline compilation.
32/// * `input`      - Input buffer of shape `[rows, cols]` (f32, f16, or bf16).
33/// * `output`     - Output buffer (same dtype and shape as input).
34/// * `params_buf` - Params buffer containing `[cols, 0]` as f32.
35/// * `rows`       - Number of rows.
36/// * `cols`       - Number of columns (softmax dimension).
37///
38/// # Errors
39///
40/// Returns `MlxError::InvalidArgument` if:
41/// - Input dtype is not f32, f16, or bf16.
42/// - Input element count does not match rows * cols.
43pub fn dispatch_softmax(
44    encoder: &mut CommandEncoder,
45    registry: &mut KernelRegistry,
46    device: &metal::DeviceRef,
47    input: &MlxBuffer,
48    output: &MlxBuffer,
49    params_buf: &MlxBuffer,
50    rows: u32,
51    cols: u32,
52) -> Result<()> {
53    if rows == 0 || cols == 0 {
54        return Err(MlxError::InvalidArgument(
55            "Softmax rows and cols must be > 0".into(),
56        ));
57    }
58
59    let expected = (rows as usize) * (cols as usize);
60    if input.element_count() != expected {
61        return Err(MlxError::InvalidArgument(format!(
62            "Softmax input element count {} != rows({}) * cols({})",
63            input.element_count(),
64            rows,
65            cols
66        )));
67    }
68    if output.element_count() != expected {
69        return Err(MlxError::InvalidArgument(format!(
70            "Softmax output element count {} != rows({}) * cols({})",
71            output.element_count(),
72            rows,
73            cols
74        )));
75    }
76    // hf2q ADR-030 iter-113 — defense-in-depth: kernel selected by
77    // input.dtype() writes output at the same stride.  Mismatched output
78    // dtype would mis-stride writes.  See ADR-030 iter-110/111/112.
79    if input.dtype() != output.dtype() {
80        return Err(MlxError::InvalidArgument(format!(
81            "Softmax dtype mismatch: input={} != output={}",
82            input.dtype(), output.dtype(),
83        )));
84    }
85
86    let kernel_name = match input.dtype() {
87        DType::F32 => "softmax_f32",
88        DType::F16 => "softmax_f16",
89        DType::BF16 => "softmax_bf16",
90        _ => {
91            return Err(MlxError::InvalidArgument(format!(
92                "Softmax unsupported dtype: {}",
93                input.dtype()
94            )));
95        }
96    };
97
98    let pipeline = registry.get_pipeline(kernel_name, device)?;
99
100    // One threadgroup per row.  Threadgroup size must be a power of 2
101    // for the tree reduction to work correctly.
102    let tg_size = std::cmp::min(256, cols.next_power_of_two()) as u64;
103
104    // Threadgroup shared memory: tg_size floats for the reduction.
105    let shared_mem_bytes = tg_size * 4; // sizeof(float) = 4
106
107    // Tag for the reorder pass (Phase 4e.3): Softmax is NOT reorderable.
108    encoder.set_op_kind(CapturedOpKind::Softmax);
109
110    encoder.encode_threadgroups_with_shared(
111        pipeline,
112        &[(0, input), (1, output), (2, params_buf)],
113        &[(0, shared_mem_bytes)],
114        MTLSize::new(rows as u64, 1, 1),
115        MTLSize::new(tg_size, 1, 1),
116    );
117
118    Ok(())
119}
mlx_native/ops/softmax.rs

mlx_native/ops/
softmax.rs
