Skip to content

Fix per-matrix CQ4 reciprocals in Metal batch transform #784

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
jakmro merged 1 commit into from
Aug 13, 2026
Merged

Fix per-matrix CQ4 reciprocals in Metal batch transform #784

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
15 changes: 9 additions & 6 deletions cactus-kernels/src/cactus_kernels.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,16 +138,19 @@ kernel void cq4_transform_simd(

kernel void cq4_transform_batch(
device const half* x [[buffer(0)]],
device const half* recip [[buffer(1)]],
device const char* lsign0 [[buffer(2)]], device const char* lsign1 [[buffer(3)]], device const char* lsign2 [[buffer(4)]],
device const char* rsign0 [[buffer(5)]], device const char* rsign1 [[buffer(6)]], device const char* rsign2 [[buffer(7)]],
device const uint* perm0 [[buffer(8)]], device const uint* perm1 [[buffer(9)]], device const uint* perm2 [[buffer(10)]],
device half* code0 [[buffer(11)]], device half* code1 [[buffer(12)]], device half* code2 [[buffer(13)]],
constant uint& ng [[buffer(14)]],
device const half* recip0 [[buffer(1)]],
device const half* recip1 [[buffer(2)]],
device const half* recip2 [[buffer(3)]],
device const char* lsign0 [[buffer(4)]], device const char* lsign1 [[buffer(5)]], device const char* lsign2 [[buffer(6)]],
device const char* rsign0 [[buffer(7)]], device const char* rsign1 [[buffer(8)]], device const char* rsign2 [[buffer(9)]],
device const uint* perm0 [[buffer(10)]], device const uint* perm1 [[buffer(11)]], device const uint* perm2 [[buffer(12)]],
device half* code0 [[buffer(13)]], device half* code1 [[buffer(14)]], device half* code2 [[buffer(15)]],
constant uint& ng [[buffer(16)]],
uint tgpos [[threadgroup_position_in_grid]], uint lane [[thread_index_in_simdgroup]],
threadgroup float* zmem [[threadgroup(0)]])
{
uint g = tgpos % ng, b = tgpos / ng;
device const half* recip = (b==0u)?recip0:((b==1u)?recip1:recip2);
device const char* lsign = (b==0u)?lsign0:((b==1u)?lsign1:lsign2);
device const char* rsign = (b==0u)?rsign0:((b==1u)?rsign1:rsign2);
device const uint* perm = (b==0u)?perm0 :((b==1u)?perm1 :perm2 );
Expand Down
12 changes: 6 additions & 6 deletions cactus-kernels/src/metal_backend.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -806,15 +806,15 @@ bool cactus_metal_encode_transform_batch(const void* x, const CactusQuantMatrix*
ensureEncoder();
[g_enc setComputePipelineState:ctx().psoTbatch];
setBufAt(x, (size_t)K*2, 0);
[g_enc setBuffer:r0.recip offset:r0.rc_off atIndex:1];
for (int b=0;b<3;b++) {
ResW* R = rw[b<B?b:0];
[g_enc setBuffer:R->lsign offset:R->ls_off atIndex:2+b];
[g_enc setBuffer:R->rsign offset:R->rs_off atIndex:5+b];
[g_enc setBuffer:R->perm offset:R->pm_off atIndex:8+b];
setBufAt(codes[b<B?b:0], (size_t)K*2, 11+b);
[g_enc setBuffer:R->recip offset:R->rc_off atIndex:1+b];
[g_enc setBuffer:R->lsign offset:R->ls_off atIndex:4+b];
[g_enc setBuffer:R->rsign offset:R->rs_off atIndex:7+b];
[g_enc setBuffer:R->perm offset:R->pm_off atIndex:10+b];
setBufAt(codes[b<B?b:0], (size_t)K*2, 13+b);
}
[g_enc setBytes:&ng length:4 atIndex:14];
[g_enc setBytes:&ng length:4 atIndex:16];
[g_enc setThreadgroupMemoryLength:gs*sizeof(float) atIndex:0];
[g_enc dispatchThreadgroups:MTLSizeMake((size_t)ng*B,1,1) threadsPerThreadgroup:MTLSizeMake(32,1,1)];
return true;
Expand Down
76 changes: 76 additions & 0 deletions cactus-kernels/tests/test_matmul.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#include "test_utils.h"
#include "metal_backend.h"
#include <vector>
#include <cmath>
#include <cstring>
Expand Down Expand Up @@ -601,6 +602,80 @@ static bool test_cq4_shared_transform_multi_projection() {
return q_ref == q_out && k_ref == k_out && v_ref == v_out;
}

static bool test_cq4_metal_transform_batch_distinct_recip() {
#if !CACTUS_HAS_METAL
return true;
#else
if (!cactus_metal_available()) return true;

constexpr uint32_t K = 256;

SyntheticCQ q(4, K, 128, 128, 1001);
SyntheticCQ k(4, K, 64, 128, 1002);
SyntheticCQ v(4, K, 64, 128, 1003);

for (uint32_t i = 0; i < K; ++i) {
q.input_scale[i] = static_cast<__fp16>(1.0f);
q.input_scale_recip[i] = static_cast<__fp16>(1.0f);

k.input_scale[i] = static_cast<__fp16>(2.0f);
k.input_scale_recip[i] = static_cast<__fp16>(0.5f);

v.input_scale[i] = static_cast<__fp16>(0.5f);
v.input_scale_recip[i] = static_cast<__fp16>(2.0f);
}

for (SyntheticCQ* matrix : {&k, &v}) {
matrix->left_signs = q.left_signs;
matrix->right_signs = q.right_signs;
matrix->permutation = q.permutation;
}

CactusQuantMatrix qm = q.matrix_interleaved();
CactusQuantMatrix km = k.matrix_interleaved();
CactusQuantMatrix vm = v.matrix_interleaved();

std::vector<__fp16> x(K);
for (size_t i = 0; i < K; ++i)
x[i] = static_cast<__fp16>(std::sin(static_cast<float>(i) * 0.031f));

std::vector<__fp16> q_ref(q.N), k_ref(k.N), v_ref(v.N);
std::vector<__fp16> q_out(q.N), k_out(k.N), v_out(v.N);
std::vector<__fp16> q_code(K), k_code(K), v_code(K);

cactus_metal_session_begin();

bool ok =
cactus_metal_encode_quant_matmul(q_ref.data(), x.data(), &qm) &&
cactus_metal_encode_quant_matmul(k_ref.data(), x.data(), &km) &&
cactus_metal_encode_quant_matmul(v_ref.data(), x.data(), &vm);

cactus_metal_session_sync();

const CactusQuantMatrix* weights[3] = {&qm, &km, &vm};
void* codes[3] = {q_code.data(), k_code.data(), v_code.data()};

ok = ok && cactus_metal_encode_transform_batch(x.data(), weights, 3, codes);

ok = ok &&
cactus_metal_encode_gemv_precoded(q_out.data(), q_code.data(), &qm) &&
cactus_metal_encode_gemv_precoded(k_out.data(), k_code.data(), &km) &&
cactus_metal_encode_gemv_precoded(v_out.data(), v_code.data(), &vm);

cactus_metal_session_sync();
cactus_metal_session_end();

bool match =
q_ref == q_out &&
k_ref == k_out &&
v_ref == v_out;

cactus_metal_invalidate_host_wraps();

return ok && match;
#endif
}

int main() {
TestRunner runner("Matrix Multiplication");
runner.run_test("matmul_f16", test_matmul_f16());
Expand All @@ -615,6 +690,7 @@ int main() {
double m_mt = 0;
runner.run_test("matmul_cq4_il_mt", test_cq4_interleaved(m_mt, 1024, 4164, 128));
runner.run_test("matmul_cq4_shared_transform", test_cq4_shared_transform_multi_projection());
runner.run_test("matmul_cq4_metal_distinct_recip", test_cq4_metal_transform_batch_distinct_recip());
}
runner.print_benchmarks_header();
runner.run_bench("benchmarks", run_benchmarks());
Expand Down
Loading