Skip to content

GatherGemv and Sparse readout #619

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
CC-Yeh merged 4 commits into from
Jul 20, 2026
Merged

GatherGemv and Sparse readout #619

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
922ccfb
naive sparse readout kernel
CC-Yeh Jul 19, 2026
6f4f5ee
Gather GEMV, 3x faster than CPU gather
CC-Yeh Jul 20, 2026
bf352e6
Merge remote-tracking branch 'origin/main' into embedding-f32-sparse-…
CC-Yeh Jul 20, 2026
45c2ddc
fix
CC-Yeh Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions crates/backend-uzu/src/backends/common/gpu_types/gemm.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ bitflags! {
const ACCUMULATE = 1 << 1;
const BIAS = 1 << 2;
const RHT = 1 << 3;
const SOFT_CAP = 1 << 4;
}
}

Expand Down
1 change: 1 addition & 0 deletions crates/backend-uzu/src/backends/common/kernel/matmul/arguments.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ pub struct MatmulArguments<'a, 'b, 'd, B: Backend, TB: BufferArg<'b, B> = &'b Al
pub b_transpose: bool,
pub d: &'d mut Allocation<B>,
pub d_transform: MatmulDOps<'d, B>,
pub gather_indices: Option<&'a Allocation<B>>,
pub m: u32,
pub n: u32,
pub k: u32,
Expand Down
10 changes: 10 additions & 0 deletions crates/backend-uzu/src/backends/common/kernel/matmul/d_ops.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ pub struct MatmulDOps<'a, B: Backend> {
pub accumulate: bool,
pub bias: Option<&'a Allocation<B>>,
pub rht_factors: Option<&'a Allocation<B>>,
pub soft_cap: Option<f32>,
}

impl<'a, B: Backend> MatmulDOps<'a, B> {
Expand All @@ -14,6 +15,7 @@ impl<'a, B: Backend> MatmulDOps<'a, B> {
accumulate: false,
bias: None,
rht_factors: None,
soft_cap: None,
}
}

Expand All @@ -31,6 +33,9 @@ impl<'a, B: Backend> MatmulDOps<'a, B> {
if self.rht_factors.is_some() {
m |= GemmDTransform::RHT;
}
if self.soft_cap.is_some() {
m |= GemmDTransform::SOFT_CAP;
Comment thread
CC-Yeh marked this conversation as resolved.
}
m
}

Expand Down Expand Up @@ -59,6 +64,11 @@ impl<'a, B: Backend> MatmulDOps<'a, B> {
} else {
self.rht_factors
},
soft_cap: if bits.contains(GemmDTransform::SOFT_CAP) {
None
} else {
self.soft_cap
},
}
}
}
33 changes: 25 additions & 8 deletions crates/backend-uzu/src/backends/cpu/kernel/matmul/kernel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ impl MatmulKernel for MatmulCpuKernel {
let accumulate = arguments.d_transform.accumulate;
let bias_alloc = arguments.d_transform.bias;
let post_rht = arguments.d_transform.rht_factors;
let soft_cap = arguments.d_transform.soft_cap;

let MatmulArguments {
a,
Expand All @@ -69,6 +70,7 @@ impl MatmulKernel for MatmulCpuKernel {
m,
n,
k,
gather_indices,
..
} = arguments;

Expand All @@ -86,6 +88,10 @@ impl MatmulKernel for MatmulCpuKernel {
let r = bias.as_buffer_range_ref();
SendPtr(unsafe { &*r.buffer().get() }.as_ptr().wrapping_byte_add(r.range().start))
});
let gather_ptr = gather_indices.map(|indices| {
let r = indices.as_buffer_range_ref();
SendPtr(unsafe { &*r.buffer().get() }.as_ptr().wrapping_byte_add(r.range().start) as *const u32)
});
let d_buffer_range = d.as_buffer_range_mut();
let d_ptr = SendPtrMut(unsafe {
(&*d_buffer_range.buffer().get()).as_ptr().wrapping_byte_add(d_buffer_range.range().start) as *mut u8
Expand Down Expand Up @@ -122,6 +128,11 @@ impl MatmulKernel for MatmulCpuKernel {
unsafe {
for row in 0..m_u {
for col in 0..n_u {
// Gather remaps output column `col` to B-row `gather_indices[row * n + col]`.
let b_col = match gather_ptr {
Some(g) => *g.as_ptr().add(row * n_u + col) as usize,
None => col,
};
let mut accumulator = 0.0f32;
for inner in 0..k_u {
let a_value = read_f32(a_ptr.as_ptr(), input_data_type, row * k_u + inner);
Expand All @@ -132,9 +143,9 @@ impl MatmulKernel for MatmulCpuKernel {
transpose,
} => {
let index = if *transpose {
col * leading_dimension + inner
b_col * leading_dimension + inner
} else {
inner * leading_dimension + col
inner * leading_dimension + b_col
};
read_f32(ptr.as_ptr(), weights_data_type, index)
},
Expand All @@ -147,7 +158,7 @@ impl MatmulKernel for MatmulCpuKernel {
group_size,
} => {
let (num_groups_k, zero_point_stride, pack_factor) = quant_layout.unwrap();
let weight_linear_index = col * k_u + inner;
let weight_linear_index = b_col * k_u + inner;
let quantized_value = if *bits == 4 {
let word_index = weight_linear_index / pack_factor;
let bit_offset = (weight_linear_index % pack_factor) * 4;
Expand All @@ -160,23 +171,26 @@ impl MatmulKernel for MatmulCpuKernel {
((w.add(word_index).read_unaligned() >> bit_offset) & 0xFF) as f32
};
let group_index = inner / group_size;
let scale =
read_f32(scales.as_ptr(), weights_data_type, col * num_groups_k + group_index);
let scale = read_f32(
scales.as_ptr(),
weights_data_type,
b_col * num_groups_k + group_index,
);
let bias_term = if let Some(zp) = zero_points {
let zero_point = if *bits == 4 {
let byte_index = col * zero_point_stride + (group_index >> 1);
let byte_index = b_col * zero_point_stride + (group_index >> 1);
let byte_value = *zp.as_ptr().add(byte_index);
if (group_index & 1) == 0 {
(byte_value & 0x0F) as f32
} else {
((byte_value >> 4) & 0x0F) as f32
}
} else {
*zp.as_ptr().add(col * zero_point_stride + group_index) as f32
*zp.as_ptr().add(b_col * zero_point_stride + group_index) as f32
};
-scale * zero_point
} else if let Some(b) = biases {
read_f32(b.as_ptr(), weights_data_type, col * num_groups_k + group_index)
read_f32(b.as_ptr(), weights_data_type, b_col * num_groups_k + group_index)
} else {
let midpoint = (1u32 << (bits - 1)) as f32;
-scale * midpoint
Expand All @@ -195,6 +209,9 @@ impl MatmulKernel for MatmulCpuKernel {
if let Some(bias) = bias_ptr {
value += read_f32(bias.as_ptr(), weights_data_type, col);
}
if let Some(cap) = soft_cap {
value = cap * (value / cap).tanh();
}
write_f32(d_ptr.as_ptr(), output_data_type, output_index, value);
}
}
Expand Down
14 changes: 14 additions & 0 deletions crates/backend-uzu/src/backends/metal/kernel/common/soft_cap.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
#pragma once

#include <metal_stdlib>

#include "defines.h"

namespace uzu {

template <typename T>
METAL_FUNC T apply_soft_cap(T value, float cap) {
return static_cast<T>(cap * metal::fast::tanh(static_cast<float>(value) / cap));
}

} // namespace uzu
1 change: 1 addition & 0 deletions crates/backend-uzu/src/backends/metal/kernel/generated/gemm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ struct GemmDTransform {
static constant constexpr uint ACCUMULATE = 1 << 1;
static constant constexpr uint BIAS = 1 << 2;
static constant constexpr uint RHT = 1 << 3;
static constant constexpr uint SOFT_CAP = 1 << 4;
constexpr bool contains(uint flag) const thread { return (raw_value & flag) != 0; }
constexpr bool contains(uint flag) const constant { return (raw_value & flag) != 0; }
constexpr uint bits() const thread { return raw_value; }
Expand Down
4 changes: 2 additions & 2 deletions crates/backend-uzu/src/backends/metal/kernel/logit_soft_cap/logit_soft_cap.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#include <metal_stdlib>
#include "../common/dsl.h"
#include "../common/soft_cap.h"

template <typename T>
VARIANTS(T, float, bfloat)
Expand All @@ -9,6 +10,5 @@ PUBLIC KERNEL(LogitSoftCap)(
constant float& soft_cap,
const uint position AXIS(length, 256)
) {
const float value = float(logits[position]);
logits[position] = T(fast::tanh(value / soft_cap) * soft_cap);
logits[position] = uzu::apply_soft_cap<T>(logits[position], soft_cap);
}
4 changes: 4 additions & 0 deletions crates/backend-uzu/src/backends/metal/kernel/matmul/gemm/kernel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,10 @@ impl GemmKernel {
&self,
arguments: &MatmulArguments<'a, 'b, 'd, Metal, TB>,
) -> bool {
if arguments.gather_indices.is_some() {
// TODO: gathered GEMM
return false;
}
match (
arguments.m,
arguments.n == arguments.k,
Expand Down
9 changes: 9 additions & 0 deletions crates/backend-uzu/src/backends/metal/kernel/matmul/gemv/common/b_source.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,10 @@ struct BSource {
const device uint8_t* zero_points,
const device BT* biases,
const device AT* a,
const device uint* gather_indices,
bool gathered,
uint in_vec_size,
uint out_vec_size,
uint out_row,
uint batch_idx,
uint simd_lane,
Expand All @@ -36,7 +39,10 @@ struct BSource {
result,
b,
a,
gather_indices,
gathered,
in_vec_size,
out_vec_size,
out_row,
batch_idx,
simd_lane,
Expand All @@ -50,7 +56,10 @@ struct BSource {
zero_points,
biases,
a,
gather_indices,
gathered,
in_vec_size,
out_vec_size,
out_row,
batch_idx,
simd_lane
Expand Down
6 changes: 6 additions & 0 deletions crates/backend-uzu/src/backends/metal/kernel/matmul/gemv/common/epilogue.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#pragma once

#include "../../../common/soft_cap.h"
#include "../../../hadamard_transform/hadamard_transform.h"

namespace uzu {
Expand All @@ -14,6 +15,7 @@ struct Epilogue {
const device int32_t* hadamard_factors,
threadgroup U* shared_results,
float ab_scale,
float soft_cap,
GemmDTransform output_transform,
uint out_row,
uint out_vec_size,
Expand All @@ -25,6 +27,7 @@ struct Epilogue {
const bool is_scale = output_transform.contains(GemmDTransform::SCALE);
const bool is_accumulate = output_transform.contains(GemmDTransform::ACCUMULATE);
const bool is_bias = output_transform.contains(GemmDTransform::BIAS);
const bool is_soft_cap = output_transform.contains(GemmDTransform::SOFT_CAP);
const bool use_hadamard = output_transform.contains(GemmDTransform::RHT);

if (writer && simd_lane == 0) {
Expand All @@ -41,6 +44,9 @@ struct Epilogue {
if (is_bias && global_row < out_vec_size) {
value += static_cast<U>(output_bias[global_row]);
}
if (is_soft_cap) {
value = apply_soft_cap(value, soft_cap);
}
result[row] = value;
}
}
Expand Down
25 changes: 18 additions & 7 deletions crates/backend-uzu/src/backends/metal/kernel/matmul/gemv/common/full_precision_b_source.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,10 @@ struct FullPrecisionBSource {
thread U (&result)[RESULTS_PER_SIMDGROUP],
const device uint32_t* b,
const device AT* a,
const device uint* gather_indices,
bool gathered,
uint in_vec_size,
uint out_vec_size,
uint out_row,
uint batch_idx,
uint simd_lane,
Expand All @@ -24,19 +27,28 @@ struct FullPrecisionBSource {

const uint k_stride = K_SPLIT * block_size;
const uint k_start = k_slice * block_size;
const uint thread_k = simd_lane * values_per_thread + k_start;
const device AT* input = a + batch_idx * in_vec_size + thread_k;

// One advancing weight pointer per output row; base row = out_row (dense) or gather index.
const uint base_row = gathered ? 0u : out_row;
const device BT* weights = reinterpret_cast<const device BT*>(b);
weights += out_row * in_vec_size + simd_lane * values_per_thread + k_start;
const device AT* input = a + batch_idx * in_vec_size + simd_lane * values_per_thread + k_start;
weights += base_row * in_vec_size + thread_k;
thread const device BT* weight_rows[RESULTS_PER_SIMDGROUP];
METAL_PRAGMA_UNROLL
for (uint row = 0; row < RESULTS_PER_SIMDGROUP; row++) {
const uint addr_row = gathered ? gather_indices[batch_idx * out_vec_size + out_row + row] : row;
weight_rows[row] = weights + addr_row * in_vec_size;
}

uint k = k_start;
for (; k + block_size <= in_vec_size; k += k_stride) {
float4 input_values = static_cast<float4>(*reinterpret_cast<const device I4*>(input));
METAL_PRAGMA_UNROLL
for (uint row = 0; row < RESULTS_PER_SIMDGROUP; row++) {
const device BT* weight_row = weights + row * in_vec_size;
result[row] += dot(static_cast<float4>(*reinterpret_cast<const device W4*>(weight_row)), input_values);
result[row] += dot(static_cast<float4>(*reinterpret_cast<const device W4*>(weight_rows[row])), input_values);
weight_rows[row] += k_stride;
}
weights += k_stride;
input += k_stride;
}

Expand All @@ -48,9 +60,8 @@ struct FullPrecisionBSource {
if (remaining > 0) {
METAL_PRAGMA_UNROLL
for (uint row = 0; row < RESULTS_PER_SIMDGROUP; row++) {
const device BT* weight_row = weights + row * in_vec_size;
for (int index = 0; index < remaining; index++) {
result[row] += static_cast<U>(weight_row[index]) * static_cast<U>(input[index]);
result[row] += static_cast<U>(weight_rows[row][index]) * static_cast<U>(input[index]);
}
}
}
Expand Down
21 changes: 15 additions & 6 deletions crates/backend-uzu/src/backends/metal/kernel/matmul/gemv/common/quantized_b_source.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,10 @@ struct QuantizedBSource {
const device uint8_t* zero_points,
const device BT* biases,
const device AT* a,
const device uint* gather_indices,
bool gathered,
uint in_vec_size,
uint out_vec_size,
uint out_row,
uint batch_idx,
uint simd_lane
Expand All @@ -41,10 +44,14 @@ struct QuantizedBSource {
const uint weights_row_stride = in_vec_size * bytes_per_pack / pack_factor;
const uint group_count = (in_vec_size + GROUP_SIZE - 1) / GROUP_SIZE;
const uint group_offset = simd_lane / scale_step_per_thread;

// Base row = out_row (dense) or 0 (gather); rows indexed relative to it.
const uint base_row = gathered ? 0u : out_row;

const device uint8_t* weights = reinterpret_cast<const device uint8_t*>(b);
weights += out_row * weights_row_stride + simd_lane * packs_per_thread * bytes_per_pack;
weights += base_row * weights_row_stride + simd_lane * packs_per_thread * bytes_per_pack;

RowState row_state(scales, zero_points, biases, out_row, group_count, group_offset);
RowState row_state(scales, zero_points, biases, base_row, group_count, group_offset);

const device AT* input = a + batch_idx * in_vec_size + simd_lane * values_per_thread;
thread U input_values[values_per_thread];
Expand All @@ -54,10 +61,11 @@ struct QuantizedBSource {
U input_sum = load_vector<AT, U, values_per_thread, BITS>(input, input_values);

RowParams row_params;
row_state.load(row_params);
row_state.load(row_params, gather_indices, gathered, batch_idx, out_vec_size, out_row);
METAL_PRAGMA_UNROLL
for (uint row = 0; row < RESULTS_PER_SIMDGROUP; row++) {
const device uint8_t* weight_row = weights + row * weights_row_stride;
const uint addr_row = gathered ? gather_indices[batch_idx * out_vec_size + out_row + row] : row;
const device uint8_t* weight_row = weights + addr_row * weights_row_stride;
result[row] += qdot<U, values_per_thread, BITS>(
weight_row,
input_values,
Expand All @@ -81,10 +89,11 @@ struct QuantizedBSource {
U input_sum = load_vector_safe<AT, U, values_per_thread>(input, input_values, remaining);

RowParams row_params;
row_state.load(row_params);
row_state.load(row_params, gather_indices, gathered, batch_idx, out_vec_size, out_row);
METAL_PRAGMA_UNROLL
for (uint row = 0; row < RESULTS_PER_SIMDGROUP; row++) {
const device uint8_t* weight_row = weights + row * weights_row_stride;
const uint addr_row = gathered ? gather_indices[batch_idx * out_vec_size + out_row + row] : row;
const device uint8_t* weight_row = weights + addr_row * weights_row_stride;
result[row] += qdot_safe<U, values_per_thread, BITS>(
weight_row,
input_values,
Expand Down
Loading
Loading