Niko1221/Strata
1
0 Fork 0
mirror of https://github.com/Niko1221/Strata.git synced 2026-11-07 14:21:06 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

[Engine] Metal: --batch N / "parallel": N runs, and the slot rows share one window

Browse Source 
- the pdl capability queries answer false instead of aborting the engine (#3981, #2715)
- each batch slot keeps its own running state in the Metal window (GDN state and conv history, QSA K/V, indexer tail and pooled keys,
  PLE history), side by side in bigger buffers; the slot copy, the checkpoints and session_zero follow the session they are for
- a batch window holds one token of each active slot: the GDN step, the QSA kernels and the PLE conv take a per-row slot, cell and
  history parity (function constants, off for every other caller), so the dense and expert weights are read once for all rows
- STRATA_METAL_BATCH_ROWS=0 runs the slots one window each instead (diagnosis)
This commit is contained in:
Niko1221 committed 2026-10-30 02:58:46 +01:00
1 parent 870aedd9e1
commit 766b54afb1
11 files changed
+488 -77
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

metal/kernels/gdn.metal
+22 -7
View File
@@ -14,6 +14,12 @@
// `hist` and the committed one written to `hist_out` (ping-pong: the q/k channels are shared by three threadgroups).
constant int GDN_NTOK [[function_constant(2)]];
// GDN_ROWS (function constant 31, with GDN_NTOK = 1): the grid is (h_v, rows) and row y is a DIFFERENT sequence (a batch slot) with its own
// state and conv history; rinfo as in qsa.metal (header [5] state stride, [6] history stride per slot and parity; row y at [8 + 4 y] =
// {slot, cell, history parity, 0}). hist / hist_out are then one buffer (hist), the parities side by side per slot. The inputs and the
// output are row y of the window's arrays. Off: exactly the plain kernel.
constant bool GDN_ROWS_C [[function_constant(31)]];
constant bool gdn_rows = is_function_constant_defined(GDN_ROWS_C) && GDN_ROWS_C;
struct GdnArgs {
uint h_k; // 16
@@ -40,6 +46,7 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b
device const float* ssm_a [[buffer(8)]], device const float* gamma [[buffer(9)]],
device float* y [[buffer(10)]], constant GdnArgs& a [[buffer(11)]],
device float* hist_out [[buffer(12)]],
device const uint* rinfo [[buffer(13), function_constant(gdn_rows)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
constexpr int S = 128;
@@ -53,6 +60,13 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b
const int cb = sg & 3, ib = sg >> 2;
const int j = 32 * cb + lane, i0 = RPT * ib;
const float scale = 1.0f / sqrt(float(S));
const int row0 = gdn_rows ? int(tg.y) : 0;
if (gdn_rows) {
const ulong sl = rinfo[8 + 4 * row0], par = rinfo[8 + 4 * row0 + 2];
state += sl * rinfo[5];
hist_out = hist + (sl * 2 + (1 - par)) * rinfo[6];
hist += (sl * 2 + par) * rinfo[6];
}
// the channel this thread convolves: 0..127 q of head hk, 128..255 k of head hk, 256..383 v of head h
int ch = -1;
@@ -70,9 +84,10 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b
const int commit_at = a.commit == 1u ? GDN_NTOK : (a.commit >> 8) ? int(a.commit & 0xFFu) : 0;
for (int t = 0; t < GDN_NTOK; ++t) {
// conv + SiLU
const int tr = t + row0;
float cv = 0.0f;
if (ch >= 0) {
const float in = qkv[t * int(a.channels) + ch];
const float in = qkv[tr * int(a.channels) + ch];
const float vals[4] = {hst[0], hst[1], hst[2], in};
float s = 0.0f;
for (int k = 0; k < 4; ++k) s = fma(vals[k], cw[k], s);
@@ -86,12 +101,12 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b
const float ssq = simd_sum(cv * cv);
if (lane == 0) part[sg] = ssq;
if (tid == 0) {
const float v = alpha[t * hv + h] + dt[h];
const float v = alpha[tr * hv + h] + dt[h];
// log1p(e^v) (Kahan: log(u) * e / (u - 1), u = 1 + e; exact where 1 + e rounds to 1)
const float e = exp(v), u = 1.0f + e;
const float sp = v > 20.0f ? v : (u == 1.0f ? e : log(u) * e / (u - 1.0f));
gb[0] = exp(sp * ssm_a[h]);
gb[1] = 1.0f / (1.0f + exp(-beta[t * hv + h]));
gb[1] = 1.0f / (1.0f + exp(-beta[tr * hv + h]));
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tid < 256) {
@@ -132,14 +147,14 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b
if (ib == 0) {
const float tot = (part[0] + part[1]) + (part[2] + part[3]);
const float rs = rsqrt(tot / float(S) + a.eps);
const float zz = z[t * hv * S + h * S + j];
const float zz = z[tr * hv * S + h * S + j];
const float yv = ((rs * o) * gamma[j]) * (1.0f / (1.0f + exp(-zz)));
y[t * hv * S + h * S + j] = yv;
y[tr * hv * S + h * S + j] = yv;
if (a.plane) {
const float q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f,
0.00006103515625f};
const float p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f};
device float* raw = y + t * hv * S + h * S + j;
device float* raw = y + tr * hv * S + h * S + j;
const int jj = j & 15;
raw[a.plane] = yv * q2s[jj & 7];
raw[2 * a.plane] = yv * p4s[jj & 3];
@@ -161,7 +176,7 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b
if (a.plane && ib == 0 && (j & 15) == 0) {
float s0 = 0.0f, s1 = 0.0f;
for (int k = 0; k < 16; k += 2) { s0 += ys[j + k]; s1 += ys[j + k + 1]; }
y[4 * a.plane + t * (hv * S / 16) + (h * S + j) / 16] = s0 + s1;
y[4 * a.plane + tr * (hv * S / 16) + (h * S + j) / 16] = s0 + s1;
}
}
}
metal/kernels/qsa.metal
+37 -12
View File
@@ -35,6 +35,13 @@ constant int QSA_NTOK [[function_constant(3)]];
// 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;
constant bool QSA_ROWS_C [[function_constant(30)]];
constant bool qsa_rows = is_function_constant_defined(QSA_ROWS_C) && QSA_ROWS_C;
// QSA_ROWS (function constant 30): the rows of the window are DIFFERENT sequences (batch slots), one token each, each at its own
// cell and in its own slot of the caches. A buffer rinfo (uint) carries the strides between the slots' caches, in elements:
// [0] kq / vq (chars), [1] ksc / vsc (halves), [2] tail, [3] dead, [4] pooled (floats), [5] GDN state (floats), [6] conv history
// (floats, per slot and parity), [7] PLE history (floats); then row t at [8 + 4 t] = {slot, cell, history parity, 0}. Off: exactly the
// plain kernels.
struct QsaArgs {
uint pos0; // the window's first cell (relative to the sequence)
@@ -92,6 +99,7 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float
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)]],
device const uint* rinfo [[buffer(21), function_constant(qsa_rows)]],
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];
@@ -102,27 +110,38 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float
// the indexer keys, every window token in order (only threadgroup (32, 0) works)
if (tg.y != 0) return;
for (int t = 0; t < int(a.n_tok); ++t) {
const uint cell = a.pos0 + uint(t);
uint cell = a.pos0 + uint(t);
device float* tail_t = tail; device float* dead_t = dead; device float* pooled_t = pooled;
if (qsa_rows) {
const ulong sl = rinfo[8 + 4 * t];
cell = rinfo[8 + 4 * t + 1];
tail_t = tail + sl * rinfo[2]; dead_t = dead + sl * rinfo[3]; pooled_t = pooled + sl * rinfo[4];
}
const float x = tid < ID ? ikraw[t * ID + tid] : 0.0f;
if (tid < ID) tail[(cell & 3) * ID + tid] = x;
if (tid < ID) tail_t[(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, mtab);
if (tid < ID) dead[tid] = o;
if (tid < ID) dead_t[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;
if (tid < ID) m = (((tail_t[tid] + tail_t[ID + tid]) + tail_t[2 * ID + tid]) + tail_t[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, mtab);
if (tid < ID) pooled[b * ID + tid] = o;
if (tid < ID) pooled_t[b * ID + tid] = o;
}
threadgroup_barrier(mem_flags::mem_device);
}
return;
}
const int t = int(tg.y);
const uint cell = a.pos0 + uint(t), absp = a.pos_base + cell;
const uint cell = qsa_rows ? rinfo[8 + 4 * t + 1] : a.pos0 + uint(t), absp = a.pos_base + cell;
device char* kqr = kq; device char* vqr = vq; device half* kscr = ksc; device half* vscr = vsc;
if (qsa_rows) {
const ulong sl = rinfo[8 + 4 * t];
kqr = kq + sl * rinfo[0]; vqr = vq + sl * rinfo[0]; kscr = ksc + sl * rinfo[1]; vscr = vsc + sl * rinfo[1];
}
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, mtab);
@@ -143,8 +162,8 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float
int q = 0;
if (sf > 0.0f) q = clamp(int(rint(x / sf)), -127, 127);
const uint rowi = cell * QKV + uint(h);
(isv ? vq : kq)[rowi * QD + tid] = char(q);
if ((tid & 63) == 0) (isv ? vsc : ksc)[rowi * 4 + g] = sh;
(isv ? vqr : kqr)[rowi * QD + tid] = char(q);
if ((tid & 63) == 0) (isv ? vscr : kscr)[rowi * 4 + g] = sh;
} else {
const int h = r - QH - 2 * QKV;
const float x = tid < ID ? iqraw[(t * IH + h) * ID + tid] : 0.0f;
@@ -157,13 +176,16 @@ kernel void qsa_prep(device const float* qfull [[buffer(0)]], device const float
kernel void qsa_score(device const float* qidx [[buffer(0)]], device const float* pooled [[buffer(1)]],
device const float* dead [[buffer(2)]], device float* scores [[buffer(3)]],
constant QsaArgs& a [[buffer(4)]],
device const uint* rinfo [[buffer(5), function_constant(qsa_rows)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]]) {
const int t = int(tg.y);
const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2;
const uint n_kv = (qsa_rows ? rinfo[8 + 4 * t + 1] : a.pos0 + uint(t)) + 1, n_bid = n_kv >> 2;
const uint b = tg.x * 8 + sg;
if (b > n_bid) return;
device const float* key = b == n_bid ? dead : pooled + b * ID;
device const float* pooled_s = pooled; device const float* dead_s = dead;
if (qsa_rows) { const ulong sl = rinfo[8 + 4 * t]; pooled_s = pooled + sl * rinfo[4]; dead_s = dead + sl * rinfo[3]; }
device const float* key = b == n_bid ? dead_s : pooled_s + b * ID;
float k[4];
for (int i = 0; i < 4; ++i) k[i] = key[lane + 32 * i];
float s = 0.0f;
@@ -275,6 +297,7 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char*
device const char* vq [[buffer(2)]], device const half* ksc [[buffer(3)]],
device const half* vsc [[buffer(4)]], device const int* ids [[buffer(5)]],
device float* part [[buffer(6)]], constant QsaArgs& a [[buffer(7)]],
device const uint* rinfo [[buffer(8), function_constant(qsa_rows)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
constexpr int HPG = QH / QKV; // 12 query heads per kv head
@@ -282,7 +305,8 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char*
threadgroup float s[HPG][QCH];
threadgroup int cid[QCH];
const int chunk = int(tg.x), g = int(tg.y), t = int(tg.z);
const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET));
const uint n_kv = (qsa_rows ? rinfo[8 + 4 * t + 1] : a.pos0 + uint(t)) + 1, width = min(n_kv, uint(QBUDGET));
if (qsa_rows) { const ulong sl = rinfo[8 + 4 * t]; kq += sl * rinfo[0]; vq += sl * rinfo[0]; ksc += sl * rinfo[1]; vsc += sl * rinfo[1]; }
const int c0 = chunk * QCH;
const int nchunks = (QBUDGET + QCH - 1) / QCH;
device float* pp = part + ((uint(t) * nchunks + chunk) * QH) * (QD + 2);
@@ -345,11 +369,12 @@ kernel void qsa_attn(device const float* qcur [[buffer(0)]], device const char*
// grid (QH, n_tok), 256 threads: out = merge * sigmoid(gate), as prepared activations of n = QH * QD per token
kernel void qsa_merge(device const float* part [[buffer(0)]], device const float* qfull [[buffer(1)]],
device float* xp [[buffer(2)]], constant QsaArgs& a [[buffer(3)]],
device const uint* rinfo [[buffer(4), function_constant(qsa_rows)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
threadgroup float msum[256];
const int h = int(tg.x), t = int(tg.y);
const int nchunks = (QBUDGET + QCH - 1) / QCH;
const uint n_kv = a.pos0 + uint(t) + 1, width = min(n_kv, uint(QBUDGET));
const uint n_kv = (qsa_rows ? rinfo[8 + 4 * t + 1] : a.pos0 + uint(t)) + 1, width = min(n_kv, uint(QBUDGET));
const int nc = int((width + QCH - 1) / QCH);
device const float* pp = part + (uint(t) * nchunks * QH + uint(h)) * (QD + 2);
float M = -INFINITY;
metal/kernels/qsa2.metal
+2 -1
View File
@@ -120,13 +120,14 @@ kernel void qsa_merge2(device const float* part [[buffer(0)]], device const floa
kernel void qsa_topk2(device const float* scores [[buffer(0)]], device int* ids [[buffer(1)]],
constant QsaArgs& a [[buffer(2)]],
device const uint* rinfo [[buffer(3), function_constant(qsa_rows)]],
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 atomic_uint hist[256];
threadgroup uint sgs[32], sgs2[32];
threadgroup uint sel[2];
const int t = int(tg.x);
const uint n_kv = a.pos0 + uint(t) + 1, n_bid = n_kv >> 2, nb = n_bid + 1;
const uint n_kv = (qsa_rows ? rinfo[8 + 4 * t + 1] : a.pos0 + uint(t)) + 1, n_bid = n_kv >> 2, nb = n_bid + 1;
device int* out = ids + t * QBUDGET;
device const float* sc = scores + uint(t) * (a.max_cells / 4 + 2);
if (n_kv <= QBUDGET) {
metal/src/engine/msl/ple.metal
+14 -2
View File
@@ -73,16 +73,28 @@ kernel void ple_gate_norm(device const float* key [[buffer(0)]], device const fl
}
// grid (hc * n_embd / 256, T), 256 threads: the conv over [history | the window's rows], SiLU, and the residual
// PLE_ROWS (function constant 32): the rows are different sequences (batch slots): row t convolves its slot's history only (slot
// stride rinfo[7] floats, row t at rinfo[8 + 4 t] = {slot, ...}); the three taps before the row itself always read the history.
constant bool PLE_ROWS_C [[function_constant(32)]];
constant bool ple_rows = is_function_constant_defined(PLE_ROWS_C) && PLE_ROWS_C;
kernel void ple_conv_add(device const float* hist [[buffer(0)]], device const float* norm_w [[buffer(1)]],
device const half* kw [[buffer(2)]], device const float* gated [[buffer(3)]],
device float* R [[buffer(4)]], constant PleArgs& a [[buffer(5)]],
device const uint* rinfo [[buffer(6), function_constant(ple_rows)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
const uint ch = tg.x * 256 + tid, t = tg.y, hcd = a.hc * a.n_embd;
if (ch >= hcd) return;
if (ple_rows) hist += ulong(rinfo[8 + 4 * t]) * rinfo[7];
float acc = 0.0f;
for (int k = 0; k < 4; ++k) {
const int p = int(t) - (3 - k) * 3; // the row this tap reads: < 0 = the committed history
const float v = p >= 0 ? norm_w[uint(p) * hcd + ch] : hist[uint(9 + p) + 9 * ch];
float v;
if (ple_rows) {
v = k == 3 ? norm_w[t * hcd + ch] : hist[uint(9 - (3 - k) * 3) + 9 * ch];
} else {
const int p = int(t) - (3 - k) * 3; // the row this tap reads: < 0 = the committed history
v = p >= 0 ? norm_w[uint(p) * hcd + ch] : hist[uint(9 + p) + 9 * ch];
}
acc += float(kw[k + 4 * ch]) * v;
}
const float conv = acc / (1.0f + exp(-acc));
metal/src/engine/stubs.cpp
+3 -2
View File
@@ -284,8 +284,9 @@ namespace strata::kernels { void native_rope_apply(const float * p0, float * p1,
namespace strata::kernels { void native_route_resident(const float * p0, int32_t * p1, float * p2, const int32_t * p3, int p4, float p5, int p6, int p7, unsigned long long * p8, void * p9) { strata::metal_engine::missing("strata::kernels::native_route_resident"); } }
namespace strata::kernels { void native_router_top10(const float * p0, int32_t * p1, float * p2, void * p3) { strata::metal_engine::missing("strata::kernels::native_router_top10"); } }
namespace strata::kernels { void native_router_top10_multi(const float * p0, int32_t * p1, float * p2, int p3, void * p4) { strata::metal_engine::missing("strata::kernels::native_router_top10_multi"); } }
namespace strata::kernels { bool & pdl_scope() { strata::metal_engine::missing("strata::kernels::pdl_scope"); static std::remove_cv_t<std::remove_reference_t<bool &>> s{}; return s; } }
namespace strata::kernels { bool pdl_supported() { strata::metal_engine::missing("strata::kernels::pdl_supported"); return {}; } }
// programmatic dependent launch is a CUDA sm_90 feature: Metal never has it (hand-edited; the batch window asks this first, #3981)
namespace strata::kernels { bool & pdl_scope() { static bool s = false; return s; } }
namespace strata::kernels { bool pdl_supported() { return false; } }
namespace strata::kernels { void ple_block(const float * p0, const float * p1, const float * p2, const PleWeights & p3, PleOut & p4, void * p5, void * p6) { strata::metal_engine::missing("strata::kernels::ple_block"); } }
namespace strata::kernels { void ple_block_projected(const float * p0, const float * p1, const float * p2, const float * p3, const PleWeights & p4, PleOut & p5, void * p6, void * p7) { strata::metal_engine::missing("strata::kernels::ple_block_projected"); } }
namespace strata::kernels { void ple_history_advance(float * p0, const float * p1, void * p2) { strata::metal_engine::missing("strata::kernels::ple_history_advance"); } }
metal/src/engine/window.cpp
+359 -40
View File
@@ -52,6 +52,8 @@ namespace strata::metal_engine {
namespace {
int g_slot_cap = 0; // --batch N: the window keeps N + 1 sequences' running states side by side (set_batch_slots)
constexpr int N = 2560, HC = 4, LR = 320, NFF = 640, TOPK = 10, NEXP = 512, MAXT = 8, NVOCAB_MAX = 248320;
constexpr int P_MAXE = 128, P_GEXP = 16, P_GSTART = P_GEXP + P_MAXE, P_ETOK = P_GSTART + P_MAXE + 4,
P_EW = P_ETOK + P_MAXE, P_TOKENT = P_EW + P_MAXE, P_INTS = P_TOKENT + P_MAXE, P_IGU = 4, P_IDN = 8;
@@ -132,6 +134,21 @@ public:
bool checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err);
bool checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err);
void reset(); // a new sequence (session_zero): the running state zeroed
// ---- batch slots (--batch N): each slot is another sequence's running state (GDN recurrence and conv history, the PLE
// history, each QSA layer's K/V, indexer tail and pooled keys). The live members always hold ONE sequence's state (the
// main one between calls); a slot's lives in a stash and is swapped in for its window (activate): the buffer objects stay
// where they are (the cached dispatches point at them), only their contents are exchanged.
bool slots_init(const std::vector<strata::core::SessionState*>& slots, std::string& err);
int n_slots() const { return sets_.empty() ? 0 : (int) sets_.size() - 1; }
int set_of(const strata::core::SessionState* s) const {
auto it = set_of_ss_.find(s);
return it == set_of_ss_.end() ? 0 : it->second;
}
// one window of S one-token rows, row t = slot rows[t] at position pos[t]: each runs and commits in its own state
bool run_slots(const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err);
// slot b's state <-> the main one (the GPU is idle between windows): the running state and the first `upto` cells of K/V
bool slot_copy(bool into_slot, int b, int64_t upto, bool kv_only, std::string& err);
void activate(int set);
void set_max_t(int t) { max_t_ = std::min(t, MAXT); }
// ---- the MTP draft layer (mtp.hpp's math on mac-kern's kernels; the MTP GGUF from tools/mtp_pack.py)
bool mtp_load(const std::string& gguf, int max_t, std::string& err);
@@ -207,7 +224,8 @@ private:
return it->second;
}
strata::core::SessionState* ss_ = nullptr;
strata::core::SessionState* ss_ = nullptr; // the session whose sequence the live state is (the main one, or a batch slot's)
strata::core::SessionState* main_ss_ = nullptr; // the main session: it owns the PLE table (read by every slot)
gguf::File f_;
std::unique_ptr<mtl::Context> cx_;
mtl::Buffer W_;
@@ -234,6 +252,40 @@ private:
int cur_inj_ = 0;
// state
std::vector<mtl::Buffer> gst_, ghist_[2], kq_, vq_, ksc_, vsc_, tail_, dead_, pooled_;
// --batch N: every one of those is a view of set 0 in a bigger buffer holding all N + 1 sets (the slots' running states side
// by side), so ONE dispatch can address any slot (rinfo_: the strides and each row's slot / cell / history parity)
bool bigs_ok_ = false;
std::vector<mtl::Buffer> gstBig_, hbBig_, kqBig_, vqBig_, kscBig_, vscBig_, tailBig_, deadBig_, pooledBig_;
mtl::Buffer pleBig_, rinfo_;
size_t st_gst_ = 0, st_hist_ = 0, st_kq_ = 0, st_sc_ = 0, st_tail_ = 0, st_dead_ = 0, st_pooled_ = 0, st_ple_ = 0; // bytes between sets
bool rows_ = false; // the window being built / run holds one token of each of rows_n_ DIFFERENT sequences (batch slots)
int rows_n_ = 0;
strata::core::SessionState* rows_ss_[MAXT] = {};
int rows_set_[MAXT] = {};
static mtl::Buffer view_of(const mtl::Buffer& big, size_t off, size_t bytes) {
mtl::Buffer v = big;
v.ptr = (char*) big.ptr + off;
v.bias = big.bias + off;
v.bytes = bytes;
return v;
}
bool run_rows(const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err);
bool finish_rows(int S, const int32_t* tokens, int32_t* out, double t0);
struct SeqSet {
std::vector<mtl::Buffer> gst, gh0, gh1, kq, vq, ksc, vsc, tail, dead, pooled;
mtl::Buffer pleHist;
std::vector<float> ple_hist;
std::vector<float> last_norm;
int hp = 0, pend_fold = 0, last_T = 0;
bool last_head = true;
int32_t last_tokens[MAXT] = {};
strata::core::SessionState* ss = nullptr;
};
std::vector<std::unique_ptr<SeqSet>> sets_; // [0] the spare the main state is exchanged with, [b + 1] slot b
std::map<const strata::core::SessionState*, int> set_of_ss_;
int active_ = 0;
void swap_state(SeqSet& o);
void alloc_set(SeqSet& o, bool real, int idx = 0);
int hp_ = 0; // which of the two history buffers holds the committed conv history (all layers)
// commit(n) of a window of several tokens: the GDN fold of the kept prefix waits for the next window and runs at
// the start of its command buffer (one submit and wait less a round). Nothing else reads the GDN state between
@@ -435,6 +487,7 @@ long repack_pid_running(const std::string& model_file) {
bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionState& ss, const std::string& shard1,
int max_t, int64_t max_cells, std::string& err) {
ss_ = &ss;
main_ss_ = &ss;
max_t_ = std::min(max_t, MAXT);
max_cells_ = max_cells;
if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert != NEXP || g.n_layers != 48) {
@@ -515,7 +568,20 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta
zL_ = A(4 * (size_t) 36 * MAXT * 6144); ydump_ = A(4 * MAXT * 6144);
keyn_ = A(4 * MAXT * 10240); qn_ = A(4 * MAXT * 10240); gated_ = A(4 * MAXT * 10240); normW_ = A(4 * MAXT * 10240);
mrope_ = A(12 * (size_t) (MAXC + 16));
pleHist_ = A(4 * 9 * 10240);
bigs_ok_ = g_slot_cap > 0;
const size_t NSET = 1 + (size_t) std::max(0, g_slot_cap);
auto R256 = [](size_t b) { return (b + 255) / 256 * 256; };
st_gst_ = R256(4 * 128 * 48 * 128); st_hist_ = R256(4 * 10240 * 3); st_kq_ = R256((size_t) MAXC * 2 * 256);
st_sc_ = R256(2 * (size_t) MAXC * 2 * 4); st_tail_ = R256(4 * 4 * 128); st_dead_ = R256(4 * 128);
st_pooled_ = R256(4 * (size_t) MB * 128); st_ple_ = R256(4 * 9 * 10240);
// the set-0 view of a new buffer of `bytes` per set (or a plain buffer when there are no batch slots)
auto SV = [&](std::vector<mtl::Buffer>& bigs, size_t bytes, size_t stride, size_t per = 1, size_t which = 0) {
if (!bigs_ok_) { bigs.push_back(mtl::Buffer()); return A(bytes); }
if (which == 0) bigs.push_back(A(stride * per * NSET));
return view_of(bigs.back(), which * stride, bytes);
};
if (bigs_ok_) { pleBig_ = A(st_ple_ * NSET); pleHist_ = view_of(pleBig_, 0, 4 * 9 * 10240); } else pleHist_ = A(4 * 9 * 10240);
rinfo_ = A(4 * (8 + 4 * MAXT));
auto act = [&](int n) { Act a; a.n = n; a.b = A(4 * ((size_t) 4 * MAXT * n + (size_t) MAXT * n / 16)); return a; };
XA_ = act(N); XF_ = act(N); XY_ = act(6144); XQ_ = act(6144); XH_ = act(NFF); XE_ = act(N);
{ // Gyro (native_experts.txt v5 + hadamard.txt): the rotation's signs and the rotated gate/up input
@@ -552,17 +618,17 @@ bool Window::init(const strata::core::ModelGeometry& g, strata::core::SessionSta
for (int c = 0; c < 2051; ++c) qids_.as<int>()[t * 2051 + c] = c;
for (int l = 0; l < 48; ++l) {
if (l % 4 != 3) {
gst_.push_back(A(4 * 128 * 48 * 128));
ghist_[0].push_back(A(4 * 10240 * 3));
ghist_[1].push_back(A(4 * 10240 * 3));
gst_.push_back(SV(gstBig_, 4 * 128 * 48 * 128, st_gst_));
ghist_[0].push_back(SV(hbBig_, 4 * 10240 * 3, st_hist_, 2, 0)); // [set][parity]: the two parities side by side
ghist_[1].push_back(SV(hbBig_, 4 * 10240 * 3, st_hist_, 2, 1));
} else {
kq_.push_back(A((size_t) MAXC * 2 * 256));
vq_.push_back(A((size_t) MAXC * 2 * 256));
ksc_.push_back(A(2 * (size_t) MAXC * 2 * 4));
vsc_.push_back(A(2 * (size_t) MAXC * 2 * 4));
tail_.push_back(A(4 * 4 * 128));
dead_.push_back(A(4 * 128));
pooled_.push_back(A(4 * (size_t) MB * 128));
kq_.push_back(SV(kqBig_, (size_t) MAXC * 2 * 256, st_kq_));
vq_.push_back(SV(vqBig_, (size_t) MAXC * 2 * 256, st_kq_));
ksc_.push_back(SV(kscBig_, 2 * (size_t) MAXC * 2 * 4, st_sc_));
vsc_.push_back(SV(vscBig_, 2 * (size_t) MAXC * 2 * 4, st_sc_));
tail_.push_back(SV(tailBig_, 4 * 4 * 128, st_tail_));
dead_.push_back(SV(deadBig_, 4 * 128, st_dead_));
pooled_.push_back(SV(pooledBig_, 4 * (size_t) MB * 128, st_pooled_));
}
}
// ---- the HC read in Q8_0 (the owner's default on Metal, as STRATA_HC_Q8 on gfx12): the BF16 down / up / inject
@@ -796,6 +862,22 @@ void Window::layer(std::vector<mtl::Dispatch>& ds, int l, int ntok, int64_t pos0
mtl::Dispatch d;
d.p = P(lgdn_, "gdn_step", mtl::Constants().i(2, ntok), "gdn_step#" + std::to_string(ntok));
const int hp = hp_;
if (rows_) { // each row is a different sequence: its own state, parity and conv history (rinfo_), one grid row each
d.p = P(lgdn_, "gdn_step", mtl::Constants().i(2, 1).b(31, true), "gdn_step#rows");
GdnArgs ar{16, 48, 10240, 1u, eps_, fuse_ ? (uint32_t) (MAXT * 6144) : 0u};
d.args = {mtl::Arg::B(gstBig_[(size_t) gi]), mtl::Arg::B(hbBig_[(size_t) gi]), mtl::Arg::B(qkvL_, o_qkv),
mtl::Arg::B(alL_, o_ab), mtl::Arg::B(beL_, o_ab), mtl::Arg::B(zL_, o_z),
mtl::Arg::B(W_, T(b + "ssm_conv1d.weight").offset),
mtl::Arg::B(W_, T(b + "ssm_dt.bias").offset), mtl::Arg::B(W_, T(b + "ssm_a").offset),
mtl::Arg::B(W_, T(b + "ssm_norm.weight").offset), mtl::Arg::B(fuse_ ? XY_.b : yg_), mtl::Arg::V(ar),
mtl::Arg::B(hbBig_[(size_t) gi]), mtl::Arg::B(rinfo_)};
d.grid[0] = 48;
d.grid[1] = (uint32_t) rows_n_;
d.tg[0] = 512;
push(ds, d);
if (!fuse_) xprep(ds, yg_, XY_, 6144);
gemv(ds, T(b + "ssm_out.weight"), XY_, bo_, N, ntok);
} else {
d.args = {mtl::Arg::B(gst_[(size_t) gi]), mtl::Arg::B(ghist_[hp][(size_t) gi]), mtl::Arg::B(qkvL_, o_qkv),
mtl::Arg::B(alL_, o_ab), mtl::Arg::B(beL_, o_ab), mtl::Arg::B(zL_, o_z),
mtl::Arg::B(W_, T(b + "ssm_conv1d.weight").offset),
@@ -808,6 +890,7 @@ void Window::layer(std::vector<mtl::Dispatch>& ds, int l, int ntok, int64_t pos0
(void) self_commit; // the history ping-pong turns after the window (run) or at commit
if (!fuse_) xprep(ds, yg_, XY_, 6144);
gemv(ds, T(b + "ssm_out.weight"), XY_, bo_, N, ntok);
}
} else {
const int qi = l / 4;
gemv(ds, T(b + "attn_q.weight"), XA_, qfull_, 12288, ntok);
@@ -821,23 +904,34 @@ void Window::layer(std::vector<mtl::Dispatch>& ds, int l, int ntok, int64_t pos0
gemv(ds, T(b + "indexer.k_proj.weight"), XA_, ikr_, 128, ntok);
QsaArgs a{(uint32_t) pos0, 0, (uint32_t) max_cells_, (uint32_t) ntok, eps_};
mtl::Dispatch d;
d.p = P(lqsa_, "qsa_prep", qsa_c(), qsa_k());
const bool rw = rows_; // (batch slots: pos0 is the highest row's cell, the rows have their own in rinfo_)
const mtl::Constants rc = mtl::Constants().b(30, true);
const mtl::Buffer& KQ = rw ? kqBig_[(size_t) qi] : kq_[(size_t) qi];
const mtl::Buffer& VQ = rw ? vqBig_[(size_t) qi] : vq_[(size_t) qi];
const mtl::Buffer& KS = rw ? kscBig_[(size_t) qi] : ksc_[(size_t) qi];
const mtl::Buffer& VS = rw ? vscBig_[(size_t) qi] : vsc_[(size_t) qi];
const mtl::Buffer& TL = rw ? tailBig_[(size_t) qi] : tail_[(size_t) qi];
const mtl::Buffer& DD = rw ? deadBig_[(size_t) qi] : dead_[(size_t) qi];
const mtl::Buffer& PL = rw ? pooledBig_[(size_t) qi] : pooled_[(size_t) qi];
d.p = rw ? P(lqsa_, "qsa_prep", rc, "qsa_prep#rows") : P(lqsa_, "qsa_prep", qsa_c(), qsa_k());
d.args = {mtl::Arg::B(qfull_), mtl::Arg::B(kcur_), mtl::Arg::B(vcur_), mtl::Arg::B(iqr_), mtl::Arg::B(ikr_),
mtl::Arg::B(W_, T(b + "attn_q_norm.weight").offset), mtl::Arg::B(W_, T(b + "attn_k_norm.weight").offset),
mtl::Arg::B(W_, T(b + "indexer.q_norm.weight").offset), mtl::Arg::B(W_, T(b + "indexer.k_norm.weight").offset),
mtl::Arg::B(cs_), mtl::Arg::B(qcur_), mtl::Arg::B(qidx_), mtl::Arg::B(kq_[(size_t) qi]),
mtl::Arg::B(vq_[(size_t) qi]), mtl::Arg::B(ksc_[(size_t) qi]), mtl::Arg::B(vsc_[(size_t) qi]),
mtl::Arg::B(tail_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), mtl::Arg::B(pooled_[(size_t) qi]),
mtl::Arg::B(cs_), mtl::Arg::B(qcur_), mtl::Arg::B(qidx_), mtl::Arg::B(KQ),
mtl::Arg::B(VQ), mtl::Arg::B(KS), mtl::Arg::B(VS),
mtl::Arg::B(TL), mtl::Arg::B(DD), mtl::Arg::B(PL),
mtl::Arg::V(a), mtl::Arg::B(mrope_)};
if (rw) d.args.push_back(mtl::Arg::B(rinfo_));
d.grid[0] = 33;
d.grid[1] = ntok;
d.tg[0] = 256;
push(ds, d);
if (rec_) rec_->push_back({(int) ds.size() - 1, 0});
mtl::Dispatch s;
s.p = P(lqsa_, "qsa_score", {}, "qsa_score");
s.args = {mtl::Arg::B(qidx_), mtl::Arg::B(pooled_[(size_t) qi]), mtl::Arg::B(dead_[(size_t) qi]), mtl::Arg::B(scores_),
s.p = rw ? P(lqsa_, "qsa_score", rc, "qsa_score#rows") : P(lqsa_, "qsa_score", {}, "qsa_score");
s.args = {mtl::Arg::B(qidx_), mtl::Arg::B(PL), mtl::Arg::B(DD), mtl::Arg::B(scores_),
mtl::Arg::V(a)};
if (rw) s.args.push_back(mtl::Arg::B(rinfo_));
s.grid[0] = (uint32_t) (((pos0 + ntok) / 4 + 1 + 7) / 8);
s.grid[1] = ntok;
s.tg[0] = 256;
@@ -847,8 +941,9 @@ void Window::layer(std::vector<mtl::Dispatch>& ds, int l, int ntok, int64_t pos0
if (rec_) rec_->push_back({(int) ds.size() - 1, 2});
}
mtl::Dispatch k;
k.p = P(lqsa_, "qsa_topk2", {}, "qsa_topk");
k.p = rw ? P(lqsa_, "qsa_topk2", rc, "qsa_topk#rows") : P(lqsa_, "qsa_topk2", {}, "qsa_topk");
k.args = {mtl::Arg::B(scores_), mtl::Arg::B(qids_), mtl::Arg::V(a)};
if (rw) k.args.push_back(mtl::Arg::B(rinfo_));
k.grid[0] = ntok;
k.tg[0] = 1024;
if (sel) {
@@ -856,9 +951,10 @@ void Window::layer(std::vector<mtl::Dispatch>& ds, int l, int ntok, int64_t pos0
if (rec_) rec_->push_back({(int) ds.size() - 1, 0});
}
mtl::Dispatch at;
at.p = P(lqsa_, "qsa_attn", {}, "qsa_attn");
at.args = {mtl::Arg::B(qcur_), mtl::Arg::B(kq_[(size_t) qi]), mtl::Arg::B(vq_[(size_t) qi]), mtl::Arg::B(ksc_[(size_t) qi]),
mtl::Arg::B(vsc_[(size_t) qi]), mtl::Arg::B(qids_), mtl::Arg::B(part_), mtl::Arg::V(a)};
at.p = rw ? P(lqsa_, "qsa_attn", rc, "qsa_attn#rows") : P(lqsa_, "qsa_attn", {}, "qsa_attn");
at.args = {mtl::Arg::B(qcur_), mtl::Arg::B(KQ), mtl::Arg::B(VQ), mtl::Arg::B(KS),
mtl::Arg::B(VS), mtl::Arg::B(qids_), mtl::Arg::B(part_), mtl::Arg::V(a)};
if (rw) at.args.push_back(mtl::Arg::B(rinfo_));
at.grid[0] = 33;
at.grid[1] = 2;
at.grid[2] = ntok;
@@ -866,10 +962,11 @@ void Window::layer(std::vector<mtl::Dispatch>& ds, int l, int ntok, int64_t pos0
push(ds, at);
if (rec_) rec_->push_back({(int) ds.size() - 1, 0});
mtl::Dispatch m;
m.p = P(lqsa_, "qsa_merge", {}, "qsa_merge");
m.p = rw ? P(lqsa_, "qsa_merge", rc, "qsa_merge#rows") : P(lqsa_, "qsa_merge", {}, "qsa_merge");
QsaArgs am = a;
am.n_tok = MAXT; // the plane stride
m.args = {mtl::Arg::B(part_), mtl::Arg::B(qfull_), mtl::Arg::B(XQ_.b), mtl::Arg::V(am)};
if (rw) m.args.push_back(mtl::Arg::B(rinfo_));
m.grid[0] = 24;
m.grid[1] = ntok;
m.tg[0] = 256;
@@ -1094,11 +1191,11 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t
uint32_t rows[MAXT * strata::kernels::PLE_N_HEADS];
int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]};
for (int t = 0; t < Tn; ++t) {
strata::kernels::ngram_rows(&tokens[t], prev, 1, ss.ple.consts, rows + t * strata::kernels::PLE_N_HEADS);
prev[0] = prev[1];
prev[1] = tokens[t];
if (rows_) { prev[0] = rows_ss_[t]->ple_prev[0]; prev[1] = rows_ss_[t]->ple_prev[1]; } // (batch slots: each row its own sequence)
strata::kernels::ngram_rows(&tokens[t], prev, 1, main_ss_->ple.consts, rows + t * strata::kernels::PLE_N_HEADS);
if (!rows_) { prev[0] = prev[1]; prev[1] = tokens[t]; }
}
if (!ss.ple.table->gather_batch(rows, (size_t) Tn, emb_.as<float>(), err)) return false;
if (!main_ss_->ple.table->gather_batch(rows, (size_t) Tn, emb_.as<float>(), err)) return false;
}
tp_pre_ = now_ms();
const bool sel = !(fuse_ && pos0 + Tn <= 2051);
@@ -1108,8 +1205,9 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t
for (int c = 0; c < 2051; ++c) qids_.as<int>()[t * 2051 + c] = c;
qdirty_ = false;
}
const int key = (((((Tn * 2 + (head ? 1 : 0)) * 2 + hp_) * (MAXT + 1)) + pend_fold_) * 2 + (sel ? 1 : 0)) * 2 +
(mrope_on() ? 1 : 0);
const int key = rows_ ? 1000000 + Tn * 2 + (sel ? 1 : 0) + (int) (pend_fold_ * 100)
: (((((Tn * 2 + (head ? 1 : 0)) * 2 + hp_) * (MAXT + 1)) + pend_fold_) * 2 + (sel ? 1 : 0)) * 2 +
(mrope_on() ? 1 : 0);
mrope_sync(pos0, Tn);
const int fold_n = pend_fold_;
pend_fold_ = 0;
@@ -1138,7 +1236,7 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t
}
tp_rg_ = now_ms() - r0;
last_n_disp_ = (int) c.ds.size();
return finish(Tn, tokens, pos0, head, out, t0);
return rows_ ? finish_rows(Tn, tokens, out, t0) : finish(Tn, tokens, pos0, head, out, t0);
}
if (stream_ && !st_->ready() && !stream_arena(err)) return false;
Cached& cw = wcache_[key];
@@ -1185,9 +1283,10 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t
mtl::Arg::B(gated_), mtl::Arg::B(normW_), mtl::Arg::V(pa)};
g3.grid[0] = HC; g3.grid[1] = Tn; g3.tg[0] = 256;
push(ds, g3);
g4.p = P(lple_, "ple_conv_add", {}, "ple_conv_add");
g4.args = {mtl::Arg::B(pleHist_), mtl::Arg::B(normW_), kw_arg(),
g4.p = rows_ ? P(lple_, "ple_conv_add", mtl::Constants().b(32, true), "ple_conv_add#rows") : P(lple_, "ple_conv_add", {}, "ple_conv_add");
g4.args = {mtl::Arg::B(rows_ ? pleBig_ : pleHist_), mtl::Arg::B(normW_), kw_arg(),
mtl::Arg::B(gated_), mtl::Arg::B(R_), mtl::Arg::V(pa)};
if (rows_) g4.args.push_back(mtl::Arg::B(rinfo_));
g4.grid[0] = HC * N / 256; g4.grid[1] = Tn; g4.tg[0] = 256;
push(ds, g4);
}
@@ -1217,7 +1316,56 @@ bool Window::run(int Tn, const int32_t* tokens, int64_t pos0, bool head, int32_t
last_gpu_ms_ = 1000.0 * run_ds(cw);
tp_rg_ = now_ms() - r0;
last_n_disp_ = (int) ds.size();
return finish(Tn, tokens, pos0, head, out, t0);
return rows_ ? finish_rows(Tn, tokens, out, t0) : finish(Tn, tokens, pos0, head, out, t0);
}
bool Window::run_rows(const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err) {
if (S < 1 || S > MAXT || S > max_t_) { err = "metal window: batch window size out of range"; return false; }
if (active_ != 0) activate(0);
uint32_t* ri = rinfo_.as<uint32_t>() + 8;
int64_t maxpos = 0;
for (int t = 0; t < S; ++t) {
if (rows[t] < 0 || rows[t] >= n_slots()) { err = "metal window: a batch row names a slot that does not exist"; return false; }
for (int u = 0; u < t; ++u)
if (rows[u] == rows[t]) { err = "metal window: two rows of one batch window are the same slot"; return false; }
const int set = rows[t] + 1;
rows_ss_[t] = sets_[(size_t) set]->ss;
rows_set_[t] = set;
ri[4 * t] = (uint32_t) set;
ri[4 * t + 1] = (uint32_t) pos[t];
ri[4 * t + 2] = (uint32_t) sets_[(size_t) set]->hp;
ri[4 * t + 3] = 0;
maxpos = std::max<int64_t>(maxpos, pos[t]);
}
rows_ = true;
rows_n_ = S;
const bool ok = run(S, tokens, maxpos, true, out, err);
rows_ = false;
return ok;
}
// the end of a window of batch slots: the picks, and each slot's commit (a one-token step committed its row: the other history
// parity is the committed one; the PLE history takes the row's normalized values; the session's token pair advances)
bool Window::finish_rows(int Tn, const int32_t* tokens, int32_t* out, double t0) {
++win_n_[Tn];
win_ms_[Tn] += now_ms() - t0;
win_gpu_[Tn] += last_gpu_ms_;
for (int t = 0; t < Tn; ++t) out[t] = tokout_.as<int32_t>()[t];
for (int t = 0; t < Tn; ++t) {
SeqSet& s = *sets_[(size_t) rows_set_[t]];
s.hp = 1 - s.hp;
if (ple_on_) {
const float* norm = normW_.as<float>() + (size_t) t * NG_HC_DIM;
for (int c = 0; c < NG_HC_DIM; ++c) {
float* col = s.pleHist.as<float>() + (size_t) c * NG_HIST;
for (int r = 0; r + 1 < NG_HIST; ++r) col[r] = col[r + 1];
col[NG_HIST - 1] = norm[c];
}
}
s.ss->ple_prev[0] = s.ss->ple_prev[1];
s.ss->ple_prev[1] = tokens[t];
}
return true;
}
double Window::run_ds(Cached& c) {
@@ -1286,7 +1434,7 @@ bool Window::stream_arena(std::string& err) {
if (sysctlbyname("vm.vm_page_filecache_min", &pages, &pl, nullptr, 0) == 0) am.filecache_min = (double) pages * (double) vm_page_size;
}
am.available = avail_now();
if (ss_ && ss_->ple.table && ss_->ple.table->mode() == strata::kernels::PleIo::Mmap)
if (main_ss_ && main_ss_->ple.table && main_ss_->ple.table->mode() == strata::kernels::PleIo::Mmap)
am.ple_mapped = (double) ss_->ple.table->rows() * 192.0; // a row is at most 192 bytes (Q8_0 + its scale)
if (const char* e = std::getenv("STRATA_METAL_KEEP_FREE_GB")) am.keep_gb = std::atof(e);
if (const char* e = std::getenv("STRATA_METAL_ARENA_GB")) am.arena_gb = std::atof(e);
@@ -1499,6 +1647,154 @@ void Window::fold(std::vector<mtl::Dispatch>& ds, int n_keep, int hp) {
nb_ = false; // what follows reads the folded state: a barrier
}
void Window::swap_state(SeqSet& o) {
auto sw = [](std::vector<mtl::Buffer>& a, std::vector<mtl::Buffer>& b) {
for (size_t i = 0; i < a.size() && i < b.size(); ++i) std::swap(a[i], b[i]);
};
sw(gst_, o.gst); sw(ghist_[0], o.gh0); sw(ghist_[1], o.gh1); sw(kq_, o.kq); sw(vq_, o.vq); sw(ksc_, o.ksc); sw(vsc_, o.vsc);
sw(tail_, o.tail); sw(dead_, o.dead); sw(pooled_, o.pooled);
std::swap(pleHist_, o.pleHist);
ple_hist_.swap(o.ple_hist);
last_norm_.swap(o.last_norm);
std::swap(hp_, o.hp); std::swap(pend_fold_, o.pend_fold); std::swap(last_T_, o.last_T); std::swap(last_head_, o.last_head);
for (int t = 0; t < MAXT; ++t) std::swap(last_tokens_[t], o.last_tokens[t]);
std::swap(ss_, o.ss);
}
// the buffers of one sequence's running state, shaped like the live ones (real: allocated and zero; else empty placeholders)
void Window::alloc_set(SeqSet& o, bool real, int idx) {
const size_t MAXC = (size_t) max_cells_, MB = MAXC / 4 + 2;
if (real && bigs_ok_) { // set idx's part of the big buffers
for (size_t i = 0; i < gst_.size(); ++i) {
o.gst.push_back(view_of(gstBig_[i], (size_t) idx * st_gst_, 4 * 128 * 48 * 128));
o.gh0.push_back(view_of(hbBig_[i], ((size_t) idx * 2 + 0) * st_hist_, 4 * 10240 * 3));
o.gh1.push_back(view_of(hbBig_[i], ((size_t) idx * 2 + 1) * st_hist_, 4 * 10240 * 3));
}
for (size_t j = 0; j < kq_.size(); ++j) {
o.kq.push_back(view_of(kqBig_[j], (size_t) idx * st_kq_, MAXC * 2 * 256));
o.vq.push_back(view_of(vqBig_[j], (size_t) idx * st_kq_, MAXC * 2 * 256));
o.ksc.push_back(view_of(kscBig_[j], (size_t) idx * st_sc_, 2 * MAXC * 2 * 4));
o.vsc.push_back(view_of(vscBig_[j], (size_t) idx * st_sc_, 2 * MAXC * 2 * 4));
o.tail.push_back(view_of(tailBig_[j], (size_t) idx * st_tail_, 4 * 4 * 128));
o.dead.push_back(view_of(deadBig_[j], (size_t) idx * st_dead_, 4 * 128));
o.pooled.push_back(view_of(pooledBig_[j], (size_t) idx * st_pooled_, 4 * MB * 128));
}
o.pleHist = view_of(pleBig_, (size_t) idx * st_ple_, 4 * 9 * 10240);
o.ple_hist.assign(ple_hist_.size(), 0.0f);
return;
}
auto A = [&](size_t bytes) { return real ? cx_->alloc(bytes) : mtl::Buffer(); };
for (size_t i = 0; i < gst_.size(); ++i) {
o.gst.push_back(A(4 * 128 * 48 * 128));
o.gh0.push_back(A(4 * 10240 * 3));
o.gh1.push_back(A(4 * 10240 * 3));
}
for (size_t j = 0; j < kq_.size(); ++j) {
o.kq.push_back(A(MAXC * 2 * 256));
o.vq.push_back(A(MAXC * 2 * 256));
o.ksc.push_back(A(2 * MAXC * 2 * 4));
o.vsc.push_back(A(2 * MAXC * 2 * 4));
o.tail.push_back(A(4 * 4 * 128));
o.dead.push_back(A(4 * 128));
o.pooled.push_back(A(4 * MB * 128));
}
o.pleHist = A(4 * 9 * 10240);
if (real) o.ple_hist.assign(ple_hist_.size(), 0.0f);
}
bool Window::slots_init(const std::vector<strata::core::SessionState*>& slots, std::string& err) {
if (!sets_.empty()) return true; // (asked once)
if (cx_ == nullptr) { err = "metal window: slots before the window is up"; return false; }
std::vector<std::unique_ptr<SeqSet>> v;
v.push_back(std::make_unique<SeqSet>());
alloc_set(*v[0], false);
v[0]->ss = ss_;
set_of_ss_[ss_] = 0;
uint64_t bytes = 0;
for (size_t b = 0; b < slots.size(); ++b) {
if ((int) b >= g_slot_cap) bigs_ok_ = false; // more slots than the window was sized for: serial windows only
auto x = std::make_unique<SeqSet>();
alloc_set(*x, true, (int) b + 1);
x->ss = slots[b];
for (auto* vec : {&x->gst, &x->gh0, &x->gh1, &x->kq, &x->vq, &x->ksc, &x->vsc, &x->tail, &x->dead, &x->pooled})
for (const auto& bf : *vec) bytes += bf.bytes;
set_of_ss_[slots[b]] = (int) b + 1;
v.push_back(std::move(x));
}
sets_ = std::move(v);
active_ = 0;
if (bigs_ok_) { // the strides the row kernels need (elements of each array's type)
uint32_t* h = rinfo_.as<uint32_t>();
h[0] = (uint32_t) st_kq_; h[1] = (uint32_t) (st_sc_ / 2); h[2] = (uint32_t) (st_tail_ / 4); h[3] = (uint32_t) (st_dead_ / 4);
h[4] = (uint32_t) (st_pooled_ / 4); h[5] = (uint32_t) (st_gst_ / 4); h[6] = (uint32_t) (st_hist_ / 4); h[7] = (uint32_t) (st_ple_ / 4);
}
std::fprintf(stderr, "strata metal: %zu batch slots, each with its own running state and K/V (%.2f GB each, %.2f GB in all)\n",
slots.size(), slots.empty() ? 0.0 : (double) bytes / (double) slots.size() / 1e9, (double) bytes / 1e9);
return true;
}
// Makes `set` the live state. The live members and the stashes exchange contents (two swaps through the active set's own stash).
void Window::activate(int set) {
if (set == active_ || set < 0 || set >= (int) sets_.size()) return;
swap_state(*sets_[(size_t) active_]); // the stash of the active one now holds its state, the live members the spare
swap_state(*sets_[(size_t) set]); // ... and now the wanted one
active_ = set;
}
bool Window::run_slots(const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err) {
if (sets_.empty()) { err = "metal window: no batch slots"; return false; }
static const bool one_window = !(std::getenv("STRATA_METAL_BATCH_ROWS") && std::string(std::getenv("STRATA_METAL_BATCH_ROWS")) == "0");
if (one_window && bigs_ok_) {
activate(0);
return run_rows(rows, S, tokens, pos, out, err);
}
bool ok = true;
for (int t = 0; t < S && ok; ++t) {
if (rows[t] < 0 || rows[t] >= n_slots()) { err = "metal window: a batch row names a slot that does not exist"; ok = false; break; }
activate(rows[t] + 1);
// a one-token window commits its row in its own step; commit(1) adds the PLE history and the session's token pair,
// which read buffers the next window rewrites
ok = run(1, tokens + t, pos[t], true, out + t, err) && commit(1, err);
}
activate(0);
return ok;
}
bool Window::slot_copy(bool into_slot, int b, int64_t upto, bool kv_only, std::string& err) {
if (b < 0 || b >= n_slots()) { err = "metal window: no such batch slot"; return false; }
activate(0);
SeqSet& o = *sets_[(size_t) b + 1];
if (pend_fold_ > 0 || o.pend_fold > 0) { err = "metal window: a slot copy with a fold pending"; return false; }
const size_t cells = (size_t) std::max<int64_t>(0, std::min<int64_t>(upto, max_cells_));
auto cp = [&](const mtl::Buffer& live, const mtl::Buffer& stash, size_t n) {
n = std::min({n, live.bytes, stash.bytes});
if (into_slot) std::memcpy(stash.ptr, live.ptr, n); else std::memcpy(live.ptr, stash.ptr, n);
};
const int hs = into_slot ? hp_ : o.hp; // the committed history parity of the source
if (!kv_only) {
if (into_slot) o.hp = hp_; else hp_ = o.hp;
for (size_t i = 0; i < gst_.size(); ++i) {
cp(gst_[i], o.gst[i], gst_[i].bytes);
cp(ghist_[hs][i], hs ? o.gh1[i] : o.gh0[i], ghist_[hs][i].bytes);
}
}
for (size_t j = 0; j < kq_.size(); ++j) {
cp(kq_[j], o.kq[j], cells * 512);
cp(vq_[j], o.vq[j], cells * 512);
cp(ksc_[j], o.ksc[j], cells * 16);
cp(vsc_[j], o.vsc[j], cells * 16);
cp(pooled_[j], o.pooled[j], ((cells + 3) / 4 + 2) * 512);
if (kv_only) continue;
cp(tail_[j], o.tail[j], tail_[j].bytes);
cp(dead_[j], o.dead[j], dead_[j].bytes);
}
if (kv_only) return true;
cp(pleHist_, o.pleHist, pleHist_.bytes);
if (into_slot) { o.ple_hist = ple_hist_; o.ss->ple_prev[0] = ss_->ple_prev[0]; o.ss->ple_prev[1] = ss_->ple_prev[1]; }
else { ple_hist_ = o.ple_hist; ss_->ple_prev[0] = o.ss->ple_prev[0]; ss_->ple_prev[1] = o.ss->ple_prev[1]; }
return true;
}
bool Window::checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err) {
(void) err;
if (pend_fold_ > 0) { // the last commit's fold first: the checkpoint is the state after it
@@ -2585,6 +2881,7 @@ std::string g_shard1;
void window_register(void* w);
bool mtp_load_pending(Window* w, std::string& err);
void set_model_shard(const char* shard1) { g_shard1 = shard1 ? shard1 : ""; }
void set_batch_slots(int n) { g_slot_cap = n > 0 ? std::min(n, 8) : 0; }
// One decode window per process: the prompt path (Prefill::init) may bring it up before the Verifier does.
namespace {
@@ -2628,16 +2925,38 @@ bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head,
static const bool logpos = std::getenv("STRATA_LOGPOS") != nullptr;
return ((Window*) w)->run(T, tokens, pos0, head || logpos, out, err);
}
void window_reset() {
if (g_the_window != nullptr && g_the_window->ready()) g_the_window->reset();
void window_reset(const strata::core::SessionState* ss) {
if (g_the_window == nullptr || !g_the_window->ready()) return;
g_the_window->activate(g_the_window->set_of(ss));
g_the_window->reset();
g_the_window->activate(0);
}
bool window_checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err) {
bool window_checkpoint_save(strata::core::ConversationCheckpoint& c, const strata::core::SessionState* ss, std::string& err) {
if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; }
return g_the_window->checkpoint_save(c, err);
g_the_window->activate(g_the_window->set_of(ss)); // a batch slot's session: that slot's running state
const bool ok = g_the_window->checkpoint_save(c, err);
g_the_window->activate(0);
return ok;
}
bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err) {
bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, const strata::core::SessionState* ss, std::string& err) {
if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; }
return g_the_window->checkpoint_restore(c, err);
g_the_window->activate(g_the_window->set_of(ss));
const bool ok = g_the_window->checkpoint_restore(c, err);
g_the_window->activate(0);
return ok;
}
bool window_init_slots(void* w, const std::vector<strata::core::SessionState*>& slots, std::string& err) {
return ((Window*) w)->slots_init(slots, err);
}
bool window_run_slots(void* w, const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err) {
const double t0 = now_ms();
const bool ok = ((Window*) w)->run_slots(rows, S, tokens, pos, out, err);
((Window*) w)->prof_add(true, now_ms() - t0);
return ok;
}
bool window_slot_copy_main(bool into_slot, int slot, int64_t upto, bool kv_only, std::string& err) {
if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; }
return g_the_window->slot_copy(into_slot, slot, upto, kv_only, err);
}
bool window_commit(void* w, int n_keep, std::string& err) {
const double t0 = now_ms();
metal/src/mtl_rt.mm
+1 -5
View File
@@ -65,11 +65,7 @@ struct Context::Impl {
static inline id<MTLBuffer> B(const Buffer& b) { return (__bridge id<MTLBuffer>)b.handle; }
// the MTLBuffer and offset that bind byte `off` of b (a chunked buffer: the chunk the offset selects)
static inline id<MTLBuffer> BO(const Buffer& b, size_t off, size_t& o) {
#if defined(STRATA_DS4_RUNTIME)
if (!b.chunks) { o = off + b.bias; return B(b); }
#else
if (!b.chunks) { o = off; return B(b); }
#endif
if (!b.chunks) { o = off + b.bias; return B(b); } // (bias: a view onto part of a bigger allocation; 0 otherwise)
const size_t k = std::min(off / b.stride, b.chunks->size() - 1);
o = off - k * b.stride;
return (__bridge id<MTLBuffer>)(*b.chunks)[k];
src/core/conversation_state.cpp
+5 -5
View File
@@ -162,8 +162,8 @@ bool conversation_checkpoint_validate(const ConversationCheckpoint& c, const Ses
} // namespace strata::core
namespace strata::metal_engine {
// the running state lives in the Metal decode window (metal/src/engine/window.cpp), not in the session's buffers
bool window_checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& error);
bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& error);
bool window_checkpoint_save(strata::core::ConversationCheckpoint& c, const strata::core::SessionState* ss, std::string& error);
bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, const strata::core::SessionState* ss, std::string& error);
} // namespace strata::metal_engine
namespace strata::core {
#endif
@@ -289,8 +289,8 @@ bool checkpoint_save(ConversationCheckpoint& c, const SessionState& ss, const Mo
if (std::getenv("STRATA_PARK_TRACE")) std::fprintf(stderr, "strata parktrace: save join %.2f ms\n",
std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - j0).count());
#if defined(STRATA_GPU_METAL)
(void) ss; (void) g;
return strata::metal_engine::window_checkpoint_save(c, error);
(void) g;
return strata::metal_engine::window_checkpoint_save(c, &ss, error);
#endif
ConversationStateSizes z;
if (!checkpoint_targets(ss, g, c.ids.size(), z, error)) return false;
@@ -401,7 +401,7 @@ bool conversation_checkpoint_restore(const ConversationCheckpoint& c, SessionSta
if (hit == nullptr) conversation_checkpoint_join(c.gdn.data());
#if defined(STRATA_GPU_METAL)
(void) g;
if (!strata::metal_engine::window_checkpoint_restore(c, error)) return false;
if (!strata::metal_engine::window_checkpoint_restore(c, &ss, error)) return false;
ss.ple_prev[0] = c.ids.size() >= 2 ? c.ids[c.ids.size() - 2] : -1;
ss.ple_prev[1] = c.ids.size() >= 1 ? c.ids[c.ids.size() - 1] : -1;
return true;
src/core/session.cpp
+2 -2
View File
@@ -136,7 +136,7 @@ void session_release(SessionState& s) {
#if defined(STRATA_GPU_METAL)
} // namespace strata::core
namespace strata::metal_engine {
void window_reset(); // metal/src/engine/window.cpp: the decode window's running state back to a new sequence's
void window_reset(const strata::core::SessionState* ss); // metal/src/engine/window.cpp: that session's running state (the main one or a batch slot's) back to a new sequence's
} // namespace strata::metal_engine
namespace strata::core {
#endif
@@ -145,7 +145,7 @@ void session_zero(SessionState& s, const ModelGeometry& g, const float* R_init,
#if defined(STRATA_GPU_METAL)
// the running state lives in the Metal decode window; the session's own device buffers are not read there
(void) g; (void) R_init; (void) stream;
strata::metal_engine::window_reset();
strata::metal_engine::window_reset(&s);
s.ple_prev[0] = -1;
s.ple_prev[1] = -1;
s.ple_token = -1;
src/core/verify.cpp
+28
View File
@@ -91,6 +91,8 @@ bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::Se
bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, std::string& err);
bool window_commit(void* w, int n_keep, std::string& err);
bool window_copy_logits(void* w, int t, float* host);
bool window_init_slots(void* w, const std::vector<strata::core::SessionState*>& slots, std::string& err);
bool window_run_slots(void* w, const int* rows, int S, const int32_t* tokens, const int64_t* pos, int32_t* out, std::string& err);
} // namespace strata::metal_engine
#endif
@@ -3704,6 +3706,17 @@ bool Verifier::init_slots(const std::vector<SessionState*>& slots, std::string&
return false;
}
}
#if defined(STRATA_GPU_METAL)
// the Metal window keeps each slot's running state itself (metal/src/engine/window.cpp): none of the CUDA batch buffers
slots_ = slots;
slot_sp_.assign(slots.size(), sampling_);
if (metal_ == nullptr || !strata::metal_engine::window_init_slots(metal_, slots, err)) {
if (err.empty()) err = "verify: the Metal window is not up";
return false;
}
std::fprintf(stderr, "strata verify: batch windows of up to %d sequences, run slot by slot in the Metal window\n", (int) slots.size());
return true;
#endif
const strata::kernels::QsaShapes s = shapes_of(*g_);
const int64_t S = (int64_t) slots.size(), CB = 2 + max_t_;
const int64_t TS = (s.idx_block - 1) * g_->idx_key_dim, nQ = g_->n_qsa_layers();
@@ -4379,6 +4392,16 @@ bool Verifier::run_slots(int S, const int32_t* tokens, const int64_t* pos, PoolM
bool Verifier::run_slot_rows(const int* rows, int S, const int32_t* tokens, const int64_t* pos, PoolMultiFn pool,
void* user, int32_t* out, std::string& err) {
#if defined(STRATA_GPU_METAL)
// each row runs in its slot's state and commits there (a one-token window); nothing is left for commit_slots
if (S < 1 || S > (int) strata::kernels::kVerifyMaxT) { err = "verify: batch window size out of range"; return false; }
(void) pool; (void) user;
if (!strata::metal_engine::window_run_slots(metal_, rows, S, tokens, pos, out, err)) return false;
last_t_ = S;
for (int t = 0; t < S; ++t) { last_rows_[t] = rows[t]; last_tokens_[t] = tokens[t]; last_pos_b_[t] = pos[t]; }
last_batch_ = true;
return true;
#endif
using namespace strata::kernels;
const OnDevice on_device(device_);
const ModelGeometry& g = *g_;
@@ -4490,6 +4513,11 @@ bool Verifier::commit_slots(std::string& err) {
}
bool Verifier::commit_slot_prefixes(const int* keep, std::string& err) {
#if defined(STRATA_GPU_METAL)
(void) keep; // (greedy one-row windows, committed by window_run_slots)
if (!last_batch_ || last_t_ < 1) { err = "verify: commit_slots without a batch window"; return false; }
return true;
#endif
const OnDevice on_device(device_);
if (!last_batch_ || last_t_ < 1) { err = "verify: commit_slots without a batch window"; return false; }
const int S = last_t_;
src/program/generate.cpp
+15 -1
View File
@@ -129,7 +129,12 @@
#include <cuda_runtime.h>
#if defined(STRATA_GPU_METAL)
namespace strata::metal_engine { void set_model_shard(const char* shard1); } // metal/src/engine/window.cpp
namespace strata::metal_engine {
void set_model_shard(const char* shard1); // metal/src/engine/window.cpp
void set_batch_slots(int n); // metal/src/engine/window.cpp: how many --batch slots the window keeps state for
// batch slot b <-> the main sequence: the running state (unless kv_only) and the first `upto` cells of K/V
bool window_slot_copy_main(bool into_slot, int slot, int64_t upto, bool kv_only, std::string& err);
}
#endif
#include <array>
@@ -3971,6 +3976,7 @@ int main(int argc, char** argv) {
const char* isa_floor = strata::kernels::cpu::isa_floor_build();
#if defined(STRATA_GPU_METAL) // Apple Silicon: every expert is on the GPU; ggml-cpu (NEON) is the only CPU expert code
strata::metal_engine::set_model_shard(o.native_preset.c_str()); // the decode window reads the model GGUF itself
strata::metal_engine::set_batch_slots(o.batch);
o.mmap_experts = true; // ... experts included: no resident copy of them (the GGUF shards in place)
if (false) {
#else
@@ -12541,6 +12547,10 @@ int main(int argc, char** argv) {
};
auto copy_to_slot = [&](int b, const std::vector<int32_t>& ids, std::string& e) -> bool {
const int64_t upto = (int64_t) ids.size();
#if defined(STRATA_GPU_METAL)
// the running state and the K/V of the Metal window, not the session's device buffers
return strata::metal_engine::window_slot_copy_main(true, b, upto, false, e);
#endif
if (slot_copy_direct) {
if (!direct_slot_copy(true, b, ids, e)) return false;
} else
@@ -12588,6 +12598,10 @@ int main(int argc, char** argv) {
conversations.take_reuse();
const int64_t upto = at != nullptr ? (int64_t) at->ids.size() : (int64_t) bs[(size_t) b].ids.size();
if (at != nullptr && at->stage_parts.size() != stages.size()) { e = "a checkpoint without its stage parts"; return false; }
#if defined(STRATA_GPU_METAL)
if (!strata::metal_engine::window_slot_copy_main(false, b, upto, at != nullptr, e)) return false;
return at == nullptr || strata::core::conversation_checkpoint_restore(*at, ss, g, e); // (the checkpoint's running state)
#endif
if (slot_copy_direct && at == nullptr) {
if (!direct_slot_copy(false, b, bs[(size_t) b].ids, e)) return false;
} else