#!/usr/bin/env python3
"""Patch vendor/strata/metal for:
1. Dynamic n_expert (K=256 / 320 / 224 / 512) in router.metal, moe.metal, prompt.metal,
stream.h, stream.cpp, window.cpp, prompt_chain.inc, and route_parity.cpp (with K=256 + MoCE parity test).
2. Zero-copy file-backed dense tensor wrapping (cx.wrap(f.map, max_dense_end, false)) in
ExpertStream::load_dense with dynamically computed locked/skipped/unpinned byte totals.
3. Support for dense-only <model>-metal.gguf paired with STRX <model>-metal.gguf.experts.
4. Race-free MoCE cache-prior reranking (STRATA_METAL_MOCE_BONUS) using a per-lane 16-bit register
mask (scaled_mask + simd_broadcast) in router.metal / moe.metal, plus top-3 unperturbed protection
and smap residency consistency in ExpertStream::ahead.
5. Transient PromptGemm scratch buffer mapped on demand, skipped memset on fresh MAP_ANON pages,
unmapped via munmap + cx_->free after prefill with measured proc_mem() footprint logging,
router GEMM over n_exp_ rows, and multi-turn hot-slot protection across prompts.
6. Standalone zero-bloat strata_window_cli (with --logits-out / --logits-ref full-vocab FP64 KLD
& PPL evaluator) and convert_k256_q2_0 targets in metal/CMakeLists.txt.
"""
import subprocess
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent / "vendor" / "strata"
def reset_git() -> None:
subprocess.run(["git", "-C", str(ROOT), "checkout", "--", "."], check=True)
def patch_mtl_rt() -> None:
h_path = ROOT / "metal" / "src" / "mtl_rt.h"
mm_path = ROOT / "metal" / "src" / "mtl_rt.mm"
h_txt = h_path.read_text(encoding="utf-8")
h_txt = h_txt.replace(
"Buffer wrap(void* page_aligned, size_t bytes); // no-copy wrap (page-aligned, size multiple of page)",
"Buffer wrap(void* page_aligned, size_t bytes, bool count_allocated = true); // no-copy wrap (page-aligned, size multiple of page)",
)
h_path.write_text(h_txt, encoding="utf-8")
mm_txt = mm_path.read_text(encoding="utf-8")
old_wrap = """Buffer Context::wrap(void* p, size_t bytes) {
id<MTLBuffer> b = [impl_->dev newBufferWithBytesNoCopy:p length:bytes options:MTLResourceStorageModeShared deallocator:nil];
if (!b) {
std::fprintf(stderr, "mtl: wrap %zu failed\\n", bytes);
std::exit(4);
}
allocated_ += bytes;"""
new_wrap = """Buffer Context::wrap(void* p, size_t bytes, bool count_allocated) {
id<MTLBuffer> b = [impl_->dev newBufferWithBytesNoCopy:p length:bytes options:MTLResourceStorageModeShared deallocator:nil];
if (!b) {
std::fprintf(stderr, "mtl: wrap %zu failed\\n", bytes);
std::exit(4);
}
if (count_allocated) allocated_ += bytes;"""
assert old_wrap in mm_txt, "wrap target not found in mtl_rt.mm"
mm_path.write_text(mm_txt.replace(old_wrap, new_wrap, 1), encoding="utf-8")
def patch_kernels() -> None:
# 1. router.metal: dynamic a.n_expert + race-free MoCE cache-prior reranking via a.pad, smap & scaled_mask
r_path = ROOT / "metal" / "kernels" / "router.metal"
r_txt = r_path.read_text(encoding="utf-8")
old_route_body = """static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan,
constant MoeArgs& a, ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp,
threadgroup int* s_ent) {
if (sg < a.n_tok) {
device const float* l = logits + sg * 512;
float v[16];
for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i];
float mx = -INFINITY;
for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
float sum = 0.0f;
for (int i = 0; i < 16; ++i) {
v[i] = exp(v[i] - mx);
sum += v[i];
}
for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
const float rcp = 1.0f / sum;
for (int i = 0; i < 16; ++i) {
v[i] *= rcp;
if (isnan(v[i])) v[i] = -FLT_MAX;
}
float selected = 0.0f, ssum = 0.0f;
for (int rank = 0; rank < 10; ++rank) {
float best = v[0];
int ex = lane;
for (int i = 1; i < 16; ++i)
if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
for (ushort m = 16; m > 0; m >>= 1) {
const float ob = simd_shuffle_xor(best, m);
const int oe = simd_shuffle_xor(ex, m);
if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; }
}
if ((ex & 31) == lane) {
v[ex / 32] = -INFINITY;
ids[sg * 10 + rank] = ex;
ssum += best;
}
if (rank == lane) selected = best;
}"""
new_route_body = """#ifndef STRATA_MOCE_FC_METAL
#define STRATA_MOCE_FC_METAL
constant int kMoceProtectFc [[function_constant(24)]];
constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3;
constant uint kMoceScaleR1BitsFc [[function_constant(25)]];
constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type<float>(kMoceScaleR1BitsFc) : 1.0f;
#endif
static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan,
constant MoeArgs& a, device const uint* smap, ushort tid, ushort nth, ushort sg, ushort lane,
threadgroup int* s_grp, threadgroup int* s_ent) {
if (sg < a.n_tok) {
device const float* l = logits + sg * 512;
float v[16];
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY;
}
float mx = -INFINITY;
for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
float sum = 0.0f;
for (int i = 0; i < 16; ++i) {
v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f;
sum += v[i];
}
for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
const float rcp = 1.0f / sum;
for (int i = 0; i < 16; ++i) {
if (lane + 32 * i >= int(a.n_expert)) {
v[i] = -INFINITY;
} else {
v[i] *= rcp;
if (isnan(v[i])) v[i] = -FLT_MAX;
}
}
const float moce_scale = a.pad != 0u ? as_type<float>(a.pad) : 1.0f;
const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f;
const int r_prot = kMoceProtect;
const float r1_scale = kMoceScaleR1;
const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f;
const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale;
ushort scaled_mask = 0;
float selected = 0.0f, ssum = 0.0f;
for (int rank = 0; rank < 10; ++rank) {
if (r1_scale > 1.0f && r_prot > 1) {
if (rank == 1 && smap != nullptr) {
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
v[i] *= r1_scale;
scaled_mask |= ushort(1u << i);
}
}
} else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) {
for (int i = 0; i < 16; ++i) {
if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) {
v[i] *= r_step_scale;
}
}
}
} else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) {
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
v[i] *= moce_scale;
scaled_mask |= ushort(1u << i);
}
}
}
float best = v[0];
int ex = lane;
for (int i = 1; i < 16; ++i)
if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
for (ushort m = 16; m > 0; m >>= 1) {
const float ob = simd_shuffle_xor(best, m);
const int oe = simd_shuffle_xor(ex, m);
if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; }
}
float true_prob = best;
const int min_scaled_rank = (r1_scale > 1.0f && r_prot > 1) ? 1 : r_prot;
if (rank >= min_scaled_rank && (moce_scale > 1.0f || r1_scale > 1.0f) && best > 0.0f) {
const ushort win_lane = ushort(ex & 31);
const ushort win_i = ushort(ex >> 5);
const ushort win_mask = simd_broadcast(scaled_mask, win_lane);
if ((win_mask & ushort(1u << win_i)) != 0u) {
true_prob = best * ((r1_scale > 1.0f && rank < r_prot) ? rcp_r1 : rcp_moce);
}
}
if ((ex & 31) == lane) {
v[ex / 32] = -INFINITY;
ids[sg * 10 + rank] = ex;
ssum += true_prob;
}
if (rank == lane) selected = true_prob;
}"""
assert old_route_body in r_txt, "route_body not found in router.metal"
r_txt = r_txt.replace(old_route_body, new_route_body, 1)
r_txt = r_txt.replace(
"device atomic_uint* counter [[buffer(9)]],",
"device atomic_uint* counter [[buffer(9)]], device const uint* smap [[buffer(10)]],",
)
r_txt = r_txt.replace(
"route_body(logits, ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent);",
"route_body(logits, ids, w, plan, a, smap, tid, nth, sg, lane, s_grp, s_ent);",
)
r_path.write_text(r_txt, encoding="utf-8")
# 2. moe.metal: dynamic a.n_expert + race-free MoCE cache-prior reranking in moe_route
m_path = ROOT / "metal" / "kernels" / "moe.metal"
m_txt = m_path.read_text(encoding="utf-8")
old_moe_route = """kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]],
ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
threadgroup int s_grp[32], s_ent[32];
if (sg < a.n_tok) {
device const float* l = logits + sg * 512;
float v[16];
for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i];
float mx = -INFINITY;
for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
float sum = 0.0f;
for (int i = 0; i < 16; ++i) {
v[i] = exp(v[i] - mx);
sum += v[i];
}
for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
const float rcp = 1.0f / sum;
for (int i = 0; i < 16; ++i) {
v[i] *= rcp;
if (isnan(v[i])) v[i] = -FLT_MAX;
}
float selected = 0.0f, ssum = 0.0f;
for (int rank = 0; rank < 10; ++rank) {
float best = v[0];
int ex = lane;
for (int i = 1; i < 16; ++i)
if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
for (ushort m = 16; m > 0; m >>= 1) {
const float ob = simd_shuffle_xor(best, m);
const int oe = simd_shuffle_xor(ex, m);
if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; }
}
if ((ex & 31) == lane) {
v[ex / 32] = -INFINITY;
ids[sg * 10 + rank] = ex;
ssum += best;
}
if (rank == lane) selected = best;
}"""
new_moe_route = """#ifndef STRATA_MOCE_FC_METAL
#define STRATA_MOCE_FC_METAL
constant int kMoceProtectFc [[function_constant(24)]];
constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3;
constant uint kMoceScaleR1BitsFc [[function_constant(25)]];
constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type<float>(kMoceScaleR1BitsFc) : 1.0f;
#endif
kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]],
device const uint* smap [[buffer(5)]],
ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
threadgroup int s_grp[32], s_ent[32];
if (sg < a.n_tok) {
device const float* l = logits + sg * 512;
float v[16];
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY;
}
float mx = -INFINITY;
for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
float sum = 0.0f;
for (int i = 0; i < 16; ++i) {
v[i] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f;
sum += v[i];
}
for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
const float rcp = 1.0f / sum;
for (int i = 0; i < 16; ++i) {
if (lane + 32 * i >= int(a.n_expert)) {
v[i] = -INFINITY;
} else {
v[i] *= rcp;
if (isnan(v[i])) v[i] = -FLT_MAX;
}
}
const float moce_scale = a.pad != 0u ? as_type<float>(a.pad) : 1.0f;
const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f;
const int r_prot = kMoceProtect;
const float r1_scale = kMoceScaleR1;
const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f;
const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale;
ushort scaled_mask = 0;
float selected = 0.0f, ssum = 0.0f;
for (int rank = 0; rank < 10; ++rank) {
if (r1_scale > 1.0f && r_prot > 1) {
if (rank == 1 && smap != nullptr) {
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
v[i] *= r1_scale;
scaled_mask |= ushort(1u << i);
}
}
} else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) {
for (int i = 0; i < 16; ++i) {
if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) {
v[i] *= r_step_scale;
}
}
}
} else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) {
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
v[i] *= moce_scale;
scaled_mask |= ushort(1u << i);
}
}
}
float best = v[0];
int ex = lane;
for (int i = 1; i < 16; ++i)
if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
for (ushort m = 16; m > 0; m >>= 1) {
const float ob = simd_shuffle_xor(best, m);
const int oe = simd_shuffle_xor(ex, m);
if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; }
}
float true_prob = best;
const int min_scaled_rank = (r1_scale > 1.0f && r_prot > 1) ? 1 : r_prot;
if (rank >= min_scaled_rank && (moce_scale > 1.0f || r1_scale > 1.0f) && best > 0.0f) {
const ushort win_lane = ushort(ex & 31);
const ushort win_i = ushort(ex >> 5);
const ushort win_mask = simd_broadcast(scaled_mask, win_lane);
if ((win_mask & ushort(1u << win_i)) != 0u) {
true_prob = best * ((r1_scale > 1.0f && rank < r_prot) ? rcp_r1 : rcp_moce);
}
}
if ((ex & 31) == lane) {
v[ex / 32] = -INFINITY;
ids[sg * 10 + rank] = ex;
ssum += true_prob;
}
if (rank == lane) selected = true_prob;
}"""
assert old_moe_route in m_txt, "moe_route target not found in moe.metal"
m_path.write_text(m_txt.replace(old_moe_route, new_moe_route, 1), encoding="utf-8")
# 3. prompt.metal: route_topk unpacks n_tok = n_tok_packed & 0xFFFF, n_exp = (n_tok_packed >> 16) ? (n_tok_packed >> 16) : 512,
# and indexes row t at logits + ulong(t) * ulong(n_exp) to match gemm(..., T, n_exp_, N)
p_path = ROOT / "metal" / "kernels" / "prompt.metal"
p_txt = p_path.read_text(encoding="utf-8")
old_ptopk = """kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
device float* w [[buffer(2)]], constant uint& n_tok [[buffer(3)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]]) {
const uint t = tg.x * 8 + sg;
if (t >= n_tok) return;
device const float* l = logits + ulong(t) * 512;
float v[16];
for (int i = 0; i < 16; ++i) v[i] = l[lane + 32 * i];
float mx = -INFINITY;
for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
float sum = 0.0f;
for (int i = 0; i < 16; ++i) {
v[i] = exp(v[i] - mx);
sum += v[i];
}
for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
const float rcp = 1.0f / sum;
for (int i = 0; i < 16; ++i) {
v[i] *= rcp;
if (isnan(v[i])) v[i] = -FLT_MAX;
}"""
new_ptopk = """kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
device float* w [[buffer(2)]], constant uint& n_tok_packed [[buffer(3)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]]) {
const uint n_tok = n_tok_packed & 0xFFFFu;
const int n_exp = (n_tok_packed >> 16) ? int(n_tok_packed >> 16) : 512;
const uint t = tg.x * 8 + sg;
if (t >= n_tok) return;
device const float* l = logits + ulong(t) * ulong(n_exp);
float v[16];
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
v[i] = idx < n_exp ? l[idx] : -INFINITY;
}
float mx = -INFINITY;
for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
float sum = 0.0f;
for (int i = 0; i < 16; ++i) {
v[i] = (lane + 32 * i < n_exp) ? exp(v[i] - mx) : 0.0f;
sum += v[i];
}
for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
const float rcp = 1.0f / sum;
for (int i = 0; i < 16; ++i) {
if (lane + 32 * i >= n_exp) {
v[i] = -INFINITY;
} else {
v[i] *= rcp;
if (isnan(v[i])) v[i] = -FLT_MAX;
}
}"""
assert old_ptopk in p_txt, "route_topk target not found in prompt.metal"
p_path.write_text(p_txt.replace(old_ptopk, new_ptopk, 1), encoding="utf-8")
def patch_stream() -> None:
sh_path = ROOT / "metal" / "src" / "engine" / "stream.h"
sc_path = ROOT / "metal" / "src" / "engine" / "stream.cpp"
sh_txt = sh_path.read_text(encoding="utf-8")
sh_txt = sh_txt.replace(
"uint8_t* region(uint64_t bytes);\n void hold_region(bool hold);",
"uint8_t* region(uint64_t bytes);\n void hold_region(bool hold);\n uint64_t release_region();\n size_t region_stolen_slots() const { return region_slots_.size(); }",
)
sh_txt = sh_txt.replace(
"struct Slot { int16_t layer = -1; int16_t e = -1; int prev = -1, next = -1; bool pre = false; int left = 0; uint64_t pin = 0; std::atomic<int> state{0}; };",
"struct Slot { int16_t layer = -1; int16_t e = -1; int prev = -1, next = -1; bool pre = false, hot = false; int left = 0; uint64_t pin = 0; std::atomic<int> state{0}; };",
)
sh_txt = sh_txt.replace(
"int nio_ = 6, pre_k_ = 10, pre_k2_ = 0, split_ = 2;",
"int ne_ = 512, nio_ = 6, pre_k_ = 10, pre_k2_ = 0, split_ = 2;\n float moce_bonus_ = 0.0f;\n bool in_prompt_ = false;\n uint64_t sep_region_bytes_ = 0;",
)
sh_path.write_text(sh_txt, encoding="utf-8")
sc_txt = sc_path.read_text(encoding="utf-8")
# 1. ExpertStream::open: dynamic ne_ from blk.0.ffn_gate_exps.weight, and optional spot check when dense-only gguf is used
old_open = """ lay_.resize(NL);
for (int l = 0; l < NL; ++l) {
const std::string b = "blk." + std::to_string(l) + ".";
const gguf::Tensor* g = f.find(b + "ffn_gate_exps.weight");
const gguf::Tensor* u = f.find(b + "ffn_up_exps.weight");
const gguf::Tensor* d = f.find(b + "ffn_down_exps.weight");
if (!g || !u || !d) { err = "metal stream: layer " + std::to_string(l) + " has no routed experts"; return false; }
lay_[(size_t) l] = {g->offset, g->bytes / NE, u->offset, u->bytes / NE, d->offset, d->bytes / NE, -1};
}"""
new_open = """ const gguf::Tensor* g0 = f.find("blk.0.ffn_gate_exps.weight");
if (g0 && g0->shape.size() >= 3 && g0->shape[2] >= 1 && g0->shape[2] <= NE) {
ne_ = (int) g0->shape[2];
}
if (const char* mb = std::getenv("STRATA_METAL_MOCE_BONUS")) moce_bonus_ = (float) std::atof(mb);
lay_.resize(NL);
for (int l = 0; l < NL; ++l) {
const std::string b = "blk." + std::to_string(l) + ".";
const gguf::Tensor* g = f.find(b + "ffn_gate_exps.weight");
const gguf::Tensor* u = f.find(b + "ffn_up_exps.weight");
const gguf::Tensor* d = f.find(b + "ffn_down_exps.weight");
if (!g || !u || !d) { err = "metal stream: layer " + std::to_string(l) + " has no routed experts"; return false; }
lay_[(size_t) l] = {g->offset, g->bytes / (uint64_t) ne_, u->offset, u->bytes / (uint64_t) ne_, d->offset, d->bytes / (uint64_t) ne_, -1};
}"""
assert old_open in sc_txt, "open target not found in stream.cpp"
sc_txt = sc_txt.replace(old_open, new_open, 1)
old_strx_chk = """ ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) NE && src == (uint64_t) st.st_size;
for (int l = 0; ok && l < NL; ++l) {
const Lay& x = lay_[(size_t) l];
ok = t[(size_t) l * 4 + 1] == x.gs && t[(size_t) l * 4 + 2] == x.us && t[(size_t) l * 4 + 3] == x.ds;
off[(size_t) l] = t[(size_t) l * 4];
}
if (ok) { // spot check: layer 0's first and layer 47's last expert, every slice
for (const int l : {0, NL - 1}) {
const Lay& x = lay_[(size_t) l];
const int e = l == 0 ? 0 : NE - 1;"""
new_strx_chk = """ ok = ok && ver == 1 && nl == (uint32_t) NL && ne == (uint32_t) ne_ && src == (uint64_t) st.st_size;
for (int l = 0; ok && l < NL; ++l) {
const Lay& x = lay_[(size_t) l];
ok = t[(size_t) l * 4 + 1] == x.gs && t[(size_t) l * 4 + 2] == x.us && t[(size_t) l * 4 + 3] == x.ds;
off[(size_t) l] = t[(size_t) l * 4];
}
const Lay& last_lay = lay_[(size_t) (NL - 1)];
const bool has_gguf_exps = (uint64_t) st.st_size >= last_lay.d + (uint64_t) ne_ * last_lay.ds;
if (ok && has_gguf_exps) { // spot check: layer 0's first and layer 47's last expert, every slice
for (const int l : {0, NL - 1}) {
const Lay& x = lay_[(size_t) l];
const int e = l == 0 ? 0 : ne_ - 1;"""
assert old_strx_chk in sc_txt, "strx check not found in stream.cpp"
sc_txt = sc_txt.replace(old_strx_chk, new_strx_chk, 1)
# 2. Zero-copy file-backed load_dense with computed byte totals (no hard-coded sizes)
old_load_dense = """bool ExpertStream::load_dense(gguf::File& f, mtl::Context& cx, const std::vector<std::string>& skip, mtl::Buffer& out,
std::string& err) {
const double t0 = now_ms();
std::vector<gguf::Tensor*> ts;
uint64_t tot = 0;
for (auto& kv : f.tensors) {
gguf::Tensor& t = kv.second;
if (t.name.find("_exps") != std::string::npos) continue;
if (std::find(skip.begin(), skip.end(), t.name) != skip.end()) { t.offset = 0; continue; }
ts.push_back(&t);
tot += (t.bytes + 255) / 256 * 256;
}"""
new_load_dense = """bool ExpertStream::load_dense(gguf::File& f, mtl::Context& cx, const std::vector<std::string>& skip, mtl::Buffer& out,
std::string& err) {
const double t0 = now_ms();
uint64_t max_dense_end = 0;
uint64_t min_exp_off = UINT64_MAX;
size_t dense_cnt = 0;
for (const auto& kv : f.tensors) {
const gguf::Tensor& t = kv.second;
if (t.name.find("_exps") != std::string::npos) {
min_exp_off = std::min(min_exp_off, t.offset);
continue;
}
++dense_cnt;
max_dense_end = std::max(max_dense_end, (t.offset + t.bytes + PAGE - 1) / PAGE * PAGE);
}
if (f.map != nullptr && max_dense_end > 0 && max_dense_end <= f.map_bytes && max_dense_end <= min_exp_off &&
env_int("STRATA_METAL_MMAP_DENSE", 1) != 0) {
const size_t alloc_before = cx.allocated();
uint64_t locked_bytes = 0, skipped_hc_bytes = 0, unpinned_embd_bytes = 0;
for (const auto& kv : f.tensors) {
const gguf::Tensor& t = kv.second;
if (t.name.find("_exps") != std::string::npos) continue;
const uintptr_t p0 = ((uintptr_t) f.map + t.offset) / PAGE * PAGE;
const uintptr_t p1 = ((uintptr_t) f.map + t.offset + t.bytes + PAGE - 1) / PAGE * PAGE;
const bool is_skip = std::find(skip.begin(), skip.end(), t.name) != skip.end();
const bool is_embd = t.name == "token_embd.weight";
if (is_skip || is_embd) {
if (is_skip) skipped_hc_bytes += t.bytes;
if (is_embd) unpinned_embd_bytes += t.bytes;
madvise((void*) p0, p1 - p0, MADV_DONTNEED);
} else if (env_int("STRATA_METAL_MLOCK_DENSE", 1) != 0) {
madvise((void*) p0, p1 - p0, MADV_WILLNEED);
if (mlock((const void*) p0, p1 - p0) == 0) locked_bytes += (p1 - p0);
}
}
if (f.map_bytes > max_dense_end) {
madvise((uint8_t*) f.map + max_dense_end, f.map_bytes - (size_t) max_dense_end, MADV_DONTNEED);
}
out = cx.wrap(f.map, (size_t) max_dense_end, /*count_allocated=*/false);
(void) alloc_before;
std::fprintf(stderr, "strata metal stream: locked %.2f GB of active GPU dense tensors (unpinned %.2f GB CPU token_embd + %.2f GB skipped BF16 HC in %.2f GB file map)\\n",
locked_bytes / 1e9, unpinned_embd_bytes / 1e9, skipped_hc_bytes / 1e9, max_dense_end / 1e9);
std::fprintf(stderr, "strata metal stream: %zu dense tensors wrapped zero-copy from file (%.2f GB clean file-backed map, %.1f ms)\\n",
dense_cnt, max_dense_end / 1e9, now_ms() - t0);
return true;
}
std::vector<gguf::Tensor*> ts;
uint64_t tot = 0;
for (auto& kv : f.tensors) {
gguf::Tensor& t = kv.second;
if (t.name.find("_exps") != std::string::npos) continue;
if (std::find(skip.begin(), skip.end(), t.name) != skip.end()) { t.offset = 0; continue; }
ts.push_back(&t);
tot += (t.bytes + 255) / 256 * 256;
}"""
assert old_load_dense in sc_txt, "load_dense target not found in stream.cpp"
sc_txt = sc_txt.replace(old_load_dense, new_load_dense, 1)
# 3. STRP profile loading & rank fallback using ne_ (while keeping where_ / smap_ indexed by l * NE + e where NE = 512)
old_prof = """ if (std::fread(magic, 1, 4, f) == 4 && std::memcmp(magic, "STRP", 4) == 0 && std::fread(h, 4, 5, f) == 5 &&
h[1] == (uint32_t) NL && h[2] == (uint32_t) NE) {
for (uint32_t i = 0; i < h[4]; ++i) {
uint16_t p[2];
if (std::fread(p, 2, 2, f) != 2) break;
if (p[0] < NL && p[1] < NE && !seen[(size_t) p[0] * NE + p[1]]) {
seen[(size_t) p[0] * NE + p[1]] = 1;
rank.push_back({p[0], p[1]});
}
}
}
std::fclose(f);
} else {
std::fprintf(stderr, "strata metal stream: no expert profile at %s: the arena starts layer-interleaved\\n", profile.c_str());
}
const size_t from_profile = rank.size();
for (int e = 0; e < NE; ++e)
for (int l = 0; l < NL; ++l)
if (!seen[(size_t) l * NE + e]) rank.push_back({l, e});"""
new_prof = """ if (std::fread(magic, 1, 4, f) == 4 && std::memcmp(magic, "STRP", 4) == 0 && std::fread(h, 4, 5, f) == 5 &&
h[1] == (uint32_t) NL && h[2] == (uint32_t) ne_) {
for (uint32_t i = 0; i < h[4]; ++i) {
uint16_t p[2];
if (std::fread(p, 2, 2, f) != 2) break;
if (p[0] < NL && p[1] < (uint16_t) ne_ && !seen[(size_t) p[0] * NE + p[1]]) {
seen[(size_t) p[0] * NE + p[1]] = 1;
rank.push_back({p[0], p[1]});
}
}
}
std::fclose(f);
} else {
std::fprintf(stderr, "strata metal stream: no expert profile at %s: the arena starts layer-interleaved\\n", profile.c_str());
}
const size_t from_profile = rank.size();
for (int e = 0; e < ne_; ++e)
for (int l = 0; l < NL; ++l)
if (!seen[(size_t) l * NE + e]) rank.push_back({l, e});"""
assert old_prof in sc_txt, "prof target not found in stream.cpp"
sc_txt = sc_txt.replace(old_prof, new_prof, 1)
# 4. ahead(): top-3 unperturbed protection + smap residency consistency with GPU moe_route
old_ahead = """void ExpertStream::ahead(int ln, int T, const float* pred, int K) {
K = std::min(K, 32);
std::pair<float, int> pw[8 * 32];
int np = 0;
for (int t = 0; t < T; ++t) {
const float* lg = pred + (size_t) t * NE;
int top[32];
int nt = 0;
for (int e = 0; e < NE; ++e) {
if (nt < K) { top[nt++] = e; continue; }
int m = 0;
for (int j = 1; j < K; ++j) if (lg[top[j]] < lg[top[m]]) m = j;
if (lg[e] > lg[top[m]]) top[m] = e;
}
for (int j = 0; j < nt; ++j) {
int q = 0;
while (q < np && pw[q].second != top[j]) ++q;
if (q == np) pw[np++] = {lg[top[j]], top[j]};
else pw[q].first = std::max(pw[q].first, lg[top[j]]);
}
}"""
new_ahead = """void ExpertStream::ahead(int ln, int T, const float* pred, int K) {
K = std::min(K, 32);
static const bool unpert_all = env_int("STRATA_METAL_PREFETCH_UNPERT", 0) != 0;
const uint32_t* sm = smap_.as<uint32_t>() + (size_t) ln * NE;
std::pair<float, int> pw[8 * 32];
int np = 0;
for (int t = 0; t < T; ++t) {
const float* lg = pred + (size_t) t * NE;
auto eff_lg = [&](int e) -> float {
return lg[e] + (moce_bonus_ > 0.0f && sm[e] != 0xFFFFFFFFu ? moce_bonus_ : 0.0f);
};
const int k_unpert = unpert_all ? K : std::min(3, K);
int top[32];
int nt = 0;
for (int r = 0; r < k_unpert; ++r) {
int best_e = -1;
float best_v = -INFINITY;
for (int e = 0; e < ne_; ++e) {
bool seen_e = false;
for (int j = 0; j < nt; ++j) if (top[j] == e) { seen_e = true; break; }
if (!seen_e && lg[e] > best_v) { best_v = lg[e]; best_e = e; }
}
if (best_e >= 0) top[nt++] = best_e;
}
for (int e = 0; e < ne_; ++e) {
if (k_unpert >= K) break;
bool in_top = false;
for (int j = 0; j < k_unpert; ++j) if (top[j] == e) { in_top = true; break; }
if (in_top) continue;
if (nt < K) { top[nt++] = e; continue; }
int m = k_unpert;
for (int j = k_unpert + 1; j < K; ++j) if (eff_lg(top[j]) < eff_lg(top[m])) m = j;
if (eff_lg(e) > eff_lg(top[m])) top[m] = e;
}
for (int j = 0; j < nt; ++j) {
const float prio = (j < k_unpert) ? (lg[top[j]] + moce_bonus_ + 1e3f) : eff_lg(top[j]);
int q = 0;
while (q < np && pw[q].second != top[j]) ++q;
if (q == np) pw[np++] = {prio, top[j]};
else pw[q].first = std::max(pw[q].first, prio);
}
}"""
assert old_ahead in sc_txt, "ahead target not found in stream.cpp"
sc_txt = sc_txt.replace(old_ahead, new_ahead, 1)
# 5. victim() & multi-turn hot-slot protection + separate prompt scratch munmap in release_region()
old_victim = """int ExpertStream::victim(int ci) {
Cls& k = cls_[(size_t) ci];
for (int s = k.tail; s >= 0; s = slots_[(size_t) s].prev)
if (slots_[(size_t) s].state.load() == 0 && slots_[(size_t) s].pin != stamp_) {"""
new_victim = """int ExpertStream::victim(int ci) {
Cls& k = cls_[(size_t) ci];
const int n_passes = in_prompt_ ? 2 : 1;
for (int pass = 0; pass < n_passes; ++pass) {
for (int s = k.tail; s >= 0; s = slots_[(size_t) s].prev)
if (slots_[(size_t) s].state.load() == 0 && slots_[(size_t) s].pin != stamp_ && (!in_prompt_ || pass == 1 || !slots_[(size_t) s].hot)) {"""
assert old_victim in sc_txt, "victim target not found in stream.cpp"
sc_txt = sc_txt.replace(old_victim, new_victim, 1)
sc_txt = sc_txt.replace(
" x.layer = x.e = -1;\n return s;\n }\n return -1;",
" x.layer = x.e = -1;\n return s;\n }\n }\n return -1;",
1,
)
old_lru_order = """ // the LRU order = the rank order (best at the front)
for (size_t i = done; i-- > 0;) lru_front(where_[(size_t) (fill[i].first * NE + fill[i].second)]);"""
new_lru_order = """ // the LRU order = the rank order (best at the front); mark top (k.n - (ne_ + 32)) STRP slots hot
std::vector<int> cls_seen(cls_.size(), 0);
for (size_t i = 0; i < done; ++i) {
const int s = where_[(size_t) (fill[i].first * NE + fill[i].second)];
if (s >= 0) {
const int ci = cls_of(s);
const int hot_cap = std::max(0, cls_[(size_t) ci].n - (ne_ + 32));
slots_[(size_t) s].hot = (cls_seen[(size_t) ci] < hot_cap);
++cls_seen[(size_t) ci];
}
}
for (size_t i = done; i-- > 0;) lru_front(where_[(size_t) (fill[i].first * NE + fill[i].second)]);"""
assert old_lru_order in sc_txt, "lru_order target not found in stream.cpp"
sc_txt = sc_txt.replace(old_lru_order, new_lru_order, 1)
old_region_fn = """uint8_t* ExpertStream::region(uint64_t bytes) {"""
new_region_fn = """uint8_t* ExpertStream::region(uint64_t bytes) {
if (env_int("STRATA_METAL_SEPARATE_PROMPT_BUF", 1) != 0) {
const uint64_t sz = (bytes + PAGE - 1) / PAGE * PAGE;
void* m = mmap(nullptr, (size_t) sz, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
if (m == MAP_FAILED) return nullptr;
region_slots_.clear();
region_ = (uint8_t*) m;
region_bytes_ = 0; // fresh zero pages from MAP_ANON; no arena slots to memset
sep_region_bytes_ = sz;
std::fprintf(stderr, "strata metal stream: prompt scratch %.2f GB mapped in separate anonymous mmap\\n",
sz / 1e9);
return region_;
}"""
assert old_region_fn in sc_txt, "region_fn not found in stream.cpp"
sc_txt = sc_txt.replace(old_region_fn, new_region_fn, 1)
sc_txt = sc_txt.replace(
"region_bytes_ = 0; // fresh zero pages; nothing in the arena to empty",
"region_bytes_ = 0; sep_region_bytes_ = sz; // fresh zero pages; unmapped in release_region()",
)
old_hold_fn = """void ExpertStream::hold_region(bool hold) {
std::unique_lock<std::mutex> g(mu_);
if (hold) {
// nothing may be reading into those slots: the queued read-aheads are dropped, the running ones waited for"""
new_hold_fn = """void ExpertStream::hold_region(bool hold) {
std::unique_lock<std::mutex> g(mu_);
in_prompt_ = hold;
if (hold) {
// Refresh hot flags from current LRU order (protecting top k.n - (ne_ + 32) slots in every prompt turn)
for (size_t ci = 0; ci < cls_.size(); ++ci) {
const int hot_cap = std::max(0, cls_[ci].n - (ne_ + 32));
int rank_idx = 0;
for (int s = cls_[ci].head; s >= 0; s = slots_[(size_t) s].next) {
slots_[(size_t) s].hot = (rank_idx < hot_cap);
++rank_idx;
}
}
// nothing may be reading into those slots: the queued read-aheads are dropped, the running ones waited for"""
assert old_hold_fn in sc_txt, "hold_fn start not found in stream.cpp"
sc_txt = sc_txt.replace(old_hold_fn, new_hold_fn, 1)
old_hold_end = """ if (region_) std::memset(region_, 0, (size_t) region_bytes_);
} else {
for (int s : region_slots_) {
slots_[(size_t) s].state = 0;
if (!retired_[(size_t) s]) lru_back(s); // a retired slot's memory serves the prompt only
}
}
}"""
new_hold_end = """ if (region_ && region_bytes_ > 0) std::memset(region_, 0, (size_t) region_bytes_);
} else {
for (int s : region_slots_) {
slots_[(size_t) s].state = 0;
if (!retired_[(size_t) s]) lru_back(s); // a retired slot's memory serves the prompt only
}
}
}
uint64_t ExpertStream::release_region() {
std::unique_lock<std::mutex> g(mu_);
if (sep_region_bytes_ > 0 && region_ != nullptr) {
const uint64_t freed = sep_region_bytes_;
munmap(region_, (size_t) sep_region_bytes_);
region_ = nullptr;
sep_region_bytes_ = 0;
return freed;
}
return 0;
}"""
assert old_hold_end in sc_txt, "hold_fn end not found in stream.cpp"
sc_txt = sc_txt.replace(old_hold_end, new_hold_end, 1)
# Also during prompt chunks in at_layer, move streamed prompt slots to lru_back so victim() cycles the unpinned tail slots!
sc_txt = sc_txt.replace(
"submit(s, ln, e, false);\n slots_[(size_t) s].pin = stamp_;\n ++p_streamed_;",
"submit(s, ln, e, false);\n slots_[(size_t) s].pin = stamp_;\n lru_back(s);\n ++p_streamed_;",
)
sc_txt = sc_txt.replace("if (e < 0 || e >= NE || used[e]) continue;", "if (e < 0 || e >= ne_ || used[e]) continue;")
sc_txt = sc_txt.replace("for (int e = 0; e < NE; ++e) {\n if (where_[(size_t) (ln * NE + e)] >= 0) continue;",
"for (int e = 0; e < ne_; ++e) {\n if (where_[(size_t) (ln * NE + e)] >= 0) continue;")
sc_txt = sc_txt.replace("tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < NE);",
"tight = tight || (cls_[i].n > 0 && cls_[i].n - stolen[i] < ne_);")
# 6. Eliminate CPU/GPU data race on smap_ during background read-ahead:
# Background read-ahead threads (x.pre == true) never write to smap_ while the GPU is running moe_route.
# Instead, smap_ is written only when !x.pre (during make_arena or urgent/promoted waits while the GPU is paused)
# or inside at_layer() when a prefetched slot is claimed by the layer!
old_io_smap = """ if (--x.left == 0) {
smap_.as<uint32_t>()[j.layer * NE + j.e] = code_[(size_t) j.slot];
x.state = 0;"""
new_io_smap = """ if (--x.left == 0) {
if (!x.pre) smap_.as<uint32_t>()[j.layer * NE + j.e] = code_[(size_t) j.slot];
x.state = 0;"""
assert old_io_smap in sc_txt, "old_io_smap not found in stream.cpp"
sc_txt = sc_txt.replace(old_io_smap, new_io_smap, 1)
old_claim = """ if (!(prompt && x.pre)) lru_front(s); // a prompt's streamed experts stay at the LRU tail
x.pre = false;"""
new_claim = """ if (!(prompt && x.pre)) lru_front(s); // a prompt's streamed experts stay at the LRU tail
x.pre = false;
if (st == 0) smap_.as<uint32_t>()[key] = code_[(size_t) s];"""
assert old_claim in sc_txt, "old_claim not found in stream.cpp"
sc_txt = sc_txt.replace(old_claim, new_claim, 1)
old_drop_block = """ for (int s : drop) cancel_slot(s);
}"""
new_drop_block = """ for (int s : drop) cancel_slot(s);
}
if (moce_bonus_ > 0.0f) {
std::unique_lock<std::mutex> g2(mu_, std::adopt_lock);
done_cv_.wait(g2, [&] {
for (int e = 0; e < ne_; ++e) {
if (used[e]) continue;
const int s = where_[(size_t) (l * NE + e)];
if (s >= 0 && slots_[(size_t) s].pre && slots_[(size_t) s].state.load() != 0) return false;
}
return true;
});
g2.release();
// Evict any unchosen speculative read-ahead of this layer that finished so where_, smap_, and LRU order stay 100% deterministic
for (int e = 0; e < ne_; ++e) {
if (used[e]) continue;
const int s = where_[(size_t) (l * NE + e)];
if (s >= 0 && slots_[(size_t) s].pre && slots_[(size_t) s].state.load() == 0) {
Slot& x = slots_[(size_t) s];
where_[(size_t) (l * NE + e)] = -1;
smap_.as<uint32_t>()[l * NE + e] = 0xFFFFFFFFu;
x.layer = x.e = -1;
x.pre = false;
lru_back(s);
}
}
}"""
assert old_drop_block in sc_txt, "old_drop_block not found in stream.cpp"
sc_txt = sc_txt.replace(old_drop_block, new_drop_block, 1)
sc_path.write_text(sc_txt, encoding="utf-8")
def patch_window() -> None:
w_path = ROOT / "metal" / "src" / "engine" / "window.cpp"
pc_path = ROOT / "metal" / "src" / "engine" / "prompt_chain.inc"
w_txt = w_path.read_text(encoding="utf-8")
# 1. Add n_exp_, moce_pad_, cached pg shader lib IDs, and release_prompt_gemm declaration to Window
old_seg_t = " int seg_T_ = 0;"
new_seg_t = " int seg_T_ = 0, n_exp_ = NEXP;\n uint32_t moce_pad_ = 0;\n int pg_lib_ = -1, pg_lgdn_ = -1, pg_lqsa_ = -1;"
assert old_seg_t in w_txt, "seg_T_ not found in window.cpp"
w_txt = w_txt.replace(old_seg_t, new_seg_t, 1)
w_txt = w_txt.replace(
"void prompt_region(bool hold) { if (stream_ && st_ && st_->ready()) st_->hold_region(hold); }",
"void prompt_region(bool hold);\n uint64_t prompt_gemm_need_bytes() const;",
)
# 2. Allow any 1 <= g.n_expert <= NEXP in Window::init, and purge unused device allocations upfront
old_init_chk = """ if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert != NEXP || g.n_layers != 48) {
err = "metal window: the geometry differs from Qwen3.8-Flash-Next's";
return false;
}"""
new_init_chk = """ if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_ff != NFF || g.n_expert < 1 || g.n_expert > NEXP || g.n_layers != 48) {
err = "metal window: the geometry differs from Qwen3.8-Flash-Next's";
return false;
}
n_exp_ = (int) g.n_expert;
if (const char* mb = std::getenv("STRATA_METAL_MOCE_BONUS")) {
const float bonus = (float) std::atof(mb);
if (bonus > 0.0f) {
const float sc = std::exp(bonus);
std::memcpy(&moce_pad_, &sc, 4);
}
}
if (!(std::getenv("STRATA_METAL_KEEP_HOST_WEIGHTS") && std::string(std::getenv("STRATA_METAL_KEEP_HOST_WEIGHTS")) == "1")) {
strata::gpu::purge_device_allocations(1u << 20);
}"""
assert old_init_chk in w_txt, "init check not found in window.cpp"
w_txt = w_txt.replace(old_init_chk, new_init_chk, 1)
# 3. Also detect n_exp_ from blk.0.ffn_gate_exps.weight in f_ right after f_.open
old_fopen = """ if (!f_.open(shard1, err)) return false;"""
new_fopen = """ if (!f_.open(shard1, err)) return false;
if (const gguf::Tensor* ge0 = f_.find("blk.0.ffn_gate_exps.weight")) {
if (ge0->shape.size() >= 3 && ge0->shape[2] >= 1 && ge0->shape[2] <= NEXP) n_exp_ = (int) ge0->shape[2];
}"""
assert old_fopen in w_txt, "f_.open not found in window.cpp"
w_txt = w_txt.replace(old_fopen, new_fopen, 1)
# 3b. Enable PLE when ss.ple.table is open (even without legacy CUDA host pointers)
w_txt = w_txt.replace(
'ple_on_ = ss.ple.ready() && f_.find("blk.1.ple_key.weight") != nullptr;',
'ple_on_ = (ss.ple.ready() || (ss.ple.table != nullptr && ss.ple.table->is_open())) && f_.find("blk.1.ple_key.weight") != nullptr;',
)
# 3c. Add reset_stream_stats() helper to Window and window_reset_stats(void* w)
w_txt = w_txt.replace(
"void stream_report(FILE* f) const {",
"void reset_stream_stats() { if (st_) st_->reset_stats(); for (int t = 0; t <= MAXT; ++t) { win_n_[t] = 0; win_ms_[t] = win_gpu_[t] = 0; } }\n void stream_report(FILE* f) const {",
)
w_txt = w_txt.replace(
"void window_reset() {",
"void window_reset_stats(void* w) { if (w) ((Window*) w)->reset_stream_stats(); }\nvoid window_reset() {",
)
# 4. In Window::layer: use n_exp_, pass moce_pad_ and smap to router_route
old_layer_route = """ a.n_tok = ntok;
a.topk = TOPK;
a.n_expert = NEXP;
const bool dq2 = td.type == ref::Q2_0;
const int gl = 32, gr = 1, gs = 4, ds_ = 4;
// the down tile: Q2_0 8,1; else 4,2 - but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there:
// 24-36 GB/s, mac-kern 513fc615)
const bool m4 = !m5_layout();
const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4,
dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2;
a.gx_gu = NFF / (gs * (32 / gl) * gr);
a.gx_dn = N / (ds_ * (32 / dl) * dr);
mtl::Dispatch r;
if (rf) {
// the router GEMV + moe_route (bitwise both), beside the shared expert's GEMVs
r.p = P(lib_, "router_route_" + std::to_string(trt.type), mtl::Constants().i(0, ntok), "router_route_" + std::to_string(trt.type) + "#" + std::to_string(ntok));
GemvArgs ga{(uint32_t) N, 512, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, 512};
r.args = {mtl::Arg::B(W_, trt.offset), mtl::Arg::B(XF_.b, 0), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4),
mtl::Arg::B(logits_), mtl::Arg::V(ga), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_),
mtl::Arg::V(a), mtl::Arg::B(rcounter_)};
r.grid[0] = 512 / 16;
r.tg[0] = 512;
nb_ = true;
} else {
r.p = P(lib_, "moe_route", {}, "moe_route");
r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)};
r.tg[0] = 512;"""
new_layer_route = """ a.n_tok = ntok;
a.topk = TOPK;
a.n_expert = (uint32_t) n_exp_;
a.pad = stream_ ? moce_pad_ : 0u;
const bool dq2 = td.type == ref::Q2_0;
const int gl = 32, gr = 1, gs = 4, ds_ = 4;
// the down tile: Q2_0 8,1; else 4,2 - but on the M4 (no M5 layout) Q5_1 4,1 and Q8_0 8,1 (4,2 spills there:
// 24-36 GB/s, mac-kern 513fc615)
const bool m4 = !m5_layout();
const int dl = dq2 || (m4 && td.type == ref::Q8_0) ? 8 : 4,
dr = dq2 || (m4 && (td.type == ref::Q8_0 || td.type == ref::Q5_1)) ? 1 : 2;
a.gx_gu = NFF / (gs * (32 / gl) * gr);
a.gx_dn = N / (ds_ * (32 / dl) * dr);
mtl::Dispatch r;
const int mp_fc = std::getenv("STRATA_METAL_MOCE_PROTECT") ? std::max(0, std::min(10, std::atoi(std::getenv("STRATA_METAL_MOCE_PROTECT")))) : 3;
uint32_t mr1_bits = 0x3F800000u; // 1.0f
if (const char* mr1_env = std::getenv("STRATA_METAL_MOCE_BONUS_R1")) {
const float mr1_b = (float) std::atof(mr1_env);
if (mr1_b > 0.0f) {
const float mr1_s = std::exp(mr1_b);
std::memcpy(&mr1_bits, &mr1_s, 4);
}
}
mtl::Constants r_fc;
std::string r_fc_tag;
if (mp_fc != 3) { r_fc.i(24, mp_fc); r_fc_tag += "#mp" + std::to_string(mp_fc); }
if (mr1_bits != 0x3F800000u) { r_fc.u(25, mr1_bits); r_fc_tag += "#mr1_" + std::to_string(mr1_bits); }
if (rf) {
// the router GEMV + moe_route (bitwise both), beside the shared expert's GEMVs
mtl::Constants rfc2 = r_fc;
rfc2.i(0, ntok);
r.p = P(lib_, "router_route_" + std::to_string(trt.type), rfc2, "router_route_" + std::to_string(trt.type) + "#" + std::to_string(ntok) + r_fc_tag);
GemvArgs ga{(uint32_t) N, (uint32_t) n_exp_, (uint32_t) ntok, (uint32_t) ref::row_bytes(trt.type, N), (uint32_t) N, 512};
r.args = {mtl::Arg::B(W_, trt.offset), mtl::Arg::B(XF_.b, 0), mtl::Arg::B(XF_.b, XF_.plane_bytes() * 4),
mtl::Arg::B(logits_), mtl::Arg::V(ga), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_),
mtl::Arg::V(a), mtl::Arg::B(rcounter_),
stream_ ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(rcounter_)};
r.grid[0] = (uint32_t) ((n_exp_ + 15) / 16);
r.tg[0] = 512;
nb_ = true;
} else {
r.p = P(lib_, "moe_route", r_fc, "moe_route" + r_fc_tag);
r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a),
stream_ ? mtl::Arg::B(st_->smap(), (size_t) l * 512 * 4) : mtl::Arg::B(rcounter_)};
r.tg[0] = 512;"""
assert old_layer_route in w_txt, "layer_route not found in window.cpp"
w_txt = w_txt.replace(old_layer_route, new_layer_route, 1)
# 5. In stream_arena: min_slots = pg_on ? (n_exp_ + 64) : 64, account for transient PromptGemm scratch in plan log,
# and only pre-init PromptGemm inside stream_arena when NOT using a separate on-demand prompt buffer
old_arena_tail = """ { // the planned total: arena + the Metal buffers (or the process's own count when it is more) + the n-gram table's
// reserve, against the cap (arena_budget.h: footprint_cap)
const FootprintPlan fp = footprint_plan(am, budget);
std::fprintf(stderr, "strata metal stream: planned resident footprint %.1f GB = arena %.1f + %.1f held before it + %.1f "
"reserved for the n-gram table; cap %.1f GB = the lesser of the GPU working set %.1f GB and the RAM "
"%.1f GB less %.1f GB left to macOS and its file cache%s\\n",
fp.total / gb, fp.arena / gb, fp.base / gb, fp.ple / gb, fp.cap / gb, cx_->recommended_working_set() / gb,
phys / gb, (am.phys - fp.cap) / gb,
am.arena_gb >= 0 && fp.over() ? " - STRATA_METAL_ARENA_GB is above the cap: kept as set (your choice), but "
"a Mac that runs short of memory can freeze (watchdog panics seen at 90+ GB "
"resident of 128)" : "");
}
const char* epg = std::getenv("STRATA_METAL_PROMPT_GEMM");
const bool pg_on = !(epg && std::string(epg) == "0");
bool minimal = false; // the arena's least layout (the min slots) is not shrunk further: it goes ahead, filled in steps
for (int tries = 0;; ++tries) {
const double planned = budget;
std::function<uint64_t()> recheck = [&]() -> uint64_t { return (uint64_t) arena_shrink(am, planned, avail_now()); };
if (st_->make_arena(*cx_, (uint64_t) budget, prof, pg_on ? 576 : 64, err, minimal ? std::function<uint64_t()>() : recheck,
[&](uint64_t step) { return arena_grow(am, avail_now(), (double) step); }))
break;
if (st_->shrink_to() == 0 || tries == 4) return false;
if ((double) st_->shrink_to() + 64.0 * (1 << 20) >= budget) minimal = true; // no smaller layout than this one
std::fprintf(stderr, "strata metal stream: less memory is free than when the arena was planned: %.1f -> %.1f GB\\n",
budget / gb, st_->shrink_to() / gb);
budget = (double) st_->shrink_to();
}
if (pg_on && !prompt_gemm_init(err)) return false;"""
new_arena_tail = """ const char* epg = std::getenv("STRATA_METAL_PROMPT_GEMM");
const bool pg_on = !(epg && std::string(epg) == "0");
const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF");
const bool sep_pg = !(esep && std::string(esep) == "0");
const double pg_scratch_need = (pg_on && sep_pg) ? (double) prompt_gemm_need_bytes() : 0.0;
if (am.arena_gb < 0 && pg_scratch_need > 0.0) {
budget = std::max(1.0 * gb, budget - pg_scratch_need);
}
{ // the planned total: arena + the Metal buffers (or the process's own count when it is more) + the n-gram table's
// reserve, against the cap (arena_budget.h: footprint_cap)
const FootprintPlan fp = footprint_plan(am, budget);
std::fprintf(stderr, "strata metal stream: planned resident footprint %.1f GB decode (%.1f GB virtual peak during prefill with %.2f GB transient prompt scratch) "
"= arena %.1f + %.1f held before it + %.1f reserved for the n-gram table; cap %.1f GB = the lesser of the GPU working set %.1f GB "
"and the RAM %.1f GB less %.1f GB left to macOS and its file cache%s\\n",
fp.total / gb, (fp.total + pg_scratch_need) / gb, pg_scratch_need / gb,
fp.arena / gb, fp.base / gb, fp.ple / gb, fp.cap / gb, cx_->recommended_working_set() / gb,
phys / gb, (am.phys - fp.cap) / gb,
am.arena_gb >= 0 && (fp.total + pg_scratch_need > fp.cap) ? " - peak prefill above cap" : "");
}
bool minimal = false; // the arena's least layout (the min slots) is not shrunk further: it goes ahead, filled in steps
for (int tries = 0;; ++tries) {
const double planned = budget;
std::function<uint64_t()> recheck = [&]() -> uint64_t { return (uint64_t) arena_shrink(am, planned, avail_now()); };
if (st_->make_arena(*cx_, (uint64_t) budget, prof, pg_on ? (n_exp_ + 64) : 64, err, minimal ? std::function<uint64_t()>() : recheck,
[&](uint64_t step) { return arena_grow(am, avail_now(), (double) step); }))
break;
if (st_->shrink_to() == 0 || tries == 4) return false;
if ((double) st_->shrink_to() + 64.0 * (1 << 20) >= budget) minimal = true; // no smaller layout than this one
std::fprintf(stderr, "strata metal stream: less memory is free than when the arena was planned: %.1f -> %.1f GB\\n",
budget / gb, st_->shrink_to() / gb);
budget = (double) st_->shrink_to();
}
if (pg_on && !sep_pg && !prompt_gemm_init(err)) return false;"""
assert old_arena_tail in w_txt, "arena_tail not found in window.cpp"
w_txt = w_txt.replace(old_arena_tail, new_arena_tail, 1)
# 6. In PromptGemm: cache compiled libraries across prompts, track all wrapped buffers in std::vector<mtl::Buffer*> bufs,
# and unmap + cx_->free them in Window::prompt_region(false) with measured proc_mem() footprint logging!
old_pg_struct = """struct Window::PromptGemm {
int MT = 0, QS = 256, TM = 16, maxrows = 0, nblk_max = 0;
int lib = -1, lgdn = -1, lqsa = -1;
bool k4 = false, gu_v3 = false, dn_v3 = false, flash = false;
mtl::Buffer R, bo, inj[2], xnh, V, loh, mixed, mixh, qkv, z, al, be, ygh, gqn, gkn, gvc, ggb, gkq, go, yg, qfull, kcur, vcur,
iqr, ikr, qcur, qcurh, qidx, scores, qids, part, XQ, xqh, logits, ids, wts, bcnt, boff, eoff, tiles, meta, stok, sw, pos,
XS, HS, YE, yr, gsh, ush, hh, ys, sgv, emb, embh, pkey, pval, keyn, qn, gated, normW;
mtl::Buffer WD, WU, Wqkv, Wz, Wal, Wbe, Wout, Wq, Wk, Wv, Wiq, Wik, Wr, Wsg, Wshg, Wshu, Wshd, Wpk, Wpv;
};"""
new_pg_struct = """struct Window::PromptGemm {
int MT = 0, QS = 256, TM = 16, maxrows = 0, nblk_max = 0;
int lib = -1, lgdn = -1, lqsa = -1;
bool k4 = false, gu_v3 = false, dn_v3 = false, flash = false;
mtl::Buffer R, bo, inj[2], xnh, V, loh, mixed, mixh, qkv, z, al, be, ygh, gqn, gkn, gvc, ggb, gkq, go, yg, qfull, kcur, vcur,
iqr, ikr, qcur, qcurh, qidx, scores, qids, part, XQ, xqh, logits, ids, wts, bcnt, boff, eoff, tiles, meta, stok, sw, pos,
XS, HS, YE, yr, gsh, ush, hh, ys, sgv, emb, embh, pkey, pval, keyn, qn, gated, normW;
mtl::Buffer WD, WU, Wqkv, Wz, Wal, Wbe, Wout, Wq, Wk, Wv, Wiq, Wik, Wr, Wsg, Wshg, Wshu, Wshd, Wpk, Wpv;
std::vector<mtl::Buffer> all_bufs;
};"""
assert old_pg_struct in w_txt, "pg_struct not found in window.cpp"
w_txt = w_txt.replace(old_pg_struct, new_pg_struct, 1)
w_txt = w_txt.replace("if (stream_ && !st_->ready()) return stream_arena(err);", "if (stream_ && !st_->ready() && !stream_arena(err)) return false;")
old_pg_init_libs = """ p.lib = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", "gemv.metal",
"prompt.metal"});
p.lgdn = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\\n");
p.lqsa = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\\n");"""
new_pg_init_libs = """ if (pg_lib_ < 0) {
pg_lib_ = cx_->compile({"quant.metal", "units.metal", "mma.metal", "dense_mm.metal", "moe_mm.metal", "gr.metal", "gemv.metal",
"prompt.metal"});
pg_lgdn_ = cx_->compile({"gdn.metal", "gdn_prompt.metal"}, "#define NO_GGML_COMMON 1\\n");
pg_lqsa_ = cx_->compile({"qsa.metal", "qsa_prompt.metal"}, "#define NO_GGML_COMMON 1\\n");
}
p.lib = pg_lib_;
p.lgdn = pg_lgdn_;
p.lqsa = pg_lqsa_;"""
assert old_pg_init_libs in w_txt, "pg_init_libs not found in window.cpp"
w_txt = w_txt.replace(old_pg_init_libs, new_pg_init_libs, 1)
old_ab = """ if (region == nullptr) return cx_->alloc(bytes);
mtl::Buffer r = cx_->wrap(region, b);
region += b;
return r;"""
new_ab = """ if (region == nullptr) {
mtl::Buffer r = cx_->alloc(bytes);
p.all_bufs.push_back(r);
return r;
}
mtl::Buffer r = cx_->wrap(region, b, /*count_allocated=*/false);
p.all_bufs.push_back(r);
region += b;
return r;"""
assert old_ab in w_txt, "old_ab not found in window.cpp"
w_txt = w_txt.replace(old_ab, new_ab, 1)
old_pg_end = """ pg_ = pg.release();
std::fprintf(stderr, "strata metal: prompts on the matrix units in chunks of %d (experts %s, QSA %s, matmul2d layout %s; "
"%.2f GB of buffers%s)\\n", C, p.gu_v3 ? "v3" : "v2", p.flash ? "flash" : "v1", p.k4 ? "K4" : "other",
need / 1e9, stream_ ? ", in the expert arena while a prompt runs" : "");
(void) err;
return true;
}"""
new_pg_end = """ if (need != prompt_gemm_need_bytes()) { err = "metal window: prompt_gemm_need_bytes() != need"; return false; }
pg_ = pg.release();
const char* esep_log = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF");
const bool sep_pg_log = !(esep_log && std::string(esep_log) == "0");
std::fprintf(stderr, "strata metal: prompts on the matrix units in chunks of %d (experts %s, QSA %s, matmul2d layout %s; "
"%.2f GB of buffers%s)\\n", C, p.gu_v3 ? "v3" : "v2", p.flash ? "flash" : "v1", p.k4 ? "K4" : "other",
need / 1e9, stream_ ? (sep_pg_log ? ", in a separate transient mmap while a prompt runs" : ", in the expert arena while a prompt runs") : "");
(void) err;
return true;
}
uint64_t Window::prompt_gemm_need_bytes() const {
const char* ec = std::getenv("STRATA_METAL_PROMPT_CHUNK");
uint64_t phys = 0;
size_t plen = sizeof phys;
sysctlbyname("hw.memsize", &phys, &plen, nullptr, 0);
const int stream_chunk = phys >= (96ull << 30) ? 4096 : 1024;
const int C = std::max(64, ec ? std::atoi(ec) : (stream_ ? stream_chunk : 2048)) / 64 * 64;
const int QS = std::min(C, 256);
const bool k4 = ::coop_k4(*cx_);
const char* eq = std::getenv("STRATA_METAL_PROMPT_QSA");
const bool flash = eq ? std::string(eq) == "flash" : k4;
PromptGemm p_tm;
const size_t MT = (size_t) C;
const size_t mr = (size_t) (C * TOPK + NEXP * p_tm.TM);
const size_t nblk_max = (size_t) ((C * TOPK + 1023) / 1024);
const size_t MB = (size_t) (max_cells_ / 4 + 2);
uint64_t need = 0;
auto AB = [&](size_t bytes) { need += (bytes + 16383) / 16384 * 16384; };
auto F32 = [&](size_t n) { AB(4 * n); };
auto F16 = [&](size_t n) { AB(2 * n); };
F32(MT * HC * N); F32(MT * N); F32(MT * HC); F32(MT * HC); F16(MT * HC * N);
F32(MT * (LR + HC)); F16(MT * LR); F32(MT * N); F16(MT * N);
F32(MT * 10240); F32(MT * 6144); F32(MT * 48); F32(MT * 48); F32(MT * 6144);
F16(MT * 6144); F32(MT * 16 * 128); F32(MT * 16 * 128); F32(MT * 6144);
F32(MT * 96); F32(MT * 16); F32(MT * 6144);
F32(MT * 12288); F32(MT * 512); F32(MT * 512); F32(MT * 512); F32(MT * 128);
F32(MT * 24 * 256); F16(MT * 24 * 256); F32(MT * 4 * 128); F32(MT * MB);
AB(4 * MT * 2051);
if (flash) F32(16); else F32((size_t) QS * 33 * 24 * 258);
F32(4 * MT * 6144 + 3 * (size_t) QS * 6144 + MT * 384); F16(MT * 6144);
F32(MT * NEXP); AB(4 * MT * TOPK); F32(MT * TOPK);
AB(4 * nblk_max * NEXP); AB(4 * nblk_max * NEXP);
AB(4 * NEXP); AB(8 * (mr / p_tm.TM + 1)); AB(64);
AB(4 * mr); F32(mr); AB(4 * MT * TOPK); F16(mr * N); F16(mr * NFF);
F32(mr * N); F32(MT * N); F32(MT * NFF); F32(MT * NFF); F16(MT * NFF);
F32(MT * N); F32(MT);
F16((size_t) (LR + HC) * HC * N); F16((size_t) HC * N * LR); F16((size_t) 10240 * N);
F16((size_t) 6144 * N); F16((size_t) 48 * N); F16((size_t) 48 * N); F16((size_t) N * 6144);
F16((size_t) 12288 * N); F16((size_t) 512 * N); F16((size_t) 512 * N); F16((size_t) 512 * N);
F16((size_t) 128 * N); F16((size_t) NEXP * N); F16((size_t) N); F16((size_t) NFF * N);
F16((size_t) NFF * N); F16((size_t) N * NFF);
if (ple_on_) {
F32(MT * N); F16(MT * N); F32(MT * 10240); F32(MT * N); F32(MT * 10240);
F32(MT * 10240); F32(MT * 10240); F32(MT * 10240);
F16((size_t) 10240 * N); F16((size_t) N * N);
}
return need;
}
void Window::prompt_region(bool hold) {
if (!stream_ || !st_ || !st_->ready()) return;
if (hold) {
st_->hold_region(true);
} else {
st_->hold_region(false);
const char* esep = std::getenv("STRATA_METAL_SEPARATE_PROMPT_BUF");
const bool sep_pg = !(esep && std::string(esep) == "0");
if (sep_pg && pg_ != nullptr) {
const ProcMem pm_before = proc_mem();
mtp_src_ = nullptr;
for (mtl::Buffer& b : pg_->all_bufs) cx_->free(b);
delete pg_;
pg_ = nullptr;
const uint64_t unmapped = st_->release_region();
strata::gpu::purge_device_allocations(1u << 20);
const ProcMem pm_after = proc_mem();
if (unmapped > 0) {
std::fprintf(stderr, "strata metal stream: prompt scratch %.2f GB unmapped after prefill (phys_footprint %.2f -> %.2f GB, resident RSS %.2f -> %.2f GB)\\n",
unmapped / 1e9, pm_before.footprint / 1e9, pm_after.footprint / 1e9, pm_before.rss / 1e9, pm_after.rss / 1e9);
}
}
}
}"""
assert old_pg_end in w_txt, "old_pg_end not found in window.cpp"
w_txt = w_txt.replace(old_pg_end, new_pg_end, 1)
old_pg_moe = """ gemm(p.mixh, p.Wr, p.logits, T, NEXP, N);
const uint32_t nt = (uint32_t) T;
D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1, 256);
const int nent = T * TOPK, nblk = (nent + 1023) / 1024;
PgPlan pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) p.TM, (uint32_t) nblk};
const uint32_t nclr = (uint32_t) p.maxrows;
D(PP(p.lib, "mm_clear"), {mtl::Arg::B(p.stok), mtl::Arg::V(nclr)}, (p.maxrows + 255) / 256, 1, 256);
D(PP(p.lib, "mm_hist"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024);
D(PP(p.lib, "mm_scan"), {mtl::Arg::B(p.bcnt), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.tiles),
mtl::Arg::B(p.meta), mtl::Arg::V(pa)}, 1, 1, NEXP);"""
new_pg_moe = """ gemm(p.mixh, p.Wr, p.logits, T, n_exp_, N);
const uint32_t nt = (uint32_t) T | ((uint32_t) n_exp_ << 16);
D(PP(p.lib, "route_topk"), {mtl::Arg::B(p.logits), mtl::Arg::B(p.ids), mtl::Arg::B(p.wts), mtl::Arg::V(nt)}, (T + 7) / 8, 1, 256);
const int nent = T * TOPK, nblk = (nent + 1023) / 1024;
PgPlan pa{(uint32_t) nent, TOPK, (uint32_t) n_exp_, (uint32_t) p.TM, (uint32_t) nblk};
const uint32_t nclr = (uint32_t) p.maxrows;
D(PP(p.lib, "mm_clear"), {mtl::Arg::B(p.stok), mtl::Arg::V(nclr)}, (p.maxrows + 255) / 256, 1, 256);
D(PP(p.lib, "mm_hist"), {mtl::Arg::B(p.ids), mtl::Arg::B(p.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024);
D(PP(p.lib, "mm_scan"), {mtl::Arg::B(p.bcnt), mtl::Arg::B(p.boff), mtl::Arg::B(p.eoff), mtl::Arg::B(p.tiles),
mtl::Arg::B(p.meta), mtl::Arg::V(pa)}, 1, 1, (uint32_t) n_exp_);"""
assert old_pg_moe in w_txt, "pg_moe target not found in window.cpp"
w_txt = w_txt.replace(old_pg_moe, new_pg_moe, 1)
w_txt = w_txt.replace(
"r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a)};",
"r.args = {mtl::Arg::B(logits_), mtl::Arg::B(ids_), mtl::Arg::B(wts_), mtl::Arg::B(plan_), mtl::Arg::V(a), mtl::Arg::B(rcounter_)};",
)
old_at_cut = """ void at_cut(int s) {
st_->at_layer(s, seg_T_, TOPK, cut_ids_ ? cut_ids_ : ids_.as<int32_t>(), st_pred_ && !cut_ids_ ? predL_.as<float>() + (size_t) s * MAXT * 512 : nullptr,
predL2_.ptr && !cut_ids_ ? predL2_.as<float>() + (size_t) s * MAXT * 512 : nullptr);
}"""
new_at_cut = """ void dump_moe_layer(const char* pfx, int l, bool dump_in, bool dump_out);
void at_cut(int s) {
if (cut_ids_ && pg_) {
if (const char* pfx = std::getenv("STRATA_METAL_DUMP_MOE")) {
if (s > 0) dump_moe_layer(pfx, s - 1, false, true);
dump_moe_layer(pfx, s, true, false);
}
}
st_->at_layer(s, seg_T_, TOPK, cut_ids_ ? cut_ids_ : ids_.as<int32_t>(), st_pred_ && !cut_ids_ ? predL_.as<float>() + (size_t) s * MAXT * 512 : nullptr,
predL2_.ptr && !cut_ids_ ? predL2_.as<float>() + (size_t) s * MAXT * 512 : nullptr);
}"""
assert old_at_cut in w_txt, "old_at_cut not found in window.cpp"
w_txt = w_txt.replace(old_at_cut, new_at_cut, 1)
old_p_rows = "int Window::prompt_chunk_rows() const { return pg_ ? pg_->MT : 0; }"
new_p_rows = """int Window::prompt_chunk_rows() const { return pg_ ? pg_->MT : 0; }
void Window::dump_moe_layer(const char* pfx, int l, bool dump_in, bool dump_out) {
if (!pfx || !pg_) return;
static std::string s_pfx;
static size_t s_prev_T = 0;
const size_t T = (size_t) seg_T_;
const bool first_chunk = (last_T_ == 0 && s_pfx != pfx);
if (l == 0 && dump_in && first_chunk) {
s_pfx = pfx;
s_prev_T = 0;
}
auto wr = [&](const char* ext, const void* ptr, size_t row_bytes) {
std::string fn = std::string(pfx) + ext;
const size_t chunk_bytes = T * row_bytes;
if (s_prev_T == 0) {
FILE* fp = std::fopen(fn.c_str(), l == 0 ? "wb" : "r+b");
if (!fp) fp = std::fopen(fn.c_str(), "wb");
if (!fp) { std::perror(("dump_moe fopen " + fn).c_str()); std::exit(2); }
if (std::fseek(fp, (long)((size_t) l * chunk_bytes), SEEK_SET) != 0 ||
std::fwrite(ptr, 1, chunk_bytes, fp) != chunk_bytes) {
std::perror(("dump_moe fwrite " + fn).c_str()); std::exit(2);
}
std::fclose(fp);
} else {
// Multi-chunk prompt: splice new T rows into [NL, s_prev_T + T, row_bytes]
const size_t old_l_bytes = s_prev_T * row_bytes;
const size_t new_l_bytes = (s_prev_T + T) * row_bytes;
static std::map<std::string, std::vector<uint8_t>> prev_files;
std::vector<uint8_t>& prev_file = prev_files[ext];
if (l == 0) {
FILE* rf = std::fopen(fn.c_str(), "rb");
if (!rf) { std::perror(("dump_moe multi-chunk r " + fn).c_str()); std::exit(2); }
prev_file.resize((size_t) 48 * old_l_bytes);
if (std::fread(prev_file.data(), 1, prev_file.size(), rf) != prev_file.size()) {
std::perror(("dump_moe multi-chunk fread " + fn).c_str()); std::exit(2);
}
std::fclose(rf);
}
FILE* wf = std::fopen(fn.c_str(), l == 0 ? "wb" : "r+b");
if (!wf) { std::perror(("dump_moe multi-chunk w " + fn).c_str()); std::exit(2); }
std::fseek(wf, (long)((size_t) l * new_l_bytes), SEEK_SET);
if (std::fwrite(prev_file.data() + (size_t) l * old_l_bytes, 1, old_l_bytes, wf) != old_l_bytes ||
std::fwrite(ptr, 1, chunk_bytes, wf) != chunk_bytes) {
std::perror(("dump_moe multi-chunk fwrite " + fn).c_str()); std::exit(2);
}
std::fclose(wf);
if (l == 47) prev_file.clear();
}
};
if (dump_in) {
wr(".mixed", pg_->mixed.ptr, N * sizeof(float));
wr(".ids", pg_->ids.ptr, TOPK * sizeof(int32_t));
wr(".wts", pg_->wts.ptr, TOPK * sizeof(float));
}
if (dump_out) {
wr(".yr", pg_->yr.ptr, N * sizeof(float));
wr(".hh", pg_->hh.ptr, NFF * sizeof(uint16_t));
wr(".sgv", pg_->sgv.ptr, sizeof(float));
if (l == 47) s_prev_T += T;
}
}"""
assert old_p_rows in w_txt, "old_p_rows not found in window.cpp"
w_txt = w_txt.replace(old_p_rows, new_p_rows, 1)
old_run_seg = """ seg_T_ = T;
cut_ids_ = p.ids.as<int32_t>();
cx_->run_graph_seg(ds, cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this);
cut_ids_ = nullptr;"""
new_run_seg = """ seg_T_ = T;
cut_ids_ = p.ids.as<int32_t>();
cx_->run_graph_seg(ds, cuts, [](void* u, int s) { ((Window*) u)->at_cut(s); }, this);
if (const char* pfx = std::getenv("STRATA_METAL_DUMP_MOE")) {
dump_moe_layer(pfx, 47, false, true);
}
cut_ids_ = nullptr;"""
assert old_run_seg in w_txt, "old_run_seg not found in window.cpp"
w_txt = w_txt.replace(old_run_seg, new_run_seg, 1)
w_txt = w_txt.replace("if (kt == Tn && head && ++seen == 20 && !stream_) kprof(c.ds, Tn);", "if (kt == Tn && head && ++seen == 20) kprof(c.ds, Tn);")
w_path.write_text(w_txt, encoding="utf-8")
# 7. prompt_chain.inc: same packed nt and n_exp_ for gemm, route_topk, PcMmPlanArgs, and mm_scan
pc_txt = pc_path.read_text(encoding="utf-8")
old_pc_moe = """ gemm(b.mixh, b.Wr, b.logits, nT, NEXP, N);
const uint32_t nt = (uint32_t) nT;
D(P_(b.lib, "route_topk"), {mtl::Arg::B(b.logits), mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::V(nt)}, (nT + 7) / 8, 1, 256);
const int nent = nT * TOPK, nblk = (nent + 1023) / 1024;
PcMmPlanArgs pa{(uint32_t) nent, TOPK, NEXP, (uint32_t) PC_TM, (uint32_t) nblk};
const uint32_t nclr = (uint32_t) b.maxrows;
D(P_(b.lib, "mm_clear"), {mtl::Arg::B(b.stok), mtl::Arg::V(nclr)}, (b.maxrows + 255) / 256, 1, 256);
D(P_(b.lib, "mm_hist"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024);
D(P_(b.lib, "mm_scan"), {mtl::Arg::B(b.bcnt), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.tiles), mtl::Arg::B(b.meta),
mtl::Arg::V(pa)}, 1, 1, NEXP);"""
new_pc_moe = """ gemm(b.mixh, b.Wr, b.logits, nT, n_exp_, N);
const uint32_t nt = (uint32_t) nT | ((uint32_t) n_exp_ << 16);
D(P_(b.lib, "route_topk"), {mtl::Arg::B(b.logits), mtl::Arg::B(b.ids), mtl::Arg::B(b.wts), mtl::Arg::V(nt)}, (nT + 7) / 8, 1, 256);
const int nent = nT * TOPK, nblk = (nent + 1023) / 1024;
PcMmPlanArgs pa{(uint32_t) nent, TOPK, (uint32_t) n_exp_, (uint32_t) PC_TM, (uint32_t) nblk};
const uint32_t nclr = (uint32_t) b.maxrows;
D(P_(b.lib, "mm_clear"), {mtl::Arg::B(b.stok), mtl::Arg::V(nclr)}, (b.maxrows + 255) / 256, 1, 256);
D(P_(b.lib, "mm_hist"), {mtl::Arg::B(b.ids), mtl::Arg::B(b.bcnt), mtl::Arg::V(pa)}, nblk, 1, 1024);
D(P_(b.lib, "mm_scan"), {mtl::Arg::B(b.bcnt), mtl::Arg::B(b.boff), mtl::Arg::B(b.eoff), mtl::Arg::B(b.tiles), mtl::Arg::B(b.meta),
mtl::Arg::V(pa)}, 1, 1, (uint32_t) n_exp_);"""
assert old_pc_moe in pc_txt, "pc_moe target not found in prompt_chain.inc"
pc_path.write_text(pc_txt.replace(old_pc_moe, new_pc_moe, 1), encoding="utf-8")
def patch_tests_and_route_parity() -> None:
# 1. Any test in metal/tests that dispatches moe_route or router_route needs the extra smap buffer argument
for p in (ROOT / "metal" / "tests").glob("*.cpp"):
txt = p.read_text(encoding="utf-8")
if "moe_route" in txt and "mtl::Arg::B(plan), mtl::Arg::V(a)};" in txt:
txt = txt.replace(
"mtl::Arg::B(plan), mtl::Arg::V(a)};",
"mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(plan)};",
)
p.write_text(txt, encoding="utf-8")
if "router_route_" in txt and "Arg::B(rcounter)};" in txt:
txt = txt.replace(
"mtl::Arg::V(a), mtl::Arg::B(rcounter)};",
"mtl::Arg::V(a), mtl::Arg::B(rcounter), mtl::Arg::B(rcounter)};",
)
p.write_text(txt, encoding="utf-8")
# 2. Extend route_parity.cpp to test BOTH K=512 (pad=0) AND K=256 + MoCE (pad=exp(0.85f), random smap, out-of-range logits in 256..511)
rp_path = ROOT / "metal" / "tests" / "route_parity.cpp"
rp_src = r"""// metal/tests/route_parity.cpp - moe_route (softmax + top-10 + the expert plan) against a CPU replica of
// native_router.cu's rule, for 1..8 tokens:
// Phase 1: 200 windows with n_expert = 512, pad = 0 (deliberate exact ties)
// Phase 2: 200 windows with n_expert = 256, pad = bit_cast(exp(0.85f)), random smap residency, and huge traps in 256..511
#include "mtl_rt.h"
#include "ref_quant.h"
#include <algorithm>
#include <cfloat>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <set>
#include <vector>
namespace {
constexpr int P_NG = 0, 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;
struct MoeArgs {
uint64_t gate_off, up_off, down_off, gu_stride, d_stride;
uint32_t gu_row, d_row, n_embd, n_ff, n_tok, topk, n_expert, gx_gu, gx_dn, pad;
};
} // namespace
int main() {
mtl::Context cx;
const int lib = cx.compile({"quant.metal", "units.metal", "moe.metal"});
mtl::Pipeline p = cx.pipeline(lib, "moe_route");
mtl::Buffer lg = cx.alloc(4 * 8 * 512), ids = cx.alloc(4 * 80), w = cx.alloc(4 * 80), plan = cx.alloc(4 * P_INTS),
smap = cx.alloc(4 * 512);
ref::Rng r(9);
int fails_512 = 0, fails_256_moce = 0;
double worst_w_512 = 0, worst_w_256_moce = 0;
// Phase 1: K=512, MoCE off (pad = 0)
for (int trial = 0; trial < 200; ++trial) {
const int nt = 1 + trial % 8;
for (int i = 0; i < nt * 512; ++i) {
float v = r.sym() * 4.0f;
if (trial % 3 == 0) v = std::round(v * 4.0f) / 4.0f; // many exact ties
lg.as<float>()[i] = v;
}
MoeArgs a{};
a.n_tok = nt;
a.topk = 10;
a.n_expert = 512;
a.gx_gu = 1;
a.gx_dn = 1;
a.pad = 0;
mtl::Dispatch d;
d.p = &p;
d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)};
d.tg[0] = 512;
cx.run({d});
for (int t = 0; t < nt; ++t) {
const float* l = lg.as<float>() + t * 512;
float mx = -INFINITY;
for (int e = 0; e < 512; ++e) mx = std::max(mx, l[e]);
std::vector<double> pr(512);
double s = 0;
for (int e = 0; e < 512; ++e) s += (pr[e] = std::exp((double)l[e] - mx));
for (int e = 0; e < 512; ++e) pr[e] /= s;
std::vector<int> order(512);
for (int e = 0; e < 512; ++e) order[e] = e;
std::stable_sort(order.begin(), order.end(), [&](int x, int y) { return l[x] > l[y]; });
double ssum = 0;
for (int k = 0; k < 10; ++k) ssum += pr[order[k]];
for (int k = 0; k < 10; ++k) {
if (ids.as<int>()[t * 10 + k] != order[k]) { ++fails_512; break; }
const double ref = pr[order[k]] / std::max(ssum, 6.103515625e-5);
worst_w_512 = std::max(worst_w_512, std::fabs(w.as<float>()[t * 10 + k] - ref) / ref);
}
}
const int* P = plan.as<int>();
std::set<int> distinct(ids.as<int>(), ids.as<int>() + nt * 10);
bool ok = P[P_NG] == (int)distinct.size();
int gi = 0;
for (int e : distinct) {
if (P[P_GEXP + gi] != e) ok = false;
int j = P[P_GSTART + gi];
for (int i = 0; i < nt * 10; ++i)
if (ids.as<int>()[i] == e) {
if (P[P_ETOK + j] != i / 10 || P[P_TOKENT + i] != j) ok = false;
++j;
}
if (j != P[P_GSTART + gi + 1]) ok = false;
++gi;
}
fails_512 += !ok;
}
// Phase 2: K=256, MoCE on (pad = bit_cast(exp(0.85f))), random smap residency, and huge trap logits in 256..511
const float moce_scale = std::exp(0.85f);
uint32_t moce_pad = 0;
std::memcpy(&moce_pad, &moce_scale, 4);
for (int trial = 0; trial < 200; ++trial) {
const int nt = 1 + trial % 8;
for (int e = 0; e < 512; ++e) {
smap.as<uint32_t>()[e] = (r.sym() > 0.0f) ? (uint32_t) e : 0xFFFFFFFFu;
}
for (int t = 0; t < nt; ++t) {
for (int e = 0; e < 256; ++e) {
float v = r.sym() * 4.0f;
if (trial % 3 == 0) v = std::round(v * 4.0f) / 4.0f;
lg.as<float>()[t * 512 + e] = v;
}
for (int e = 256; e < 512; ++e) {
lg.as<float>()[t * 512 + e] = 50.0f + (float) e; // trap logits past n_expert = 256
}
}
MoeArgs a{};
a.n_tok = nt;
a.topk = 10;
a.n_expert = 256;
a.gx_gu = 1;
a.gx_dn = 1;
a.pad = moce_pad;
mtl::Dispatch d;
d.p = &p;
d.args = {mtl::Arg::B(lg), mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(plan), mtl::Arg::V(a), mtl::Arg::B(smap)};
d.tg[0] = 512;
cx.run({d});
const uint32_t* sm = smap.as<uint32_t>();
for (int t = 0; t < nt; ++t) {
const float* l = lg.as<float>() + t * 512;
float mx = -INFINITY;
for (int e = 0; e < 256; ++e) mx = std::max(mx, l[e]);
std::vector<float> pr(256);
float s = 0.0f;
for (int e = 0; e < 256; ++e) s += (pr[e] = std::exp(l[e] - mx));
const float rcp = 1.0f / s;
for (int e = 0; e < 256; ++e) pr[e] *= rcp;
std::vector<float> work = pr;
std::vector<int> picked(10, -1);
double ssum = 0.0;
for (int rank = 0; rank < 10; ++rank) {
if (rank == 3) {
for (int e = 0; e < 256; ++e) {
if (work[e] > 0.0f && sm[e] != 0xFFFFFFFFu) work[e] *= moce_scale;
}
}
int best_e = 0;
float best_v = work[0];
for (int e = 1; e < 256; ++e) {
if (work[e] > best_v) { best_v = work[e]; best_e = e; }
}
picked[rank] = best_e;
work[best_e] = -INFINITY;
ssum += pr[best_e];
}
for (int k = 0; k < 10; ++k) {
const int got_id = ids.as<int>()[t * 10 + k];
if (got_id != picked[k] || got_id >= 256) { ++fails_256_moce; break; }
const double ref = (double) pr[picked[k]] / std::max(ssum, 6.103515625e-5);
worst_w_256_moce = std::max(worst_w_256_moce, std::fabs(w.as<float>()[t * 10 + k] - ref) / ref);
}
}
const int* P = plan.as<int>();
std::set<int> distinct(ids.as<int>(), ids.as<int>() + nt * 10);
bool ok = P[P_NG] == (int)distinct.size();
int gi = 0;
for (int e : distinct) {
if (e >= 256 || P[P_GEXP + gi] != e) ok = false;
int j = P[P_GSTART + gi];
for (int i = 0; i < nt * 10; ++i)
if (ids.as<int>()[i] == e) {
if (P[P_ETOK + j] != i / 10 || P[P_TOKENT + i] != j) ok = false;
++j;
}
if (j != P[P_GSTART + gi + 1]) ok = false;
++gi;
}
fails_256_moce += !ok;
}
const int total_fails = fails_512 + fails_256_moce;
const double worst_w = std::max(worst_w_512, worst_w_256_moce);
std::printf("route parity K=512 (pad=0, 200 windows): mismatches %d, worst weight rel err %.2e\n",
fails_512, worst_w_512);
std::printf("route parity K=256 + MoCE=0.85 (random smap + 256..511 traps, 200 windows): mismatches %d, worst weight rel err %.2e %s\n",
fails_256_moce, worst_w_256_moce, total_fails == 0 && worst_w < 2e-6 ? "PASS" : "FAIL");
return total_fails == 0 && worst_w < 2e-6 ? 0 : 1;
}
"""
rp_path.write_text(rp_src, encoding="utf-8")
def write_strata_window_cli_and_converter() -> None:
cli_src = r"""// metal/src/strata_window_cli.cpp - Direct zero-bloat Metal runner for Strata's Window + ExpertStream + PleTable.
// Supports both:
// 1. Prefill (PromptGemm) + decode (Window::run + ExpertStream F_NOCACHE pread + speculative router lookahead)
// 2. Teacher-forced full-vocab (248,320) FP64 KLD, PPL, and top-1 agreement evaluation (--logits-out / --logits-ref)
#include "strata/core/layer.hpp"
#include "strata/core/session.hpp"
#include "strata/kernels/ngram.hpp"
#include "strata/spec/suffix_drafter.hpp"
#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fcntl.h>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
namespace strata::metal_engine {
void set_model_shard(const char* shard1);
void* window_create();
void window_destroy(void* w);
bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::SessionState& ss, int max_t, int64_t max_cells,
int64_t* n_vocab, std::string& err);
bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0,
const std::function<bool(const float*, int64_t, int64_t, std::string&)>& on_chunk, std::string& err);
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);
void window_reset_stats(void* w);
} // namespace strata::metal_engine
static double now_ms() {
return std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now().time_since_epoch()).count();
}
static double logsumexp_f64(const float* lg, int64_t n, int32_t* argmax_out) {
float mx = lg[0];
int32_t am = 0;
for (int64_t i = 1; i < n; ++i) {
if (lg[i] > mx) {
mx = lg[i];
am = (int32_t) i;
}
}
if (argmax_out) *argmax_out = am;
double sum = 0.0;
for (int64_t i = 0; i < n; ++i) {
sum += std::exp((double) lg[i] - (double) mx);
}
return (double) mx + std::log(sum);
}
int main(int argc, char** argv) {
std::string metal_gguf, ple_gguf, tokens_file, logits_out, logits_ref, logits_ref2;
int n_expert = 256;
int64_t max_cells = 4096;
int n_prompt = 22;
int n_gen = 64;
int suffix_draft = 0;
for (int i = 1; i < argc; ++i) {
std::string a = argv[i];
if (a == "--metal-gguf" && i + 1 < argc) metal_gguf = argv[++i];
else if (a == "--ple-gguf" && i + 1 < argc) ple_gguf = argv[++i];
else if (a == "--tokens-file" && i + 1 < argc) tokens_file = argv[++i];
else if (a == "--n-expert" && i + 1 < argc) n_expert = std::atoi(argv[++i]);
else if (a == "--ctx" && i + 1 < argc) max_cells = std::atoll(argv[++i]);
else if (a == "--n-prompt" && i + 1 < argc) n_prompt = std::atoi(argv[++i]);
else if (a == "--n-gen" && i + 1 < argc) n_gen = std::atoi(argv[++i]);
else if (a == "--suffix-draft" && i + 1 < argc) suffix_draft = std::max(0, std::min(7, std::atoi(argv[++i])));
else if (a == "--logits-out" && i + 1 < argc) logits_out = argv[++i];
else if (a == "--logits-ref" && i + 1 < argc) logits_ref = argv[++i];
else if (a == "--logits-ref2" && i + 1 < argc) logits_ref2 = argv[++i];
}
if (metal_gguf.empty()) {
std::fprintf(stderr, "usage: strata_window_cli --metal-gguf <K256-metal.gguf> [--ple-gguf <shard1.gguf>] "
"[--n-expert 256] [--ctx 4096] [--n-prompt 22] [--n-gen 64] [--tokens-file <ids.txt>] "
"[--logits-out <out.bin>] [--logits-ref <ref.bin>]\n");
return 2;
}
const bool tf_mode = !logits_out.empty() || !logits_ref.empty();
std::vector<int64_t> all_ids;
if (!tokens_file.empty()) {
std::ifstream f(tokens_file);
int64_t tid;
while (f >> tid) all_ids.push_back(tid);
}
if (all_ids.empty()) {
const int64_t def_ids[] = {
248045, 8678, 198, 2523, 513, 264, 10631, 17313, 13, 248046, 198, 248045, 846, 198, 814, 20139, 1204, 27413, 71284, 3728, 13081, 32907, 12361, 321, 264, 3074, 6700, 11, 321, 948, 264, 28336, 628, 6561, 571, 13, 248046, 198, 248045, 74455, 198, 248068, 198
};
const size_t nd = sizeof(def_ids) / sizeof(def_ids[0]);
for (int i = 0; i < n_prompt + (tf_mode ? n_gen : 0); ++i) all_ids.push_back(def_ids[i % nd]);
}
if (tf_mode && (int) all_ids.size() < n_prompt + n_gen) {
std::fprintf(stderr, "strata_window_cli: teacher-forced mode needs at least n_prompt + n_gen (%d) tokens in --tokens-file, got %zu\n",
n_prompt + n_gen, all_ids.size());
return 1;
}
std::vector<int64_t> prompt_ids(all_ids.begin(), all_ids.begin() + std::min<size_t>(all_ids.size(), (size_t) n_prompt));
const double t_init0 = now_ms();
std::string err;
strata::kernels::PleTable ple_table;
strata::core::SessionState ss{};
ss.ple_prev[0] = -1;
ss.ple_prev[1] = -1;
ss.ple.consts = strata::kernels::ple_artifact_consts();
if (!ple_gguf.empty()) {
strata::kernels::PleIoOptions io{};
io.mode = strata::kernels::PleIo::Direct;
io.max_inflight = 256;
io.cache_rows = 1u << 18;
if (!ple_table.open(ple_gguf, err, io)) {
std::fprintf(stderr, "strata_window_cli: PleTable::open failed: %s\n", err.c_str());
return 1;
}
ss.ple.table = &ple_table;
}
strata::core::ModelGeometry g{};
g.n_embd = 2560;
g.hc = 4;
g.hc_lr = 320;
g.n_ff = 640;
g.n_expert = n_expert;
g.n_layers = 48;
strata::metal_engine::set_model_shard(metal_gguf.c_str());
void* w = strata::metal_engine::window_create();
int64_t n_vocab = 0;
if (!strata::metal_engine::window_init(w, g, ss, 1 + suffix_draft, max_cells, &n_vocab, err)) {
std::fprintf(stderr, "strata_window_cli: window_init failed: %s\n", err.c_str());
return 1;
}
const double init_ms = now_ms() - t_init0;
FILE* fout_lg = nullptr;
if (!logits_out.empty()) {
fout_lg = std::fopen(logits_out.c_str(), "wb");
if (!fout_lg) {
std::perror("fopen logits_out");
return 1;
}
fcntl(fileno(fout_lg), F_NOCACHE, 1);
}
FILE* fref_lg = nullptr;
if (!logits_ref.empty()) {
fref_lg = std::fopen(logits_ref.c_str(), "rb");
if (!fref_lg) {
std::perror("fopen logits_ref");
return 1;
}
fcntl(fileno(fref_lg), F_NOCACHE, 1);
}
FILE* fref2_lg = nullptr;
if (!logits_ref2.empty()) {
fref2_lg = std::fopen(logits_ref2.c_str(), "rb");
if (!fref2_lg) {
std::perror("fopen logits_ref2");
return 1;
}
fcntl(fileno(fref2_lg), F_NOCACHE, 1);
}
strata::spec::SuffixDrafter suffix(3, 32);
for (int64_t tid : prompt_ids) suffix.append((int32_t) tid);
// Prefill all but the last prompt token via prompt_run_rows
const int64_t n_pre = std::max<int64_t>(0, (int64_t) prompt_ids.size() - 1);
const double t_pre0 = now_ms();
if (n_pre > 0) {
if (!strata::metal_engine::prompt_run_rows(prompt_ids.data(), n_pre, 0, nullptr, err)) {
std::fprintf(stderr, "strata_window_cli: prompt_run_rows failed: %s\n", err.c_str());
return 1;
}
}
// Run the last prompt token with head=true to get the first output distribution
int32_t cur_tok = (int32_t) prompt_ids.back();
int32_t next_tok = 0;
int64_t pos = n_pre;
if (!strata::metal_engine::window_run(w, 1, &cur_tok, pos, true, &next_tok, err) ||
!strata::metal_engine::window_commit(w, 1, err)) {
std::fprintf(stderr, "strata_window_cli: first window_run failed: %s\n", err.c_str());
return 1;
}
++pos;
const double prefill_ms = now_ms() - t_pre0;
std::vector<float> cur_lg(tf_mode ? (size_t) n_vocab : 0);
std::vector<float> ref_lg((fref_lg || fref2_lg) ? (size_t) n_vocab : 0);
double sum_nll_cand = 0.0, sum_nll_ref = 0.0, sum_nll_ref2 = 0.0;
std::vector<double> klds, klds2;
int top1_matches = 0, top1_matches2 = 0;
auto eval_tf_step = [&](int step_idx) -> bool {
if (!strata::metal_engine::window_copy_logits(w, 0, cur_lg.data())) return false;
if (fout_lg) {
if (std::fwrite(cur_lg.data(), sizeof(float), (size_t) n_vocab, fout_lg) != (size_t) n_vocab) return false;
}
const int32_t tgt = (int32_t) all_ids[(size_t) n_prompt + (size_t) step_idx];
int32_t cand_top1 = 0;
const double lse_cand = logsumexp_f64(cur_lg.data(), n_vocab, &cand_top1);
sum_nll_cand += -((double) cur_lg[(size_t) tgt] - lse_cand);
if (fref_lg) {
if (std::fread(ref_lg.data(), sizeof(float), (size_t) n_vocab, fref_lg) != (size_t) n_vocab) return false;
int32_t ref_top1 = 0;
const double lse_ref = logsumexp_f64(ref_lg.data(), n_vocab, &ref_top1);
sum_nll_ref += -((double) ref_lg[(size_t) tgt] - lse_ref);
if (cand_top1 == ref_top1) ++top1_matches;
double kld = 0.0;
for (int64_t v = 0; v < n_vocab; ++v) {
const double log_p = (double) ref_lg[(size_t) v] - lse_ref;
const double log_q = (double) cur_lg[(size_t) v] - lse_cand;
const double p = std::exp(log_p);
if (p > 0.0) kld += p * (log_p - log_q);
}
klds.push_back(std::max(0.0, kld));
}
if (fref2_lg) {
if (std::fread(ref_lg.data(), sizeof(float), (size_t) n_vocab, fref2_lg) != (size_t) n_vocab) return false;
int32_t ref2_top1 = 0;
const double lse_ref2 = logsumexp_f64(ref_lg.data(), n_vocab, &ref2_top1);
sum_nll_ref2 += -((double) ref_lg[(size_t) tgt] - lse_ref2);
if (cand_top1 == ref2_top1) ++top1_matches2;
double kld2 = 0.0;
for (int64_t v = 0; v < n_vocab; ++v) {
const double log_p = (double) ref_lg[(size_t) v] - lse_ref2;
const double log_q = (double) cur_lg[(size_t) v] - lse_cand;
const double p = std::exp(log_p);
if (p > 0.0) kld2 += p * (log_p - log_q);
}
klds2.push_back(std::max(0.0, kld2));
}
return true;
};
if (tf_mode && n_gen > 0) {
if (!eval_tf_step(0)) {
std::fprintf(stderr, "strata_window_cli: eval_tf_step(0) failed\n");
return 1;
}
}
std::vector<int32_t> gen_tokens;
gen_tokens.push_back(next_tok);
suffix.append(next_tok);
cur_tok = tf_mode ? (int32_t) all_ids[(size_t) n_prompt] : next_tok;
strata::metal_engine::window_reset_stats(w);
const int decode_steps = std::max(0, n_gen - 1);
const double t_dec0 = now_ms();
int produced = 0;
while (produced < decode_steps) {
int32_t win_in[8] = {cur_tok};
int32_t win_out[8] = {};
int nd = 0;
if (!tf_mode && suffix_draft > 0) {
const int rem = decode_steps - produced;
const int max_d = std::min(suffix_draft, std::max(0, rem - 1));
if (max_d > 0) nd = suffix.propose(max_d, win_in + 1);
}
const int T = 1 + nd;
if (!strata::metal_engine::window_run(w, T, win_in, pos, true, win_out, err)) {
std::fprintf(stderr, "strata_window_cli: decode step failed: %s\n", err.c_str());
return 1;
}
int n_keep = 1;
for (int i = 0; i < nd; ++i) {
if (win_out[i] == win_in[i + 1]) ++n_keep;
else break;
}
if (!strata::metal_engine::window_commit(w, n_keep, err)) {
std::fprintf(stderr, "strata_window_cli: window_commit failed: %s\n", err.c_str());
return 1;
}
if (tf_mode) {
if (!eval_tf_step(produced + 1)) {
std::fprintf(stderr, "strata_window_cli: eval_tf_step(%d) failed\n", produced + 1);
return 1;
}
gen_tokens.push_back(win_out[0]);
++produced;
if (produced < decode_steps) {
cur_tok = (int32_t) all_ids[(size_t) n_prompt + (size_t) produced];
}
pos += 1;
} else {
for (int i = 0; i < n_keep && produced < decode_steps; ++i) {
gen_tokens.push_back(win_out[i]);
suffix.append(win_out[i]);
cur_tok = win_out[i];
++produced;
}
pos += n_keep;
}
}
const double decode_ms = now_ms() - t_dec0;
const double decode_tok_s = decode_steps > 0 ? (decode_steps * 1000.0 / decode_ms) : 0.0;
const double decode_ms_tok = decode_steps > 0 ? (decode_ms / decode_steps) : 0.0;
const double prefill_tok_s = prompt_ids.size() * 1000.0 / prefill_ms;
if (fout_lg) std::fclose(fout_lg);
if (fref_lg) std::fclose(fref_lg);
if (fref2_lg) std::fclose(fref2_lg);
std::printf("STRATA_RESULT init_ms=%.2f prompt_tokens=%zu prefill_ms=%.2f prefill_tok_s=%.2f "
"decode_tokens=%d decode_ms=%.2f decode_tok_s=%.2f decode_ms_per_tok=%.2f\n",
init_ms, prompt_ids.size(), prefill_ms, prefill_tok_s,
decode_steps, decode_ms, decode_tok_s, decode_ms_tok);
if (tf_mode && n_gen > 0) {
const double cand_ppl = std::exp(sum_nll_cand / n_gen);
if (!klds.empty()) {
const double ref_ppl = std::exp(sum_nll_ref / n_gen);
double mean_kld = 0.0;
for (double x : klds) mean_kld += x;
mean_kld /= klds.size();
std::vector<double> sk = klds;
std::sort(sk.begin(), sk.end());
const double med_kld = sk[sk.size() / 2];
const double p90_kld = sk[(size_t) (0.90 * (sk.size() - 1))];
const double top1_pct = 100.0 * (double) top1_matches / (double) n_gen;
std::printf("STRATA_KLD eval_tokens=%d ref_ppl=%.4f cand_ppl=%.4f delta_ppl=%.4f mean_kld=%.6f median_kld=%.6f p90_kld=%.6f top1_agree_pct=%.2f\n",
n_gen, ref_ppl, cand_ppl, cand_ppl - ref_ppl, mean_kld, med_kld, p90_kld, top1_pct);
if (!klds2.empty()) {
const double ref2_ppl = std::exp(sum_nll_ref2 / n_gen);
double mean_kld2 = 0.0;
for (double x : klds2) mean_kld2 += x;
mean_kld2 /= klds2.size();
std::vector<double> sk2 = klds2;
std::sort(sk2.begin(), sk2.end());
const double med_kld2 = sk2[sk2.size() / 2];
const double p90_kld2 = sk2[(size_t) (0.90 * (sk2.size() - 1))];
const double top1_pct2 = 100.0 * (double) top1_matches2 / (double) n_gen;
std::printf("STRATA_KLD_REF2 eval_tokens=%d ref_ppl=%.4f cand_ppl=%.4f delta_ppl=%.4f mean_kld=%.6f median_kld=%.6f p90_kld=%.6f top1_agree_pct=%.2f\n",
n_gen, ref2_ppl, cand_ppl, cand_ppl - ref2_ppl, mean_kld2, med_kld2, p90_kld2, top1_pct2);
}
} else {
std::printf("STRATA_PPL eval_tokens=%d cand_ppl=%.4f\n", n_gen, cand_ppl);
}
}
std::printf("GEN_IDS:");
for (size_t i = 0; i < gen_tokens.size(); ++i) {
std::printf(" %d", gen_tokens[i]);
}
std::printf("\n");
std::fflush(stdout);
strata::metal_engine::window_destroy(w);
return 0;
}
"""
(ROOT / "metal" / "src" / "strata_window_cli.cpp").write_text(cli_src, encoding="utf-8")
# Also update convert_k256_q2_0.cpp so SSE search evaluates the fp16-rounded scale d
conv_path = ROOT / "metal" / "src" / "convert_k256_q2_0.cpp"
if conv_path.exists():
conv_txt = conv_path.read_text(encoding="utf-8")
conv_txt = conv_txt.replace(
"const float d = amax * f;\n const float inv = 1.0f / d;",
"const float d = ref::f16_to_f32(ref::f32_to_f16(amax * f));\n if (d <= 0.0f) continue;\n const float inv = 1.0f / d;",
)
conv_path.write_text(conv_txt, encoding="utf-8")
heal_src = Path(__file__).resolve().parent / "heal_k256_q2_0.cpp"
if heal_src.exists():
(ROOT / "metal" / "src" / "heal_k256_q2_0.cpp").write_text(heal_src.read_text(encoding="utf-8"), encoding="utf-8")
cm_path = ROOT / "metal" / "CMakeLists.txt"
cm_txt = cm_path.read_text(encoding="utf-8")
if "strata_window_cli" not in cm_txt:
cm_txt += """
add_executable(strata_window_cli src/strata_window_cli.cpp)
target_link_libraries(strata_window_cli PRIVATE strata_metal_kernels strata_metal_host strata_metal_kernels)
add_executable(convert_k256_q2_0 src/convert_k256_q2_0.cpp)
target_link_libraries(convert_k256_q2_0 PRIVATE mtl_rt)
add_executable(heal_k256_q2_0 src/heal_k256_q2_0.cpp)
target_compile_definitions(heal_k256_q2_0 PRIVATE ACCELERATE_NEW_LAPACK)
target_link_libraries(heal_k256_q2_0 PRIVATE mtl_rt "-framework Accelerate")
"""
cm_path.write_text(cm_txt, encoding="utf-8")
if __name__ == "__main__":
reset_git()
patch_mtl_rt()
patch_kernels()
patch_stream()
patch_window()
patch_tests_and_route_parity()
write_strata_window_cli_and_converter()
print("Patched vendor/strata/metal successfully!")