Skip to content

Commit 7113f64

Browse files
bri-prismbri-prism
authored
metal: PTQ1_0 exact-float trit extraction (+30-38% decode) and the missing per-expert mat-vec (#157)
* metal: extract PTQ1_0 trits in exact float instead of the integer recurrence The mat-vec pulled each trit out of its byte with the base-3 remainder recurrence, w = v*3, t = w >> 8, v = w & 0xFF: four integer ops and a convert per trit, on an ISA that cannot co-issue integer and floating-point work. The kernel was already known to be instruction-bound rather than bandwidth-bound. A packed byte is a base-3 fraction of 256. With u = b/256, trit n is g_{n+1} - 3*g_n for g_k = floor(3^k*u), and every 3^k*b is below 2^24, so the floors are exact and this matches the recurrence bit for bit over all 256 bytes and all five positions. Summing t_n*y_n over a byte's trits collapses to sum_k g_k*(y_{k-1} - 3*y_k) + g_5*y_4, whose coefficients depend only on the activations, so the caller stages those in place of the raw y and reuses them across every row. The inner loop is one floor and one fma per trit with no integer arithmetic. The qh trit uses the two-floor single-trit form, replacing a lane-divergent loop that stepped the recurrence up to three times. Verified against the recurrence for all bytes before the kernel was touched. test-backend-ops MUL_MAT: 45 PTQ1_0 cases pass, zero failures. Output is not bit-identical: the coefficient staging reassociates the fp32 sum. Measured on an M5 Pro, llama-bench r=3, two interleaved passes, dense models: small decode 221 -> 287 tok/s (+29.5%) large decode 59 -> 82 tok/s (+37.8%) Prefill unchanged; the mat-mul path does not use this routine. * metal: instantiate the per-expert PTQ1_0 mat-vec that supports_op already claims supports_op accepts PTQ1_0 for MUL_MAT_ID and the dispatcher carries an nsg/nr0 entry for it, but the kernel_mul_mv_id_ptq1_0_f32 pipeline was never instantiated, so the first per-expert PTQ1_0 mat-vec on Metal fails the library lookup and the process segfaults. test-backend-ops hit this on its first PTQ1_0 MUL_MAT_ID case and died there, which truncated the Metal MUL_MAT_ID run to 24 cases while still printing a clean tail; nothing after that case, for any type, was being exercised. Add the instantiation next to the other low-bit per-expert kernels. It reuses kernel_mul_mv_ptq1_0_f32_impl, so it inherits the exact-float extraction from the previous commit. The full Metal MUL_MAT_ID run now completes: 1022 cases, 75 of them PTQ1_0, zero failures.
1 parent d0ce941 commit 7113f64

1 file changed

Lines changed: 63 additions & 44 deletions

File tree

‎ggml/src/ggml-metal/kernels/mul_mv.metal‎

Lines changed: 63 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -862,46 +862,48 @@ kernel void kernel_mul_mv_q2_0_f32(
862862
// qs[2*it], qs[2*it+1] -> elements n*16 + m for n in 0..4
863863
// qs[16 + it] -> elements 80 + n*8 + it for n in 0..4
864864
// qh[it] (it < 2 only) -> elements 120 + n*2 + it for n in 0..3
865-
// Dot against y already staged in registers by the caller and reused across all nr0
866-
// rows. Trits come out of the byte by the base-3 remainder recurrence
867-
// t = (v*3) >> 8, v = (v*3) & 0xFF, v starting at the byte
868-
// which is two integer ops per trit with no table and no pow3 lookup. A 256-entry
869-
// lookup was measurably worse here: it adds a dependent load per byte, and decode on
870-
// this kernel is instruction-bound, not bandwidth-bound. Verified against the
871-
// reference decode for all 256 bytes and all five positions.
872-
// Accumulates the raw 0/1/2 trit and subtracts the staged y sum once, so there is no
873-
// per-element offset correction.
865+
// Dot against coefficients already staged in registers by the caller and reused across
866+
// all nr0 rows. A packed byte is a base-3 fraction of 256: with u = b/256, trit n is
867+
// t_n = g_{n+1} - 3*g_n, g_k = floor(3^k * u)
868+
// and 3^k*b <= 61965 is exact in fp32, so the floors are exact and this matches the
869+
// integer recurrence bit for bit over all 256 bytes and all five positions. Summing
870+
// t_n*y_n over a byte's trits then collapses to
871+
// sum_{k=1..4} g_k*(y_{k-1} - 3*y_k) + g_5*y_4
872+
// whose coefficients depend only on the activations, so the caller stages those in
873+
// place of the raw y. The inner loop is one floor and one fma per trit and never
874+
// leaves the float pipe, which matters because this ISA cannot co-issue integer and
875+
// floating-point work; the recurrence spent four integer ops and a convert per trit.
876+
// sumy is subtracted once for the -1 offset, exactly as before.
874877
inline float ptq1_0_dot_reg(device const block_ptq1_0 * qb, thread const float * yl, float sumy, short it) {
875878
float acc = 0.f;
876-
short c = 0;
877879

878880
FOR_UNROLL (short k = 0; k < 2; ++k) {
879-
ushort v = qb->qs[2*it + k];
880-
FOR_UNROLL (short n = 0; n < 5; ++n) {
881-
const ushort w = v * 3;
882-
acc += (float) (w >> 8) * yl[c++];
883-
v = w & 0xFF;
884-
}
881+
const float u = (float) qb->qs[2*it + k] * (1.0f/256.0f);
882+
thread const float * c = yl + 5*k;
883+
acc += floor( 3.0f*u)*c[0];
884+
acc += floor( 9.0f*u)*c[1];
885+
acc += floor( 27.0f*u)*c[2];
886+
acc += floor( 81.0f*u)*c[3];
887+
acc += floor(243.0f*u)*c[4];
885888
}
886889

887890
{
888-
ushort v = qb->qs[16 + it];
889-
FOR_UNROLL (short n = 0; n < 5; ++n) {
890-
const ushort w = v * 3;
891-
acc += (float) (w >> 8) * yl[c++];
892-
v = w & 0xFF;
893-
}
894-
}
895-
896-
// qh holds 8 elements; give every thread exactly one so all eight do 16 elements.
897-
// Element 120+it sits at trit it>>1 of byte qh[it&1], so step the recurrence to it.
891+
const float u = (float) qb->qs[16 + it] * (1.0f/256.0f);
892+
thread const float * c = yl + 10;
893+
acc += floor( 3.0f*u)*c[0];
894+
acc += floor( 9.0f*u)*c[1];
895+
acc += floor( 27.0f*u)*c[2];
896+
acc += floor( 81.0f*u)*c[3];
897+
acc += floor(243.0f*u)*c[4];
898+
}
899+
900+
// qh holds 8 elements; every thread takes exactly one, trit it>>1 of byte qh[it&1].
901+
// The single-trit form is two floors and replaces a lane-divergent loop that
902+
// stepped the recurrence up to three times.
898903
{
899-
ushort v = qb->qh[it & 1];
900-
const short n = it >> 1;
901-
for (short i = 0; i < n; ++i) {
902-
v = (v * 3) & 0xFF;
903-
}
904-
acc += (float) ((v * 3) >> 8) * yl[c++];
904+
const float u = (float) qb->qh[it & 1] * (1.0f/256.0f);
905+
const float p0 = yl[16]; // 3^n for this thread's trit
906+
acc += (floor(3.0f*p0*u) - 3.0f*floor(p0*u)) * yl[15];
905907
}
906908

907909
return (acc - sumy) * (float) qb->d;
@@ -940,7 +942,8 @@ void kernel_mul_mv_ptq1_0_f32_impl(
940942
ax[row] = (device const block_ptq1_0 *) ((device char *) src0 + offset0);
941943
}
942944

943-
float yl[16];
945+
// 15 collapse coefficients, the qh activation, and the qh trit's 3^n
946+
float yl[17];
944947
float sumf[nr0] = {0.f};
945948

946949
// eight threads cover one 128-weight block; each owns whole bytes, not a
@@ -950,28 +953,43 @@ void kernel_mul_mv_ptq1_0_f32_impl(
950953

951954
device const float * yb = y + ix*QK_PTQ1_0;
952955

956+
{
957+
const float pow3f[4] = {1.0f, 3.0f, 9.0f, 27.0f};
958+
yl[16] = pow3f[it >> 1];
959+
}
960+
953961
for (int ib = ix; ib < nb; ib += N_SIMDWIDTH/8) {
954-
// stage this thread's y values once, then reuse them for every row
962+
// stage this thread's activations once as collapse coefficients, then reuse
963+
// them for every row: c[k-1] = y_{k-1} - 3*y_k for k = 1..4, c[4] = y_4
955964
float sumy = 0.f;
956-
short c = 0;
957965

958966
FOR_UNROLL (short k = 0; k < 2; ++k) {
959967
const short m = 2*it + k;
968+
float y[5];
960969
FOR_UNROLL (short n = 0; n < 5; ++n) {
961-
const float v = yb[n*16 + m];
962-
yl[c++] = v;
963-
sumy += v;
970+
y[n] = yb[n*16 + m];
971+
sumy += y[n];
964972
}
973+
FOR_UNROLL (short n = 0; n < 4; ++n) {
974+
yl[5*k + n] = y[n] - 3.0f*y[n+1];
975+
}
976+
yl[5*k + 4] = y[4];
965977
}
966-
FOR_UNROLL (short n = 0; n < 5; ++n) {
967-
const float v = yb[80 + n*8 + it];
968-
yl[c++] = v;
969-
sumy += v;
978+
{
979+
float y[5];
980+
FOR_UNROLL (short n = 0; n < 5; ++n) {
981+
y[n] = yb[80 + n*8 + it];
982+
sumy += y[n];
983+
}
984+
FOR_UNROLL (short n = 0; n < 4; ++n) {
985+
yl[10 + n] = y[n] - 3.0f*y[n+1];
986+
}
987+
yl[14] = y[4];
970988
}
971989
{
972990
const float v = yb[120 + it];
973-
yl[c++] = v;
974-
sumy += v;
991+
yl[15] = v;
992+
sumy += v;
975993
}
976994

977995
FOR_UNROLL (short row = 0; row < nr0; row++) {
@@ -3834,6 +3852,7 @@ template [[host_name("kernel_mul_mv_id_q8_0_f32")]] kernel kernel_mul_mv_id_t
38343852
template [[host_name("kernel_mul_mv_id_q1_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q1_0_f32_impl<N_R0_Q1_0>>>;
38353853
template [[host_name("kernel_mul_mv_id_q2_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q2_0_f32_impl<N_R0_Q2_0>>>;
38363854
template [[host_name("kernel_mul_mv_id_pq2_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_pq2_0_f32_impl<N_R0_PQ2_0>>>;
3855+
template [[host_name("kernel_mul_mv_id_ptq1_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_ptq1_0_f32_impl<N_R0_PTQ1_0>>>;
38373856
template [[host_name("kernel_mul_mv_id_q4_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q4_0, N_R0_Q4_0>>>;
38383857
template [[host_name("kernel_mul_mv_id_q4_1_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q4_1, N_R0_Q4_1>>>;
38393858
template [[host_name("kernel_mul_mv_id_q5_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q5_0, N_R0_Q5_0>>>;

0 commit comments

Comments
 (0)