// metal/kernels/dense_mm.metal - dense projections for a PROMPT chunk (after quant/units/mma.metal):
// act_f16 X[t][k] -> half (the matrix units' operand)
// dq_f16<F> a quantized matrix -> half rows (ref_quant's values rounded to fp16), once per chunk
// gemm_f16 Y[t][n] = Xh[t] . Wh[n] (fp32) by MetalPerformancePrimitives matmul2d, TM x TN tiles, K dynamic;
// rows of Y past n_out / t past n_tok are not written (the tensors carry the extents)
#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
using namespace mpp::tensor_ops;
kernel void act_f16(device const float* x [[buffer(0)]], device half* xh [[buffer(1)]], constant uint& n [[buffer(2)]],
uint gid [[thread_position_in_grid]]) {
const uint n4 = n >> 2;
if (gid < n4) {
((device half4*)xh)[gid] = half4(((device const float4*)x)[gid]);
} else if (gid == n4) {
for (uint i = n4 << 2; i < n; ++i) xh[i] = half(x[i]);
}
}
struct DqArgs {
uint n_in; // values per row
uint n_out; // rows
uint row_bytes;
uint perm; // 0: row r -> r; p > 0: row r = c * (n_out / p) + d -> d * p + c (the HC up matrix with the p
// streams of an output interleaved, for gemm_hc_up's epilogue)
};
// one thread per unit: grid ((units per row + nth.x - 1) / nth.x, n_out), nth.x threads
template <int F>
kernel void dq_f16(device const uchar* W [[buffer(0)]], device half* Wh [[buffer(1)]], constant DqArgs& a [[buffer(2)]],
uint3 gid [[thread_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
ushort3 nth [[threads_per_threadgroup]]) {
constexpr int UNIT = fmt<F>::UNIT, H = UNIT / 2;
threadgroup uint tab[fmt<F>::TAB];
if (F == T_IQ4_XS || F == T_IQ4_NL) {
if (tid < 64) {
const uint r0 = tid & 15u, sh = (tid >> 4) * 8u;
const half h0 = half(as_type<char>(uchar(((r0 & 8u) ? 0x7159453526190d01ul : 0xf6eaddcfbfad9881ul) >> ((r0 & 7u) * 8u))));
tab[tid] = as_type<uint>(half2(h0, half(as_type<char>(uchar(0xbfad9881u >> sh)))));
tab[tid + 64] = as_type<uint>(half2(h0, half(as_type<char>(uchar(0xf6eaddcfu >> sh)))));
tab[tid + 128] = as_type<uint>(half2(h0, half(as_type<char>(uchar(0x26190d01u >> sh)))));
tab[tid + 192] = as_type<uint>(half2(h0, half(as_type<char>(uchar(0x71594535u >> sh)))));
}
threadgroup_barrier(mem_flags::mem_threadgroup);
} else if (fmt<F>::TAB > 1) {
load_tab<F>(tab, tid, nth.x);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const int u = int(gid.x);
if (u * UNIT >= int(a.n_in)) return;
uint r = gid.y, ro = r;
if (a.perm) {
if (fmt<F>::BB * 16 < fmt<F>::QK * 17) {
r = a.perm == 4 ? (ro & 3u) * (a.n_out >> 2) + (ro >> 2) : (ro % a.perm) * (a.n_out / a.perm) + ro / a.perm;
} else {
const uint q = a.n_out / a.perm;
ro = (r % q) * a.perm + r / q;
}
}
float w[UNIT];
int p0, p1;
DW<F>::run(W + ulong(r) * a.row_bytes, u, 0, w, p0, p1, tab);
device half* o = Wh + ulong(ro) * a.n_in;
#pragma clang loop unroll(full)
for (int j = 0; j < H; j += 4) {
*(device half4*)(o + p0 + j) = half4(half(w[j]), half(w[j + 1]), half(w[j + 2]), half(w[j + 3]));
*(device half4*)(o + p1 + j) = half4(half(w[H + j]), half(w[H + j + 1]), half(w[H + j + 2]), half(w[H + j + 3]));
}
}
struct GemmArgs { int M, N, K; int pad; }; // Y is M rows of N floats
template <int TM, int TN, int SG>
kernel void gemm_f16(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device float* Y [[buffer(2)]],
constant GemmArgs& a [[buffer(3)]], uint3 tg [[threadgroup_position_in_grid]]) {
tensor<device half, dextents<int32_t, 2>, tensor_inline> tA(A, dextents<int32_t, 2>(a.K, a.M));
tensor<device half, dextents<int32_t, 2>, tensor_inline> tB(B, dextents<int32_t, 2>(a.K, a.N));
tensor<device float, dextents<int32_t, 2>, tensor_inline> tY(Y, dextents<int32_t, 2>(a.N, a.M));
constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v<int>, false, true, false);
matmul2d<desc, execution_simdgroups<SG>> op;
auto mA = tA.slice(0, int(tg.y) * TM);
auto mB = tB.slice(0, int(tg.x) * TN);
auto mY = tY.slice(int(tg.x) * TN, int(tg.y) * TM);
op.run(mA, mB, mY);
}
#define DQ(F) template [[host_name("dq_f16_" #F)]] kernel decltype(dq_f16<F>) dq_f16<F>;
#define DQ_(F) DQ(F)
DQ_(T_F32) DQ_(T_BF16) DQ_(T_Q8_0) DQ_(T_Q4_0) DQ_(T_Q5_0) DQ_(T_Q5_1) DQ_(T_IQ4_NL) DQ_(T_IQ4_XS) DQ_(T_Q2_0) DQ_(T_IQ3_S) DQ_(T_Q3_K)
DQ_(T_Q4_K) DQ_(T_Q5_K) DQ_(T_Q6_K)
template [[host_name("gemm_f16_64_64_4")]] kernel decltype(gemm_f16<64, 64, 4>) gemm_f16<64, 64, 4>;
template [[host_name("gemm_f16_128_64_4")]] kernel decltype(gemm_f16<128, 64, 4>) gemm_f16<128, 64, 4>;
template [[host_name("gemm_f16_64_128_4")]] kernel decltype(gemm_f16<64, 128, 4>) gemm_f16<64, 128, 4>;
template [[host_name("gemm_f16_64_32_4")]] kernel decltype(gemm_f16<64, 32, 4>) gemm_f16<64, 32, 4>;
template [[host_name("gemm_f16_128_128_4")]] kernel decltype(gemm_f16<128, 128, 4>) gemm_f16<128, 128, 4>;
template [[host_name("gemm_f16_64_64_2")]] kernel decltype(gemm_f16<64, 64, 2>) gemm_f16<64, 64, 2>;
template [[host_name("gemm_f16_64_128_2")]] kernel decltype(gemm_f16<64, 128, 2>) gemm_f16<64, 128, 2>;
template [[host_name("gemm_f16_16_64_2")]] kernel decltype(gemm_f16<16, 64, 2>) gemm_f16<16, 64, 2>;
template [[host_name("gemm_f16_32_64_2")]] kernel decltype(gemm_f16<32, 64, 2>) gemm_f16<32, 64, 2>;
template [[host_name("gemm_f16_32_128_2")]] kernel decltype(gemm_f16<32, 128, 2>) gemm_f16<32, 128, 2>;
template [[host_name("gemm_f16_64_32_2")]] kernel decltype(gemm_f16<64, 32, 2>) gemm_f16<64, 32, 2>;
template [[host_name("gemm_f16_16_96_2")]] kernel decltype(gemm_f16<16, 96, 2>) gemm_f16<16, 96, 2>;