refactor(sampling): migrate sampling family into kernels/sampling/

crates/metaltile-codegen/src/spirv/mod.rs

@@ -1057,7 +1057,7 @@ impl GlslGenerator {
                 // step on top of the hardware reciprocal gives correctly-
                 // rounded f32 divide assuming hardware fma is correctly
                 // rounded (true on AMD GFX1xxx V_FMAC_F32). Without this
-                // `test_logits_repetition_penalty` misses bit-exact tol=0
+                // `test_mt_logits_repetition_penalty` misses bit-exact tol=0
                 // by 1 ULP on 2 elements out of 8192.
                 if matches!(bop, BinOpKind::Div) && ty == "float" {
                     writeln!(out, "{pad}{ty} {v} = mt_fdiv({l}, {r});").ok();

crates/metaltile-std/src/ffai/dsv4_indexer_topk.rs

@@ -14,7 +14,7 @@
 //! shared memory. We do a parallel bitonic sort **descending** over
 //! `(score, original_index)` pairs and emit the first `k` indices.
 //! Same Batcher-pattern compare-and-swap as
-//! [`crate::mlx::sort::mt_sort`] but with parallel storage for an
+//! [`crate::kernels::sampling::sort::mt_sort`] but with parallel storage for an
 //! `u32` index tag so the original cache-position survives the swaps.
 //!
 //! Scores below `n_kv` are padded with `-INFINITY` so the sort

crates/metaltile-std/src/ffai/mod.rs

@@ -79,10 +79,6 @@ pub mod im2col_patch_interleaved;
 pub mod kv_cache;
 pub mod kv_cache_update_many;
 pub mod leaky_relu;
-pub mod logits_min_p;
-pub mod logits_processors;
-pub mod logits_top_p;
-pub mod logits_topk;
 pub mod lstm;
 pub mod mel_spectrogram;
 pub mod moe;
@@ -124,7 +120,6 @@ pub mod patch_embed_mma_block_scaled;
 pub mod patch_unfold_qwen;
 pub mod pos_emb_2d_add;
 pub mod resize_normalize;
-pub mod sampling;
 pub mod sdpa_bidirectional;
 pub mod sdpa_bidirectional_d128_relpos;
 pub mod sdpa_bidirectional_windowed;

crates/metaltile-std/src/kernels/mod.rs

@@ -10,3 +10,4 @@
 pub mod convolution;
 pub mod norm;
 pub mod rope;
+pub mod sampling;

crates/metaltile-std/src/ffai/sampling.rs → crates/metaltile-std/src/kernels/sampling/categorical_sample.rs

@@ -40,7 +40,7 @@ use metaltile::kernel;
 // replaced by 1 × n/lsize chunk-traverse per lane + an 8-stage scan +
 // 1 × n/lsize finalizing walk on the winning lane.
 #[kernel]
-pub fn softmax_categorical_sample<T>(
+pub fn mt_softmax_categorical_sample<T>(
     inp: Tensor<T>,
     out: Tensor<u32>,
     temperature_in: Tensor<f32>,
@@ -136,7 +136,7 @@ pub fn softmax_categorical_sample<T>(
 pub mod kernel_tests {
     use metaltile::{test::*, test_kernel};
 
-    use super::softmax_categorical_sample;
+    use super::mt_softmax_categorical_sample;
     use crate::utils::pack_f32;
 
     #[test_kernel(dtypes = [f32, f16, bf16], tol = 0.5)]
@@ -148,7 +148,7 @@ pub mod kernel_tests {
         let (n, idx, spike) = (1024usize, 813usize, 30.0f32);
         let mut logits = vec![0.0f32; n];
         logits[idx] = spike;
-        TestSetup::new(softmax_categorical_sample::kernel_ir_for(dt))
+        TestSetup::new(mt_softmax_categorical_sample::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .input(TestBuffer::from_vec("inp", pack_f32(&logits, dt), dt))
             .input(TestBuffer::zeros("out", 1, DType::U32))
@@ -160,19 +160,19 @@ pub mod kernel_tests {
     }
 }
 
-/// New-syntax benchmark for `softmax_categorical_sample` at Qwen3 vocab
+/// New-syntax benchmark for `mt_softmax_categorical_sample` at Qwen3 vocab
 /// scale (Reduction, single threadgroup; cooperative max + scan + CDF
 /// walk). Random logits, fixed temperature and uniform draw.
 pub mod kernel_benches {
     use metaltile::{bench, test::*};
 
-    use super::softmax_categorical_sample;
+    use super::mt_softmax_categorical_sample;
     use crate::utils::pack_f32;
 
     #[bench(dtypes = [f32, f16, bf16])]
     fn bench_softmax_categorical_sample(dt: DType) -> BenchSetup {
         let n = 152_064usize;
-        BenchSetup::new(softmax_categorical_sample::kernel_ir_for(dt))
+        BenchSetup::new(mt_softmax_categorical_sample::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .buffer(BenchBuffer::random("inp", n, dt))
             .buffer(BenchBuffer::zeros("out", 1, DType::U32).output())

crates/metaltile-std/src/ffai/logits_min_p.rs → crates/metaltile-std/src/kernels/sampling/logits_min_p.rs

@@ -33,7 +33,7 @@
 use metaltile::kernel;
 
 #[kernel]
-pub fn logits_min_p_mask<T>(
+pub fn mt_logits_min_p_mask<T>(
     inp: Tensor<T>,
     out: Tensor<T>,
     #[constexpr] n: u32,
@@ -59,7 +59,7 @@ pub fn logits_min_p_mask<T>(
 pub mod kernel_tests {
     use metaltile::{test::*, test_kernel};
 
-    use super::logits_min_p_mask;
+    use super::mt_logits_min_p_mask;
     use crate::utils::{pack_f32, unpack_f32};
 
     /// CPU oracle: per row, keep `v` iff `exp(v − row_max) ≥ min_p`.
@@ -85,7 +85,7 @@ pub mod kernel_tests {
         let logits: Vec<f32> = (0..n * rows).map(|i| (i % 53) as f32 * 0.2 - 5.0).collect();
         let rounded = unpack_f32(&pack_f32(&logits, dt), dt);
         let expected = cpu_min_p_mask(&rounded, n, rows, min_p);
-        TestSetup::new(logits_min_p_mask::kernel_ir_for(dt))
+        TestSetup::new(mt_logits_min_p_mask::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .input(TestBuffer::from_vec("inp", pack_f32(&logits, dt), dt))
             .input(TestBuffer::zeros("out", n * rows, dt))
@@ -96,17 +96,17 @@ pub mod kernel_tests {
     }
 }
 
-/// New-syntax benchmark for `logits_min_p_mask` at Qwen3 vocab scale
+/// New-syntax benchmark for `mt_logits_min_p_mask` at Qwen3 vocab scale
 /// (Reduction, one TG per row).
 pub mod kernel_benches {
     use metaltile::{bench, test::*};
 
-    use super::logits_min_p_mask;
+    use super::mt_logits_min_p_mask;
 
     #[bench(dtypes = [f32, f16, bf16])]
     fn bench_logits_min_p_mask(dt: DType) -> BenchSetup {
         let (n, rows) = (152_064usize, 2usize);
-        BenchSetup::new(logits_min_p_mask::kernel_ir_for(dt))
+        BenchSetup::new(mt_logits_min_p_mask::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .buffer(BenchBuffer::random("inp", n * rows, dt))
             .buffer(BenchBuffer::zeros("out", n * rows, dt).output())

crates/metaltile-std/src/ffai/logits_processors.rs → crates/metaltile-std/src/kernels/sampling/logits_processors.rs

@@ -52,7 +52,7 @@ use metaltile::kernel;
 //   read/write out of bounds; the runtime should size the dispatch so the
 //   total thread count exactly matches the logits length.
 #[kernel]
-pub fn logits_temperature<T>(inp: Tensor<T>, out: Tensor<T>, #[constexpr] temperature: f32) {
+pub fn mt_logits_temperature<T>(inp: Tensor<T>, out: Tensor<T>, #[constexpr] temperature: f32) {
     let i = program_id::<0>();
     let inv_t = 1.0f32 / temperature;
     let v = load(inp[i]).cast::<f32>();
@@ -89,7 +89,7 @@ pub fn logits_temperature<T>(inp: Tensor<T>, out: Tensor<T>, #[constexpr] temper
 // - **No `threadgroup_*` / `simd_*` cooperation** — every thread is
 //   independent. The only invariant is the dedupe contract above.
 #[kernel]
-pub fn logits_repetition_penalty<T>(
+pub fn mt_logits_repetition_penalty<T>(
     mut logits: Tensor<T>,
     token_ids: Tensor<u32>,
     #[constexpr] penalty: f32,
@@ -104,7 +104,7 @@ pub fn logits_repetition_penalty<T>(
 pub mod kernel_tests {
     use metaltile::{test::*, test_kernel};
 
-    use super::{logits_repetition_penalty, logits_temperature};
+    use super::{mt_logits_repetition_penalty, mt_logits_temperature};
     use crate::utils::{pack_f32, unpack_f32};
 
     fn u32_bytes(v: &[u32]) -> Vec<u8> { v.iter().flat_map(|x| x.to_le_bytes()).collect() }
@@ -118,7 +118,7 @@ pub mod kernel_tests {
         // Kernel computes `v * (1/T)` in f32 then casts back; oracle mirrors.
         let inv_t = 1.0f32 / temperature;
         let expected: Vec<f32> = rounded.iter().map(|&v| v * inv_t).collect();
-        TestSetup::new(logits_temperature::kernel_ir_for(dt))
+        TestSetup::new(mt_logits_temperature::kernel_ir_for(dt))
             .mode(KernelMode::Grid3D)
             .input(TestBuffer::from_vec("inp", pack_f32(&logits, dt), dt))
             .input(TestBuffer::zeros("out", n, dt))
@@ -145,7 +145,7 @@ pub mod kernel_tests {
             let v = expected[tok as usize];
             expected[tok as usize] = if v > 0.0 { v / penalty } else { v * penalty };
         }
-        TestSetup::new(logits_repetition_penalty::kernel_ir_for(dt))
+        TestSetup::new(mt_logits_repetition_penalty::kernel_ir_for(dt))
             .mode(KernelMode::Grid3D)
             .input(TestBuffer::from_vec("logits", pack_f32(&logits, dt), dt))
             .input(TestBuffer::from_vec("token_ids", u32_bytes(&token_ids), DType::U32))
@@ -159,7 +159,7 @@ pub mod kernel_tests {
 pub mod kernel_benches {
     use metaltile::{bench, test::*};
 
-    use super::{logits_repetition_penalty, logits_temperature};
+    use super::{mt_logits_repetition_penalty, mt_logits_temperature};
 
     fn u32_bytes(v: impl Iterator<Item = u32>) -> Vec<u8> {
         v.flat_map(|x| x.to_le_bytes()).collect()
@@ -168,7 +168,7 @@ pub mod kernel_benches {
     #[bench(dtypes = [f32, f16, bf16])]
     fn bench_logits_temperature(dt: DType) -> BenchSetup {
         let n = 152_064usize;
-        BenchSetup::new(logits_temperature::kernel_ir_for(dt))
+        BenchSetup::new(mt_logits_temperature::kernel_ir_for(dt))
             .mode(KernelMode::Grid3D)
             .buffer(BenchBuffer::random("inp", n, dt))
             .buffer(BenchBuffer::zeros("out", n, dt).output())
@@ -182,7 +182,7 @@ pub mod kernel_benches {
         // A modest context window of distinct token ids over a Qwen-scale
         // vocab — one thread per token id.
         let (vocab, n_tokens) = (152_064usize, 2048usize);
-        BenchSetup::new(logits_repetition_penalty::kernel_ir_for(dt))
+        BenchSetup::new(mt_logits_repetition_penalty::kernel_ir_for(dt))
             .mode(KernelMode::Grid3D)
             .buffer(BenchBuffer::random("logits", vocab, dt).output())
             .buffer(BenchBuffer::from_vec(

crates/metaltile-std/src/ffai/logits_top_p.rs → crates/metaltile-std/src/kernels/sampling/logits_top_p.rs

@@ -42,7 +42,7 @@
 use metaltile::kernel;
 
 #[kernel]
-pub fn logits_top_p_mask<T>(
+pub fn mt_logits_top_p_mask<T>(
     inp: Tensor<T>,
     out: Tensor<T>,
     #[constexpr] n: u32,
@@ -102,7 +102,7 @@ pub fn logits_top_p_mask<T>(
 pub mod kernel_tests {
     use metaltile::{test::*, test_kernel};
 
-    use super::logits_top_p_mask;
+    use super::mt_logits_top_p_mask;
     use crate::utils::{pack_f32, unpack_f32};
 
     /// Bisection halvings — must match the kernel's loop bound.
@@ -142,7 +142,7 @@ pub mod kernel_tests {
         let logits: Vec<f32> = (0..n * rows).map(|i| (i % 53) as f32 * 0.2 - 5.0).collect();
         let rounded = unpack_f32(&pack_f32(&logits, dt), dt);
         let expected = cpu_top_p_mask(&rounded, n, rows, top_p);
-        TestSetup::new(logits_top_p_mask::kernel_ir_for(dt))
+        TestSetup::new(mt_logits_top_p_mask::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .input(TestBuffer::from_vec("inp", pack_f32(&logits, dt), dt))
             .input(TestBuffer::zeros("out", n * rows, dt))
@@ -153,17 +153,17 @@ pub mod kernel_tests {
     }
 }
 
-/// New-syntax benchmark for `logits_top_p_mask` at Qwen3 vocab scale
+/// New-syntax benchmark for `mt_logits_top_p_mask` at Qwen3 vocab scale
 /// (Reduction, one TG per row; 24 bisection passes over the row).
 pub mod kernel_benches {
     use metaltile::{bench, test::*};
 
-    use super::logits_top_p_mask;
+    use super::mt_logits_top_p_mask;
 
     #[bench(dtypes = [f32, f16, bf16])]
     fn bench_logits_top_p_mask(dt: DType) -> BenchSetup {
         let (n, rows) = (152_064usize, 2usize);
-        BenchSetup::new(logits_top_p_mask::kernel_ir_for(dt))
+        BenchSetup::new(mt_logits_top_p_mask::kernel_ir_for(dt))
             .mode(KernelMode::Reduction)
             .buffer(BenchBuffer::random("inp", n * rows, dt))
             .buffer(BenchBuffer::zeros("out", n * rows, dt).output())

crates/metaltile-std/src/ffai/logits_topk.rs → crates/metaltile-std/src/kernels/sampling/logits_topk.rs

@@ -38,7 +38,7 @@
 use metaltile::kernel;
 
 #[kernel]
-pub fn logits_topk_mask<T>(inp: Tensor<T>, out: Tensor<T>, #[constexpr] threshold: f32) {
+pub fn mt_logits_topk_mask<T>(inp: Tensor<T>, out: Tensor<T>, #[constexpr] threshold: f32) {
     let i = program_id::<0>();
     let v = load(inp[i]).cast::<f32>();
     // `select(cond, lhs, rhs)` returns lhs when cond is true.
@@ -51,7 +51,7 @@ pub fn logits_topk_mask<T>(inp: Tensor<T>, out: Tensor<T>, #[constexpr] threshol
 pub mod kernel_tests {
     use metaltile::{test::*, test_kernel};
 
-    use super::logits_topk_mask;
+    use super::mt_logits_topk_mask;
     use crate::utils::{pack_f32, unpack_f32};
 
     /// K-th largest value (descending) — how callers pre-compute the cutoff.
@@ -72,7 +72,7 @@ pub mod kernel_tests {
         let threshold = kth_largest(&rounded, k);
         let expected: Vec<f32> =
             rounded.iter().map(|&v| if v >= threshold { v } else { f32::NEG_INFINITY }).collect();
-        TestSetup::new(logits_topk_mask::kernel_ir_for(dt))
+        TestSetup::new(mt_logits_topk_mask::kernel_ir_for(dt))
             .mode(KernelMode::Grid3D)
             .input(TestBuffer::from_vec("inp", pack_f32(&logits, dt), dt))
             .input(TestBuffer::zeros("out", n, dt))
@@ -82,17 +82,17 @@ pub mod kernel_tests {
     }
 }
 
-/// New-syntax benchmark for `logits_topk_mask` at Qwen3 vocab scale
+/// New-syntax benchmark for `mt_logits_topk_mask` at Qwen3 vocab scale
 /// (Grid3D, one thread per vocab position).
 pub mod kernel_benches {
     use metaltile::{bench, test::*};
 
-    use super::logits_topk_mask;
+    use super::mt_logits_topk_mask;
 
     #[bench(dtypes = [f32, f16, bf16])]
     fn bench_logits_topk_mask(dt: DType) -> BenchSetup {
         let n = 152_064usize;
-        BenchSetup::new(logits_topk_mask::kernel_ir_for(dt))
+        BenchSetup::new(mt_logits_topk_mask::kernel_ir_for(dt))
             .mode(KernelMode::Grid3D)
             .buffer(BenchBuffer::random("inp", n, dt))
             .buffer(BenchBuffer::zeros("out", n, dt).output())

crates/metaltile-std/src/kernels/sampling/mod.rs

@@ -0,0 +1,15 @@
+//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric
+//! SPDX-License-Identifier: Apache-2.0
+//! Sampling kernels — the sampling family (see
+//! `docs/specs/KERNEL_CONSOLIDATION_PLAN.md`): the logits→token pipeline.
+//! Softmax and sort (from the legacy `mlx/`), the fused categorical sampler,
+//! and the logits processors / masks (top-k, top-p, min-p, temperature,
+//! repetition penalty) from the legacy `ffai/`.
+
+pub mod categorical_sample;
+pub mod logits_min_p;
+pub mod logits_processors;
+pub mod logits_top_p;
+pub mod logits_topk;
+pub mod softmax;
+pub mod sort;

crates/metaltile-std/src/mlx/mod.rs

@@ -52,8 +52,6 @@ pub mod scan;
 pub mod scatter_axis;
 pub mod sdpa_vector;
 pub mod sgload_smoke;
-pub mod softmax;
-pub mod sort;
 pub mod steel;
 pub mod strided;
 pub mod swiglu;

crates/metaltile-std/tests/kernel_registry_consistency.rs

@@ -322,7 +322,7 @@ fn no_undefined_mt_symbols_in_emitted_msl() {
 #[test]
 fn kernel_annotations_have_matching_inventory_submit() {
     // Collect kernel-named registry entries (kernel_name → matched-via-name)
-    // PLUS kernel_ir function paths (e.g. `softmax_categorical_sample::kernel_ir_for`)
+    // PLUS kernel_ir function paths (e.g. `mt_softmax_categorical_sample::kernel_ir_for`)
     // so we can match either form. The macro expands `#[kernel]` to
     // give each function a `kernel_ir_for(dt) -> Kernel` associated
     // function inside a module of the same name; `inventory::submit!`

crates/metaltile-std/tests/softmax_categorical_sample_matrix.rs

@@ -1,6 +1,6 @@
 //! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric
 //! SPDX-License-Identifier: Apache-2.0
-//! Matrix coverage for `ffai::softmax_categorical_sample` —
+//! Matrix coverage for `ffai::mt_softmax_categorical_sample` —
 //! dtype × vocab × temperature × distribution × uniform-draw, plus
 //! invariant property assertions that catch silent codegen bugs
 //! numeric drift alone can't detect.
@@ -61,7 +61,7 @@ use std::collections::BTreeMap;
 
 use common::{Dt, gpu_lock, pack_bytes};
 use metaltile::{Context, core::ir::KernelMode};
-use metaltile_std::ffai::sampling::softmax_categorical_sample;
+use metaltile_std::kernels::sampling::categorical_sample::mt_softmax_categorical_sample;
 
 // ── CPU oracle (mirrors the kernel's three-pass shape) ────────────────
 
@@ -102,7 +102,7 @@ fn run_sample(ctx: &Context, dt: Dt, logits: &[f32], temperature: f32, uniform:
     buffers.insert("uniform_in".into(), uniform.to_le_bytes().to_vec());
     buffers.insert("n".into(), (n as u32).to_le_bytes().to_vec());
 
-    let mut kernel = softmax_categorical_sample::kernel_ir_for(dt.to_dtype());
+    let mut kernel = mt_softmax_categorical_sample::kernel_ir_for(dt.to_dtype());
     kernel.mode = KernelMode::Reduction;
 
     // Fixed TPG = 256 per kernel invariant (tg_max + tg_sum both 256-wide,