RadixTopKSmall

crates/backend-uzu/src/backends/common/kernel/mod.rs

@@ -4,6 +4,7 @@ pub mod attention_gemm;
 pub mod delta_net_chunked_prefill;
 pub mod delta_net_tree_verify;
 pub mod matmul;
+pub mod radix_top_k_small;
 
 include!(concat!(env!("OUT_DIR"), "/traits.rs"));
 
@@ -15,6 +16,7 @@ pub trait Kernels: Sized {
     type DeltaNetChunkedPrefill: delta_net_chunked_prefill::DeltaNetChunkedPrefill<Self::Backend>;
     type DeltaNetTreeVerify: delta_net_tree_verify::DeltaNetTreeVerify<Self::Backend>;
     type MatmulKernel: matmul::MatmulKernel<Backend = Self::Backend>;
+    type RadixTopKSmall: radix_top_k_small::RadixTopKSmall<Self::Backend>;
 }
 
 #[cfg(test)]

crates/backend-uzu/src/backends/common/kernel/radix_top_k_small.rs

@@ -0,0 +1,20 @@
+use crate::backends::common::{Allocation, Backend, Encoder};
+
+pub const MAX_K: u32 = 512;
+
+pub trait RadixTopKSmall<B: Backend>: Sized {
+    fn new(
+        context: &B::Context,
+        columns: u32,
+    ) -> Result<Self, B::Error>;
+
+    fn encode(
+        &self,
+        input: &Allocation<B>,
+        output_ids: &mut Allocation<B>,
+        output_scores: &mut Allocation<B>,
+        rows: u32,
+        k: u32,
+        encoder: &mut Encoder<B>,
+    ) -> Result<(), B::Error>;
+}

crates/backend-uzu/src/backends/cpu/kernel/mod.rs

@@ -13,6 +13,7 @@ mod matmul;
 mod moe;
 mod normalization;
 mod pooling;
+mod radix_top_k_small;
 mod sampling;
 mod short_conv;
 mod softmax;
@@ -34,4 +35,5 @@ impl Kernels for CpuKernels {
     type DeltaNetChunkedPrefill = Infallible;
     type DeltaNetTreeVerify = Infallible;
     type MatmulKernel = matmul::MatmulCpuKernel;
+    type RadixTopKSmall = radix_top_k_small::CpuRadixTopKSmall;
 }

crates/backend-uzu/src/backends/cpu/kernel/radix_top_k_small.rs

@@ -0,0 +1,72 @@
+use crate::{
+    backends::{
+        common::{
+            Allocation, AsBufferRangeMut, AsBufferRangeRef, Encoder,
+            kernel::radix_top_k_small::{MAX_K, RadixTopKSmall},
+        },
+        cpu::{Cpu, context::CpuContext, error::CpuError},
+    },
+    utils::pointers::{SendPtr, SendPtrMut},
+};
+
+pub struct CpuRadixTopKSmall {
+    columns: usize,
+}
+
+impl RadixTopKSmall<Cpu> for CpuRadixTopKSmall {
+    fn new(
+        _context: &CpuContext,
+        columns: u32,
+    ) -> Result<Self, CpuError> {
+        assert!(columns > 0);
+        Ok(Self {
+            columns: columns as usize,
+        })
+    }
+
+    fn encode(
+        &self,
+        input: &Allocation<Cpu>,
+        output_ids: &mut Allocation<Cpu>,
+        output_scores: &mut Allocation<Cpu>,
+        rows: u32,
+        k: u32,
+        encoder: &mut Encoder<Cpu>,
+    ) -> Result<(), CpuError> {
+        let rows = rows as usize;
+        let columns = self.columns;
+        let k = k as usize;
+        assert!(rows > 0 && k > 0 && k <= MAX_K as usize && k <= columns);
+        let input = input.as_buffer_range_ref();
+        let input = SendPtr(unsafe { (&*input.buffer().get()).as_ptr().add(input.range().start).cast::<f32>() });
+        let output_ids = output_ids.as_buffer_range_mut();
+        let output_ids = SendPtrMut(unsafe {
+            (&mut *output_ids.buffer().get()).as_mut_ptr().add(output_ids.range().start).cast::<u32>()
+        });
+        let output_scores = output_scores.as_buffer_range_mut();
+        let output_scores = SendPtrMut(unsafe {
+            (&mut *output_scores.buffer().get()).as_mut_ptr().add(output_scores.range().start).cast::<f32>()
+        });
+        encoder.as_command_buffer_mut().push_command(move || {
+            let values = unsafe { std::slice::from_raw_parts(input.as_ptr(), rows * columns) };
+            let output_ids = unsafe { std::slice::from_raw_parts_mut(output_ids.as_ptr(), rows * k) };
+            let output_scores = unsafe { std::slice::from_raw_parts_mut(output_scores.as_ptr(), rows * k) };
+            for (row, values) in values.chunks_exact(columns).enumerate() {
+                let mut indices = (0..columns).collect::<Vec<_>>();
+                let compare = |&left: &usize, &right: &usize| {
+                    values[right].total_cmp(&values[left]).then_with(|| left.cmp(&right))
+                };
+                if k < columns {
+                    indices.select_nth_unstable_by(k, compare);
+                    indices.truncate(k);
+                }
+                indices.sort_unstable_by(compare);
+                for (rank, index) in indices.into_iter().enumerate() {
+                    output_ids[row * k + rank] = index as u32;
+                    output_scores[row * k + rank] = values[index];
+                }
+            }
+        });
+        Ok(())
+    }
+}

crates/backend-uzu/src/backends/metal/kernel/common/top_k.h

@@ -0,0 +1,22 @@
+#pragma once
+
+#include <metal_stdlib>
+using namespace metal;
+
+constant uint TOP_K_SIGN_BIT = 0x80000000u;
+
+static inline uint top_k_score_key(float score) {
+  const uint bits = as_type<uint>(score);
+  return (bits & TOP_K_SIGN_BIT) ? ~bits : (bits ^ TOP_K_SIGN_BIT);
+}
+
+static inline float top_k_score_from_key(uint key) {
+  const uint bits = (key & TOP_K_SIGN_BIT) ? (key ^ TOP_K_SIGN_BIT) : ~key;
+  return as_type<float>(bits);
+}
+
+// Larger keys follow the f32 total order; equal scores prefer smaller indices.
+static inline ulong top_k_ordered_key(float score, uint index, uint index_bits) {
+  const ulong index_mask = (1ul << index_bits) - 1ul;
+  return (ulong(top_k_score_key(score)) << index_bits) | (index_mask - ulong(index));
+}

crates/backend-uzu/src/backends/metal/kernel/mod.rs

@@ -9,6 +9,7 @@ use crate::backends::{
 pub mod attention;
 pub mod gdn;
 pub mod matmul;
+mod radix_top_k_small;
 
 pub const MTLB: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/default.metallib"));
 
@@ -24,4 +25,5 @@ impl Kernels for MetalKernels {
     type DeltaNetChunkedPrefill = gdn::chunked::MetalDeltaNetChunkedPrefill;
     type DeltaNetTreeVerify = gdn::tree_verify::MetalDeltaNetTreeVerify;
     type MatmulKernel = matmul::MatmulMetalKernel;
+    type RadixTopKSmall = radix_top_k_small::MetalRadixTopKSmall;
 }

crates/backend-uzu/src/backends/metal/kernel/radix_top_k_small.metal

@@ -0,0 +1,216 @@
+#include <metal_stdlib>
+#include <metal_atomic>
+#include "common/defines.h"
+#include "common/dsl.h"
+#include "common/top_k.h"
+
+using namespace metal;
+
+#define THREADS_PER_TG 256
+constant uint MAX_K = 512;
+constant uint RADIX_BITS = 10;
+constant uint BUCKETS = 1 << RADIX_BITS;
+constant uint BUCKETS_PER_THREAD = BUCKETS / THREADS_PER_TG;
+constant uint VECTOR_WIDTH = 4;
+
+METAL_FUNC bool arrive_last(device atomic_uint* count, uint expected) {
+  atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);
+  const bool last = atomic_fetch_add_explicit(count, 1u, memory_order_relaxed) == expected - 1;
+  if (last)
+    atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);
+  return last;
+}
+
+METAL_FUNC void reset_arrival(device atomic_uint* count) {
+  atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);
+  atomic_store_explicit(count, 0u, memory_order_relaxed);
+}
+
+template <typename Visitor>
+METAL_FUNC void visit_partition_keys(
+    const device float* input,
+    uint columns,
+    uint index_bits,
+    uint partition,
+    uint partitions,
+    uint lid,
+    Visitor visit
+) {
+  if (columns % VECTOR_WIDTH != 0) {
+    const uint begin = columns * partition / partitions;
+    const uint end = columns * (partition + 1u) / partitions;
+    for (uint column = begin + lid; column < end; column += THREADS_PER_TG)
+      visit(top_k_ordered_key(input[column], column, index_bits));
+    return;
+  }
+  const device float4* input4 = reinterpret_cast<const device float4*>(input);
+  const uint vector_columns = columns / VECTOR_WIDTH;
+  const uint vector_begin = vector_columns * partition / partitions;
+  const uint vector_end = vector_columns * (partition + 1u) / partitions;
+  for (uint vector = vector_begin + lid; vector < vector_end; vector += THREADS_PER_TG) {
+    const float4 values = input4[vector];
+    for (uint lane = 0; lane < VECTOR_WIDTH; ++lane)
+      visit(top_k_ordered_key(values[lane], vector * VECTOR_WIDTH + lane, index_bits));
+  }
+  if (partition + 1 == partitions) {
+    for (uint column = vector_columns * VECTOR_WIDTH + lid; column < columns; column += THREADS_PER_TG)
+      visit(top_k_ordered_key(input[column], column, index_bits));
+  }
+}
+
+KERNEL(RadixTopKSmallPass)(
+    const device float* input,
+    device atomic_uint* partial_histograms,
+    device ulong* prefixes,
+    device ulong* prefix_masks,
+    device uint* ranks,
+    device atomic_uint* done_counts,
+    constant uint& rows,
+    constant uint& k,
+    constant uint& pass,
+    const uint columns SPECIALIZE,
+    const uint partitions SPECIALIZE,
+    threadgroup uint histogram[BUCKETS],
+    threadgroup uint simd_prefixes[THREADS_PER_TG / METAL_SIMD_SIZE],
+    threadgroup uint& is_last,
+    const uint group GROUPS(rows* partitions),
+    const uint lid THREADS(THREADS_PER_TG)
+) {
+  const uint row = group / partitions;
+  const uint partition = group % partitions;
+  const uint index_bits = columns <= 1 ? 1u : 32u - clz(columns - 1u);
+  const uint passes = (32u + index_bits + RADIX_BITS - 1u) / RADIX_BITS;
+  const uint shift = (passes - pass - 1u) * RADIX_BITS;
+  const device float* row_input = input + ulong(row) * columns;
+  const ulong prefix = prefixes[row];
+  const ulong mask = prefix_masks[row];
+
+  // Build this partition's histogram.
+  device atomic_uint* partition_histogram = partial_histograms + (row * partitions + partition) * BUCKETS;
+  for (uint bucket = lid; bucket < BUCKETS; bucket += THREADS_PER_TG)
+    atomic_store_explicit(&partition_histogram[bucket], 0u, memory_order_relaxed);
+  // Order this partition's clears before its histogram updates.
+  threadgroup_barrier(mem_flags::mem_device);
+  visit_partition_keys(row_input, columns, index_bits, partition, partitions, lid, [&](ulong key) {
+    if ((key & mask) == prefix)
+      atomic_fetch_add_explicit(&partition_histogram[(key >> shift) & (BUCKETS - 1)], 1u, memory_order_relaxed);
+  });
+  threadgroup_barrier(mem_flags::mem_device);
+
+  // The last partition merges the row.
+  if (lid == 0)
+    is_last = arrive_last(&done_counts[row], partitions);
+  threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device);
+  if (!is_last)
+    return;
+
+  for (uint bucket = lid; bucket < BUCKETS; bucket += THREADS_PER_TG) {
+    uint count = 0;
+    for (uint p = 0; p < partitions; ++p)
+      count +=
+          atomic_load_explicit(&partial_histograms[(row * partitions + p) * BUCKETS + bucket], memory_order_relaxed);
+    histogram[bucket] = count;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  uint sum = 0;
+  const uint begin = BUCKETS - (lid + 1) * BUCKETS_PER_THREAD;
+  for (uint offset = 0; offset < BUCKETS_PER_THREAD; ++offset)
+    sum += histogram[begin + offset];
+  uint higher = simd_prefix_exclusive_sum(sum);
+  if ((lid + 1) % METAL_SIMD_SIZE == 0)
+    simd_prefixes[lid / METAL_SIMD_SIZE] = higher + sum;
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (lid < METAL_SIMD_SIZE) {
+    const uint simdgroups = THREADS_PER_TG / METAL_SIMD_SIZE;
+    const uint simd_sum = lid < simdgroups ? simd_prefixes[lid] : 0;
+    const uint simd_prefix = simd_prefix_exclusive_sum(simd_sum);
+    if (lid < simdgroups)
+      simd_prefixes[lid] = simd_prefix;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  higher += simd_prefixes[lid / METAL_SIMD_SIZE];
+
+  uint rank = pass == 0 ? k - 1 : ranks[row];
+  if (higher <= rank && rank < higher + sum) {
+    rank -= higher;
+    const uint group_begin = BUCKETS - (lid + 1) * BUCKETS_PER_THREAD;
+    for (uint offset = 0; offset < BUCKETS_PER_THREAD; ++offset) {
+      const uint bucket = group_begin + BUCKETS_PER_THREAD - offset - 1;
+      const uint count = histogram[bucket];
+      if (rank < count) {
+        prefixes[row] |= ulong(bucket) << shift;
+        prefix_masks[row] |= ulong(BUCKETS - 1) << shift;
+        break;
+      }
+      rank -= count;
+    }
+    ranks[row] = rank;
+  }
+  threadgroup_barrier(mem_flags::mem_device);
+  if (lid == 0)
+    reset_arrival(&done_counts[row]);
+}
+
+KERNEL(RadixTopKSmallCollect)(
+    const device float* input,
+    device uint* output_ids,
+    device float* output_scores,
+    const device ulong* prefixes,
+    device atomic_uint* selected_counts,
+    device atomic_uint* done_counts,
+    device ulong* selected_keys,
+    constant uint& rows,
+    constant uint& k,
+    const uint columns SPECIALIZE,
+    const uint partitions SPECIALIZE,
+    threadgroup ulong sorted_keys[MAX_K],
+    threadgroup uint& is_last,
+    const uint group GROUPS(rows* partitions),
+    const uint lid THREADS(THREADS_PER_TG)
+) {
+  const uint row = group / partitions;
+  const uint partition = group % partitions;
+  const uint index_bits = columns <= 1 ? 1u : 32u - clz(columns - 1u);
+  const device float* row_input = input + ulong(row) * columns;
+  const ulong prefix = prefixes[row];
+
+  visit_partition_keys(row_input, columns, index_bits, partition, partitions, lid, [&](ulong key) {
+    if (key >= prefix) {
+      const uint index = atomic_fetch_add_explicit(&selected_counts[row], 1u, memory_order_relaxed);
+      // The final radix prefix selects exactly k unique keys.
+      selected_keys[row * k + index] = key;
+    }
+  });
+  threadgroup_barrier(mem_flags::mem_device);
+  if (lid == 0)
+    is_last = arrive_last(&done_counts[row], partitions);
+  threadgroup_barrier(mem_flags::mem_threadgroup | mem_flags::mem_device);
+  if (!is_last)
+    return;
+
+  for (uint index = lid; index < MAX_K; index += THREADS_PER_TG)
+    sorted_keys[index] = index < k ? selected_keys[row * k + index] : 0;
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  for (uint size = 2; size <= MAX_K; size <<= 1) {
+    for (uint stride = size >> 1; stride; stride >>= 1) {
+      for (uint index = lid; index < MAX_K; index += THREADS_PER_TG) {
+        const uint other = index ^ stride;
+        if (other > index) {
+          const bool descending = (index & size) == 0;
+          if (descending ? sorted_keys[index] < sorted_keys[other] : sorted_keys[index] > sorted_keys[other]) {
+            const ulong key = sorted_keys[index];
+            sorted_keys[index] = sorted_keys[other];
+            sorted_keys[other] = key;
+          }
+        }
+      }
+      threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+  }
+  for (uint index = lid; index < k; index += THREADS_PER_TG) {
+    const ulong key = sorted_keys[index];
+    const ulong mask = (1ul << index_bits) - 1;
+    output_ids[row * k + index] = uint(mask - (key & mask));
+    output_scores[row * k + index] = top_k_score_from_key(uint(key >> index_bits));
+  }
+}