mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-07 22:59:00 -10:00
metal: K-quant / Q5_1 paired-run prompt decoders (W8) and paired mm_down3; qsa_prep M-RoPE (function constant 24)
- moe_mm.metal: W8<Q4_K / Q5_K / Q5_1 / Q6_K> (the lane's runs k and k + 16 share the sub-block's scale, min and high bits; 4-aligned code runs read as one uint); mm_down3 takes the paired-run decode on the COOP_K4 layout (TK 32). Prompt per layer at T 2048: Q4_K gate/up 10.6 -> 6.0 ms, Q5_K 15.4 -> 9.0, Q6_K 17.2 -> 14.0, Q5_1 down 5.0 -> 3.3; the same values (moe_mm_parity v3 incl. Q2_0 and IQ3_XXS / IQ4_NL) - qsa.metal: QSA_MROPE (function constant 24) - interleaved M-RoPE (ggml IMROPE, sections 11/11/10): buffer 20 of qsa_prep = int32 [cells][3] (t, h, w); rotary pair i of the row at pos_base + cell rotates by the cos/sin row tab[(pos_base + cell) * 3 + i % 3] (pooled indexer keys at pos_base + 4 b too). qsa_parity QSA_MROPE=1 (random table vs the CPU reference) / QSA_MROPE=id (identity table: the same numbers as the plain kernel)
This commit is contained in:
Niko1221 committed
1 parent
945c0d2424
commit
29724e5ea9
3 files changed
+113
-10
No files matched your search
@@ -691,6 +691,77 @@ template <> struct W8<T_IQ3_S> {
|
||||
}
|
||||
};
|
||||
|
||||
// W8 for the K-quants and Q5_1: the lane's runs k and k + 16 (k % 32 < 16) share the block's scales / mins / high
|
||||
// bits; the 4-byte code runs read as one uint where the block layout keeps them 4-aligned (Q4_K 144-, Q5_K 176-,
|
||||
// Q5_1 24-byte blocks: rows of whole blocks start 4-aligned). The same per-value arithmetic as W4.
|
||||
static inline float4 nib4(uint v, int hi) { const uint s = hi ? (v >> 4) : v; return float4(as_type<uchar4>(s & 0x0F0F0F0Fu)); }
|
||||
template <> struct W8<T_Q4_K> {
|
||||
static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) {
|
||||
device const uchar* bl = row + (k >> 8) * 144;
|
||||
const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31;
|
||||
float s, mm;
|
||||
scale_min_k4(2 * c + hi, bl + 4, s, mm);
|
||||
const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm;
|
||||
device const uint* q = (device const uint*)(bl + 16 + 32 * c + l);
|
||||
const float4 qa = nib4(q[0], hi), qb = nib4(q[4], hi);
|
||||
a = half4(d1 * qa - m1);
|
||||
b = half4(d1 * qb - m1);
|
||||
}
|
||||
};
|
||||
template <> struct W8<T_Q5_K> {
|
||||
static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) {
|
||||
device const uchar* bl = row + (k >> 8) * 176;
|
||||
const int j = k & 255, c = j >> 6, hi = (j >> 5) & 1, l = j & 31;
|
||||
float s, mm;
|
||||
scale_min_k4(2 * c + hi, bl + 4, s, mm);
|
||||
const float d1 = rdh(bl) * s, m1 = rdh(bl + 2) * mm;
|
||||
device const uint* qh = (device const uint*)(bl + 16 + l);
|
||||
device const uint* ql = (device const uint*)(bl + 48 + 32 * c + l);
|
||||
const uint sh = uint(2 * c + hi);
|
||||
const float4 ha = float4(as_type<uchar4>((qh[0] >> sh) & 0x01010101u)) * 16.0f;
|
||||
const float4 hb = float4(as_type<uchar4>((qh[4] >> sh) & 0x01010101u)) * 16.0f;
|
||||
a = half4(d1 * (nib4(ql[0], hi) + ha) - m1);
|
||||
b = half4(d1 * (nib4(ql[4], hi) + hb) - m1);
|
||||
}
|
||||
};
|
||||
template <> struct W8<T_Q5_1> {
|
||||
static inline void run(device const uchar* row, int k, threadgroup const uint*, thread half4& a, thread half4& b) {
|
||||
device const uchar* bl = row + (k >> 5) * 24;
|
||||
const float d = rdh(bl), m = rdh(bl + 2);
|
||||
const uint hb = rd32(bl + 4);
|
||||
const int j = k & 15; // runs at j and j + 16: the low nibbles of qs[j..j+3], then the high ones
|
||||
const uint v = *(device const uint*)(bl + 8 + j);
|
||||
float4 wa, wb;
|
||||
for (int c = 0; c < 4; ++c) {
|
||||
const uint qv = (v >> (8 * c)) & 0xFFu;
|
||||
wa[c] = float(int(qv & 15) | int(((hb >> (j + c)) & 1) << 4)) * d + m;
|
||||
wb[c] = float(int(qv >> 4) | int(((hb >> (j + 16 + c)) & 1) << 4)) * d + m;
|
||||
}
|
||||
a = half4(wa);
|
||||
b = half4(wb);
|
||||
}
|
||||
};
|
||||
template <> struct W8<T_Q6_K> {
|
||||
static inline void run(device const uchar* row, int k, threadgroup const uint* tab, thread half4& a, thread half4& b) {
|
||||
device const uchar* bl = row + (k >> 8) * 210;
|
||||
const int j = k & 255, n = j >> 7, jj = j & 127, qd = jj >> 5, l = jj & 31; // l < 16; the pair at l + 16
|
||||
device const uchar* ql = bl + 64 * n + l + ((qd & 1) ? 32 : 0);
|
||||
device const uchar* qh = bl + 128 + 32 * n + l;
|
||||
const float d = rdh(bl + 208);
|
||||
device const char* sc = (device const char*)(bl + 192 + 8 * n);
|
||||
const float dsa = d * float(sc[2 * qd]), dsb = d * float(sc[1 + 2 * qd]);
|
||||
float4 wa, wb;
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
const int qa = int(((qd >> 1) ? (ql[i] >> 4) : (ql[i] & 15)) | (((qh[i] >> (2 * qd)) & 3) << 4)) - 32;
|
||||
const int qb = int(((qd >> 1) ? (ql[16 + i] >> 4) : (ql[16 + i] & 15)) | (((qh[16 + i] >> (2 * qd)) & 3) << 4)) - 32;
|
||||
wa[i] = dsa * float(qa);
|
||||
wb[i] = dsb * float(qb);
|
||||
}
|
||||
a = half4(wa);
|
||||
b = half4(wb);
|
||||
}
|
||||
};
|
||||
|
||||
template <int F, int TM, int TN, int TK, int SG>
|
||||
kernel void mm_gu3(device const uchar* W [[buffer(0)]], device const half* XS [[buffer(1)]],
|
||||
device half* HS [[buffer(2)]], device const int2* tiles [[buffer(3)]], constant MmArgs& a [[buffer(4)]],
|
||||
@@ -796,6 +867,18 @@ kernel void mm_down3(device const uchar* W [[buffer(0)]], device const half* HS
|
||||
const int k = k0 + ix[0];
|
||||
rw[i] = W4<F>::run(w0 + ulong(ix[1]) * a.d_row, k & ~3, tab)[k & 3];
|
||||
}
|
||||
} else if (TK == 32) {
|
||||
// the paired runs (element i + cap / 2 is (k + 16, n), the COOP_K4 layout): one W8 a pair
|
||||
const uint16_t half_cap = rw.get_capacity() / 2;
|
||||
#pragma clang loop unroll(full)
|
||||
for (uint16_t i = 0; i < half_cap; i += 4) {
|
||||
const auto ix = rw.get_multidimensional_index(i);
|
||||
half4 w, w2;
|
||||
W8<F>::run(w0 + ulong(ix[1]) * a.d_row, k0 + ix[0], tab, w, w2);
|
||||
rw[i] = w.x; rw[i + 1] = w.y; rw[i + 2] = w.z; rw[i + 3] = w.w;
|
||||
const uint16_t j = i + half_cap;
|
||||
rw[j] = w2.x; rw[j + 1] = w2.y; rw[j + 2] = w2.z; rw[j + 3] = w2.w;
|
||||
}
|
||||
} else
|
||||
#pragma clang loop unroll(full)
|
||||
for (uint16_t i = 0; i < rw.get_capacity(); i += 4) {
|
||||
|
||||
+15
-7
@@ -29,6 +29,12 @@
|
||||
#define QCH 64
|
||||
|
||||
constant int QSA_NTOK [[function_constant(3)]];
|
||||
// QSA_MROPE (function constant 24): interleaved M-RoPE (the vision path; strata/kernels/mrope.hpp, ggml IMROPE,
|
||||
// sections 11/11/10): buffer 20 of qsa_prep is a table int32 [cells][3] (t, h, w), and rotary pair i (0..31) of the row
|
||||
// at (pos_base + cell) rotates by the cos/sin table's row tab[(pos_base + cell) * 3 + i % 3] instead of its own -
|
||||
// the pooled indexer keys too (position pos_base + 4 b of a block). Off: exactly the plain RoPE.
|
||||
constant bool QSA_MROPE [[function_constant(24)]];
|
||||
constant bool qsa_mrope = is_function_constant_defined(QSA_MROPE) && QSA_MROPE;
|
||||
|
||||
struct QsaArgs {
|
||||
uint pos0; // the window's first cell (relative to the sequence)
|
||||
@@ -55,7 +61,7 @@ static inline float tg_sum256(float v, threadgroup float* part, ushort sg, ushor
|
||||
// one row (width w <= 256 = threads) : out = rope((rs * x) * g) (RoPE on the first 64 dims, NEOX pairs i, i+32)
|
||||
static inline float norm_rope(float x, float xhi_unused, device const float* g, threadgroup float* row,
|
||||
threadgroup float* part, ushort tid, ushort sg, ushort lane, int w, float eps,
|
||||
device const float2* cs, uint absp, bool rope) {
|
||||
device const float2* cs, uint absp, bool rope, device const int* mtab) {
|
||||
const float ss = tg_sum256(tid < w ? x * x : 0.0f, part, sg, lane, 8);
|
||||
const float rs = rsqrt(ss / float(w) + eps);
|
||||
const float y = tid < w ? (rs * x) * g[tid] : 0.0f;
|
||||
@@ -64,7 +70,8 @@ static inline float norm_rope(float x, float xhi_unused, device const float* g,
|
||||
float o = y;
|
||||
if (rope && tid < QROT) {
|
||||
const int p = tid & 31;
|
||||
const float2 c = cs[absp * 32 + p];
|
||||
const uint pp = qsa_mrope ? uint(mtab[absp * 3 + uint(p % 3)]) : absp;
|
||||
const float2 c = cs[pp * 32 + p];
|
||||
const float a = row[p], b = row[p + 32];
|
||||
o = tid < 32 ? a * c.x - b * c.y : a * c.y + b * c.x;
|
||||
}
|
||||
@@ -84,6 +91,7 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float
|
||||
device char* kq [[buffer(12)]], device char* vq [[buffer(13)]], device half* ksc [[buffer(14)]],
|
||||
device half* vsc [[buffer(15)]], device float* tail [[buffer(16)]], device float* dead [[buffer(17)]],
|
||||
device float* pooled [[buffer(18)]], constant QsaArgs& a [[buffer(19)]],
|
||||
device const int* mtab [[buffer(20)]],
|
||||
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
|
||||
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
|
||||
threadgroup float part[8];
|
||||
@@ -99,14 +107,14 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float
|
||||
if (tid < ID) tail[(cell & 3) * ID + tid] = x;
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
if (cell == 0) {
|
||||
const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false);
|
||||
const float o = norm_rope(x, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, 0, false, mtab);
|
||||
if (tid < ID) dead[tid] = o;
|
||||
}
|
||||
if ((cell & 3) == 3) {
|
||||
float m = 0.0f;
|
||||
if (tid < ID) m = (((tail[tid] + tail[ID + tid]) + tail[2 * ID + tid]) + tail[3 * ID + tid]) * 0.25f;
|
||||
const uint b = cell >> 2;
|
||||
const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true);
|
||||
const float o = norm_rope(m, 0.0f, ikn, row, part, tid, sg, lane, ID, a.eps, cs, a.pos_base + 4 * b, true, mtab);
|
||||
if (tid < ID) pooled[b * ID + tid] = o;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_device);
|
||||
@@ -117,13 +125,13 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float
|
||||
const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell;
|
||||
if (r < QH) {
|
||||
const float x = qfull[(t * QH + r) * 2 * QD + tid];
|
||||
const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true);
|
||||
const float o = norm_rope(x, 0.0f, qn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab);
|
||||
qcur[(t * QH + r) * QD + tid] = o;
|
||||
} else if (r < QH + 2 * QKV) {
|
||||
const bool isv = r >= QH + QKV;
|
||||
const int h = r - QH - (isv ? QKV : 0);
|
||||
float x = (isv ? vcur : kcur)[(t * QKV + h) * QD + tid];
|
||||
if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true);
|
||||
if (!isv) x = norm_rope(x, 0.0f, kn, row, part, tid, sg, lane, QD, a.eps, cs, absp, true, mtab);
|
||||
// int8 with an fp16 scale per 64 values: groups = simdgroup pairs (64 threads)
|
||||
float m = simd_max(fabs(x));
|
||||
if (lane == 0) amax[sg] = m;
|
||||
@@ -140,7 +148,7 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float
|
||||
} else {
|
||||
const int h = r - QH - 2 * QKV;
|
||||
const float x = tid < ID ? iqraw[(t * IH + h) * ID + tid] : 0.0f;
|
||||
const float o = norm_rope(x, 0.0f, iqn, row, part, tid, sg, lane, ID, a.eps, cs, absp, true);
|
||||
const float o = norm_rope(x, 0.0f, iqn, row, part, tid, sg, lane, ID, a.eps, cs, absp, true, mtab);
|
||||
if (tid < ID) qidx[(t * IH + h) * ID + tid] = o;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -41,7 +41,7 @@ int main(int argc, char** argv) {
|
||||
// QSA_TOPK=1: the v1 top-k; QSA_ATTN=1|2|3 (default 3): which attention (2 has its own chunking and merge)
|
||||
const int av = std::getenv("QSA_ATTN") ? std::atoi(std::getenv("QSA_ATTN")) : 1;
|
||||
const bool v2 = av == 2;
|
||||
mtl::Pipeline pprep = cx.pipeline(lib, "qsa_prep"), pscore = cx.pipeline(lib, "qsa_score"),
|
||||
mtl::Pipeline pprep = cx.pipeline(lib, "qsa_prep", mtl::Constants().b(24, std::getenv("QSA_MROPE") != nullptr)), pscore = cx.pipeline(lib, "qsa_score"),
|
||||
ptopk = cx.pipeline(lib, std::getenv("QSA_TOPK") ? "qsa_topk" : "qsa_topk2"),
|
||||
pattn = cx.pipeline(lib, av == 1 ? "qsa_attn" : av == 2 ? "qsa_attn2" : "qsa_attn3"),
|
||||
pmerge = cx.pipeline(lib, v2 ? "qsa_merge2" : "qsa_merge");
|
||||
@@ -66,6 +66,17 @@ int main(int argc, char** argv) {
|
||||
for (int i = 0; i < QD; ++i) { qn.as<float>()[i] = 1.0f + 0.2f * r.sym(); kn.as<float>()[i] = 1.0f + 0.2f * r.sym(); }
|
||||
for (int i = 0; i < ID; ++i) { iqn.as<float>()[i] = 1.0f + 0.2f * r.sym(); ikn.as<float>()[i] = 1.0f + 0.2f * r.sym(); }
|
||||
const uint32_t pos_base = 0;
|
||||
// QSA_MROPE=1: qsa_prep with the M-RoPE table (function constant 24): random (t, h, w) positions per cell, the CPU
|
||||
// reference rotating pair i by row tab[cell * 3 + i % 3]; QSA_MROPE=id: the identity table, checked bitwise
|
||||
// against the plain kernel
|
||||
const char* mre = std::getenv("QSA_MROPE");
|
||||
const bool mrope = mre != nullptr, mid = mrope && std::string(mre) == "id";
|
||||
mtl::Buffer mtab = cx.alloc(4 * 3 * (size_t)(MAXC + 16));
|
||||
{
|
||||
ref::Rng rm(77);
|
||||
for (int c = 0; c < MAXC + 16; ++c)
|
||||
for (int k = 0; k < 3; ++k) mtab.as<int>()[c * 3 + k] = mid ? c : (int)(rm.next() % (uint64_t)(MAXC + 16));
|
||||
}
|
||||
for (int p = 0; p < MAXC + 16; ++p)
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
const double inv = std::pow(1e7, -2.0 * i / 64.0), th = (double)p * inv;
|
||||
@@ -81,7 +92,7 @@ int main(int argc, char** argv) {
|
||||
mtl::Arg::B(ik, 4 * (size_t)pos0 * ID), mtl::Arg::B(qn), mtl::Arg::B(kn), mtl::Arg::B(iqn),
|
||||
mtl::Arg::B(ikn), mtl::Arg::B(cs), mtl::Arg::B(qcur), mtl::Arg::B(qidx), mtl::Arg::B(kq),
|
||||
mtl::Arg::B(vq), mtl::Arg::B(ksc), mtl::Arg::B(vsc), mtl::Arg::B(tail), mtl::Arg::B(dead),
|
||||
mtl::Arg::B(pooled), mtl::Arg::V(a)};
|
||||
mtl::Arg::B(pooled), mtl::Arg::V(a), mtl::Arg::B(mtab)};
|
||||
d.grid[0] = 33;
|
||||
d.grid[1] = nt;
|
||||
d.tg[0] = 256;
|
||||
@@ -165,7 +176,8 @@ int main(int argc, char** argv) {
|
||||
for (int i = 0; i < w; ++i) out[i] = rs * x[i] * g[i];
|
||||
if (rope)
|
||||
for (int i = 0; i < 32; ++i) {
|
||||
const double c = cs.as<float>()[(absp * 32 + i) * 2], s = cs.as<float>()[(absp * 32 + i) * 2 + 1];
|
||||
const int pp = mrope ? mtab.as<int>()[absp * 3 + i % 3] : absp;
|
||||
const double c = cs.as<float>()[(pp * 32 + i) * 2], s = cs.as<float>()[(pp * 32 + i) * 2 + 1];
|
||||
const double a = out[i], b = out[i + 32];
|
||||
out[i] = a * c - b * s;
|
||||
out[i + 32] = a * s + b * c;
|
||||
|
||||
Reference in new issue
Block a user