From 6742fead6f72f65a0025c7c77ab5ada80f840adb Mon Sep 17 00:00:00 2001 From: Bhavyansh Sharma <66562679+clumsylad21@users.noreply.github.com> Date: Wed, 12 Aug 2026 09:24:05 -0500 Subject: [PATCH] fix: use per-matrix CQ4 reciprocals in Metal batch transform Signed-off-by: Bhavyansh Sharma <66562679+clumsylad21@users.noreply.github.com> --- cactus-kernels/src/cactus_kernels.metal | 15 +++-- cactus-kernels/src/metal_backend.mm | 12 ++-- cactus-kernels/tests/test_matmul.cpp | 76 +++++++++++++++++++++++++ 3 files changed, 91 insertions(+), 12 deletions(-) diff --git a/cactus-kernels/src/cactus_kernels.metal b/cactus-kernels/src/cactus_kernels.metal index 1c4a724c5..1ef064202 100644 --- a/cactus-kernels/src/cactus_kernels.metal +++ b/cactus-kernels/src/cactus_kernels.metal @@ -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 ); diff --git a/cactus-kernels/src/metal_backend.mm b/cactus-kernels/src/metal_backend.mm index 829a11fb6..cb55e36e1 100644 --- a/cactus-kernels/src/metal_backend.mm +++ b/cactus-kernels/src/metal_backend.mm @@ -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[blsign 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[brecip 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 #include #include @@ -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(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()); @@ -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());