diff --git a/Cargo.lock b/Cargo.lock index 2650c74..d744705 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1133,6 +1133,7 @@ dependencies = [ name = "ds4-server" version = "0.1.0" dependencies = [ + "cc", "diesel", "diesel_migrations", "iced", diff --git a/Cargo.toml b/Cargo.toml index e6f9542..6b54a98 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,10 @@ rust-version = "1.97" description = "A native macOS coding-agent GUI for DwarfStar" license = "MIT" publish = false +build = "build.rs" + +[build-dependencies] +cc = "1.3.0" [dependencies] diesel = { version = "2.3.11", features = ["sqlite", "returning_clauses_for_sqlite_3_35", "64-column-tables"] } @@ -26,6 +30,7 @@ identifier = "DS4Server.rfc1437.de" description = "A native macOS coding-agent GUI for DwarfStar" binaries = [{ path = "ds4-server", main = true }] icons = ["assets/DS4Server.icns", "assets/app-icon.png"] +resources = ["metal"] [package.metadata.packager.macos] minimum-system-version = "13.0" diff --git a/PLAN.md b/PLAN.md index 85369b3..d16b45b 100644 --- a/PLAN.md +++ b/PLAN.md @@ -12,23 +12,23 @@ sessions, tools, and turn orchestration. | Area | Status | Available now | Still open | | --- | --- | --- | --- | -| App shell | Implemented | Native multi-window Iced app, Codex-inspired layout, native Application/Edit/Window menus, Preferences panel, and Model Manager window | Functional chat content and native app packaging | +| App shell | Implemented | Native multi-window Iced app, Codex-inspired layout, native Application/Edit/Window menus, Preferences panel, Model Manager, and streaming composer | Native app release packaging | | Projects and sessions | Implemented for metadata | Native folder picker, project-name dialog, create/select/delete projects and sessions | Transcripts, KV payloads, rename/archive, and model binding | | Persistence | Implemented for metadata and preferences | SQLite in Application Support, Diesel ORM, embedded migrations, constrained defaults and CRUD tests | Messages, downloaded-artifact metadata, and session-runtime migrations | -| Model runtime | Model intake implemented | Typed DS4 preferences, managed downloads, mmap-backed GGUF v3 loading, complete DS4 shape/quantization validation, DeepSeek/GLM tokenization, and prompt rendering | Metal inference, sessions/KV state, lifecycle coordination, and idle unloading | +| Model runtime | DeepSeek Flash vertical implemented | Rust-owned mmap/model/session graph, reused Metal kernels, 2K compressed-attention generation, DS4 sampling, lazy background loading, cancellation, and idle unloading | Ratio-4 sparse indexer beyond 2K, DSpark/SSD/steering, GLM/Pro, KV checkpoint persistence, and shared acquisition for HTTP | | Local endpoint | Not started | Nothing listening | OpenAI-compatible HTTP and streaming surfaces | -| Agent | UI shell only | Empty conversation pane and composer placeholder | Messages, generation, tools, approvals, compaction, and cancellation | +| Agent | Basic local chat implemented | Multi-turn role rendering, streamed DeepSeek Flash text, composer send/stop, and in-memory transcript | Durable transcripts, reasoning presentation, tools, approvals, compaction, statistics, and KV-backed resume | | A2UI | Future extension | bDS2 provides the reference structured render-tool contract | Native inline surfaces for local chat after core chat and tools are stable | | Dev brain | Not started | No vault access | Obsidian vault selection, durable access, and agent knowledge/memory tools | | Release | Not started | Development binary builds and tests | `.app` packaging, signing, entitlements, and notarization | -The application currently manages durable project/session metadata and model -preferences. Its separate Model Manager can download, fully validate, and -delete main and DSpark GGUF artifacts in managed Application Support storage -without blocking the UI. The engine can mmap a configured model, validate its -metadata and complete tensor layout, load its tokenizer, and render a prompt; -the app does not yet acquire that engine for inference, serve HTTP, save chat -messages, or run an agent. +The application manages durable project/session metadata and model preferences. +Its Model Manager downloads, validates, and deletes managed GGUF artifacts +without blocking the UI. Local chat now lazily acquires the Rust DeepSeek Flash +executor, streams generated text, supports cancellation and multi-turn role +replay, and unloads model resources after the configured idle timeout. The +current Rust graph supports DS4 compressed attention through 2K context; the sparse ratio-4 indexer, +durable transcripts/KV checkpoints, HTTP service, and agent tools remain open. ## Phase 0 — project and session shell @@ -47,27 +47,27 @@ Exit criterion: restart the app and see the same project/session tree. ## Phase 1 — model library and on-demand loading -Status: **partially implemented**. The typed DS4 preference contract, shared -effective-settings builder, background download workflow, and model-intake -boundary are implemented; Metal inference, session state, and idle lifecycle -coordination remain open. +Status: **partially implemented**. The typed preference/download/intake path and +a DeepSeek Flash Rust/Metal vertical are implemented, including compressed KV +through 2K context, DS4 sampling, a single background owner, cancellation, and +idle unload. Sparse indexed attention, persistent KV/session sync, other model +families, and optional execution modes remain open. 1. **Implemented:** the typed model/runtime preference contract described below is complete before chat, the HTTP endpoint, or the engine lifecycle. -2. **Partially implemented:** the mmap-backed GGUF loader, DeepSeek/GLM - tokenizer, prompt renderer, and model `open`/`close` ownership boundary are - complete. The Metal session API modeled on `ds4.h` remains: - `create_session`, `sync`, `sample`, and KV save/load. +2. **Partially implemented:** mmap-backed loading, tokenizer/prompt rendering, + Rust-owned Flash session creation/evaluation/sampling, and model ownership + are complete. Prefix sync plus KV save/load remain. 3. **Implemented:** retain the tested model restrictions from DwarfStar. Main and DSpark artifacts reject invalid GGUF versions, catalog mismatches, and incompatible metadata, tensor shapes, offsets, or quantization before promotion and again before model use. -4. Reuse the existing `.metal` kernels and preserve mmap-backed resident +4. **Partially implemented:** reuse the existing `.metal` kernels and preserve mmap-backed resident loading plus the explicit SSD-streaming path. Keep Objective-C only at the Metal interop boundary. -5. Keep one engine resident at most, matching DwarfStar's instance-lock and +5. **Implemented for local chat:** keep one engine resident at most, matching DwarfStar's instance-lock and memory assumptions. -6. Move engine work off the UI thread and support cooperative cancellation at +6. **Implemented for local chat:** move engine work off the UI thread and support cooperative cancellation at the safe session boundaries already defined by `ds4_session_sync`. ### DS4 KV-cache port and shared session core @@ -328,10 +328,11 @@ the same conversation without prefill when its KV payload is compatible. ## Phase 4 — agent chat and tools -Status: **UI shell only**. The conversation area, composer, and icons are -visual placeholders with no message or agent runtime behind them. +Status: **basic local chat implemented**. The conversation area streams +multi-turn DeepSeek Flash output and supports Stop; durable messages, tool +turns, and the coding-agent runtime remain open. -1. Build the classic chat surface: transcript, reasoning disclosure, tool-call +1. **Partially implemented:** build the classic chat surface: transcript, reasoning disclosure, tool-call cards, composer, queued input, stop button, prefill progress, and generation statistics. Local turns must call the same session registry and KV path used by the HTTP endpoint. @@ -481,9 +482,10 @@ tests, and release packaging remain open. ## Recommended next milestones -1. Implement the single-engine lifecycle and local generation path: lazy load, - shared concurrent acquisition, activity tracking, and idle unload. -2. Put the OpenAI-compatible endpoint on that same engine lifecycle so local +1. Finish the Flash session core as one larger slice: port the ratio-4 indexer + and indexed attention for full configured context, then add prefix sync and + compatible KV checkpoint save/load. +2. Put the OpenAI-compatible endpoint on the same engine lifecycle so local chats and HTTP requests cannot load duplicate engines. 3. Add transcript/KV persistence, then connect the existing conversation UI and finally port the agent tools. diff --git a/README.md b/README.md index c9aa2da..9017838 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # DS4Server -DS4Server is a native macOS coding-agent application for the DwarfStar (`ds4`) -inference engine. It uses Rust and Iced and will combine local model loading, an +DS4Server is a native macOS coding-agent application that rewrites the +DwarfStar (`ds4`) inference engine in Rust. It uses Rust and Iced and will combine local model loading, an OpenAI-compatible localhost endpoint, and project-scoped agent chat in one app. The current milestone provides a Codex-inspired project/session layout. A native @@ -17,6 +17,14 @@ bar. Stopping or quitting keeps the partial file; the next Download/Resume action continues from that exact byte after relaunch. Exact size and SHA-256 verification happen before an artifact becomes usable. +The selected DeepSeek V4 Flash model can run directly from a project session. +The model, KV/compressor state, 43-layer graph, sampling, and lifecycle are +owned by Rust; a fixed snapshot of the Objective-C Metal boundary and unchanged +Metal kernels is vendored and built inside this repository. Tokens stream into the chat UI, Stop cancels generation, +follow-up turns replay their role history, and the model unloads after the idle +timeout. The current graph supports 2K context while the sparse long-context +indexer and persistent transcript/KV restore remain the next engine slice. + ```sh cargo install cargo-packager --locked --version 0.11.8 make bundle diff --git a/build.rs b/build.rs new file mode 100644 index 0000000..0b8c2e0 --- /dev/null +++ b/build.rs @@ -0,0 +1,18 @@ +use std::path::Path; + +fn main() { + if std::env::var("CARGO_CFG_TARGET_OS").as_deref() != Ok("macos") { + return; + } + + let metal = Path::new("native/metal/ds4_metal.m"); + println!("cargo:rerun-if-changed={}", metal.display()); + cc::Build::new() + .include("native/metal") + .file(metal) + .flag("-fobjc-arc") + .opt_level(3) + .compile("ds4_metal"); + println!("cargo:rustc-link-lib=framework=Foundation"); + println!("cargo:rustc-link-lib=framework=Metal"); +} diff --git a/metal/argsort.metal b/metal/argsort.metal new file mode 100644 index 0000000..77c473f --- /dev/null +++ b/metal/argsort.metal @@ -0,0 +1,275 @@ +struct ds4_metal_args_argsort { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + int32_t top_k; +}; + +struct ds4_metal_args_argsort_merge { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + int32_t top_k; + int32_t len; +}; + +typedef void (argsort_t)( + constant ds4_metal_args_argsort & args, + device const char * src0, + device int32_t * dst, + threadgroup int32_t * shmem_i32 [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]); + +// Sort one float row into an index row. DS4 only exports the descending +// instance because router and indexer selection both need top-k order. +template +kernel void kernel_argsort_f32_i32( + constant ds4_metal_args_argsort & args, + device const char * src0, + device int32_t * dst, + threadgroup int32_t * shmem_i32 [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + // bitonic sort + const int col = tpitg[0]; + const int ib = tgpig[0] / args.ne01; + + const int i00 = ib*ntg.x; + const int i01 = tgpig[0] % args.ne01; + const int i02 = tgpig[1]; + const int i03 = tgpig[2]; + + device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03); + + // initialize indices + shmem_i32[col] = i00 + col; + + // Stage this block's score slice in threadgroup memory (indices stay in + // [i00, i00+ntg.x), so shmem_f32[idx - i00] replaces the device gather). + // The host allocates ntg.x extra floats after the index array. Values and + // the comparison network are unchanged, so the permutation is identical. + threadgroup float * shmem_f32 = (threadgroup float *) (shmem_i32 + ntg.x); + if (i00 + col < args.ne00) { + shmem_f32[col] = src0_row[i00 + col]; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int k = 2; k <= ntg.x; k *= 2) { + for (int j = k / 2; j > 0; j /= 2) { + int ixj = col ^ j; + if (ixj > col) { + if ((col & k) == 0) { + if (shmem_i32[col] >= args.ne00 || + (shmem_i32[ixj] < args.ne00 && (order == DS4_SORT_ORDER_ASC ? + shmem_f32[shmem_i32[col] - i00] > shmem_f32[shmem_i32[ixj] - i00] : + shmem_f32[shmem_i32[col] - i00] < shmem_f32[shmem_i32[ixj] - i00])) + ) { + SWAP(shmem_i32[col], shmem_i32[ixj]); + } + } else { + if (shmem_i32[ixj] >= args.ne00 || + (shmem_i32[col] < args.ne00 && (order == DS4_SORT_ORDER_ASC ? + shmem_f32[shmem_i32[col] - i00] < shmem_f32[shmem_i32[ixj] - i00] : + shmem_f32[shmem_i32[col] - i00] > shmem_f32[shmem_i32[ixj] - i00])) + ) { + SWAP(shmem_i32[col], shmem_i32[ixj]); + } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + const int64_t i0 = ib*args.top_k; + + // copy the result to dst without the padding + if (i0 + col < args.ne0 && col < args.top_k) { + dst += i0 + args.ne0*i01 + args.ne0*args.ne1*i02 + args.ne0*args.ne1*args.ne2*i03; + + dst[col] = shmem_i32[col]; + } +} + +// Host-visible sort variant used by DS4 top-k selection. +template [[host_name("kernel_argsort_f32_i32_desc")]] kernel argsort_t kernel_argsort_f32_i32; + +typedef void (argsort_merge_t)( + constant ds4_metal_args_argsort_merge & args, + device const char * src0, + device const int32_t * tmp, + device int32_t * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]); + +// Merges sorted index runs produced by kernel_argsort_f32_i32. In the DS4 graph +// this finishes top-k over router or compressed-attention score rows. +template +kernel void kernel_argsort_merge_f32_i32( + constant ds4_metal_args_argsort_merge & args, + device const char * src0, + device const int32_t * tmp, + device int32_t * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + + const int im = tgpig[0] / args.ne01; + const int i01 = tgpig[0] % args.ne01; + const int i02 = tgpig[1]; + const int i03 = tgpig[2]; + + const int start = im * (2 * args.len); + + const int len0 = MIN(args.len, MAX(0, args.ne0 - (int)(start))); + const int len1 = MIN(args.len, MAX(0, args.ne0 - (int)(start + args.len))); + + const int total = len0 + len1; + + device const int32_t * tmp0 = tmp + start + + i01*args.ne0 + + i02*args.ne0*args.ne01 + + i03*args.ne0*args.ne01*args.ne02; + + device const int32_t * tmp1 = tmp0 + args.len; + + dst += start + + i01*args.top_k + + i02*args.top_k*args.ne01 + + i03*args.top_k*args.ne01*args.ne02; + + device const float * src0_row = (device const float *)(src0 + + args.nb01*i01 + + args.nb02*i02 + + args.nb03*i03); + + if (total == 0) { + return; + } + + const int chunk = (total + ntg.x - 1) / ntg.x; + + const int k0 = tpitg.x * chunk; + const int k1 = MIN(MIN(k0 + chunk, total), args.top_k); + + if (k0 >= args.top_k) { + return; + } + + if (k0 >= total) { + return; + } + + int low = k0 > len1 ? k0 - len1 : 0; + int high = MIN(k0, len0); + + // binary-search partition (i, j) such that i + j = k + while (low < high) { + const int mid = (low + high) >> 1; + + const int32_t idx0 = tmp0[mid]; + const int32_t idx1 = tmp1[k0 - mid - 1]; + + const float val0 = src0_row[idx0]; + const float val1 = src0_row[idx1]; + + bool take_left; + if (order == DS4_SORT_ORDER_ASC) { + take_left = (val0 <= val1); + } else { + take_left = (val0 >= val1); + } + + if (take_left) { + low = mid + 1; + } else { + high = mid; + } + } + + int i = low; + int j = k0 - i; + + // keep the merge fronts into registers + int32_t idx0 = 0; + float val0 = 0.0f; + if (i < len0) { + idx0 = tmp0[i]; + val0 = src0_row[idx0]; + } + + int32_t idx1 = 0; + float val1 = 0.0f; + if (j < len1) { + idx1 = tmp1[j]; + val1 = src0_row[idx1]; + } + + for (int k = k0; k < k1; ++k) { + int32_t out_idx; + + if (i >= len0) { + while (k < k1) { + dst[k++] = tmp1[j++]; + } + break; + } else if (j >= len1) { + while (k < k1) { + dst[k++] = tmp0[i++]; + } + break; + } else { + bool take_left; + + if (order == DS4_SORT_ORDER_ASC) { + take_left = (val0 <= val1); + } else { + take_left = (val0 >= val1); + } + + if (take_left) { + out_idx = idx0; + ++i; + if (i < len0) { + idx0 = tmp0[i]; + val0 = src0_row[idx0]; + } + } else { + out_idx = idx1; + ++j; + if (j < len1) { + idx1 = tmp1[j]; + val1 = src0_row[idx1]; + } + } + } + + dst[k] = out_idx; + } +} + +// Host-visible merge variant used by DS4 top-k selection. +template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32; diff --git a/metal/bin.metal b/metal/bin.metal new file mode 100644 index 0000000..ca7def1 --- /dev/null +++ b/metal/bin.metal @@ -0,0 +1,217 @@ +struct ds4_metal_args_bin { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + uint64_t offs; + uint64_t o1[8]; +}; + +struct ds4_metal_args_add3 { + uint32_t n; +}; + +constant short FC_bin_op [[function_constant(FC_BIN + 0)]]; +constant short FC_bin_f [[function_constant(FC_BIN + 1)]]; +constant bool FC_bin_rb [[function_constant(FC_BIN + 2)]]; +constant bool FC_bin_cb [[function_constant(FC_BIN + 3)]]; + +// Generic binary elementwise op with compile-time operation and broadcast +// modes. DS4 currently instantiates this as add, multiply, scalar multiply, and +// row division in the static graph. +template +kernel void kernel_bin_fuse_impl( + constant ds4_metal_args_bin & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { +#define FC_OP FC_bin_op +#define FC_F FC_bin_f +#define FC_RB FC_bin_rb +#define FC_CB FC_bin_cb + + if (FC_RB) { + const uint i0 = tgpig.y*args.ne00 + tgpig.x; + const uint i1 = FC_CB ? tgpig.x%args.ne10 : tgpig.x; + + device const T0 * src0_row = (device const T0 *) (src0); + device T * dst_row = (device T *) (dst); + + if (FC_F == 1) { + device const T1 * src1_row = (device const T1 *) (src1 + args.o1[0]); + + if (FC_OP == 0) { + dst_row[i0] = src0_row[i0] + src1_row[i1]; + } + + if (FC_OP == 1) { + dst_row[i0] = src0_row[i0] - src1_row[i1]; + } + + if (FC_OP == 2) { + dst_row[i0] = src0_row[i0] * src1_row[i1]; + } + + if (FC_OP == 3) { + dst_row[i0] = src0_row[i0] / src1_row[i1]; + } + } else { + T0 res = src0_row[i0]; + + if (FC_OP == 0) { + FOR_UNROLL (short j = 0; j < FC_F; ++j) { + res += ((device const T1 *) (src1 + args.o1[j]))[i1]; + } + } + + if (FC_OP == 1) { + FOR_UNROLL (short j = 0; j < FC_F; ++j) { + res -= ((device const T1 *) (src1 + args.o1[j]))[i1]; + } + } + + if (FC_OP == 2) { + FOR_UNROLL (short j = 0; j < FC_F; ++j) { + res *= ((device const T1 *) (src1 + args.o1[j]))[i1]; + } + } + + if (FC_OP == 3) { + FOR_UNROLL (short j = 0; j < FC_F; ++j) { + res /= ((device const T1 *) (src1 + args.o1[j]))[i1]; + } + } + + dst_row[i0] = res; + } + } else { + const int i03 = tgpig.z; + const int i02 = tgpig.y; + const int i01 = tgpig.x; + + if (i01 >= args.ne01) { + return; + } + + const int i13 = i03%args.ne13; + const int i12 = i02%args.ne12; + const int i11 = i01%args.ne11; + + device const T0 * src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + args.offs); + device T * dst_ptr = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1 + args.offs); + + if (FC_F == 1) { + device const T1 * src1_ptr = (device const T1 *) (src1 + args.o1[0] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11); + + for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) { + const int i10 = FC_CB ? i0%args.ne10 : i0; + + if (FC_OP == 0) { + dst_ptr[i0] = src0_ptr[i0] + src1_ptr[i10]; + } + + if (FC_OP == 1) { + dst_ptr[i0] = src0_ptr[i0] - src1_ptr[i10]; + } + + if (FC_OP == 2) { + dst_ptr[i0] = src0_ptr[i0] * src1_ptr[i10]; + } + + if (FC_OP == 3) { + dst_ptr[i0] = src0_ptr[i0] / src1_ptr[i10]; + } + } + } else { + device const T1 * src1_ptr[8]; + FOR_UNROLL (short j = 0; j < FC_F; ++j) { + src1_ptr[j] = (device const T1 *) (src1 + args.o1[j] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11); + } + + for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) { + const int i10 = FC_CB ? i0%args.ne10 : i0; + + T res = src0_ptr[i0]; + + if (FC_OP == 0) { + FOR_UNROLL (short j = 0; j < FC_F; ++j) { + res += src1_ptr[j][i10]; + } + } + + if (FC_OP == 1) { + FOR_UNROLL (short j = 0; j < FC_F; ++j) { + res -= src1_ptr[j][i10]; + } + } + + if (FC_OP == 2) { + FOR_UNROLL (short j = 0; j < FC_F; ++j) { + res *= src1_ptr[j][i10]; + } + } + + if (FC_OP == 3) { + FOR_UNROLL (short j = 0; j < FC_F; ++j) { + res /= src1_ptr[j][i10]; + } + } + + dst_ptr[i0] = res; + } + } + } + +#undef FC_OP +#undef FC_F +#undef FC_RB +#undef FC_CB +} + +typedef decltype(kernel_bin_fuse_impl) kernel_bin_fuse_t; +// Host-visible F32 binary op; function constants specialize it per use site. +template [[host_name("kernel_bin_fuse_f32_f32_f32")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl; + +kernel void kernel_add2_f32( + constant ds4_metal_args_add3 &args, + device const float *a, + device const float *b, + device float *out, + uint i [[thread_position_in_grid]]) { + if (i >= args.n) return; + out[i] = a[i] + b[i]; +} + +kernel void kernel_add3_f32( + constant ds4_metal_args_add3 &args, + device const float *a, + device const float *b, + device const float *c, + device float *out, + uint i [[thread_position_in_grid]]) { + if (i >= args.n) return; + out[i] = a[i] + b[i] + c[i]; +} diff --git a/metal/concat.metal b/metal/concat.metal new file mode 100644 index 0000000..bfe957f --- /dev/null +++ b/metal/concat.metal @@ -0,0 +1,62 @@ +// DS4 Metal concat kernel used by the graph. + +struct ds4_metal_args_concat { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + int32_t dim; +}; + +// Concatenates two float tensors along one dimension. In DS4 this is a graph +// utility for assembling attention inputs with exactly the same tensor layout +// expected by the downstream kernels. +kernel void kernel_concat( + constant ds4_metal_args_concat & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + const int i3 = tgpig.z; + const int i2 = tgpig.y; + const int i1 = tgpig.x; + + int o[4] = {0, 0, 0, 0}; + o[args.dim] = args.dim == 0 ? args.ne00 : (args.dim == 1 ? args.ne01 : (args.dim == 2 ? args.ne02 : args.ne03)); + + device const float * x; + + for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) { + if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) { + x = (device const float *)(src0 + (i3 )*args.nb03 + (i2 )*args.nb02 + (i1 )*args.nb01 + (i0 )*args.nb00); + } else { + x = (device const float *)(src1 + (i3 - o[3])*args.nb13 + (i2 - o[2])*args.nb12 + (i1 - o[1])*args.nb11 + (i0 - o[0])*args.nb10); + } + + device float * y = (device float *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0); + + *y = *x; + } +} diff --git a/metal/cpy.metal b/metal/cpy.metal new file mode 100644 index 0000000..e8127d8 --- /dev/null +++ b/metal/cpy.metal @@ -0,0 +1,247 @@ +struct ds4_metal_args_cpy { + int64_t nk0; + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int64_t ne0; + int64_t ne1; + int64_t ne2; + int64_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +}; + +// Typed copy/conversion between graph tensors. DS4 uses this for layout +// materialization and F32/F16 conversions at graph boundaries such as KV/cache +// packing and compressor pooling. +template +kernel void kernel_cpy_t_t( + constant ds4_metal_args_cpy & args, + device const char * src0, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + const int i03 = tgpig[2]; + const int i02 = tgpig[1]; + const int i01 = ntg[1] == 1 ? tgpig[0]%args.ne01 : tgpig[0]*ntg[1] + tiitg/ntg[0]; + const int iw0 = ntg[1] == 1 ? tgpig[0]/args.ne01 : 0; + + const int64_t n = i03*args.ne02*args.ne01*args.ne00 + i02*args.ne01*args.ne00 + i01*args.ne00; + + const int64_t i3 = n/(args.ne2*args.ne1*args.ne0); + const int64_t i2 = (n - i3*args.ne2*args.ne1*args.ne0)/(args.ne1*args.ne0); + const int64_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0; + const int64_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0); + + device T1 * dst_data = (device T1 *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0); + + for (int64_t i00 = iw0*ntg[0] + tiitg%ntg[0]; i00 < args.ne00; ) { + device const T0 * src = (device T0 *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00); + dst_data[i00] = (T1) src[0]; + break; + } +} + +typedef decltype(kernel_cpy_t_t) kernel_cpy_t; +// Host-visible copy/conversion variants used by the DS4 graph. +template [[host_name("kernel_cpy_f32_f32")]] kernel kernel_cpy_t kernel_cpy_t_t; +template [[host_name("kernel_cpy_f32_f16")]] kernel kernel_cpy_t kernel_cpy_t_t; +template [[host_name("kernel_cpy_f16_f32")]] kernel kernel_cpy_t kernel_cpy_t_t; +template [[host_name("kernel_cpy_f16_f16")]] kernel kernel_cpy_t kernel_cpy_t_t; + +// Contiguous 1D conversions avoid the generic tensor-index reconstruction +// above. Packed vector types retain scalar alignment, so tensor views whose +// offsets are float/half aligned do not need additional 16/8-byte alignment. +// The final vector is converted element-by-element when n is not divisible by +// four, preserving the generic kernel's bounds and conversion semantics. +kernel void kernel_cpy_contig_f32_f16_4( + constant uint & n, + device const packed_float4 * src, + device packed_half4 * dst, + uint gid [[thread_position_in_grid]]) { + const uint i = gid * 4u; + if (i >= n) { + return; + } + + const uint remaining = n - i; + if (remaining >= 4u) { + const float4 value = float4(src[gid]); + dst[gid] = packed_half4(half4(value)); + return; + } + + device const float * src_scalar = (device const float *)src; + device half * dst_scalar = (device half *)dst; + for (uint lane = 0; lane < remaining; ++lane) { + dst_scalar[i + lane] = half(src_scalar[i + lane]); + } +} + +kernel void kernel_cpy_contig_f16_f32_4( + constant uint & n, + device const packed_half4 * src, + device packed_float4 * dst, + uint gid [[thread_position_in_grid]]) { + const uint i = gid * 4u; + if (i >= n) { + return; + } + + const uint remaining = n - i; + if (remaining >= 4u) { + const half4 value = half4(src[gid]); + dst[gid] = packed_float4(float4(value)); + return; + } + + device const half * src_scalar = (device const half *)src; + device float * dst_scalar = (device float *)dst; + for (uint lane = 0; lane < remaining; ++lane) { + dst_scalar[i + lane] = float(src_scalar[i + lane]); + } +} + +// Bitwise F16 transport for cache staging. Use ushort rather than half so NaN +// payloads and every other binary16 encoding pass through unchanged. +kernel void kernel_cpy_contig_f16_f16_bits_4( + constant uint & n, + device const packed_ushort4 * src, + device packed_ushort4 * dst, + uint gid [[thread_position_in_grid]]) { + const uint i = gid * 4u; + if (i >= n) { + return; + } + + const uint remaining = n - i; + if (remaining >= 4u) { + dst[gid] = src[gid]; + return; + } + + device const ushort * src_scalar = (device const ushort *)src; + device ushort * dst_scalar = (device ushort *)dst; + for (uint lane = 0; lane < remaining; ++lane) { + dst_scalar[i + lane] = src_scalar[i + lane]; + } +} + +struct ds4_metal_args_flash_kv_stage_f16 { + uint raw_cap; + uint raw_start; + uint n_raw; + uint n_comp; + uint pad_rows; + uint shared_pad; +}; + +// Decode-time gathered attention consumes a logical raw-cache ring followed +// by an already-F16 compressed cache. Pack both regions into the contiguous +// F16 FlashAttention scratch in one dispatch. The raw conversion expression +// and compressed ushort4 transport exactly match the standalone copy kernels. +kernel void kernel_dsv4_flash_kv_stage_f16( + constant ds4_metal_args_flash_kv_stage_f16 & args, + device const char * raw_src, + device const char * comp_src, + device char * dst, + device const char * mask_src, + device char * pad_dst, + uint gid [[thread_position_in_grid]]) { + constexpr uint row_vecs = 128; + const uint raw_vecs = args.n_raw * row_vecs; + const uint n_keys = args.n_raw + args.n_comp; + const uint total_vecs = n_keys * row_vecs; + + if (gid < raw_vecs) { + const uint logical_row = gid >> 7; + const uint col = gid & 127u; + uint physical_row = args.raw_start + logical_row; + if (physical_row >= args.raw_cap) { + physical_row -= args.raw_cap; + } + device const packed_float4 *raw = + (device const packed_float4 *)raw_src; + device packed_half4 *dst_half = (device packed_half4 *)dst; + const float4 value = + float4(raw[physical_row * row_vecs + col]); + dst_half[gid] = packed_half4(half4(value)); + return; + } + + if (gid < total_vecs) { + device const packed_ushort4 *comp = + (device const packed_ushort4 *)comp_src; + device packed_ushort4 *dst_bits = (device packed_ushort4 *)dst; + dst_bits[gid] = comp[gid - raw_vecs]; + return; + } + + // The vector FlashAttention kernel redirects its final partial block to + // a compact K/V/mask buffer. When requested, append those writes to this + // dispatch so gathered decode does not need a standalone pad dispatch. + const uint pad_rows = args.pad_rows; + const uint pad_vecs = pad_rows * row_vecs; + uint pad_gid = gid - total_vecs; + if (pad_gid < pad_vecs) { + const uint row = pad_gid >> 7; + const uint col = pad_gid & 127u; + const uint valid_rows = n_keys % pad_rows; + device packed_half4 *pad_half = (device packed_half4 *)pad_dst; + device packed_ushort4 *pad_bits = (device packed_ushort4 *)pad_dst; + if (row >= valid_rows) { + const packed_half4 zero = packed_half4(half4(0.0h)); + pad_half[pad_gid] = zero; + if (!args.shared_pad) { + pad_half[pad_vecs + pad_gid] = zero; + } + return; + } + + const uint logical_row = n_keys - valid_rows + row; + if (logical_row < args.n_raw) { + uint physical_row = args.raw_start + logical_row; + if (physical_row >= args.raw_cap) { + physical_row -= args.raw_cap; + } + device const packed_float4 *raw = + (device const packed_float4 *)raw_src; + const float4 value = + float4(raw[physical_row * row_vecs + col]); + const packed_half4 value_half = packed_half4(half4(value)); + pad_half[pad_gid] = value_half; + if (!args.shared_pad) { + pad_half[pad_vecs + pad_gid] = value_half; + } + } else { + device const packed_ushort4 *comp = + (device const packed_ushort4 *)comp_src; + const packed_ushort4 value_bits = + comp[(logical_row - args.n_raw) * row_vecs + col]; + pad_bits[pad_gid] = value_bits; + if (!args.shared_pad) { + pad_bits[pad_vecs + pad_gid] = value_bits; + } + } + return; + } + + pad_gid -= pad_vecs; + if (pad_gid < pad_rows) { + const uint valid_rows = n_keys % pad_rows; + device const ushort *mask_bits = (device const ushort *)mask_src; + device ushort *pad_mask_bits = + (device ushort *)pad_dst + 2u * pad_vecs * 4u; + pad_mask_bits[pad_gid] = pad_gid < valid_rows + ? mask_bits[n_keys - valid_rows + pad_gid] + : 0xfbffu; + } +} diff --git a/metal/dense.metal b/metal/dense.metal new file mode 100644 index 0000000..ffcdb36 --- /dev/null +++ b/metal/dense.metal @@ -0,0 +1,2389 @@ +// DS4 Metal matvec kernels used by generation. + +constant short FC_mul_mv_nsg [[function_constant(FC_MUL_MV + 0)]]; +constant short FC_mul_mv_nxpsg [[function_constant(FC_MUL_MV + 1)]]; + +struct ds4_metal_args_mul_mv { + int ne00; + int ne01; + int ne02; + ulong nb00; + ulong nb01; + ulong nb02; + ulong nb03; + int ne10; + int ne11; + int ne12; + ulong nb10; + ulong nb11; + ulong nb12; + ulong nb13; + int ne0; + int ne1; + int nr0; + short r2; + short r3; +}; + +struct ds4_metal_args_compressor_pair_store { + uint32_t width; + uint32_t ratio; + uint32_t pos; + uint32_t ape_type; +}; + +struct ds4_metal_args_mul_mm { + int32_t ne00; + int32_t ne02; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int16_t r2; + int16_t r3; +}; + +struct ds4_metal_args_mul_mv_ext { + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int16_t r2; + int16_t r3; +}; + +template +static inline void helper_mv_reduce_and_write( + device float * dst_f32, + float sumf[NR0], + const int r0, + const int ne01, + ushort tiisg, + ushort sgitg, + threadgroup char * shmem) { + constexpr short NW = N_SIMDWIDTH; + + threadgroup float * shmem_f32[NR0]; + + for (short row = 0; row < NR0; ++row) { + shmem_f32[row] = (threadgroup float *) shmem + NW*row; + + if (sgitg == 0) { + shmem_f32[row][tiisg] = 0.0f; + } + + sumf[row] = simd_sum(sumf[row]); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + shmem_f32[row][sgitg] = sumf[row]; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short row = 0; row < NR0 && r0 + row < ne01; ++row) { + float tot = simd_sum(shmem_f32[row][tiisg]); + + if (tiisg == 0 && sgitg == 0) { + dst_f32[r0 + row] = tot; + } + } +} + +template +void kernel_mul_mv_q8_0_f32_impl( + args_t args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + + const int nb = args.ne00/QK8_0; + + const int r0 = tgpig.x*NR0; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + + const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; + + device const float * y = (device const float *) (src1 + offset1); + + device const block_q8_0 * ax[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + + ax[row] = (device const block_q8_0 *) ((device char *) src0 + offset0); + } + + float sumf[NR0] = { 0.f }; + + const short ix = tiisg/(NW/NQ); + const short il = tiisg%(NW/NQ); + + const int ib0 = sgitg*NQ + ix; + + float yl[NQ]; + + device const float * yb = y + ib0*QK8_0 + il*NQ; + + for (int ib = ib0; ib < nb; ib += NSG*NQ) { + for (short i = 0; i < NQ; ++i) { + yl[i] = yb[i]; + } + + for (short row = 0; row < NR0; row++) { + device const int8_t * qs = ax[row][ib].qs + il*NQ; + + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sumq += qs[i] * yl[i]; + } + + sumf[row] += sumq*ax[row][ib].d; + } + + yb += NSG*NQ*QK8_0; + } + + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + + helper_mv_reduce_and_write(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem); +} + +// Decode-time Q8_0 matrix-vector multiply. DS4 uses this for Q8_0 dense +// projections such as shared experts and output-side small matvecs. +[[host_name("kernel_mul_mv_q8_0_f32")]] +kernel void kernel_mul_mv_q8_0_f32( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_q8_0_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +} + +[[host_name("kernel_mul_mv_q8_0_f32_r4")]] +kernel void kernel_mul_mv_q8_0_f32_r4( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_q8_0_f32_impl<4, constant ds4_metal_args_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +} + +// Output projection alias used by the optimized host dispatch. +[[host_name("kernel_mul_mv_q8_0_f32_nr4")]] +kernel void kernel_mul_mv_q8_0_f32_nr4( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_q8_0_f32_impl<4, constant ds4_metal_args_mul_mv &>( + args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +} + +// Decode Q-A/KV pair. Both projections consume the same activation row but +// have independent weight ranges and output extents. Keep the standalone Q8_0 +// lane/block traversal and two-stage reduction verbatim for each bank; only +// the activation load and threadgroup scheduling are shared. +[[host_name("kernel_mul_mv_q8_0_f32_pair")]] +kernel void kernel_mul_mv_q8_0_f32_pair( + constant ds4_metal_args_mul_mv & args0, + constant ds4_metal_args_mul_mv & args1, + device const char * src0_a, + device const char * src0_b, + device const char * src1, + device char * dst_a, + device char * dst_b, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + constexpr short NR0 = N_R0_Q8_0; + + const int nb = args0.ne00 / QK8_0; + const int r0 = tgpig.x * NR0; + const bool active_a = r0 < args0.ne01; + const bool active_b = r0 < args1.ne01; + + device const float *y = (device const float *)src1; + device const block_q8_0 *ax_a[NR0]; + device const block_q8_0 *ax_b[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const int out_row = r0 + row; + ax_a[row] = active_a && out_row < args0.ne01 + ? (device const block_q8_0 *)(src0_a + (uint64_t)out_row * args0.nb01) + : (device const block_q8_0 *)src0_a; + ax_b[row] = active_b && out_row < args1.ne01 + ? (device const block_q8_0 *)(src0_b + (uint64_t)out_row * args1.nb01) + : (device const block_q8_0 *)src0_b; + } + + float suma[NR0] = { 0.f }; + float sumb[NR0] = { 0.f }; + + const short ix = tiisg / (NW / NQ); + const short il = tiisg % (NW / NQ); + const int ib0 = sgitg * NQ + ix; + float yl[NQ]; + device const float *yb = y + ib0 * QK8_0 + il * NQ; + + for (int ib = ib0; ib < nb; ib += NSG * NQ) { + FOR_UNROLL (short i = 0; i < NQ; ++i) { + yl[i] = yb[i]; + } + + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const int out_row = r0 + row; + if (active_a && out_row < args0.ne01) { + device const int8_t *qs = ax_a[row][ib].qs + il * NQ; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sumq += qs[i] * yl[i]; + } + suma[row] += sumq * ax_a[row][ib].d; + } + if (active_b && out_row < args1.ne01) { + device const int8_t *qs = ax_b[row][ib].qs + il * NQ; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sumq += qs[i] * yl[i]; + } + sumb[row] += sumq * ax_b[row][ib].d; + } + } + + yb += NSG * NQ * QK8_0; + } + + threadgroup float *shared = (threadgroup float *)shmem; + threadgroup float *sha[NR0]; + threadgroup float *shb[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + sha[row] = shared + NW * row; + shb[row] = shared + NW * (NR0 + row); + if (sgitg == 0) { + sha[row][tiisg] = 0.0f; + if (active_b) shb[row][tiisg] = 0.0f; + } + suma[row] = simd_sum(suma[row]); + if (active_b) sumb[row] = simd_sum(sumb[row]); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + sha[row][sgitg] = suma[row]; + if (active_b) shb[row][sgitg] = sumb[row]; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *out_a = (device float *)dst_a; + device float *out_b = (device float *)dst_b; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const float total_a = simd_sum(sha[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + const int out_row = r0 + row; + if (active_a && out_row < args0.ne01) out_a[out_row] = total_a; + } + if (active_b) { + const float total_b = simd_sum(shb[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + const int out_row = r0 + row; + if (out_row < args1.ne01) out_b[out_row] = total_b; + } + } + } +} + +// Decode shared-expert gate/up projections followed by SwiGLU: +// +// mid = silu(min(gate, limit)) * clamp(up, -limit, limit) +// +// DS4's shared expert uses two Q8_0 matrices with the same input row. This +// kernel preserves the exact Q8_0 dot-product reduction shape for both +// projections, still writes gate/up for diagnostics, and derives `mid` in the +// same lane that owns the reduced output row. The point is not to fuse two +// independent weight streams into one matmul; it is to remove the separate +// activation pass and its reread of the two 2048-wide rows. +template +void kernel_dsv4_shared_gate_up_swiglu_q8_0_impl( + constant ds4_metal_args_mul_mv & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + constant float &clamp_value, + threadgroup char * shmem, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + + const int nb = args.ne00 / QK8_0; + const int r0 = tgpig.x * NR0; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const uint i12 = im % args.ne12; + const uint i13 = im / args.ne12; + const uint64_t offset1 = r1 * args.nb11 + i12 * args.nb12 + i13 * args.nb13; + device const float *y = (device const float *)(src1 + offset1); + + device const block_q8_0 *ag[NR0]; + device const block_q8_0 *au[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row) * args.nb01 + + (i12 / args.r2) * args.nb02 + + (i13 / args.r3) * args.nb03; + ag[row] = (device const block_q8_0 *)((device const char *)src0_gate + offset0); + au[row] = (device const block_q8_0 *)((device const char *)src0_up + offset0); + } + + float sumg[NR0] = { 0.f }; + float sumu[NR0] = { 0.f }; + + const short ix = tiisg / (NW / NQ); + const short il = tiisg % (NW / NQ); + const int ib0 = sgitg * NQ + ix; + float yl[NQ]; + device const float *yb = y + ib0 * QK8_0 + il * NQ; + + for (int ib = ib0; ib < nb; ib += NSG * NQ) { + FOR_UNROLL (short i = 0; i < NQ; ++i) { + yl[i] = yb[i]; + } + + FOR_UNROLL (short row = 0; row < NR0; ++row) { + device const int8_t *qg = ag[row][ib].qs + il * NQ; + device const int8_t *qu = au[row][ib].qs + il * NQ; + + float sg = 0.f; + float su = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sg += qg[i] * yl[i]; + su += qu[i] * yl[i]; + } + + sumg[row] += sg * ag[row][ib].d; + sumu[row] += su * au[row][ib].d; + } + + yb += NSG * NQ * QK8_0; + } + + threadgroup float *shmem_f32 = (threadgroup float *)shmem; + threadgroup float *sh_gate[NR0]; + threadgroup float *sh_up[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + sh_gate[row] = shmem_f32 + NW * row; + sh_up[row] = shmem_f32 + NW * (NR0 + row); + if (sgitg == 0) { + sh_gate[row][tiisg] = 0.0f; + sh_up[row][tiisg] = 0.0f; + } + sumg[row] = simd_sum(sumg[row]); + sumu[row] = simd_sum(sumu[row]); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + sh_gate[row][sgitg] = sumg[row]; + sh_up[row][sgitg] = sumu[row]; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *gate_f32 = (device float *)dst_gate + + (uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0; + device float *up_f32 = (device float *)dst_up + + (uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0; + device float *mid_f32 = (device float *)dst_mid + + (uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0; + + FOR_UNROLL (short row = 0; row < NR0 && r0 + row < args.ne01; ++row) { + const float gate = simd_sum(sh_gate[row][tiisg]); + const float up = simd_sum(sh_up[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + const uint out_row = r0 + row; + if (STORE_GATE_UP) { + gate_f32[out_row] = gate; + up_f32[out_row] = up; + } + float g = gate; + float u = up; + if (clamp_value > 1.0e-6f) { + g = min(g, clamp_value); + u = clamp(u, -clamp_value, clamp_value); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u; + } + } +} + +[[host_name("kernel_dsv4_shared_gate_up_swiglu_q8_0")]] +kernel void kernel_dsv4_shared_gate_up_swiglu_q8_0( + constant ds4_metal_args_mul_mv & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + constant float &clamp_value, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_dsv4_shared_gate_up_swiglu_q8_0_impl( + args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid, + clamp_value, shmem, tgpig, tiisg, sgitg); +} + +[[host_name("kernel_dsv4_shared_gate_up_swiglu_q8_0_r4")]] +kernel void kernel_dsv4_shared_gate_up_swiglu_q8_0_r4( + constant ds4_metal_args_mul_mv & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + constant float &clamp_value, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<4, true>( + args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid, + clamp_value, shmem, tgpig, tiisg, sgitg); +} + +[[host_name("kernel_dsv4_shared_mid_swiglu_q8_0")]] +kernel void kernel_dsv4_shared_mid_swiglu_q8_0( + constant ds4_metal_args_mul_mv & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + constant float &clamp_value, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_dsv4_shared_gate_up_swiglu_q8_0_impl( + args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid, + clamp_value, shmem, tgpig, tiisg, sgitg); +} + +[[host_name("kernel_dsv4_shared_mid_swiglu_q8_0_r4")]] +kernel void kernel_dsv4_shared_mid_swiglu_q8_0_r4( + constant ds4_metal_args_mul_mv & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + constant float &clamp_value, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_dsv4_shared_gate_up_swiglu_q8_0_impl<4, false>( + args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid, + clamp_value, shmem, tgpig, tiisg, sgitg); +} + +template +void kernel_mul_mv_t_t_impl( + args_t args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + + constexpr short NW = N_SIMDWIDTH; + constexpr short NB = 32; + constexpr short NF = 8; + + const int nb = args.ne00/NB; + + const int r0 = tgpig.x*NR0; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + + const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; + + device const T1 * y = (device const T1 *) (src1 + offset1); + + device const T0 * ax[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + + ax[row] = (device const T0 *) ((device char *) src0 + offset0); + } + + float sumf[NR0] = { 0.f }; + + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + + const int ib0 = sgitg*NF + ix; + + T1 yl[NF]; + + device const T1 * yb = y + (ib0*NB + il*NF); + + for (int ib = ib0; ib < nb; ib += NSG*NF) { + for (short i = 0; i < NF; ++i) { + yl[i] = yb[i]; + } + + for (short row = 0; row < NR0; row++) { + device const T0 * xb = ax[row] + (ib*NB + il*NF); + + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NF; ++i) { + sumq += xb[i] * yl[i]; + } + + sumf[row] += sumq; + } + + yb += NSG*NF*NW; + } + + for (int i = nb*NB + sgitg*NW + tiisg; i < args.ne00; i += NW*NSG) { + for (short row = 0; row < NR0; row++) { + sumf[row] += ax[row][i] * y[i]; + } + } + + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + + helper_mv_reduce_and_write(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem); +} + +template +void kernel_mul_mv_t_t_disp( + args_t args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + switch (args.nr0) { + case 2: kernel_mul_mv_t_t_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break; + case 4: kernel_mul_mv_t_t_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break; + } +} + +// Decode-time dense F32/F16 matrix-vector multiply. The instantiated kernels +// handle unquantized DS4 weights and activations that are already float rows. +template +kernel void kernel_mul_mv_t_t( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_t_t_disp(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +} + +typedef decltype(kernel_mul_mv_t_t) mul_mv_t_t; + +// Host-visible dense matvec variants used by the graph for F32 and F16 weights. +template [[host_name("kernel_mul_mv_f32_f32")]] kernel mul_mv_t_t kernel_mul_mv_t_t; +template [[host_name("kernel_mul_mv_f16_f32")]] kernel mul_mv_t_t kernel_mul_mv_t_t; + +template +void kernel_mul_mv_t_t_4_impl( + args_t args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + + constexpr short NW = N_SIMDWIDTH; + constexpr short NB = 32; + constexpr short NF = 16; + constexpr short NF4 = NF/4; + + const int nb = args.ne00/NB; + + const int r0 = tgpig.x*NR0; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + + const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; + + device const T1 * y = (device const T1 *) (src1 + offset1); + device const T14 * y4 = (device const T14 *) (src1 + offset1); + + device const T0 * ax [NR0]; + device const T04 * ax4[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + + ax [row] = (device const T0 *) ((device char *) src0 + offset0); + ax4[row] = (device const T04 *) ((device char *) src0 + offset0); + } + + float sumf[NR0] = { 0.f }; + + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + + const int ib0 = sgitg*NF + ix; + + T14 yl4[NF4]; + + device const T14 * yb4 = y4 + (ib0*NB + il*NF)/4; + + for (int ib = ib0; ib < nb; ib += NSG*NF) { + for (short i = 0; i < NF4; ++i) { + yl4[i] = yb4[i]; + } + + for (short row = 0; row < NR0; row++) { + device const T04 * xb4 = ax4[row] + (ib*NB + il*NF)/4; + + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NF4; ++i) { + sumq += dot(float4(xb4[i]), float4(yl4[i])); + } + + sumf[row] += sumq; + } + + yb4 += NSG*NF*NW/4; + } + + for (int i = nb*NB + sgitg*NW + tiisg; i < args.ne00; i += NW*NSG) { + for (short row = 0; row < NR0; row++) { + sumf[row] += ax[row][i] * y[i]; + } + } + + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + + helper_mv_reduce_and_write(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem); +} + +template +void kernel_mul_mv_t_t_4_disp( + args_t args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + switch (args.nr0) { + case 2: kernel_mul_mv_t_t_4_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break; + case 4: kernel_mul_mv_t_t_4_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break; + }; +} + +// Vectorized dense matvec using float4/half4 loads. DS4 uses this where the +// inner dimension and alignment make vector loads cheaper than scalar lanes. +template +kernel void kernel_mul_mv_t_t_4( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_t_t_4_disp(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +} + +typedef decltype(kernel_mul_mv_t_t_4) mul_mv_t_t_4; + +// Host-visible vectorized dense matvec variants for F32 and F16 weights. +template [[host_name("kernel_mul_mv_f32_f32_4")]] kernel mul_mv_t_t_4 kernel_mul_mv_t_t_4; +template [[host_name("kernel_mul_mv_f16_f32_4")]] kernel mul_mv_t_t_4 kernel_mul_mv_t_t_4; + +// DS4 compressor projections always compute two same-shaped F16 matvecs from +// the same normalized activation: one for projected KV and one for pooling +// scores. This paired variant keeps the exact dense F16 row-reduction shape +// for each matrix, but shares one dispatch and one activation stream. +template +void kernel_mul_mv_f16_f32_pair_4_impl( + args_t args, + device const char * src0_a, + device const char * src0_b, + device const char * src1, + device char * dst_a, + device char * dst_b, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + + constexpr short NW = N_SIMDWIDTH; + constexpr short NB = 32; + constexpr short NF = 16; + constexpr short NF4 = NF/4; + + const int nb = args.ne00/NB; + + const int r0 = tgpig.x*NR0; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + + const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; + + device const float * y = (device const float *) (src1 + offset1); + device const float4 * y4 = (device const float4 *) (src1 + offset1); + + device const half * ax_a [NR0]; + device const half4 * ax4_a[NR0]; + device const half * ax_b [NR0]; + device const half4 * ax4_b[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + + ax_a [row] = (device const half *) ((device char *) src0_a + offset0); + ax4_a[row] = (device const half4 *) ((device char *) src0_a + offset0); + ax_b [row] = (device const half *) ((device char *) src0_b + offset0); + ax4_b[row] = (device const half4 *) ((device char *) src0_b + offset0); + } + + float sum_a[NR0] = { 0.f }; + float sum_b[NR0] = { 0.f }; + + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + + const int ib0 = sgitg*NF + ix; + + float4 yl4[NF4]; + + device const float4 * yb4 = y4 + (ib0*NB + il*NF)/4; + + for (int ib = ib0; ib < nb; ib += NSG*NF) { + for (short i = 0; i < NF4; ++i) { + yl4[i] = yb4[i]; + } + + for (short row = 0; row < NR0; row++) { + device const half4 * xb4_a = ax4_a[row] + (ib*NB + il*NF)/4; + device const half4 * xb4_b = ax4_b[row] + (ib*NB + il*NF)/4; + + float suma = 0.f; + float sumb = 0.f; + FOR_UNROLL (short i = 0; i < NF4; ++i) { + const float4 yv = float4(yl4[i]); + suma += dot(float4(xb4_a[i]), yv); + sumb += dot(float4(xb4_b[i]), yv); + } + + sum_a[row] += suma; + sum_b[row] += sumb; + } + + yb4 += NSG*NF*NW/4; + } + + for (int i = nb*NB + sgitg*NW + tiisg; i < args.ne00; i += NW*NSG) { + for (short row = 0; row < NR0; row++) { + const float yi = y[i]; + sum_a[row] += ax_a[row][i] * yi; + sum_b[row] += ax_b[row][i] * yi; + } + } + + device float * dst_a_f32 = (device float *) dst_a + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + device float * dst_b_f32 = (device float *) dst_b + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + + helper_mv_reduce_and_write(dst_a_f32, sum_a, r0, args.ne01, tiisg, sgitg, shmem); + threadgroup_barrier(mem_flags::mem_threadgroup); + helper_mv_reduce_and_write(dst_b_f32, sum_b, r0, args.ne01, tiisg, sgitg, shmem); +} + +template +void kernel_mul_mv_f16_f32_pair_4_disp( + args_t args, + device const char * src0_a, + device const char * src0_b, + device const char * src1, + device char * dst_a, + device char * dst_b, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + switch (args.nr0) { + case 2: kernel_mul_mv_f16_f32_pair_4_impl<2>(args, src0_a, src0_b, src1, dst_a, dst_b, shmem, tgpig, tiisg, sgitg); break; + case 4: kernel_mul_mv_f16_f32_pair_4_impl<4>(args, src0_a, src0_b, src1, dst_a, dst_b, shmem, tgpig, tiisg, sgitg); break; + } +} + +kernel void kernel_mul_mv_f16_f32_pair_4( + constant ds4_metal_args_mul_mv & args, + device const char * src0_a, + device const char * src0_b, + device const char * src1, + device char * dst_a, + device char * dst_b, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_f16_f32_pair_4_disp( + args, src0_a, src0_b, src1, dst_a, dst_b, shmem, tgpig, tiisg, sgitg); +} + +// Decode compressor projection plus recurrent-state append. The paired +// matvec remains unchanged and still materializes both F32 outputs. After a +// device-memory barrier, the first NR0 threads reload those exact stored bits +// and perform the same state write and score+APE addition as +// kernel_dsv4_compressor_store_one. +kernel void kernel_mul_mv_f16_f32_pair_compressor_store_4( + constant ds4_metal_args_mul_mv & args, + constant ds4_metal_args_compressor_pair_store & store, + device const char * src0_a, + device const char * src0_b, + device const char * src1, + device char * dst_a, + device char * dst_b, + device const char * ape, + device float * state_kv, + device float * state_score, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_f16_f32_pair_4_disp( + args, src0_a, src0_b, src1, dst_a, dst_b, + shmem, tgpig, tiisg, sgitg); + + threadgroup_barrier(mem_flags::mem_device); + + if (tiitg >= args.nr0 || store.width == 0u || store.ratio == 0u) { + return; + } + const uint col = tgpig.x * (uint)args.nr0 + tiitg; + if (col >= store.width) return; + + const uint pos_mod = store.pos % store.ratio; + const uint dst_row = store.ratio == 4u ? store.ratio + pos_mod : pos_mod; + const uint dst = dst_row * store.width + col; + const uint ape_i = pos_mod * store.width + col; + + device volatile const float * projected_kv = + (device volatile const float *)dst_a; + device volatile const float * projected_score = + (device volatile const float *)dst_b; + float ape_v; + if (store.ape_type == 1u) { + ape_v = (float)(((device const half *)ape)[ape_i]); + } else { + ape_v = ((device const float *)ape)[ape_i]; + } + + state_kv[dst] = projected_kv[col]; + state_score[dst] = projected_score[col] + ape_v; +} + +template +void kernel_mul_mv_t_t_short_impl( + args_t args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig, + ushort tiisg) { + const int r0 = tgpig.x*32 + tiisg; + const int r1 = tgpig.y; + const int im = tgpig.z; + + if (r0 >= args.ne01) { + return; + } + + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + + const uint64_t offset0 = r0*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + + device const T0 * x = (device const T0 *) (src0 + offset0); + + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1; + + const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; + + device const T1 * y = (device const T1 *) (src1 + offset1); + + float res = 0.0f; + + for (int i = 0; i < args.ne00; ++i) { + res += (float) x[i] * (float) y[i]; + } + + dst_f32[(uint64_t)r1*args.ne0 + r0] = res; +} + +// Scalar fallback for short rows. It trades parallelism for lower dispatch and +// reduction overhead when DS4 asks for tiny dense matvecs. +template +kernel void kernel_mul_mv_t_t_short( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]]) { + kernel_mul_mv_t_t_short_impl( + args, + src0, + src1, + dst, + tgpig, + tiisg); +} + +typedef decltype(kernel_mul_mv_t_t_short) mul_mv_t_t_short_t; + +// Host-visible short-row dense matvec variants. +template [[host_name("kernel_mul_mv_f32_f32_short")]] kernel mul_mv_t_t_short_t kernel_mul_mv_t_t_short; +template [[host_name("kernel_mul_mv_f16_f32_short")]] kernel mul_mv_t_t_short_t kernel_mul_mv_t_t_short; + +template +void dequantize_f32(device const float4x4 * src, short il, thread type4x4 & reg) { + reg = (type4x4)(*src); +} + +template +void dequantize_f16(device const half4x4 * src, short il, thread type4x4 & reg) { + reg = (type4x4)(*src); +} + +template +void dequantize_q8_0(device const block_q8_0 *xb, short il, thread type4x4 & reg) { + device const int8_t * qs = ((device const int8_t *)xb->qs); + const float d = xb->d; + + float4x4 reg_f; + + for (int i = 0; i < 16; i++) { + reg_f[i/4][i%4] = (qs[i + 16*il] * d); + } + + reg = (type4x4) reg_f; +} + +struct ds4_dense_block_q4_0 { + half d; + uchar qs[16]; +}; + +struct ds4_dense_block_q4_K { + half d; + half dmin; + uchar scales[12]; + uchar qs[128]; +}; + +static inline uchar2 ds4_dense_q4_K_scale_min(int j, int k, device const uchar *q) { + return j < 4 ? uchar2{uchar(q[j + 0 + k] & 63), uchar(q[j + 4 + k] & 63)} + : uchar2{uchar((q[j + 4 + k] & 0x0f) | ((q[j - 4 + k] & 0xc0) >> 2)), + uchar((q[j + 4 + k] >> 4) | ((q[j - 0 + k] & 0xc0) >> 2))}; +} + +template +void dequantize_dense_q4_0(device const ds4_dense_block_q4_0 *xb, short il, thread type4x4 ®) { + float4x4 reg_f; + const float d = (float)xb->d; + const int base = 16 * (int)il; + for (int i = 0; i < 16; i++) { + const int k = base + i; + /* ggml Q4_0: elems 0..15 = low nibbles of qs[0..15], 16..31 = high. */ + const uchar packed = xb->qs[(uint)(k & 15)]; + const uchar q = (k < 16) ? (packed & 0x0f) : (packed >> 4); + reg_f[i / 4][i % 4] = d * ((float)q - 8.0f); + } + reg = (type4x4)reg_f; +} + +template +void dequantize_dense_q4_K(device const ds4_dense_block_q4_K *xb, short il, thread type4x4 ®) { + device const uchar *q = xb->qs; + short is = (il / 4) * 2; + q = q + (il / 4) * 32 + 16 * (il & 1); + il = il & 3; + const uchar2 sc = ds4_dense_q4_K_scale_min(is, il / 2, xb->scales); + const float d = il < 2 ? (float)xb->d : (float)xb->d * (1.0f / 16.0f); + const float min = (float)xb->dmin; + const float dl = d * sc[0]; + const float ml = min * sc[1]; + const ushort mask = il < 2 ? 0x0F : 0xF0; + + float4x4 reg_f; + for (int i = 0; i < 16; ++i) { + reg_f[i / 4][i % 4] = dl * (q[i] & mask) - ml; + } + reg = (type4x4)reg_f; +} + +/* + * Bit-identical twin of dequantize_q8_0 for the MPP staging loop: same + * half(float(qs[i]) * d) per element, but the 16 consecutive int8 lanes are + * fetched as eight aligned 16-bit loads instead of sixteen byte loads. + * xb->qs is always 2-byte aligned (block_q8_0.d is a half). + */ +void dequantize_q8_0_pairs(device const block_q8_0 *xb, short il, thread half4x4 & reg) { + device const ushort *qs16 = (device const ushort *)(xb->qs + 16*il); + const float d = xb->d; + + float4x4 reg_f; + + FOR_UNROLL (short i = 0; i < 8; i++) { + const ushort u = qs16[i]; + reg_f[i/2][(i%2)*2 + 0] = ((float)(int8_t)(u & 0xFF)) * d; + reg_f[i/2][(i%2)*2 + 1] = ((float)(int8_t)(u >> 8)) * d; + } + + reg = (half4x4) reg_f; +} + +template +void dequantize_q8_0_t4(device const block_q8_0 *xb, short il, thread type4 & reg) { + device const int8_t * qs = ((device const int8_t *)xb->qs); + const float d = xb->d; + + for (int i = 0; i < 4; i++) { + reg[i] = (qs[4*(il%4) + i + 16*(il/4)] * d); + } +} + +template +void dequantize_dense_q4_0_t4(device const ds4_dense_block_q4_0 *xb, short il, thread type4 ®) { + const float d = (float)xb->d; + const int base = 4 * (int)il; + for (int i = 0; i < 4; i++) { + const int k = base + i; + /* ggml Q4_0: elems 0..15 = low nibbles of qs[0..15], 16..31 = high. */ + const uchar packed = xb->qs[(uint)(k & 15)]; + const uchar q = (k < 16) ? (packed & 0x0f) : (packed >> 4); + reg[i] = d * ((float)q - 8.0f); + } +} + +template +void dequantize_dense_q4_K_t4(device const ds4_dense_block_q4_K *xb, short il, thread type4 ®) { + float4x4 tmp; + dequantize_dense_q4_K(xb, il / 4, tmp); + const short row = il & 3; + for (int i = 0; i < 4; i++) { + reg[i] = tmp[row][i]; + } +} + +// DS4 small-batch mat-vec kernel used for 2..8 prompt tokens. +template +void kernel_mul_mv_ext_q4_f32_impl( + constant ds4_metal_args_mul_mv_ext & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nxpsg = FC_mul_mv_nxpsg; + + const short chpt = 4; // chunks per thread + + const short nypsg = (32/nxpsg); + + const short tx = tiisg%nxpsg; + const short ty = tiisg/nxpsg; + + const int i01 = tgpig.x*(nypsg*NSG) + nypsg*sgitg + ty; + const int i11 = tgpig.y*r1ptg; + const int i1m = tgpig.z; + + const int i12 = i1m%args.ne12; + const int i13 = i1m/args.ne12; + + const uint64_t offset0 = i01*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + const uint64_t offset1 = i11*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; + + device const q_t * xq = (i01 < args.ne01) ? (device const q_t *) (src0 + offset0) + tx/chpb : (device const q_t *) src0; + + device const float4 * y4[r1ptg]; + + for (int ir1 = 0; ir1 < r1ptg; ++ir1) { + y4[ir1] = (i11 + ir1 < args.ne11) ? (device const float4 *) (src1 + offset1 + ir1*args.nb11) + tx : (device const float4 *) src1; + } + + float sumf[r1ptg] = { [ 0 ... r1ptg - 1 ] = 0.0f }; + + short cch = tx%chpb; // current chunk index + + for (int ich = tx; 4*ich < args.ne00; ich += chpt*nxpsg) { + float4 lx[chpt]; + +#pragma unroll(chpt) + for (short ch = 0; ch < chpt; ++ch) { + deq_t4(xq, cch, lx[ch]); + + cch += nxpsg; + if (cch >= chpb) { + xq += cch/chpb; + cch %= chpb; + } + } + +#pragma unroll(chpt) + for (short ch = 0; ch < chpt; ++ch) { +#pragma unroll(r1ptg) + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + sumf[ir1] += dot(lx[ch], y4[ir1][ch*nxpsg]); + } + } + +#pragma unroll(r1ptg) + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + y4[ir1] += chpt*nxpsg; + } + } + + // reduce only the threads in each row + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + if (nxpsg >= 32) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 16); + } + if (nxpsg >= 16) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 8); + } + if (nxpsg >= 8) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 4); + } + if (nxpsg >= 4) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 2); + } + if (nxpsg >= 2) { + sumf[ir1] += simd_shuffle_down(sumf[ir1], 1); + } + } + + if (tx == 0) { + for (short ir1 = 0; ir1 < r1ptg && i11 + ir1 < args.ne11; ++ir1) { + device float * dst_f32 = (device float *) dst + (uint64_t)i1m*args.ne0*args.ne1 + (uint64_t)(i11 + ir1)*args.ne0; + + if (i01 < args.ne01) { + dst_f32[i01] = sumf[ir1]; + } + } + } +} + +// Small-batch prompt matvec for 2..5 tokens. It bridges decode-style matvec and +// full matmul when DS4 prefill chunks are too small to amortize matrix tiles. +template +kernel void kernel_mul_mv_ext_q4_f32_disp( + constant ds4_metal_args_mul_mv_ext & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_ext_q4_f32_impl(args, src0, src1, dst, tgpig, tiisg, sgitg); +} + +typedef decltype(kernel_mul_mv_ext_q4_f32_disp<2, block_q8_0, 32, dequantize_q8_0_t4>) mul_mv_ext_q4_f32_t; + +// Host-visible small-batch variants for r1=2..5 during tiny prompt/support +// paths. +template [[host_name("kernel_mul_mv_ext_f32_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, float4, 4, dequantize_f32_t4>; +template [[host_name("kernel_mul_mv_ext_f32_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, float4, 4, dequantize_f32_t4>; +template [[host_name("kernel_mul_mv_ext_f32_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, float4, 4, dequantize_f32_t4>; +template [[host_name("kernel_mul_mv_ext_f32_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, float4, 4, dequantize_f32_t4>; + +template +void kernel_mul_mv_ext_q8_0_pair_swiglu_f32_impl( + constant ds4_metal_args_mul_mv_ext & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + constant float &clamp_value, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nxpsg = FC_mul_mv_nxpsg; + + const short chpt = 4; + const short chpb = 8; + const short nypsg = (32/nxpsg); + + const short tx = tiisg%nxpsg; + const short ty = tiisg/nxpsg; + + const int i01 = tgpig.x*(nypsg*NSG) + nypsg*sgitg + ty; + const int i11 = tgpig.y*r1ptg; + const int i1m = tgpig.z; + + const int i12 = i1m%args.ne12; + const int i13 = i1m/args.ne12; + + const uint64_t offset0 = i01*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + const uint64_t offset1 = i11*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; + + device const block_q8_0 * xq_gate = + (i01 < args.ne01) ? (device const block_q8_0 *)(src0_gate + offset0) + tx/chpb + : (device const block_q8_0 *)src0_gate; + device const block_q8_0 * xq_up = + (i01 < args.ne01) ? (device const block_q8_0 *)(src0_up + offset0) + tx/chpb + : (device const block_q8_0 *)src0_up; + + device const float4 * y4[r1ptg]; + for (int ir1 = 0; ir1 < r1ptg; ++ir1) { + y4[ir1] = (i11 + ir1 < args.ne11) + ? (device const float4 *)(src1 + offset1 + ir1*args.nb11) + tx + : (device const float4 *)src1; + } + + float sum_gate[r1ptg] = { [ 0 ... r1ptg - 1 ] = 0.0f }; + float sum_up[r1ptg] = { [ 0 ... r1ptg - 1 ] = 0.0f }; + + short cch = tx%chpb; + for (int ich = tx; 4*ich < args.ne00; ich += chpt*nxpsg) { + float4 lg[chpt]; + float4 lu[chpt]; + +#pragma unroll(chpt) + for (short ch = 0; ch < chpt; ++ch) { + dequantize_q8_0_t4(xq_gate, cch, lg[ch]); + dequantize_q8_0_t4(xq_up, cch, lu[ch]); + + cch += nxpsg; + if (cch >= chpb) { + xq_gate += cch/chpb; + xq_up += cch/chpb; + cch %= chpb; + } + } + +#pragma unroll(chpt) + for (short ch = 0; ch < chpt; ++ch) { +#pragma unroll(r1ptg) + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + const float4 y = y4[ir1][ch*nxpsg]; + sum_gate[ir1] += dot(lg[ch], y); + sum_up[ir1] += dot(lu[ch], y); + } + } + +#pragma unroll(r1ptg) + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + y4[ir1] += chpt*nxpsg; + } + } + + for (short ir1 = 0; ir1 < r1ptg; ++ir1) { + if (nxpsg >= 32) { + sum_gate[ir1] += simd_shuffle_down(sum_gate[ir1], 16); + sum_up[ir1] += simd_shuffle_down(sum_up[ir1], 16); + } + if (nxpsg >= 16) { + sum_gate[ir1] += simd_shuffle_down(sum_gate[ir1], 8); + sum_up[ir1] += simd_shuffle_down(sum_up[ir1], 8); + } + if (nxpsg >= 8) { + sum_gate[ir1] += simd_shuffle_down(sum_gate[ir1], 4); + sum_up[ir1] += simd_shuffle_down(sum_up[ir1], 4); + } + if (nxpsg >= 4) { + sum_gate[ir1] += simd_shuffle_down(sum_gate[ir1], 2); + sum_up[ir1] += simd_shuffle_down(sum_up[ir1], 2); + } + if (nxpsg >= 2) { + sum_gate[ir1] += simd_shuffle_down(sum_gate[ir1], 1); + sum_up[ir1] += simd_shuffle_down(sum_up[ir1], 1); + } + } + + if (tx == 0 && i01 < args.ne01) { + for (short ir1 = 0; ir1 < r1ptg && i11 + ir1 < args.ne11; ++ir1) { + const uint64_t dst_base = + (uint64_t)i1m*args.ne0*args.ne1 + (uint64_t)(i11 + ir1)*args.ne0; + device float * gate_f32 = (device float *)dst_gate + dst_base; + device float * up_f32 = (device float *)dst_up + dst_base; + device float * mid_f32 = (device float *)dst_mid + dst_base; + + const float gate = sum_gate[ir1]; + const float up = sum_up[ir1]; + gate_f32[i01] = gate; + up_f32[i01] = up; + + float g = gate; + float u = up; + if (clamp_value > 1.0e-6f) { + g = min(g, clamp_value); + u = clamp(u, -clamp_value, clamp_value); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[i01] = silu * u; + } + } +} + +template +kernel void kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp( + constant ds4_metal_args_mul_mv_ext & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + constant float &clamp_value, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_ext_q8_0_pair_swiglu_f32_impl( + args, src0_gate, src0_up, src1, dst_gate, dst_up, dst_mid, clamp_value, + tgpig, tiisg, sgitg); +} + +typedef decltype(kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp<2>) mul_mv_ext_q8_0_pair_swiglu_f32_t; + +// Host-visible small-batch variants. DS4 currently needs F16 and Q8_0 weights +// for r1=2..5 during the prompt path. +template [[host_name("kernel_mul_mv_ext_f16_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, half4, 4, dequantize_f16_t4>; +template [[host_name("kernel_mul_mv_ext_f16_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, half4, 4, dequantize_f16_t4>; +template [[host_name("kernel_mul_mv_ext_f16_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, half4, 4, dequantize_f16_t4>; +template [[host_name("kernel_mul_mv_ext_f16_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, half4, 4, dequantize_f16_t4>; + +template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q8_0, 32, dequantize_q8_0_t4>; +template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q8_0, 32, dequantize_q8_0_t4>; +template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q8_0, 32, dequantize_q8_0_t4>; +template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q8_0, 32, dequantize_q8_0_t4>; + +template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_1")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<1, ds4_dense_block_q4_0, 32, dequantize_dense_q4_0_t4>; +template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, ds4_dense_block_q4_0, 32, dequantize_dense_q4_0_t4>; +template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, ds4_dense_block_q4_0, 32, dequantize_dense_q4_0_t4>; +template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, ds4_dense_block_q4_0, 32, dequantize_dense_q4_0_t4>; +template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, ds4_dense_block_q4_0, 32, dequantize_dense_q4_0_t4>; + +template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_1")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<1, ds4_dense_block_q4_K, 256, dequantize_dense_q4_K_t4>; +template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, ds4_dense_block_q4_K, 256, dequantize_dense_q4_K_t4>; +template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, ds4_dense_block_q4_K, 256, dequantize_dense_q4_K_t4>; +template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, ds4_dense_block_q4_K, 256, dequantize_dense_q4_K_t4>; +template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, ds4_dense_block_q4_K, 256, dequantize_dense_q4_K_t4>; + +template [[host_name("kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_2")]] kernel mul_mv_ext_q8_0_pair_swiglu_f32_t kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp<2>; +template [[host_name("kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_3")]] kernel mul_mv_ext_q8_0_pair_swiglu_f32_t kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp<3>; +template [[host_name("kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_4")]] kernel mul_mv_ext_q8_0_pair_swiglu_f32_t kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp<4>; +template [[host_name("kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_5")]] kernel mul_mv_ext_q8_0_pair_swiglu_f32_t kernel_mul_mv_ext_q8_0_pair_swiglu_f32_disp<5>; + +constant bool FC_mul_mm_bc_inp [[function_constant(FC_MUL_MM + 0)]]; +constant bool FC_mul_mm_bc_out [[function_constant(FC_MUL_MM + 1)]]; + +#ifdef DS4_METAL_HAS_TENSOR +template< + short NR0, short NR1, + typename SA, typename SA_4x4, typename block_q, short nl, + void (*dequantize_func)(device const block_q *, short, thread SA_4x4 &), + typename T0, typename T0_4x4, typename T1> +kernel void kernel_mul_mm_mpp( + constant ds4_metal_args_mul_mm & args, + device const char * srcA, + device const char * srcB, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + (void) sgitg; + + constexpr int NK = 32; + constexpr int NL = NK/16; + constexpr int NUM_THREADS = 128; + + const int K = args.ne00; + const int M = args.ne0; + const int N = args.ne1; + const int im = tgpig.z; + const int i12 = im%args.ne12; + const int i13 = im/args.ne12; + const int r0 = tgpig.y*NR0; + const int r1 = tgpig.x*NR1; + + const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + + threadgroup SA *sa = (threadgroup SA *)shmem; + threadgroup SA *sb = sa + NR0*NK; + auto tA = tensor(sa, dextents(NK, NR0)); + auto tB = tensor(sb, dextents(NK, NR1)); + + device const T1 *ptrB = (device const T1 *)(srcB + args.nb12*i12 + args.nb13*i13); + const int strideB = args.nb11/sizeof(T1); + + matmul2d< + matmul2d_descriptor(NR1, NR0, NK, false, true, false, + matmul2d_descriptor::mode::multiply_accumulate), + execution_simdgroups<4>> mm; + + auto cT = mm.template get_destination_cooperative_tensor(); + + #pragma unroll + for (uint16_t i = 0; i < cT.get_capacity(); ++i) { + if (cT.is_valid_element(i)) { + cT[i] = 0.0f; + } + } + + for (int loop_k = 0; loop_k < K; loop_k += NK) { + for (int work = tiitg; work < NR0*NL; work += NUM_THREADS) { + const int row = work/NL; + const int k_chunk = work%NL; + const int k_pos = loop_k + k_chunk*16; + const short k_base = k_chunk*16; + + if (!FC_mul_mm_bc_out || r0 + row < M) { + if (is_same::value && FC_mul_mm_bc_inp) { + device const T0 *row_ptr = (device const T0 *)(srcA + args.nb01*(r0 + row) + offset0); + FOR_UNROLL (short i = 0; i < 16; i++) { + sa[row*NK + k_base + i] = (k_pos + i < K) ? (SA)row_ptr[k_pos + i] : (SA)0; + } + } else { + const int block_idx = k_pos/(16*nl); + const short il = (k_pos/16)%nl; + device const block_q *row_ptr = (device const block_q *)(srcA + args.nb01*(r0 + row) + offset0); + + SA_4x4 temp_a; + dequantize_func(row_ptr + block_idx, il, temp_a); + FOR_UNROLL (short i = 0; i < 16; i++) { + sa[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (SA)0; + } + } + } else { + FOR_UNROLL (short i = 0; i < 16; i++) { + sa[row*NK + k_base + i] = (SA)0; + } + } + } + for (int work = tiitg; work < NK*NR1; work += NUM_THREADS) { + const int col = work/NK; + const int k = work%NK; + if ((!FC_mul_mm_bc_out && !FC_mul_mm_bc_inp) || + (r1 + col < N && loop_k + k < K)) { + sb[col*NK + k] = (SA)ptrB[(uint64_t)(r1 + col)*strideB + loop_k + k]; + } else { + sb[col*NK + k] = (SA)0; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + auto mA = tA.slice(0, 0); + auto mB = tB.slice(0, 0); + mm.run(mB, mA, cT); + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + device float *dst_batch = (device float *)dst + im*N*M; + if (!FC_mul_mm_bc_out) { + device float *dst_tile = dst_batch + r0 + (uint64_t)r1*M; + auto tD = tensor(dst_tile, dextents(NR0, NR1), array({1, M})); + cT.store(tD); + } else { + auto tD = tensor(dst_batch, dextents(M, N), array({1, M})); + auto mD = tD.slice(r0, r1); + cT.store(mD); + } +} + +typedef decltype(kernel_mul_mm_mpp<64, 32, half, half4x4, float4x4, 1, dequantize_f32, float, float4x4, float>) mul_mm_mpp_t; + +template [[host_name("kernel_mul_mm_f16_f32_mpp")]] kernel mul_mm_mpp_t kernel_mul_mm_mpp<64, 32, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>; + +// Retained Metal4/TensorOps dense prefill kernel. The legacy MPP prototype +// staged both operands in threadgroup memory; this version stages only the +// model weight tile and lets MPP read the dense RHS activation matrix directly +// from device memory. That direct-RHS shape was the clear win for DS4's large +// aligned F16/Q8_0 prompt matmuls. The host selects the widest token tile that +// evenly divides the batch, with 128-token tiles retained after the 64-token +// retest was neutral or slower. +// +// The host dispatch guarantees M % NR0 == 0 and K % NK == 0, so the dequant +// stage does no bounds work. The weight tile is double-buffered: the next +// k-step's dequant overlaps the current cooperative matmul, so the k-loop +// needs one threadgroup barrier per step instead of two. +template< + short NR1, + typename SA, typename SA_4x4, typename block_q, short nl, + void (*dequantize_func)(device const block_q *, short, thread SA_4x4 &), + typename T0, typename T0_4x4, typename T1> +kernel void kernel_mul_mm_mpp_direct_rhs( + constant ds4_metal_args_mul_mm & args, + device const char * srcA, + device const char * srcB, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + (void) sgitg; + + constexpr int NR0 = 64; + constexpr int NK = 32; + constexpr int NL = NK/16; + constexpr int NUM_THREADS = 128; + + const int K = args.ne00; + const int M = args.ne0; + const int N = args.ne1; + const int im = tgpig.z; + const int i12 = im%args.ne12; + const int i13 = im/args.ne12; + const int r0 = tgpig.y*NR0; + const int r1 = tgpig.x*NR1; + + const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + + threadgroup SA *sa = (threadgroup SA *)shmem; + auto tA0 = tensor(sa, dextents(NK, NR0)); + auto tA1 = tensor(sa + NR0*NK, dextents(NK, NR0)); + + device T1 *ptrB = (device T1 *)(srcB + args.nb12*i12 + args.nb13*i13); + const int strideB = args.nb11/sizeof(T1); + auto tB = tensor(ptrB, dextents(K, N), array({1, strideB})); + + matmul2d< + matmul2d_descriptor(NR1, NR0, NK, false, true, true, + matmul2d_descriptor::mode::multiply_accumulate), + execution_simdgroups<4>> mm; + + auto cT = mm.template get_destination_cooperative_tensor(); + + #pragma unroll + for (uint16_t i = 0; i < cT.get_capacity(); ++i) { + if (cT.is_valid_element(i)) { + cT[i] = 0.0f; + } + } + + // NR0*NL/NUM_THREADS 16-value weight chunks per thread (1 at NR0=64). + auto stage_tile = [&](const int loop_k, threadgroup SA *buf) { + FOR_UNROLL (int work = tiitg; work < NR0*NL; work += NUM_THREADS) { + const int row = work / NL; + const int k_chunk = work % NL; + const int k_pos = loop_k + k_chunk*16; + const short k_base = k_chunk*16; + if (is_same::value && FC_mul_mm_bc_inp) { + device const T0 *row_ptr_f = + (device const T0 *)(srcA + args.nb01*(r0 + row) + offset0); + FOR_UNROLL (short i = 0; i < 16; i++) { + buf[row*NK + k_base + i] = (SA)row_ptr_f[k_pos + i]; + } + } else { + device const block_q *row_ptr = + (device const block_q *)(srcA + args.nb01*(r0 + row) + offset0); + SA_4x4 temp_a; + dequantize_func(row_ptr + k_pos/(16*nl), (k_pos/16)%nl, temp_a); + typedef vec SA4; + threadgroup SA4 *dst4 = (threadgroup SA4 *)(buf + row*NK + k_base); + dst4[0] = temp_a[0]; + dst4[1] = temp_a[1]; + dst4[2] = temp_a[2]; + dst4[3] = temp_a[3]; + } + } + }; + + stage_tile(0, sa); + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint buf_sel = 0; + for (int loop_k = 0; loop_k < K; loop_k += NK) { + auto mA = (buf_sel ? tA1 : tA0).slice(0, 0); + auto mB = tB.slice(loop_k, r1); + mm.run(mB, mA, cT); + + const int next_k = loop_k + NK; + if (next_k < K) { + buf_sel ^= 1u; + stage_tile(next_k, buf_sel ? sa + NR0*NK : sa); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + device float *dst_batch = (device float *)dst + im*N*M; + auto tD = tensor(dst_batch, dextents(M, N), array({1, M})); + auto mD = tD.slice(r0, r1); + cT.store(mD); +} + +typedef decltype(kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, float4x4, 1, dequantize_f32, float, float4x4, float>) mul_mm_mpp_direct_rhs_t; + +template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>; +template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>; +template [[host_name("kernel_mul_mm_f16_f32_mpp_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, half4x4, 1, dequantize_f16, half, half4x4, float>; +template [[host_name("kernel_mul_mm_q4_0_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, ds4_dense_block_q4_0, 2, dequantize_dense_q4_0, float, float4x4, float>; +template [[host_name("kernel_mul_mm_q4_0_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, ds4_dense_block_q4_0, 2, dequantize_dense_q4_0, float, float4x4, float>; +template [[host_name("kernel_mul_mm_q4_0_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, ds4_dense_block_q4_0, 2, dequantize_dense_q4_0, float, float4x4, float>; +template [[host_name("kernel_mul_mm_q4_K_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, ds4_dense_block_q4_K, 16, dequantize_dense_q4_K, float, float4x4, float>; +template [[host_name("kernel_mul_mm_q4_K_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, ds4_dense_block_q4_K, 16, dequantize_dense_q4_K, float, float4x4, float>; +template [[host_name("kernel_mul_mm_q4_K_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, ds4_dense_block_q4_K, 16, dequantize_dense_q4_K, float, float4x4, float>; + +template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<32, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>; +template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<64, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>; +template [[host_name("kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128")]] kernel mul_mm_mpp_direct_rhs_t kernel_mul_mm_mpp_direct_rhs<128, half, half4x4, block_q8_0, 2, dequantize_q8_0_pairs, float, float4x4, float>; +#endif + +// Tiled matrix-matrix kernel used for prompt batches larger than 8. DS4 uses +// this to turn prefill into large simdgroup matrix operations; each block_q +// contains 16*nl weights. +template +kernel void kernel_mul_mm( + constant ds4_metal_args_mul_mm & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + + threadgroup S0 * sa = (threadgroup S0 *)(shmem); + threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096); + + constexpr int NR0 = 64; + constexpr int NR1 = 32; + + constexpr int NK = 32; + constexpr int NL0 = NK/16; + constexpr int NL1 = NK/8; + + const int im = tgpig.z; + const int r0 = tgpig.y*NR0; + const int r1 = tgpig.x*NR1; + + // if this block is of 64x32 shape or smaller + const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; + const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1; + + // a thread shouldn't load data outside of the matrix + const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63 + const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31 + + const short il0 = (tiitg % NL0); + + short il = il0; + + const int i12 = im%args.ne12; + const int i13 = im/args.ne12; + + const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + const short offset1 = il0/nl; + + device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1; + + const short iy = 8*(tiitg % NL1); + + device const T1 * y = (device const T1 *)(src1 + + args.nb13*i13 + + args.nb12*i12 + + args.nb11*(r1 + lr1) + + args.nb10*iy); + + S0_8x8 ma[4]; + S1_8x8 mb[2]; + + simdgroup_float8x8 mc[8]; + + for (short i = 0; i < 8; i++){ + mc[i] = make_filled_simdgroup_matrix(0.f); + } + + for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { + // load data and store to threadgroup memory + if (is_same::value && FC_mul_mm_bc_inp) { + threadgroup_barrier(mem_flags::mem_threadgroup); + + // no need for dequantization + for (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + *(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0; + } + } else { + S0_4x4 temp_a; + dequantize_func(x, il, temp_a); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + // Pointer-form store avoids a slower address-lowering path in + // current Apple Metal compilers for this dequantized tile write. + *(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4]; + } + } + + if (FC_mul_mm_bc_inp) { + for (short i = 0; i < 8; ++i) { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short lx = i; + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + + *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0; + } + } else { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + + *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y)); + } + + il = (il + 2 < nl) ? il + 2 : il % 2; + x = (il < 2) ? x + (2 + nl - 1)/nl : x; + + y += NK; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // load matrices from threadgroup memory and conduct outer products + threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2)); + threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2)); + + FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma[i], lsma + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++){ + simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]); + } + + lsma += 8*64; + lsmb += 4*64; + } + } + + if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) { + // if no bounds checks on the output are needed, we can directly write to device memory + device float * C = (device float *) dst + + (r0 + 32*(sgitg & 1)) + \ + (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false); + } + } else { + // block is smaller than 64x32, we should avoid writing data outside of the matrix + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sgitg == 0) { + for (int j = tiitg; j < nr1; j += NR1) { + device float * D = (device float *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0; + device float4 * D4 = (device float4 *) D; + + threadgroup float * C = temp_str + (j*NR0); + threadgroup float4 * C4 = (threadgroup float4 *) C; + + int i = 0; + for (; i < nr0/4; i++) { + *(D4 + i) = *(C4 + i); + } + + i *= 4; + for (; i < nr0; i++) { + *(D + i) = *(C + i); + } + } + } + } +} + +// Legacy F16-weight/F32-RHS prefill matmul with a per-row RMS scale applied at +// the existing F32-to-F16 RHS staging boundary. The tile layout, half inputs, +// float accumulators, and output path intentionally mirror kernel_mul_mm so the +// only arithmetic change versus materializing RMSNorm first is where x*scale is +// rounded from F32 to F16. +kernel void kernel_mul_mm_f16_f32_scaled( + constant ds4_metal_args_mul_mm & args, + device const char * src0, + device const char * src1, + device char * dst, + device const float * scales, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + + threadgroup half * sa = (threadgroup half *)(shmem); + threadgroup half * sb = (threadgroup half *)(shmem + 4096); + + constexpr int NR0 = 64; + constexpr int NR1 = 32; + + constexpr int NK = 32; + constexpr int NL0 = NK/16; + constexpr int NL1 = NK/8; + + const int im = tgpig.z; + const int r0 = tgpig.y*NR0; + const int r1 = tgpig.x*NR1; + + // if this block is of 64x32 shape or smaller + const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; + const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1; + + // a thread shouldn't load data outside of the matrix + const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63 + const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31 + + const short il0 = (tiitg % NL0); + + short il = il0; + + const int i12 = im%args.ne12; + const int i13 = im/args.ne12; + + const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + const short offset1 = il0; + + device const half4x4 * x = (device const half4x4 *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1; + + const short iy = 8*(tiitg % NL1); + + device const float * y = (device const float *)(src1 + + args.nb13*i13 + + args.nb12*i12 + + args.nb11*(r1 + lr1) + + args.nb10*iy); + const float row_scale = scales[r1 + lr1]; + + simdgroup_half8x8 ma[4]; + simdgroup_half8x8 mb[2]; + + simdgroup_float8x8 mc[8]; + + for (short i = 0; i < 8; i++){ + mc[i] = make_filled_simdgroup_matrix(0.f); + } + + for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { + // load data and store to threadgroup memory + if (FC_mul_mm_bc_inp) { + threadgroup_barrier(mem_flags::mem_threadgroup); + + // no need for dequantization + for (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + *(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device half *) x + i) : 0; + } + } else { + half4x4 temp_a; + dequantize_f16(x, il, temp_a); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + // Pointer-form store matches the legacy F16 tile layout. + *(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4]; + } + } + + if (FC_mul_mm_bc_inp) { + for (short i = 0; i < 8; ++i) { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short lx = i; + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + const float scaled = loop_k + iy + i < args.ne00 ? y[i] * row_scale : 0.0f; + + *(sb + 64*ib + 8*ly + lx) = (half)scaled; + } + } else { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + + const float2x4 raw = *((device const float2x4 *) y); + float2x4 scaled; + scaled[0] = raw[0] * row_scale; + scaled[1] = raw[1] * row_scale; + *(threadgroup half2x4 *)(sb + 64*ib + 8*ly) = (half2x4)scaled; + } + + il = (il + 2 < 1) ? il + 2 : il % 2; + x = (il < 2) ? x + 2 : x; + + y += NK; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // load matrices from threadgroup memory and conduct outer products + threadgroup const half * lsma = (sa + 4*64*(sgitg%2)); + threadgroup const half * lsmb = (sb + 2*64*(sgitg/2)); + + FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma[i], lsma + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++){ + simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]); + } + + lsma += 8*64; + lsmb += 4*64; + } + } + + if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) { + // if no bounds checks on the output are needed, we can directly write to device memory + device float * C = (device float *) dst + + (r0 + 32*(sgitg & 1)) + \ + (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false); + } + } else { + // block is smaller than 64x32, we should avoid writing data outside of the matrix + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sgitg == 0) { + for (int j = tiitg; j < nr1; j += NR1) { + device float * D = (device float *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0; + device float4 * D4 = (device float4 *) D; + + threadgroup float * C = temp_str + (j*NR0); + threadgroup float4 * C4 = (threadgroup float4 *) C; + + int i = 0; + for (; i < nr0/4; i++) { + *(D4 + i) = *(C4 + i); + } + + i *= 4; + for (; i < nr0; i++) { + *(D + i) = *(C + i); + } + } + } + } +} + +kernel void kernel_mul_mm_f16_f32_pair( + constant ds4_metal_args_mul_mm & args, + device const char * src0_a, + device const char * src0_b, + device const char * src1, + device char * dst_a, + device char * dst_b, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + threadgroup half * sa_a = (threadgroup half *)(shmem); + threadgroup half * sa_b = (threadgroup half *)(shmem + 4096); + threadgroup half * sb = (threadgroup half *)(shmem + 8192); + + constexpr int NR0 = 64; + constexpr int NR1 = 32; + constexpr int NK = 32; + constexpr int NL0 = NK/16; + constexpr int NL1 = NK/8; + + const int im = tgpig.z; + const int r0 = tgpig.y*NR0; + const int r1 = tgpig.x*NR1; + + const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; + const short nr1 = (args.ne1 - r1 < NR1) ? (args.ne1 - r1) : NR1; + + const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; + const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; + + const short il0 = (tiitg % NL0); + short il = il0; + + const int i12 = im%args.ne12; + const int i13 = im/args.ne12; + + const uint64_t offset0 = (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + const short offset1 = il0; + + device const half4x4 * xa = (device const half4x4 *)(src0_a + args.nb01*(r0 + lr0) + offset0) + offset1; + device const half4x4 * xb = (device const half4x4 *)(src0_b + args.nb01*(r0 + lr0) + offset0) + offset1; + + const short iy = 8*(tiitg % NL1); + + device const float * y = (device const float *)(src1 + + args.nb13*i13 + + args.nb12*i12 + + args.nb11*(r1 + lr1) + + args.nb10*iy); + + simdgroup_half8x8 ma[4]; + simdgroup_half8x8 mb[2]; + + simdgroup_float8x8 mc_a[8]; + simdgroup_float8x8 mc_b[8]; + + for (short i = 0; i < 8; i++) { + mc_a[i] = make_filled_simdgroup_matrix(0.f); + mc_b[i] = make_filled_simdgroup_matrix(0.f); + } + + for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { + half4x4 temp_a; + half4x4 temp_b; + dequantize_f16(xa, il, temp_a); + dequantize_f16(xb, il, temp_b); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + *(sa_a + 64*ib + 8*ly + lx) = temp_a[i/4][i%4]; + *(sa_b + 64*ib + 8*ly + lx) = temp_b[i/4][i%4]; + } + + if (FC_mul_mm_bc_inp) { + for (short i = 0; i < 8; ++i) { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short lx = i; + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + + *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (half) *((device float *) y + i) : 0; + } + } else { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + + *(threadgroup half2x4 *)(sb + 64*ib + 8*ly) = (half2x4)(*((device float2x4 *) y)); + } + + il = (il + 2 < 1) ? il + 2 : il % 2; + xa = (il < 2) ? xa + 2 : xa; + xb = (il < 2) ? xb + 2 : xb; + + y += NK; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup const half * lsma_a = (sa_a + 4*64*(sgitg%2)); + threadgroup const half * lsma_b = (sa_b + 4*64*(sgitg%2)); + threadgroup const half * lsmb = (sb + 2*64*(sgitg/2)); + + FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma[i], lsma_a + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++) { + simdgroup_multiply_accumulate(mc_a[i], mb[i/4], ma[i%4], mc_a[i]); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma[i], lsma_b + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++) { + simdgroup_multiply_accumulate(mc_b[i], mb[i/4], ma[i%4], mc_b[i]); + } + + lsma_a += 8*64; + lsma_b += 8*64; + lsmb += 4*64; + } + } + + if (!FC_mul_mm_bc_out || (r0 + NR0 <= args.ne0 && r1 + NR1 <= args.ne1)) { + device float * C_a = (device float *) dst_a + + (r0 + 32*(sgitg & 1)) + + (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0; + device float * C_b = (device float *) dst_b + + (r0 + 32*(sgitg & 1)) + + (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc_a[i], C_a + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false); + simdgroup_store(mc_b[i], C_b + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false); + } + } else { + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup float * temp_str = (threadgroup float *) shmem; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc_a[i], + temp_str + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0 + 8*(i%4) + 8*NR0*(i/4), + NR0, + 0, + false); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sgitg == 0) { + for (int j = tiitg; j < nr1; j += NR1) { + device float * D = (device float *) dst_a + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0; + device float4 * D4 = (device float4 *) D; + + threadgroup float * C = temp_str + (j*NR0); + threadgroup float4 * C4 = (threadgroup float4 *) C; + + int i = 0; + for (; i < nr0/4; i++) { + *(D4 + i) = *(C4 + i); + } + + i *= 4; + for (; i < nr0; i++) { + *(D + i) = *(C + i); + } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc_b[i], + temp_str + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0 + 8*(i%4) + 8*NR0*(i/4), + NR0, + 0, + false); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sgitg == 0) { + for (int j = tiitg; j < nr1; j += NR1) { + device float * D = (device float *) dst_b + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0; + device float4 * D4 = (device float4 *) D; + + threadgroup float * C = temp_str + (j*NR0); + threadgroup float4 * C4 = (threadgroup float4 *) C; + + int i = 0; + for (; i < nr0/4; i++) { + *(D4 + i) = *(C4 + i); + } + + i *= 4; + for (; i < nr0; i++) { + *(D + i) = *(C + i); + } + } + } + } +} + +typedef decltype(kernel_mul_mm) mul_mm_t; + +// Host-visible prefill matmul variants for F16 and Q8_0 weights. +template [[host_name("kernel_mul_mm_f16_f32")]] kernel mul_mm_t kernel_mul_mm; +template [[host_name("kernel_mul_mm_q8_0_f32")]] kernel mul_mm_t kernel_mul_mm; +template [[host_name("kernel_mul_mm_q4_0_f32")]] kernel mul_mm_t kernel_mul_mm; +template [[host_name("kernel_mul_mm_q4_K_f32")]] kernel mul_mm_t kernel_mul_mm; diff --git a/metal/dsv4_hc.metal b/metal/dsv4_hc.metal new file mode 100644 index 0000000..af5c0e9 --- /dev/null +++ b/metal/dsv4_hc.metal @@ -0,0 +1,1017 @@ +struct ds4_metal_args_dsv4_hc_split_sinkhorn { + int32_t n_hc; + int32_t sinkhorn_iters; + int64_t n_rows; + int64_t mix_hc; + uint64_t nb01; + uint64_t nb1; + float eps; +}; + +struct ds4_metal_args_dsv4_hc_weighted_sum { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb_w0; + uint64_t nb_w1; + uint64_t nb0; + uint64_t nb1; +}; + +struct ds4_metal_args_dsv4_hc_weighted_sum_norm { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb_w0; + uint64_t nb_w1; + uint64_t nb0; + uint64_t nb1; + uint64_t nb_norm1; + float norm_eps; +}; + +struct ds4_metal_args_dsv4_output_hc_weights4 { + float post_scale; + float eps; +}; + +struct ds4_metal_args_dsv4_hc_split_weighted_sum { + int64_t n_embd; + int32_t n_hc; + int32_t sinkhorn_iters; + int64_t n_rows; + int64_t mix_hc; + uint64_t nb_mix1; + uint64_t nb_split1; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb0; + uint64_t nb1; + float eps; +}; + +struct ds4_metal_args_dsv4_hc_split_weighted_sum_norm { + int64_t n_embd; + int32_t n_hc; + int32_t sinkhorn_iters; + int64_t n_rows; + int64_t mix_hc; + uint64_t nb_mix1; + uint64_t nb_split1; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb0; + uint64_t nb1; + uint64_t nb_norm1; + float eps; + float norm_eps; +}; + +struct ds4_metal_args_dsv4_hc_expand { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_block0; + uint64_t nb_block1; + uint64_t nb_add0; + uint64_t nb_add1; + uint64_t nb_res0; + uint64_t nb_res1; + uint64_t nb_res2; + uint64_t nb_post0; + uint64_t nb_post1; + uint64_t nb_comb0; + uint64_t nb_comb1; + uint64_t nb_comb2; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + int32_t has_add; +}; + +// Numerically stable sigmoid for the standalone split/sinkhorn path. The naive +// form 1/(1+exp(-z)) overflows for large negative z (exp(-z) blows up); +// replacing it with the 0.5*(tanh(z/2)+1) identity keeps the value bounded in +// [0, 1] across the entire float range. Gated by DS4_METAL_HC_STABLE so we can +// A/B vs the historical form on M5 Max where the faster ALU is more likely to +// push HC mixer inputs into the unstable regime. +// +// Do not automatically use these helpers in the fused HC decode kernels below: +// routing the fused vector sites through the tanh form produced non-finite +// logits on M5 Max, while the historical inline exp form remains finite and is +// the decode throughput baseline. +#ifdef DS4_METAL_HC_STABLE +static inline float ds4_hc_sigmoid(float z) { return 0.5f * tanh(0.5f * z) + 0.5f; } +static inline float4 ds4_hc_sigmoid(float4 z) { return 0.5f * tanh(0.5f * z) + 0.5f; } +// 2 * sigmoid(z) == 1 + tanh(z/2). +static inline float ds4_hc_twice_sigmoid(float z) { return 1.0f + tanh(0.5f * z); } +static inline float4 ds4_hc_twice_sigmoid(float4 z) { return 1.0f + tanh(0.5f * z); } +#else +static inline float ds4_hc_sigmoid(float z) { return 1.0f / (1.0f + exp(-z)); } +static inline float4 ds4_hc_sigmoid(float4 z) { return 1.0f / (1.0f + exp(-z)); } +static inline float ds4_hc_twice_sigmoid(float z) { return 2.0f / (1.0f + exp(-z)); } +static inline float4 ds4_hc_twice_sigmoid(float4 z) { return 2.0f / (1.0f + exp(-z)); } +#endif + +// Splits an HC mixer row into pre weights, post gates, and the HC-to-HC +// combination matrix. The 4-channel path is specialized because DS4 Flash uses +// HC=4 in normal inference, while the scalar fallback keeps diagnostics usable. +kernel void kernel_dsv4_hc_split_sinkhorn( + constant ds4_metal_args_dsv4_hc_split_sinkhorn & args, + device const float * mixes, + device const float * scale, + device const float * base, + device float * dst, + uint tid [[thread_position_in_grid]]) { + if ((int64_t) tid >= args.n_rows) { + return; + } + + constexpr int HC_MAX = 16; + const int HC = args.n_hc; + if (HC <= 0 || HC > HC_MAX) { + return; + } + + device const float * mix = mixes + ((int64_t) tid)*args.mix_hc; + device float * out = dst + ((int64_t) tid)*args.mix_hc; + + const float epsv = args.eps; + const float pre_scale = scale[0]; + const float post_scale = scale[1]; + const float comb_scale = scale[2]; + + if (HC == 4) { + const float4 pre_z = + *((device const float4 *) mix) * pre_scale + + *((device const float4 *) base); + *((device float4 *) out) = ds4_hc_sigmoid(pre_z) + epsv; + + const float4 post_z = + *((device const float4 *) (mix + 4)) * post_scale + + *((device const float4 *) (base + 4)); + *((device float4 *) (out + 4)) = ds4_hc_twice_sigmoid(post_z); + + float4 r0 = + *((device const float4 *) (mix + 8)) * comb_scale + + *((device const float4 *) (base + 8)); + float4 r1 = + *((device const float4 *) (mix + 12)) * comb_scale + + *((device const float4 *) (base + 12)); + float4 r2 = + *((device const float4 *) (mix + 16)) * comb_scale + + *((device const float4 *) (base + 16)); + float4 r3 = + *((device const float4 *) (mix + 20)) * comb_scale + + *((device const float4 *) (base + 20)); + + const float m0 = max(max(r0.x, r0.y), max(r0.z, r0.w)); + const float m1 = max(max(r1.x, r1.y), max(r1.z, r1.w)); + const float m2 = max(max(r2.x, r2.y), max(r2.z, r2.w)); + const float m3 = max(max(r3.x, r3.y), max(r3.z, r3.w)); + + r0 = exp(r0 - m0); + r1 = exp(r1 - m1); + r2 = exp(r2 - m2); + r3 = exp(r3 - m3); + + r0 = r0 * (1.0f / (r0.x + r0.y + r0.z + r0.w)) + epsv; + r1 = r1 * (1.0f / (r1.x + r1.y + r1.z + r1.w)) + epsv; + r2 = r2 * (1.0f / (r2.x + r2.y + r2.z + r2.w)) + epsv; + r3 = r3 * (1.0f / (r3.x + r3.y + r3.z + r3.w)) + epsv; + + float4 col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col_inv; + r1 *= col_inv; + r2 *= col_inv; + r3 *= col_inv; + + for (int iter = 1; iter < args.sinkhorn_iters; ++iter) { + r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv); + r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv); + r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv); + r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv); + + col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col_inv; + r1 *= col_inv; + r2 *= col_inv; + r3 *= col_inv; + } + + *((device float4 *) (out + 8)) = r0; + *((device float4 *) (out + 12)) = r1; + *((device float4 *) (out + 16)) = r2; + *((device float4 *) (out + 20)) = r3; + return; + } + + for (int i = 0; i < HC; ++i) { + const float z = mix[i] * pre_scale + base[i]; + out[i] = ds4_hc_sigmoid(z) + epsv; + } + + for (int i = 0; i < HC; ++i) { + const int off = HC + i; + const float z = mix[off] * post_scale + base[off]; + out[off] = ds4_hc_twice_sigmoid(z); + } + + float c[HC_MAX*HC_MAX]; + + for (int dst_hc = 0; dst_hc < HC; ++dst_hc) { + float row_max = -INFINITY; + for (int src_hc = 0; src_hc < HC; ++src_hc) { + const int idx = src_hc + dst_hc*HC; + const int off = 2*HC + idx; + const float v = mix[off] * comb_scale + base[off]; + c[idx] = v; + row_max = max(row_max, v); + } + + float row_sum = 0.0f; + for (int src_hc = 0; src_hc < HC; ++src_hc) { + const int idx = src_hc + dst_hc*HC; + const float v = exp(c[idx] - row_max); + c[idx] = v; + row_sum += v; + } + + const float inv_sum = 1.0f / row_sum; + for (int src_hc = 0; src_hc < HC; ++src_hc) { + const int idx = src_hc + dst_hc*HC; + c[idx] = c[idx] * inv_sum + epsv; + } + } + + for (int src_hc = 0; src_hc < HC; ++src_hc) { + float sum = 0.0f; + for (int dst_hc = 0; dst_hc < HC; ++dst_hc) { + sum += c[src_hc + dst_hc*HC]; + } + + const float inv_denom = 1.0f / (sum + epsv); + for (int dst_hc = 0; dst_hc < HC; ++dst_hc) { + c[src_hc + dst_hc*HC] *= inv_denom; + } + } + + for (int iter = 1; iter < args.sinkhorn_iters; ++iter) { + for (int dst_hc = 0; dst_hc < HC; ++dst_hc) { + float sum = 0.0f; + for (int src_hc = 0; src_hc < HC; ++src_hc) { + sum += c[src_hc + dst_hc*HC]; + } + + const float inv_denom = 1.0f / (sum + epsv); + for (int src_hc = 0; src_hc < HC; ++src_hc) { + c[src_hc + dst_hc*HC] *= inv_denom; + } + } + + for (int src_hc = 0; src_hc < HC; ++src_hc) { + float sum = 0.0f; + for (int dst_hc = 0; dst_hc < HC; ++dst_hc) { + sum += c[src_hc + dst_hc*HC]; + } + + const float inv_denom = 1.0f / (sum + epsv); + for (int dst_hc = 0; dst_hc < HC; ++dst_hc) { + c[src_hc + dst_hc*HC] *= inv_denom; + } + } + } + + for (int i = 0; i < HC*HC; ++i) { + out[2*HC + i] = c[i]; + } +} + +// Decode-side fusion of HC split and pre-weighted HC reduction. One threadgroup +// handles one token row: lane 0 computes the HC=4 mixer split once, stores the +// post/comb data for the following HC expand, and all lanes reuse the pre +// weights from threadgroup memory to produce the embedding row. +kernel void kernel_dsv4_hc_split_weighted_sum( + constant ds4_metal_args_dsv4_hc_split_weighted_sum & args, + device const char * mixes, + device const float * scale, + device const float * base, + device const char * x, + device char * split, + device char * dst, + threadgroup float * pre_shmem [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint ntg [[threads_per_threadgroup]]) { + if ((int64_t) row >= args.n_rows || args.n_hc != 4) { + return; + } + + device const float * mix = (device const float *) (mixes + (uint64_t)row*args.nb_mix1); + device float * out = (device float *) (split + (uint64_t)row*args.nb_split1); + + if (tid == 0) { + const float epsv = args.eps; + const float pre_scale = scale[0]; + const float post_scale = scale[1]; + const float comb_scale = scale[2]; + + const float4 pre_z = + *((device const float4 *) mix) * pre_scale + + *((device const float4 *) base); + const float4 pre = 1.0f / (1.0f + exp(-pre_z)) + epsv; + *((device float4 *) out) = pre; + pre_shmem[0] = pre.x; + pre_shmem[1] = pre.y; + pre_shmem[2] = pre.z; + pre_shmem[3] = pre.w; + + const float4 post_z = + *((device const float4 *) (mix + 4)) * post_scale + + *((device const float4 *) (base + 4)); + *((device float4 *) (out + 4)) = 2.0f / (1.0f + exp(-post_z)); + + float4 r0 = + *((device const float4 *) (mix + 8)) * comb_scale + + *((device const float4 *) (base + 8)); + float4 r1 = + *((device const float4 *) (mix + 12)) * comb_scale + + *((device const float4 *) (base + 12)); + float4 r2 = + *((device const float4 *) (mix + 16)) * comb_scale + + *((device const float4 *) (base + 16)); + float4 r3 = + *((device const float4 *) (mix + 20)) * comb_scale + + *((device const float4 *) (base + 20)); + + const float m0 = max(max(r0.x, r0.y), max(r0.z, r0.w)); + const float m1 = max(max(r1.x, r1.y), max(r1.z, r1.w)); + const float m2 = max(max(r2.x, r2.y), max(r2.z, r2.w)); + const float m3 = max(max(r3.x, r3.y), max(r3.z, r3.w)); + + r0 = exp(r0 - m0); + r1 = exp(r1 - m1); + r2 = exp(r2 - m2); + r3 = exp(r3 - m3); + + r0 = r0 * (1.0f / (r0.x + r0.y + r0.z + r0.w)) + epsv; + r1 = r1 * (1.0f / (r1.x + r1.y + r1.z + r1.w)) + epsv; + r2 = r2 * (1.0f / (r2.x + r2.y + r2.z + r2.w)) + epsv; + r3 = r3 * (1.0f / (r3.x + r3.y + r3.z + r3.w)) + epsv; + + float4 col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col_inv; + r1 *= col_inv; + r2 *= col_inv; + r3 *= col_inv; + + for (int iter = 1; iter < args.sinkhorn_iters; ++iter) { + r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv); + r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv); + r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv); + r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv); + + col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col_inv; + r1 *= col_inv; + r2 *= col_inv; + r3 *= col_inv; + } + + *((device float4 *) (out + 8)) = r0; + *((device float4 *) (out + 12)) = r1; + *((device float4 *) (out + 16)) = r2; + *((device float4 *) (out + 20)) = r3; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int64_t d = tid; d < args.n_embd; d += ntg) { + float acc = 0.0f; + acc += *((device const float *) (x + d*args.nb_x0 + 0*args.nb_x1 + (uint64_t)row*args.nb_x2)) * pre_shmem[0]; + acc += *((device const float *) (x + d*args.nb_x0 + 1*args.nb_x1 + (uint64_t)row*args.nb_x2)) * pre_shmem[1]; + acc += *((device const float *) (x + d*args.nb_x0 + 2*args.nb_x1 + (uint64_t)row*args.nb_x2)) * pre_shmem[2]; + acc += *((device const float *) (x + d*args.nb_x0 + 3*args.nb_x1 + (uint64_t)row*args.nb_x2)) * pre_shmem[3]; + *((device float *) (dst + d*args.nb0 + (uint64_t)row*args.nb1)) = acc; + } +} + +// Decode HC-pre plus the following RMSNorm. DS4 uses HC=4 here. The normal +// release path computes HC coefficients, collapses four residual streams into +// the model row, then immediately launches a weighted RMSNorm over the row. +// This kernel keeps the HC split math identical to +// kernel_dsv4_hc_split_weighted_sum, stores the HC-pre row for diagnostics, and +// reuses the just-collapsed values from threadgroup memory for the RMSNorm +// reduction. +kernel void kernel_dsv4_hc_split_weighted_sum_norm4( + constant ds4_metal_args_dsv4_hc_split_weighted_sum_norm & args, + device const char * mixes, + device const float * scale, + device const float * base, + device const char * x, + device char * split, + device char * dst, + device const char * norm_weight, + device char * norm_dst, + threadgroup float * shared [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + ushort tid [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort ntg [[threads_per_threadgroup]]) { + if ((int64_t)row >= args.n_rows || args.n_hc != 4 || (args.n_embd & 3) != 0) { + return; + } + + const uint n_embd = uint(args.n_embd); + const uint n4 = n_embd >> 2; + threadgroup float4 *row_shmem = (threadgroup float4 *)shared; + threadgroup float *pre_shmem = shared + n_embd; + threadgroup float *sum_shmem = pre_shmem + 4; + + device const float *mix = (device const float *)(mixes + (uint64_t)row * args.nb_mix1); + device float *out = (device float *)(split + (uint64_t)row * args.nb_split1); + + if (sgitg == 0) { + sum_shmem[tiisg] = 0.0f; + } + + if (tid == 0) { + const float epsv = args.eps; + const float pre_scale = scale[0]; + const float post_scale = scale[1]; + const float comb_scale = scale[2]; + + const float4 pre_z = + *((device const float4 *)mix) * pre_scale + + *((device const float4 *)base); + const float4 pre = 1.0f / (1.0f + exp(-pre_z)) + epsv; + *((device float4 *)out) = pre; + pre_shmem[0] = pre.x; + pre_shmem[1] = pre.y; + pre_shmem[2] = pre.z; + pre_shmem[3] = pre.w; + + const float4 post_z = + *((device const float4 *)(mix + 4)) * post_scale + + *((device const float4 *)(base + 4)); + *((device float4 *)(out + 4)) = 2.0f / (1.0f + exp(-post_z)); + + float4 r0 = + *((device const float4 *)(mix + 8)) * comb_scale + + *((device const float4 *)(base + 8)); + float4 r1 = + *((device const float4 *)(mix + 12)) * comb_scale + + *((device const float4 *)(base + 12)); + float4 r2 = + *((device const float4 *)(mix + 16)) * comb_scale + + *((device const float4 *)(base + 16)); + float4 r3 = + *((device const float4 *)(mix + 20)) * comb_scale + + *((device const float4 *)(base + 20)); + + const float m0 = max(max(r0.x, r0.y), max(r0.z, r0.w)); + const float m1 = max(max(r1.x, r1.y), max(r1.z, r1.w)); + const float m2 = max(max(r2.x, r2.y), max(r2.z, r2.w)); + const float m3 = max(max(r3.x, r3.y), max(r3.z, r3.w)); + + r0 = exp(r0 - m0); + r1 = exp(r1 - m1); + r2 = exp(r2 - m2); + r3 = exp(r3 - m3); + + r0 = r0 * (1.0f / (r0.x + r0.y + r0.z + r0.w)) + epsv; + r1 = r1 * (1.0f / (r1.x + r1.y + r1.z + r1.w)) + epsv; + r2 = r2 * (1.0f / (r2.x + r2.y + r2.z + r2.w)) + epsv; + r3 = r3 * (1.0f / (r3.x + r3.y + r3.z + r3.w)) + epsv; + + float4 col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col_inv; + r1 *= col_inv; + r2 *= col_inv; + r3 *= col_inv; + + for (int iter = 1; iter < args.sinkhorn_iters; ++iter) { + r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv); + r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv); + r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv); + r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv); + + col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col_inv; + r1 *= col_inv; + r2 *= col_inv; + r3 *= col_inv; + } + + *((device float4 *)(out + 8)) = r0; + *((device float4 *)(out + 12)) = r1; + *((device float4 *)(out + 16)) = r2; + *((device float4 *)(out + 20)) = r3; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + float sumf = 0.0f; + for (uint i = tid; i < n4; i += ntg) { + device const float4 *x0 = (device const float4 *)(x + 0 * args.nb_x1 + (uint64_t)row * args.nb_x2); + device const float4 *x1 = (device const float4 *)(x + 1 * args.nb_x1 + (uint64_t)row * args.nb_x2); + device const float4 *x2 = (device const float4 *)(x + 2 * args.nb_x1 + (uint64_t)row * args.nb_x2); + device const float4 *x3 = (device const float4 *)(x + 3 * args.nb_x1 + (uint64_t)row * args.nb_x2); + // Preserve the standalone HC collapse's explicit accumulation order. + float4 v = 0.0f; + v += x0[i] * pre_shmem[0]; + v += x1[i] * pre_shmem[1]; + v += x2[i] * pre_shmem[2]; + v += x3[i] * pre_shmem[3]; + row_shmem[i] = v; + sumf += dot(v, v); + } + + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) { + sum_shmem[sgitg] = sumf; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + sumf = sum_shmem[tiisg]; + sumf = simd_sum(sumf); + // Batched prefill must match kernel_rms_norm_fuse_impl so enabling this + // fusion does not change its scale by an ULP. Keep the established + // single-row decode result unchanged; that path historically used rsqrt. + const float norm_arg = sumf / float(n_embd) + args.norm_eps; + const float norm_scale = args.n_rows > 1 ? 1.0f / sqrt(norm_arg) : rsqrt(norm_arg); + + device float4 *dst4 = (device float4 *)(dst + (uint64_t)row * args.nb1); + device const float4 *w4 = (device const float4 *)norm_weight; + device float4 *norm4 = (device float4 *)(norm_dst + (uint64_t)row * args.nb_norm1); + for (uint i = tid; i < n4; i += ntg) { + const float4 v = row_shmem[i]; + dst4[i] = v; + norm4[i] = (v * norm_scale) * w4[i]; + } +} + +// Expands an embedding-sized block back into HC channels after attention/FFN. +// The post gate scales the current block, while the Sinkhorn combination matrix +// mixes residual HC channels from the previous state. +kernel void kernel_dsv4_hc_expand( + constant ds4_metal_args_dsv4_hc_expand & args, + device const char * block_out, + device const char * residual, + device const char * post, + device const char * comb, + device const char * block_add, + device char * dst, + uint gid [[thread_position_in_grid]]) { + const int64_t n_elem = args.n_embd * args.n_hc * args.n_tokens; + if ((int64_t) gid >= n_elem) { + return; + } + + const int64_t d = ((int64_t) gid) % args.n_embd; + const int64_t tmp = ((int64_t) gid) / args.n_embd; + const int64_t dst_hc = tmp % args.n_hc; + const int64_t t = tmp / args.n_hc; + + float block_v = *((device const float *) (block_out + d*args.nb_block0 + t*args.nb_block1)); + if (args.has_add) { + block_v += *((device const float *) (block_add + d*args.nb_add0 + t*args.nb_add1)); + } + const float post_v = *((device const float *) (post + dst_hc*args.nb_post0 + t*args.nb_post1)); + + float acc = block_v * post_v; + for (int64_t src_hc = 0; src_hc < args.n_hc; ++src_hc) { + const float comb_v = *((device const float *) (comb + dst_hc*args.nb_comb0 + src_hc*args.nb_comb1 + t*args.nb_comb2)); + const float res_v = *((device const float *) (residual + d*args.nb_res0 + src_hc*args.nb_res1 + t*args.nb_res2)); + acc += comb_v * res_v; + } + + *((device float *) (dst + d*args.nb0 + dst_hc*args.nb1 + t*args.nb2)) = acc; +} + +// HC=4 specialization of the post/expand step. One thread computes all four +// destination HC streams for one token/dimension, reusing the same block output +// and residual HC values while preserving the per-stream accumulation order. +kernel void kernel_dsv4_hc_expand4( + constant ds4_metal_args_dsv4_hc_expand & args, + device const char * block_out, + device const char * residual, + device const char * post, + device const char * comb, + device const char * block_add, + device char * dst, + uint gid [[thread_position_in_grid]]) { + if (args.n_hc != 4) { + return; + } + + const int64_t n_elem = args.n_embd * args.n_tokens; + if ((int64_t) gid >= n_elem) { + return; + } + + const int64_t d = ((int64_t) gid) % args.n_embd; + const int64_t t = ((int64_t) gid) / args.n_embd; + + float block_v = *((device const float *) (block_out + d*args.nb_block0 + t*args.nb_block1)); + if (args.has_add) { + block_v += *((device const float *) (block_add + d*args.nb_add0 + t*args.nb_add1)); + } + + const float r0 = *((device const float *) (residual + d*args.nb_res0 + 0*args.nb_res1 + t*args.nb_res2)); + const float r1 = *((device const float *) (residual + d*args.nb_res0 + 1*args.nb_res1 + t*args.nb_res2)); + const float r2 = *((device const float *) (residual + d*args.nb_res0 + 2*args.nb_res1 + t*args.nb_res2)); + const float r3 = *((device const float *) (residual + d*args.nb_res0 + 3*args.nb_res1 + t*args.nb_res2)); + + for (int64_t dst_hc = 0; dst_hc < 4; ++dst_hc) { + float acc = block_v * *((device const float *) (post + dst_hc*args.nb_post0 + t*args.nb_post1)); + + acc += *((device const float *) (comb + dst_hc*args.nb_comb0 + 0*args.nb_comb1 + t*args.nb_comb2)) * r0; + acc += *((device const float *) (comb + dst_hc*args.nb_comb0 + 1*args.nb_comb1 + t*args.nb_comb2)) * r1; + acc += *((device const float *) (comb + dst_hc*args.nb_comb0 + 2*args.nb_comb1 + t*args.nb_comb2)) * r2; + acc += *((device const float *) (comb + dst_hc*args.nb_comb0 + 3*args.nb_comb1 + t*args.nb_comb2)) * r3; + + *((device float *) (dst + d*args.nb0 + dst_hc*args.nb1 + t*args.nb2)) = acc; + } +} + +// Decode-time FFN tail fusion: +// +// shared_out = shared_mid @ Wshared_down +// after_ffn_hc = HCPost(routed_out + shared_out, residual_hc, split) +// +// The Q8_0 dot reduction is intentionally copied from the normal matvec shape +// so the shared expert result is bit-identical. The only specialization is +// that DS4 decode has one token and HC=4, so the thread that finishes each +// shared-down output row can immediately expand it into the four HC streams. +kernel void kernel_dsv4_shared_down_hc_expand4_q8_0( + constant ds4_metal_args_mul_mv & mv, + constant ds4_metal_args_dsv4_hc_expand & hc, + device const char * weight, + device const char * shared_mid, + device char * shared_out, + device const char * routed_out, + device const char * residual, + device const char * post, + device const char * comb, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + if (hc.n_hc != 4 || hc.n_tokens != 1) { + return; + } + + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + constexpr short NR0 = N_R0_Q8_0; + + const int nb = mv.ne00 / QK8_0; + const int row0 = tgpig.x * NR0; + + const short ix = tiisg / (NW / NQ); + const short il = tiisg % (NW / NQ); + const int ib0 = sgitg * NQ + ix; + + device const float *y = (device const float *)(shared_mid); + device const float *yb = y + ib0 * QK8_0 + il * NQ; + + device const block_q8_0 *ax[NR0]; + FOR_UNROLL(short row = 0; row < NR0; ++row) { + const uint64_t off0 = (uint64_t)(row0 + row) * mv.nb01; + ax[row] = (device const block_q8_0 *)(weight + off0); + } + + float sumf[NR0] = { 0.0f }; + float yl[NQ]; + + for (int ib = ib0; ib < nb; ib += NSG * NQ) { + FOR_UNROLL(short i = 0; i < NQ; ++i) { + yl[i] = yb[i]; + } + + FOR_UNROLL(short row = 0; row < NR0; ++row) { + device const int8_t *qs = ax[row][ib].qs + il * NQ; + + float sumq = 0.0f; + FOR_UNROLL(short i = 0; i < NQ; ++i) { + sumq += qs[i] * yl[i]; + } + + sumf[row] += sumq * ax[row][ib].d; + } + + yb += NSG * NQ * QK8_0; + } + + threadgroup float *shmem_f32[NR0]; + FOR_UNROLL(short row = 0; row < NR0; ++row) { + shmem_f32[row] = (threadgroup float *)shmem + NW * row; + if (sgitg == 0) { + shmem_f32[row][tiisg] = 0.0f; + } + sumf[row] = simd_sum(sumf[row]); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL(short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + shmem_f32[row][sgitg] = sumf[row]; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL(short row = 0; row < NR0; ++row) { + const int d = row0 + row; + if (d >= mv.ne01) { + continue; + } + + const float shared_v = simd_sum(shmem_f32[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + *((device float *)(shared_out + (uint64_t)d * sizeof(float))) = shared_v; + + float block_v = *((device const float *)(routed_out + (uint64_t)d * hc.nb_block0)); + block_v += shared_v; + + const float r0 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 0 * hc.nb_res1)); + const float r1 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 1 * hc.nb_res1)); + const float r2 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 2 * hc.nb_res1)); + const float r3 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 3 * hc.nb_res1)); + + for (int64_t dst_hc = 0; dst_hc < 4; ++dst_hc) { + float acc = block_v * *((device const float *)(post + dst_hc * hc.nb_post0)); + + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 0 * hc.nb_comb1)) * r0; + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 1 * hc.nb_comb1)) * r1; + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 2 * hc.nb_comb1)) * r2; + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 3 * hc.nb_comb1)) * r3; + + *((device float *)(dst + (uint64_t)d * hc.nb0 + dst_hc * hc.nb1)) = acc; + } + } + } +} + +// Decode-time attention output tail fusion: +// +// attn_out = attn_low @ Wob +// after_attn_hc = HCPost(attn_out, residual_hc, split) +// +// This is the no-add sibling of the shared-down/FFN fusion above. It preserves +// the exact Q8_0 matvec reduction, stores `attn_out` for diagnostics, and then +// writes the four HC streams for the same embedding dimension. +kernel void kernel_dsv4_q8_hc_expand4_q8_0( + constant ds4_metal_args_mul_mv & mv, + constant ds4_metal_args_dsv4_hc_expand & hc, + device const char * weight, + device const char * input, + device char * block_out, + device const char * residual, + device const char * post, + device const char * comb, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + if (hc.n_hc != 4 || hc.n_tokens != 1) { + return; + } + + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + constexpr short NR0 = N_R0_Q8_0; + + const int nb = mv.ne00 / QK8_0; + const int row0 = tgpig.x * NR0; + + const short ix = tiisg / (NW / NQ); + const short il = tiisg % (NW / NQ); + const int ib0 = sgitg * NQ + ix; + + device const float *y = (device const float *)(input); + device const float *yb = y + ib0 * QK8_0 + il * NQ; + + device const block_q8_0 *ax[NR0]; + FOR_UNROLL(short row = 0; row < NR0; ++row) { + const uint64_t off0 = (uint64_t)(row0 + row) * mv.nb01; + ax[row] = (device const block_q8_0 *)(weight + off0); + } + + float sumf[NR0] = { 0.0f }; + float yl[NQ]; + + for (int ib = ib0; ib < nb; ib += NSG * NQ) { + FOR_UNROLL(short i = 0; i < NQ; ++i) { + yl[i] = yb[i]; + } + + FOR_UNROLL(short row = 0; row < NR0; ++row) { + device const int8_t *qs = ax[row][ib].qs + il * NQ; + + float sumq = 0.0f; + FOR_UNROLL(short i = 0; i < NQ; ++i) { + sumq += qs[i] * yl[i]; + } + + sumf[row] += sumq * ax[row][ib].d; + } + + yb += NSG * NQ * QK8_0; + } + + threadgroup float *shmem_f32[NR0]; + FOR_UNROLL(short row = 0; row < NR0; ++row) { + shmem_f32[row] = (threadgroup float *)shmem + NW * row; + if (sgitg == 0) { + shmem_f32[row][tiisg] = 0.0f; + } + sumf[row] = simd_sum(sumf[row]); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL(short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + shmem_f32[row][sgitg] = sumf[row]; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL(short row = 0; row < NR0; ++row) { + const int d = row0 + row; + if (d >= mv.ne01) { + continue; + } + + const float block_v = simd_sum(shmem_f32[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + *((device float *)(block_out + (uint64_t)d * sizeof(float))) = block_v; + + const float r0 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 0 * hc.nb_res1)); + const float r1 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 1 * hc.nb_res1)); + const float r2 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 2 * hc.nb_res1)); + const float r3 = *((device const float *)(residual + (uint64_t)d * hc.nb_res0 + 3 * hc.nb_res1)); + + for (int64_t dst_hc = 0; dst_hc < 4; ++dst_hc) { + float acc = block_v * *((device const float *)(post + dst_hc * hc.nb_post0)); + + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 0 * hc.nb_comb1)) * r0; + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 1 * hc.nb_comb1)) * r1; + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 2 * hc.nb_comb1)) * r2; + acc += *((device const float *)(comb + dst_hc * hc.nb_comb0 + 3 * hc.nb_comb1)) * r3; + + *((device float *)(dst + (uint64_t)d * hc.nb0 + dst_hc * hc.nb1)) = acc; + } + } + } +} + +// Reduces HC channels to a normal embedding row with the learned pre weights. +// This is the input adapter before the attention block and before the FFN block. +kernel void kernel_dsv4_hc_weighted_sum( + constant ds4_metal_args_dsv4_hc_weighted_sum & args, + device const char * x, + device const char * weights, + device char * dst, + uint gid [[thread_position_in_grid]]) { + const int64_t n_elem = args.n_embd * args.n_tokens; + if ((int64_t) gid >= n_elem) { + return; + } + + const int64_t d = ((int64_t) gid) % args.n_embd; + const int64_t t = ((int64_t) gid) / args.n_embd; + + float acc = 0.0f; + for (int64_t h = 0; h < args.n_hc; ++h) { + const float xv = *((device const float *) (x + d*args.nb_x0 + h*args.nb_x1 + t*args.nb_x2)); + const float wv = *((device const float *) (weights + h*args.nb_w0 + t*args.nb_w1)); + acc += xv * wv; + } + + *((device float *) (dst + d*args.nb0 + t*args.nb1)) = acc; +} + +// The one-row output head immediately applies a learned RMSNorm after reducing +// its four HC streams. Preserve the standalone scalar HC accumulation, write +// the collapsed row for diagnostics, then reload its F32 values from +// threadgroup memory using the standalone RMSNorm's float4 reduction mapping. +kernel void kernel_dsv4_hc_weighted_sum_norm4( + constant ds4_metal_args_dsv4_hc_weighted_sum_norm & args, + device const char * x, + device const char * weights, + device char * dst, + device const char * norm_weight, + device char * norm_dst, + threadgroup float * shared [[threadgroup(0)]], + ushort tid [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort ntg [[threads_per_threadgroup]]) { + if (args.n_tokens != 1 || args.n_hc != 4 || + args.n_embd <= 0 || (args.n_embd & 3) != 0) { + return; + } + + const uint n_embd = uint(args.n_embd); + const uint n4 = n_embd >> 2; + threadgroup float *row_shmem = shared; + threadgroup float *sum_shmem = shared + n_embd; + + if (sgitg == 0) { + sum_shmem[tiisg] = 0.0f; + } + + for (uint d = tid; d < n_embd; d += ntg) { + float acc = 0.0f; + for (int64_t h = 0; h < args.n_hc; ++h) { + const float xv = *((device const float *)( + x + (uint64_t)d*args.nb_x0 + (uint64_t)h*args.nb_x1)); + const float wv = *((device const float *)( + weights + (uint64_t)h*args.nb_w0)); + acc += xv * wv; + } + row_shmem[d] = acc; + *((device float *)(dst + (uint64_t)d*args.nb0)) = acc; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup const float4 *row4 = + (threadgroup const float4 *)row_shmem; + float sumf = 0.0f; + for (uint i = tid; i < n4; i += ntg) { + sumf += dot(row4[i], row4[i]); + } + sumf = simd_sum(sumf); + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) { + sum_shmem[sgitg] = sumf; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + sumf = sum_shmem[tiisg]; + sumf = simd_sum(sumf); + + const float mean = sumf/args.n_embd; + const float scale = 1.0f/sqrt(mean + args.norm_eps); + device const float4 *w4 = (device const float4 *)norm_weight; + device float4 *norm4 = (device float4 *)norm_dst; + for (uint i = tid; i < n4; i += ntg) { + norm4[i] = (row4[i]*scale)*w4[i]; + } +} + +// The one-row HC=4 output head historically materializes four device-F32 +// stages across separate launches. Collapse those launches into one tiny +// two-thread group while preserving the scalar/vector lane mapping and every +// global rounding boundary. +kernel void kernel_dsv4_output_hc_weights4( + constant ds4_metal_args_dsv4_output_hc_weights4 & args, + device const float * pre, + device const float * hc_scale, + device const float * hc_base, + device float * dst, + ushort tid [[thread_position_in_threadgroup]]) { + device volatile float *stage = (device volatile float *)dst; + + for (uint i = tid; i < 4; i += 2) { + stage[i] = pre[i] * hc_scale[0]; + } + threadgroup_barrier(mem_flags::mem_device); + + for (uint i = tid; i < 4; i += 2) { + stage[i] = stage[i] + hc_base[i]; + } + threadgroup_barrier(mem_flags::mem_device); + + if (tid == 0) { + const float4 x = *((device volatile float4 *)stage); + *((device volatile float4 *)stage) = 1 / (1 + exp(-x)); + } + threadgroup_barrier(mem_flags::mem_device); + + if (tid == 0) { + const float4 x = *((device volatile float4 *)stage); + *((device volatile float4 *)stage) = + args.post_scale * x + args.eps; + } +} diff --git a/metal/dsv4_kv.metal b/metal/dsv4_kv.metal new file mode 100644 index 0000000..4f44a43 --- /dev/null +++ b/metal/dsv4_kv.metal @@ -0,0 +1,369 @@ +constant float dsv4_e4m3fn_exp_scale[16] = { + 0.0f, 0.015625f, 0.03125f, 0.0625f, + 0.125f, 0.25f, 0.5f, 1.0f, + 2.0f, 4.0f, 8.0f, 16.0f, + 32.0f, 64.0f, 128.0f, 256.0f, +}; + +constant float dsv4_e2m1fn_values[8] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, +}; + +struct ds4_metal_args_dsv4_fp8_kv_quantize { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + ulong nb00; + ulong nb01; + ulong nb02; + ulong nb03; + ulong nb0; + ulong nb1; + ulong nb2; + ulong nb3; + int n_rot; +}; + +struct ds4_metal_args_dsv4_kv_fp8_store { + int32_t head_dim; + int32_t n_rot; + int32_t raw_row; +}; + +struct ds4_metal_args_dsv4_indexer_qat { + uint32_t n_rows; + uint32_t head_dim; + uint64_t row_stride; +}; + +struct ds4_metal_args_dsv4_ratio4_shift { + uint32_t width; +}; + +struct ds4_metal_args_dsv4_compressor_pack_ratio4 { + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t n_threads; +}; + +struct ds4_metal_args_dsv4_compressor_store_one { + uint32_t width; + uint32_t ratio; + uint32_t pos; + uint32_t ape_type; +}; + +static inline float dsv4_e4m3fn_value(int i) { + const int exp = (i >> 3) & 0x0f; + const int mant = i & 0x07; + return exp == 0 + ? float(mant) * 0.001953125f + : (1.0f + float(mant) * 0.125f) * dsv4_e4m3fn_exp_scale[exp]; +} + +static inline float dsv4_e4m3fn_dequant(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float ax = min(abs(x), 448.0f); + + int lo = 0; + int hi = 126; + while (lo < hi) { + const int mid = (lo + hi + 1) >> 1; + if (dsv4_e4m3fn_value(mid) <= ax) { + lo = mid; + } else { + hi = mid - 1; + } + } + + int best = lo; + if (best < 126) { + const float best_diff = abs(ax - dsv4_e4m3fn_value(best)); + const float next_diff = abs(ax - dsv4_e4m3fn_value(best + 1)); + if (next_diff < best_diff || (next_diff == best_diff && ((best + 1) & 1) == 0 && (best & 1) != 0)) { + best = best + 1; + } + } + + return sign * dsv4_e4m3fn_value(best); +} + +static inline float dsv4_e2m1fn_dequant(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float ax = min(abs(x), 6.0f); + int best = 0; + float best_diff = abs(ax - dsv4_e2m1fn_values[0]); + for (int i = 1; i < 8; i++) { + const float diff = abs(ax - dsv4_e2m1fn_values[i]); + if (diff < best_diff || (diff == best_diff && ((i & 1) == 0) && ((best & 1) != 0))) { + best = i; + best_diff = diff; + } + } + return sign * dsv4_e2m1fn_values[best]; +} + +// Quantizes the non-RoPE part of a KV row through E4M3FN and writes the +// dequantized value back as float. DS4 uses this to match the FP8 KV-cache +// semantics while keeping the Metal graph's cache buffers float-addressable. +kernel void kernel_dsv4_fp8_kv_quantize_f32( + constant ds4_metal_args_dsv4_fp8_kv_quantize & args, + device const char * src0, + device char * dst, + threadgroup float * scratch [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + const int64_t n_rows = args.ne01 * args.ne02 * args.ne03; + if ((int64_t) row >= n_rows) { + return; + } + + const int64_t i1 = row % args.ne01; + const int64_t i2 = (row / args.ne01) % args.ne02; + const int64_t i3 = row / (args.ne01 * args.ne02); + + device const char * src_base = src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03; + device char * dst_base = dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3; + + const int64_t n_nope = args.ne00 - args.n_rot; + + for (int64_t off = 0; off < n_nope; off += 64) { + float v = 0.0f; + if (tid < 64) { + v = *((device const float *) (src_base + (off + tid)*args.nb00)); + scratch[tid] = abs(v); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint stride = 32; stride > 0; stride >>= 1) { + if (tid < stride) { + scratch[tid] = max(scratch[tid], scratch[tid + stride]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float amax = max(scratch[0], 1.0e-4f); + const float scale = exp2(ceil(log2(amax / 448.0f))); + if (tid < 64) { + const float q = dsv4_e4m3fn_dequant(clamp(v / scale, -448.0f, 448.0f)) * scale; + *((device float *) (dst_base + (off + tid)*args.nb0)) = q; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (int64_t i = n_nope + tid; i < args.ne00; i += 64) { + *((device float *) (dst_base + i*args.nb0)) = *((device const float *) (src_base + i*args.nb00)); + } +} + +// The official DS4 indexer applies a 128-wide Hadamard rotation and then an +// inplace FP4 activation-simulation pass to both indexer Q and indexer KV. +kernel void kernel_dsv4_indexer_hadamard_fp4_f32( + constant ds4_metal_args_dsv4_indexer_qat & args, + device char * x, + threadgroup float * scratch [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + if (row >= args.n_rows || args.head_dim != 128u || tid >= 128u) { + return; + } + + threadgroup float *vals = scratch; + threadgroup float *absbuf = scratch + 128; + device float *xr = (device float *)(x + (uint64_t)row * args.row_stride); + + vals[tid] = xr[tid]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint stride = 1u; stride < 128u; stride <<= 1u) { + if ((tid & stride) == 0u) { + const uint base = (tid & ~(2u * stride - 1u)) + (tid & (stride - 1u)); + const float a = vals[base]; + const float b = vals[base + stride]; + vals[base] = a + b; + vals[base + stride] = a - b; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + float v = vals[tid] * 0.08838834764831845f; + const uint block = tid >> 5u; + const uint lane = tid & 31u; + const uint block_base = block * 32u; + absbuf[tid] = abs(v); + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint stride = 16u; stride > 0u; stride >>= 1u) { + if (lane < stride) { + absbuf[block_base + lane] = max(absbuf[block_base + lane], + absbuf[block_base + lane + stride]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float amax = max(absbuf[block_base], 7.052966104933725e-38f); + const float scale = exp2(ceil(log2(amax / 6.0f))); + xr[tid] = dsv4_e2m1fn_dequant(clamp(v / scale, -6.0f, 6.0f)) * scale; +} + +// Decode-side KV finalizer after RoPE. The normal RoPE kernel intentionally +// remains separate because tiny trigonometric codegen changes can flip later +// sampled tokens. This kernel only fuses the FP8 round-trip for the non-RoPE +// prefix with the F16-rounded raw-cache row used by FlashAttention. +kernel void kernel_dsv4_kv_fp8_store_f32( + constant ds4_metal_args_dsv4_kv_fp8_store & args, + device float * kv, + device float * raw_cache, + threadgroup float * scratch [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]]) { + const int head_dim = args.head_dim; + const int n_rot = args.n_rot; + const int n_nope = head_dim - n_rot; + if (head_dim <= 0 || n_rot < 0 || n_nope < 0 || tid >= 64) { + return; + } + + device float * raw = raw_cache + (int64_t)args.raw_row * head_dim; + + for (int off = 0; off < n_nope; off += 64) { + float v = 0.0f; + if (off + (int)tid < n_nope) { + v = kv[off + tid]; + scratch[tid] = abs(v); + } else { + scratch[tid] = 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint stride = 32; stride > 0; stride >>= 1) { + if (tid < stride) { + scratch[tid] = max(scratch[tid], scratch[tid + stride]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float amax = max(scratch[0], 1.0e-4f); + const float fp8_scale = exp2(ceil(log2(amax / 448.0f))); + if (off + (int)tid < n_nope) { + const float q = dsv4_e4m3fn_dequant(clamp(v / fp8_scale, -448.0f, 448.0f)) * fp8_scale; + kv[off + tid] = q; + // Diagnostic only: skip the FP16 round-trip that normally matches the + // half-typed FlashAttention KV buffer's precision. With this enabled the + // indexer will see higher-precision raw values than FlashAttention does, + // which is informative but not a production-ready setting. +#ifdef DS4_METAL_KV_RAW_F32 + raw[off + tid] = q; +#else + raw[off + tid] = (float)((half)q); +#endif + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (int i = n_nope + tid; i < head_dim; i += 64) { +#ifdef DS4_METAL_KV_RAW_F32 + raw[i] = kv[i]; +#else + raw[i] = (float)((half)kv[i]); +#endif + } +} + +// Builds the two ratio-4 softmax-pool packs together. Each output plane holds +// four previous-half rows followed by four current-half rows. Normal prefill +// seeds plane zero with 0/-inf; replay seeds it from the previous compressor +// state. Later planes take their previous half from the prior input group. +kernel void kernel_dsv4_compressor_pack_ratio4( + constant ds4_metal_args_dsv4_compressor_pack_ratio4 & args, + device const uint * kv, + device const uint * score, + device const uint * state_kv, + device const uint * state_score, + device uint * packed_kv, + device uint * packed_score, + uint2 group [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + if (group.x >= args.n_comp || group.y >= 8u || + args.head_dim == 0u || args.n_threads == 0u) { + return; + } + + const uint plane = group.x; + const uint row = group.y; + const uint64_t input_row_stride = 2ull * args.head_dim; + const uint64_t dst_row = ((uint64_t)plane * 8u + row) * args.head_dim; + + for (uint col = tid; col < args.head_dim; col += args.n_threads) { + const uint64_t dst = dst_row + col; + if (row >= 4u) { + const uint token = plane * 4u + (row - 4u); + const uint64_t src = (uint64_t)token * input_row_stride + + args.head_dim + col; + packed_kv[dst] = kv[src]; + packed_score[dst] = score[src]; + } else if (plane != 0u) { + const uint token = (plane - 1u) * 4u + row; + const uint64_t src = (uint64_t)token * input_row_stride + col; + packed_kv[dst] = kv[src]; + packed_score[dst] = score[src]; + } else if (args.replay != 0u) { + const uint64_t src = (uint64_t)row * input_row_stride + col; + packed_kv[dst] = state_kv[src]; + packed_score[dst] = state_score[src]; + } else { + packed_kv[dst] = 0u; + packed_score[dst] = 0xff800000u; + } + } +} + +// Ratio-4 compression keeps two 4-row halves of recurrent state. After an +// emitted compressed row, the second half becomes the next window's previous +// half. The old encoder expressed this as four generic copies; this DS4-specific +// kernel performs the KV and score copies together. +kernel void kernel_dsv4_ratio4_shift_f32( + constant ds4_metal_args_dsv4_ratio4_shift & args, + device float * state_kv, + device float * state_score, + uint gid [[thread_position_in_grid]]) { + const uint n = 4u * args.width; + if (gid >= n) return; + + state_kv[gid] = state_kv[n + gid]; + state_score[gid] = state_score[n + gid]; +} + +// One-token compressor frontier update. Decode appends exactly one projected KV +// row and one score row into a small recurrent state. The generic batch helper +// expresses this as APE copy, score add, and two set_rows operations; this +// kernel writes both state tensors directly while preserving the same +// score + APE arithmetic. +kernel void kernel_dsv4_compressor_store_one( + constant ds4_metal_args_dsv4_compressor_store_one & args, + device const float * kv, + device const float * score, + device const char * ape, + device float * state_kv, + device float * state_score, + uint gid [[thread_position_in_grid]]) { + if (gid >= args.width || args.width == 0 || args.ratio == 0) { + return; + } + + const uint pos_mod = args.pos % args.ratio; + const uint dst_row = args.ratio == 4u ? args.ratio + pos_mod : pos_mod; + const uint dst = dst_row * args.width + gid; + const uint ape_i = pos_mod * args.width + gid; + + float ape_v; + if (args.ape_type == 1u) { + ape_v = (float)(((device const half *)ape)[ape_i]); + } else { + ape_v = ((device const float *)ape)[ape_i]; + } + + state_kv[dst] = kv[gid]; + state_score[dst] = score[gid] + ape_v; +} diff --git a/metal/dsv4_misc.metal b/metal/dsv4_misc.metal new file mode 100644 index 0000000..b24167a --- /dev/null +++ b/metal/dsv4_misc.metal @@ -0,0 +1,6157 @@ +struct ds4_metal_args_dsv4_topk_mask { + int64_t ne00; + int64_t ne01; + uint64_t nb00; + uint64_t nb01; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; +}; + +struct ds4_metal_args_dsv4_indexer_weighted_sum { + int64_t ne00; + int64_t ne01; + int64_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + int64_t ne10; + int64_t ne11; + uint64_t nb10; + uint64_t nb11; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; + float scale; +}; + +struct ds4_metal_args_dsv4_softmax_pool { + int64_t ne00; + int64_t ne01; + int64_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; +}; + +struct ds4_metal_args_dsv4_softmax_pool_ratio4_direct { + int64_t n_rows; + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t pad; +}; + +struct ds4_metal_args_dsv4_compressor_score_ape { + uint32_t width; + uint32_t ratio; + uint32_t pos0; + uint32_t n_tokens; +}; + +struct ds4_metal_args_dsv4_indexed_attention { + uint32_t n_tokens; + uint32_t n_head; + uint32_t n_raw; + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_comp; + uint32_t top_k; + uint32_t pos0; + uint32_t window; + uint32_t ratio; + uint32_t comp_kv_f16; + uint32_t pad0; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t raw_row_stride; + uint64_t comp_row_stride; + uint64_t topk_token_stride; + uint64_t dst_token_stride; + uint64_t dst_head_stride; + float scale; +}; + +struct ds4_metal_args_dsv4_indexer_scores_fused { + uint32_t n_comp; + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t pos0; + uint32_t ratio; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t weights_token_stride; + uint64_t index_row_stride; + uint64_t score_token_stride; + float scale; +}; + +struct ds4_metal_args_dsv4_router_select_one { + uint32_t has_bias; + uint32_t hash_mode; + uint32_t use_token_buffer; + uint32_t token; + uint32_t hash_rows; +}; + +struct ds4_metal_args_glm_router_select_one { + uint32_t n_expert; + uint32_t n_expert_used; + float expert_weight_scale; + uint32_t pad0; +}; + +struct ds4_metal_args_glm_kv_lora_rms_norm { + uint32_t n_tokens; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + float eps; +}; + +struct ds4_metal_args_glm_k_b_project { + uint32_t n_tokens; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t n_head; + uint32_t row_bytes; + uint32_t weight_type; + uint32_t pad1; + uint32_t pad2; +}; + +struct ds4_metal_args_glm_build_kv_cache { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t n_head; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_rope; + uint32_t value_dim; + uint32_t n_ctx_orig; + uint32_t cache_f16; + uint32_t pad0; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +}; + +struct ds4_metal_args_glm_store_compact_kv { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t qk_rope; + uint32_t cache_f16; + uint32_t pad1; +}; + +struct ds4_metal_args_glm_qkv_norm_store_compact_kv { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t q_n; + uint32_t q_n4; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t kv_lora_n4; + uint32_t qk_rope; + uint32_t cache_f16; + float eps; + uint32_t pad0; +}; + +struct ds4_metal_args_glm_store_indexer_k { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t head_dim; + uint32_t rot_dim; + uint32_t n_ctx_orig; + uint32_t cache_f16; + uint32_t pad0; + float eps; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + float pad1; +}; + +struct ds4_metal_args_glm_attention_full { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_len; + uint32_t cache_cap; + uint32_t n_head; + uint32_t qk_dim; + uint32_t value_dim; + uint32_t pad0; + uint32_t cache_f16; + uint32_t pad1; + uint32_t pad2; + float scale; +}; + +struct ds4_metal_args_glm_fill_selected_range { + uint32_t n_selected; +}; + +struct ds4_metal_args_glm_fill_selected_range_batch { + uint32_t n_tokens; + uint32_t pos0; + uint32_t n_selected; + uint32_t pad_row; +}; + +struct ds4_metal_args_glm_indexer_rope_tail { + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t rot_dim; + uint32_t rot_offset; + uint32_t pos0; + uint32_t n_ctx_orig; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +}; + +struct ds4_metal_args_glm_indexer_score_one { + uint32_t n_rows; + uint32_t n_head; + uint32_t head_dim; + uint32_t cache_f16; + float scale; +}; + +struct ds4_metal_args_glm_indexer_scores_batch { + uint32_t n_rows; + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t pos0; + uint32_t cache_f16; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t weights_token_stride; + uint64_t score_token_stride; + float scale; +}; + +struct ds4_metal_args_glm_qk_lowrank { + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_dim; + uint32_t row_bytes; + uint32_t weight_type; + uint32_t pad1; + uint32_t pad2; +}; + +struct ds4_metal_args_glm_qk_lowrank_batch { + uint32_t n_tokens; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_dim; + uint32_t row_bytes; + uint32_t weight_type; + /* First head this dispatch computes: under tensor-parallel head split + * each rank covers a contiguous half of the heads; buffers and weights + * keep full-model layout and are indexed by absolute head. */ + uint32_t head_base; +}; + +struct ds4_metal_args_glm_attention_indexed_decode { + uint32_t n_selected; + uint32_t cache_cap; + uint32_t cache_f16; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_rope; + uint32_t value_dim; + uint32_t n_ctx_orig; + uint32_t value_row_bytes; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t value_type; +}; + +struct ds4_metal_args_glm_attention_indexed_decode_split { + uint32_t n_selected; + uint32_t cache_cap; + uint32_t cache_f16; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_rope; + uint32_t value_dim; + uint32_t n_ctx_orig; + uint32_t value_row_bytes; + uint32_t block_rows; + uint32_t n_blocks; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t value_type; +}; + +struct ds4_metal_args_glm_attention_indexed_batch { + uint32_t n_tokens; + uint32_t n_selected; + uint32_t cache_cap; + uint32_t cache_f16; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_rope; + uint32_t value_dim; + uint32_t n_ctx_orig; + uint32_t value_row_bytes; + uint32_t value_type; + uint32_t pos0; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t head_base; +}; + +struct ds4_metal_args_dsv4_directional_steering_project { + uint32_t width; + uint32_t rows; + uint32_t layer; + uint32_t n_threads; + float scale; +}; + +// Optional directional steering projection. +// +// Each threadgroup owns one 4096-wide token row, computes +// dot(row, direction[layer]), then subtracts scale * direction * dot in-place. +// Positive scales remove a concept direction; negative scales amplify it. The +// kernel is not used unless a steering file and nonzero scale are provided. +kernel void kernel_dsv4_directional_steering_project_f32( + constant ds4_metal_args_dsv4_directional_steering_project & args, + device float *x, + device const float *directions, + threadgroup float *scratch [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + if (row >= args.rows || args.width == 0) return; + + device float *xr = x + (uint64_t)row * args.width; + device const float *dir = directions + (uint64_t)args.layer * args.width; + const uint nth = args.n_threads; + + float sum = 0.0f; + for (uint i = tid; i < args.width; i += nth) { + sum += xr[i] * dir[i]; + } + scratch[tid] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) scratch[tid] += scratch[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float coeff = args.scale * scratch[0]; + for (uint i = tid; i < args.width; i += nth) { + xr[i] -= coeff * dir[i]; + } +} + +// Decode-only DS4 ratio-4 indexer score builder. One threadgroup owns one +// compressed row for the current token, stages that 128-wide row once, then +// walks the 64 indexer heads in four-head groups. This avoids materializing the +// intermediate [compressed rows x heads] score matrix used by the generic +// matvec + weighted-sum path. +kernel void kernel_dsv4_indexer_score_one_direct( + constant ds4_metal_args_dsv4_indexer_scores_fused & args, + device const char *q, + device const char *weights, + device const char *index_comp, + device char *scores, + threadgroup float *shared [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + if (row >= args.n_comp || args.n_head != 64u || args.head_dim != 128u) { + return; + } + + threadgroup float *ktg = shared; // [128] + threadgroup float *psum = ktg + 128u; // [4] + + if (tid < 128u) { + device const float *krow = (device const float *)(index_comp + + (uint64_t)row * args.index_row_stride); + ktg[tid] = krow[tid]; + } + + float acc = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint head0 = 0; head0 < 64u; head0 += 4u) { + const uint head = head0 + (uint)sg; + device const float4 *q4 = (device const float4 *)(q + + (uint64_t)head * args.q_head_stride); + threadgroup const float4 *k4 = (threadgroup const float4 *)ktg; + + float s = dot(q4[lane], k4[lane]); + s = simd_sum(s); + if (lane == 0) { + device const float *w = (device const float *)weights; + psum[sg] = max(s, 0.0f) * (w[head] * args.scale); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + acc += psum[0]; + acc += psum[1]; + acc += psum[2]; + acc += psum[3]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (tid == 0) { + device float *dst = (device float *)scores; + dst[row] = acc; + } +} + +// Decode router post-processing for one token. The selected expert ids are +// already known; this gathers their probabilities, normalizes by the selected +// sum, clamps the denominator like the reference path, and applies DS4's 1.5 +// expert-weight scale in one tiny dispatch. +kernel void kernel_dsv4_router_weights_one( + device const char *probs, + device const char *selected, + device char *weights, + uint tid [[thread_position_in_grid]]) { + if (tid >= 6) return; + + device const float *p = (device const float *)probs; + device const int *s = (device const int *)selected; + + float sum = 0.0f; + for (uint i = 0; i < 6; i++) { + sum += p[s[i]]; + } + sum = max(sum, 6.103515625e-5f); + + device float *w = (device float *)weights; + w[tid] = p[s[tid]] / sum * 1.5f; +} + +static inline float ds4_glm_router_sigmoid(float x) { + if (x >= 0.0f) { + const float e = exp(-x); + return 1.0f / (1.0f + e); + } else { + const float e = exp(x); + return e / (1.0f + e); + } +} + +static inline bool ds4_glm_router_better( + threadgroup const float *scores, + int32_t a, + int32_t b) { + const float sa = scores[(uint)a]; + const float sb = scores[(uint)b]; + return sa > sb || (sa == sb && a < b); +} + +static float glm_rope_yarn_ramp(const float low, const float high, const int i0) { + const float y = (i0 / 2 - low) / max(0.001f, high - low); + return 1.0f - min(1.0f, max(0.0f, y)); +} + +static void glm_rope_yarn( + float theta_extrap, + float freq_scale, + float corr_dims[2], + int i0, + float ext_factor, + float mscale, + thread float *cos_theta, + thread float *sin_theta) { + float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + if (ext_factor != 0.0f) { + float ramp_mix = glm_rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * ext_factor; + theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix; + mscale *= 1.0f + 0.1f * log(1.0f / freq_scale); + } + *cos_theta = cos(theta) * mscale; + *sin_theta = sin(theta) * mscale; +} + +static float glm_rope_yarn_corr_factor(int n_dims, int n_ctx_orig, float n_rot, float base) { + return n_dims * log(n_ctx_orig / (n_rot * 2 * M_PI_F)) / (2 * log(base)); +} + +static void glm_rope_yarn_corr_dims( + int n_dims, + int n_ctx_orig, + float freq_base, + float beta_fast, + float beta_slow, + float dims[2]) { + dims[0] = max(0.0f, + floor(glm_rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_fast, freq_base))); + dims[1] = min(n_dims - 1.0f, + ceil(glm_rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_slow, freq_base))); +} + +kernel void kernel_glm_kv_lora_rms_norm( + constant ds4_metal_args_glm_kv_lora_rms_norm & args, + device const char *src, + device const char *weight, + device char *dst, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid_u [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]]) { + const uint row = tgpig.x; + if (row >= args.n_tokens) return; + + const uint tid = tid_u; + const uint nth = ntg_u.x; + device const float *x = (device const float *)(src + (uint64_t)row * args.kv_raw_dim * sizeof(float)); + device const float *w = (device const float *)weight; + device float *out = (device float *)(dst + (uint64_t)row * args.kv_lora_dim * sizeof(float)); + + float ss = 0.0f; + for (uint i = tid; i < args.kv_lora_dim; i += nth) { + const float v = x[i]; + ss += v * v; + } + scratch[tid] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) scratch[tid] += scratch[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float inv = rsqrt(scratch[0] / (float)args.kv_lora_dim + args.eps); + for (uint i = tid; i < args.kv_lora_dim; i += nth) { + out[i] = x[i] * inv * w[i]; + } +} + +static inline float glm_quant_weight_at( + uint weight_type, + device const char *row, + uint col); + +kernel void kernel_glm_k_b_project_q8_0( + constant ds4_metal_args_glm_k_b_project & args, + device const char *weight, + device const char *kv_norm, + device char *dst, + threadgroup float *kv_scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint token = tgpig.x; + const uint head = tgpig.y; + if (token >= args.n_tokens || head >= args.n_head) return; + + const uint nth = (uint)ntg_u.x * (uint)ntg_u.y; + device const float *kv = + (device const float *)(kv_norm + (uint64_t)token * args.kv_lora_dim * sizeof(float)); + device float *out = + (device float *)(dst + + ((uint64_t)token * args.n_head + head) * args.qk_nope * sizeof(float)); + + for (uint j = tid; j < args.kv_lora_dim; j += nth) { + kv_scratch[j] = kv[j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint block = (uint)sgitg; + const uint q = (block << 5) + (uint)tiisg; + if (q < args.qk_nope) { + float acc = 0.0f; + for (uint j = 0; j < args.kv_lora_dim; j++) { + device const char *row = + weight + ((uint64_t)head * args.kv_lora_dim + j) * args.row_bytes; + acc += glm_quant_weight_at(args.weight_type, row, q) * kv_scratch[j]; + } + out[q] = acc; + } +} + +kernel void kernel_glm_store_compact_kv( + constant ds4_metal_args_glm_store_compact_kv & args, + device const char *kv_norm, + device const char *kv_raw, + device char *kv_lora_cache, + device char *k_rope_cache, + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint token = tgpig.x; + const uint part = tgpig.y; + if (token >= args.n_tokens || part > 1u) return; + + const uint pos = args.pos0 + token; + if (pos >= args.cache_cap) return; + + const uint nth = ntg_u.x; + if (part == 0) { + device const float *src = + (device const float *)(kv_norm + + (uint64_t)token * args.kv_lora_dim * sizeof(float)); + if (args.cache_f16 != 0u) { + device half *dst = + (device half *)(kv_lora_cache + + (uint64_t)pos * args.kv_lora_dim * sizeof(half)); + for (uint i = tid; i < args.kv_lora_dim; i += nth) { + dst[i] = (half)src[i]; + } + } else { + device float *dst = + (device float *)(kv_lora_cache + + (uint64_t)pos * args.kv_lora_dim * sizeof(float)); + for (uint i = tid; i < args.kv_lora_dim; i += nth) { + dst[i] = src[i]; + } + } + } else { + device const float *src = + (device const float *)(kv_raw + + ((uint64_t)token * args.kv_raw_dim + args.kv_lora_dim) * sizeof(float)); + if (args.cache_f16 != 0u) { + device half *dst = + (device half *)(k_rope_cache + + (uint64_t)pos * args.qk_rope * sizeof(half)); + for (uint i = tid; i < args.qk_rope; i += nth) { + dst[i] = (half)src[i]; + } + } else { + device float *dst = + (device float *)(k_rope_cache + + (uint64_t)pos * args.qk_rope * sizeof(float)); + for (uint i = tid; i < args.qk_rope; i += nth) { + dst[i] = src[i]; + } + } + } +} + +kernel void kernel_glm_qkv_norm_store_compact_kv( + constant ds4_metal_args_glm_qkv_norm_store_compact_kv & args, + device const char *q_src, + device const char *q_weight, + device char *q_dst, + device const char *kv_raw, + device const char *kv_weight, + device char *kv_lora_cache, + device char *k_rope_cache, + threadgroup float *shmem_f32 [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint token = tgpig.x; + const uint part = tgpig.y; + if (token >= args.n_tokens || part > 2u) return; + + const uint nth = ntg_u.x; + if (part == 2u) { + const uint pos = args.pos0 + token; + if (pos >= args.cache_cap) return; + device const float *src = + (device const float *)(kv_raw + + ((uint64_t)token * args.kv_raw_dim + args.kv_lora_dim) * sizeof(float)); + if (args.cache_f16 != 0u) { + device half *dst = + (device half *)(k_rope_cache + + (uint64_t)pos * args.qk_rope * sizeof(half)); + for (uint i = tid; i < args.qk_rope; i += nth) { + dst[i] = (half)src[i]; + } + } else { + device float *dst = + (device float *)(k_rope_cache + + (uint64_t)pos * args.qk_rope * sizeof(float)); + for (uint i = tid; i < args.qk_rope; i += nth) { + dst[i] = src[i]; + } + } + return; + } + + if (sgitg == 0) { + shmem_f32[tiisg] = 0.0f; + } + + const bool kv_task = part != 0u; + const uint n = kv_task ? args.kv_lora_dim : args.q_n; + const uint n4 = kv_task ? args.kv_lora_n4 : args.q_n4; + device const float4 *x = + kv_task + ? (device const float4 *)(kv_raw + + (uint64_t)token * args.kv_raw_dim * sizeof(float)) + : (device const float4 *)(q_src + + (uint64_t)token * args.q_n * sizeof(float)); + device const float4 *w = + kv_task ? (device const float4 *)kv_weight + : (device const float4 *)q_weight; + + float sumf = 0.0f; + for (uint i = tid; i < n4; i += nth) { + const float4 v = x[i]; + sumf += dot(v, v); + } + sumf = simd_sum(sumf); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiisg == 0) { + shmem_f32[sgitg] = sumf; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + sumf = shmem_f32[tiisg]; + sumf = simd_sum(sumf); + +#ifdef DS4_METAL_NORM_RSQRT_DISABLE + const float scale = 1.0f / sqrt(sumf / float(n) + args.eps); +#else + const float scale = rsqrt(sumf / float(n) + args.eps); +#endif + + if (!kv_task) { + device float4 *y = + (device float4 *)(q_dst + + (uint64_t)token * args.q_n * sizeof(float)); + for (uint i = tid; i < n4; i += nth) { + y[i] = (x[i] * scale) * w[i]; + } + return; + } + + const uint pos = args.pos0 + token; + if (pos >= args.cache_cap) return; + device const float *x1 = + (device const float *)(kv_raw + + (uint64_t)token * args.kv_raw_dim * sizeof(float)); + device const float *w1 = (device const float *)kv_weight; + if (args.cache_f16 != 0u) { + device half *dst = + (device half *)(kv_lora_cache + + (uint64_t)pos * args.kv_lora_dim * sizeof(half)); + for (uint i = tid; i < args.kv_lora_dim; i += nth) { + dst[i] = (half)((x1[i] * scale) * w1[i]); + } + } else { + device float *dst = + (device float *)(kv_lora_cache + + (uint64_t)pos * args.kv_lora_dim * sizeof(float)); + for (uint i = tid; i < args.kv_lora_dim; i += nth) { + dst[i] = (x1[i] * scale) * w1[i]; + } + } +} + +kernel void kernel_glm_store_indexer_k( + constant ds4_metal_args_glm_store_indexer_k & args, + device const char *raw_k, + device const char *weight, + device const char *bias, + device char *indexer_key_cache, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint token = tgpig.x; + if (token >= args.n_tokens) return; + + const uint pos = args.pos0 + token; + if (pos >= args.cache_cap) return; + + const uint nth = ntg_u.x; + const uint head_dim = args.head_dim; + const uint rot_dim = args.rot_dim; + + device const float *src = + (device const float *)(raw_k + (uint64_t)token * head_dim * sizeof(float)); + device const float *w = (device const float *)weight; + device const float *b = (device const float *)bias; + + float sum = 0.0f; + for (uint i = tid; i < head_dim; i += nth) { + sum += src[i]; + } + scratch[tid] = sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) scratch[tid] += scratch[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float mean = scratch[0] / (float)head_dim; + + float ss = 0.0f; + for (uint i = tid; i < head_dim; i += nth) { + const float d = src[i] - mean; + ss += d * d; + } + scratch[tid] = ss; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) scratch[tid] += scratch[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float inv = rsqrt(scratch[0] / (float)head_dim + args.eps); + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)rot_dim, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + const float theta_base = (float)pos; + const float inv_ndims = -1.0f / (float)rot_dim; + + if (args.cache_f16 != 0u) { + device half *dst = + (device half *)(indexer_key_cache + + (uint64_t)pos * head_dim * sizeof(half)); + for (uint i = tid; i < head_dim; i += nth) { + if (i < rot_dim) { + if ((i & 1u) != 0u) continue; + const uint rel_i0 = i; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)rel_i0 * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims * (float)rel_i0); +#endif + float cos_theta; + float sin_theta; + glm_rope_yarn(theta, + args.freq_scale, + corr_dims, + (int)rel_i0, + args.ext_factor, + args.attn_factor, + &cos_theta, + &sin_theta); + const float x0 = (src[i] - mean) * inv * w[i] + b[i]; + const uint j = i + 1u; + const float x1 = (src[j] - mean) * inv * w[j] + b[j]; + dst[i] = (half)(x0 * cos_theta - x1 * sin_theta); + dst[j] = (half)(x0 * sin_theta + x1 * cos_theta); + } else if (i >= rot_dim) { + const float x = (src[i] - mean) * inv * w[i] + b[i]; + dst[i] = (half)x; + } + } + } else { + device float *dst = + (device float *)(indexer_key_cache + + (uint64_t)pos * head_dim * sizeof(float)); + for (uint i = tid; i < head_dim; i += nth) { + if (i < rot_dim) { + if ((i & 1u) != 0u) continue; + const uint rel_i0 = i; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)rel_i0 * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims * (float)rel_i0); +#endif + float cos_theta; + float sin_theta; + glm_rope_yarn(theta, + args.freq_scale, + corr_dims, + (int)rel_i0, + args.ext_factor, + args.attn_factor, + &cos_theta, + &sin_theta); + const float x0 = (src[i] - mean) * inv * w[i] + b[i]; + const uint j = i + 1u; + const float x1 = (src[j] - mean) * inv * w[j] + b[j]; + dst[i] = x0 * cos_theta - x1 * sin_theta; + dst[j] = x0 * sin_theta + x1 * cos_theta; + } else if (i >= rot_dim) { + const float x = (src[i] - mean) * inv * w[i] + b[i]; + dst[i] = x; + } + } + } +} + +static inline void glm_dense_cache_store_f32_or_f16( + device char *base, + uint64_t index, + uint cache_f16, + float x) { + if (cache_f16 != 0u) { + ((device half *)base)[index] = (half)x; + } else { + ((device float *)base)[index] = x; + } +} + +static inline float glm_dense_cache_load_f32_or_f16( + device const char *base, + uint64_t index, + uint cache_f16) { + if (cache_f16 != 0u) { + return (float)((device const half *)base)[index]; + } + return ((device const float *)base)[index]; +} + +static inline float4 glm_dense_cache_load4_f32_or_f16( + device const char *base, + uint64_t index, + uint cache_f16) { + if (cache_f16 != 0u) { + device const half *h = (device const half *)base; + return float4((float)h[index + 0u], + (float)h[index + 1u], + (float)h[index + 2u], + (float)h[index + 3u]); + } + return ((device const float4 *)base)[index >> 2u]; +} + +kernel void kernel_glm_build_kv_cache( + constant ds4_metal_args_glm_build_kv_cache & args, + device const char *kv_raw, + device const char *k_nope, + device const char *value, + device char *key_cache, + device char *value_cache, + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint token = tgpig.x; + const uint head = tgpig.y; + if (token >= args.n_tokens || head >= args.n_head) return; + + const uint nth = ntg_u.x; + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint pos = args.pos0 + token; + device const float *raw = + (device const float *)(kv_raw + (uint64_t)token * args.kv_raw_dim * sizeof(float)); + device const float *kn = + (device const float *)(k_nope + + ((uint64_t)token * args.n_head + head) * args.qk_nope * sizeof(float)); + device const float *val = + (device const float *)(value + + ((uint64_t)token * args.n_head + head) * args.value_dim * sizeof(float)); + const uint64_t kbase = ((uint64_t)pos * args.n_head + head) * qk_dim; + const uint64_t vbase = ((uint64_t)pos * args.n_head + head) * args.value_dim; + + for (uint i = tid; i < args.qk_nope; i += nth) { + glm_dense_cache_store_f32_or_f16(key_cache, kbase + i, args.cache_f16, kn[i]); + } + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + const float theta_base = (float)pos; + const float inv_ndims = -1.0f / (float)args.qk_rope; + for (uint r = tid * 2u; r < args.qk_rope; r += nth * 2u) { +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims * (float)r); +#endif + float cos_theta; + float sin_theta; + glm_rope_yarn(theta, + args.freq_scale, + corr_dims, + (int)r, + args.ext_factor, + args.attn_factor, + &cos_theta, + &sin_theta); + const uint src0 = args.kv_lora_dim + r; + const float x0 = raw[src0]; + const float x1 = raw[src0 + 1u]; + const uint dst0 = args.qk_nope + r; + glm_dense_cache_store_f32_or_f16(key_cache, + kbase + dst0, + args.cache_f16, + x0 * cos_theta - x1 * sin_theta); + glm_dense_cache_store_f32_or_f16(key_cache, + kbase + dst0 + 1u, + args.cache_f16, + x0 * sin_theta + x1 * cos_theta); + } + + for (uint i = tid; i < args.value_dim; i += nth) { + glm_dense_cache_store_f32_or_f16(value_cache, vbase + i, args.cache_f16, val[i]); + } +} + +kernel void kernel_glm_build_kv_cache_decode_group4( + constant ds4_metal_args_glm_build_kv_cache & args, + device const char *kv_raw, + device const char *k_nope, + device const char *value, + device char *key_cache, + device char *value_cache, + uint tid [[thread_index_in_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint token = tgpig.x; + const uint group_head0 = tgpig.y * 4u; + if (token >= args.n_tokens || group_head0 >= args.n_head) return; + + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint pos = args.pos0 + token; + const uint lane = tid & 63u; + const uint slot = tid >> 6; + device const float *raw = + (device const float *)(kv_raw + (uint64_t)token * args.kv_raw_dim * sizeof(float)); + + const uint head = group_head0 + slot; + if (slot < 4u && head < args.n_head) { + device const float *kn = + (device const float *)(k_nope + + ((uint64_t)token * args.n_head + head) * args.qk_nope * sizeof(float)); + device const float *val = + (device const float *)(value + + ((uint64_t)token * args.n_head + head) * args.value_dim * sizeof(float)); + const uint64_t kbase = ((uint64_t)pos * args.n_head + head) * qk_dim; + const uint64_t vbase = ((uint64_t)pos * args.n_head + head) * args.value_dim; + + for (uint i = lane; i < args.qk_nope; i += 64u) { + glm_dense_cache_store_f32_or_f16(key_cache, kbase + i, args.cache_f16, kn[i]); + } + for (uint i = lane; i < args.value_dim; i += 64u) { + glm_dense_cache_store_f32_or_f16(value_cache, vbase + i, args.cache_f16, val[i]); + } + } + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + const float theta_base = (float)pos; + const float inv_ndims = -1.0f / (float)args.qk_rope; + for (uint r = tid * 2u; r < args.qk_rope; r += 512u) { +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims * (float)r); +#endif + float cos_theta; + float sin_theta; + glm_rope_yarn(theta, + args.freq_scale, + corr_dims, + (int)r, + args.ext_factor, + args.attn_factor, + &cos_theta, + &sin_theta); + const uint src0 = args.kv_lora_dim + r; + const float x0 = raw[src0]; + const float x1 = raw[src0 + 1u]; + const uint dst0 = args.qk_nope + r; + const float y0 = x0 * cos_theta - x1 * sin_theta; + const float y1 = x0 * sin_theta + x1 * cos_theta; + for (uint h = group_head0; h < min(group_head0 + 4u, args.n_head); h++) { + const uint64_t kbase = ((uint64_t)pos * args.n_head + h) * qk_dim; + glm_dense_cache_store_f32_or_f16(key_cache, kbase + dst0, args.cache_f16, y0); + glm_dense_cache_store_f32_or_f16(key_cache, kbase + dst0 + 1u, args.cache_f16, y1); + } + } +} + +kernel void kernel_glm_build_kv_cache_flash( + constant ds4_metal_args_glm_build_kv_cache & args, + device const char *kv_raw, + device const char *k_nope, + device const char *value, + device char *key_cache, + device char *value_cache, + device char *key_f16, + device char *value_f16, + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint token = tgpig.x; + const uint head = tgpig.y; + if (token >= args.n_tokens || head >= args.n_head) return; + + const uint nth = ntg_u.x; + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint pos = args.pos0 + token; + device const float *raw = + (device const float *)(kv_raw + (uint64_t)token * args.kv_raw_dim * sizeof(float)); + device const float *kn = + (device const float *)(k_nope + + ((uint64_t)token * args.n_head + head) * args.qk_nope * sizeof(float)); + device const float *val = + (device const float *)(value + + ((uint64_t)token * args.n_head + head) * args.value_dim * sizeof(float)); + const uint64_t kbase = ((uint64_t)pos * args.n_head + head) * qk_dim; + const uint64_t vbase = ((uint64_t)pos * args.n_head + head) * args.value_dim; + device half *kdst_f16 = + (device half *)(key_f16 + + ((uint64_t)head * args.n_tokens + token) * qk_dim * sizeof(half)); + device half *vdst_f16 = + (device half *)(value_f16 + + ((uint64_t)head * args.n_tokens + token) * args.value_dim * sizeof(half)); + + for (uint i = tid; i < args.qk_nope; i += nth) { + const float x = kn[i]; + glm_dense_cache_store_f32_or_f16(key_cache, kbase + i, args.cache_f16, x); + kdst_f16[i] = (half)x; + } + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + const float theta_base = (float)pos; + const float inv_ndims = -1.0f / (float)args.qk_rope; + for (uint r = tid * 2u; r < args.qk_rope; r += nth * 2u) { +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims * (float)r); +#endif + float cos_theta; + float sin_theta; + glm_rope_yarn(theta, + args.freq_scale, + corr_dims, + (int)r, + args.ext_factor, + args.attn_factor, + &cos_theta, + &sin_theta); + const uint src0 = args.kv_lora_dim + r; + const float x0 = raw[src0]; + const float x1 = raw[src0 + 1u]; + const uint dst0 = args.qk_nope + r; + const float y0 = x0 * cos_theta - x1 * sin_theta; + const float y1 = x0 * sin_theta + x1 * cos_theta; + glm_dense_cache_store_f32_or_f16(key_cache, kbase + dst0, args.cache_f16, y0); + glm_dense_cache_store_f32_or_f16(key_cache, kbase + dst0 + 1u, args.cache_f16, y1); + kdst_f16[dst0] = (half)y0; + kdst_f16[dst0 + 1u] = (half)y1; + } + + for (uint i = tid; i < args.value_dim; i += nth) { + const float x = val[i]; + glm_dense_cache_store_f32_or_f16(value_cache, vbase + i, args.cache_f16, x); + vdst_f16[i] = (half)x; + } +} + +kernel void kernel_glm_attention_full( + constant ds4_metal_args_glm_attention_full & args, + device const char *q, + device const char *key_cache, + device const char *value_cache, + device char *heads, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint token = tgpig.x; + const uint head = tgpig.y; + if (token >= args.n_tokens || head >= args.n_head) return; + + const uint nth = ntg_u.x; + const uint qk4 = args.qk_dim / 4u; + const uint visible = min(args.cache_len, args.pos0 + token + 1u); + threadgroup float *red = scratch; + threadgroup float *scores = scratch + 256u; + + device const float4 *q4 = (device const float4 *)(q + + ((uint64_t)token * args.n_head + head) * args.qk_dim * sizeof(float)); + + if (args.pad0 == 2u) { + for (uint s = tid; s < visible; s += nth) { + const uint64_t kbase = ((uint64_t)s * args.n_head + head) * args.qk_dim; + float dotv = 0.0f; + for (uint i = 0; i < qk4; i++) { + dotv += dot(q4[i], + glm_dense_cache_load4_f32_or_f16(key_cache, + kbase + 4u * (uint64_t)i, + args.cache_f16)); + } + scores[s] = dotv * args.scale; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tid == 0u) { + float max_score = -INFINITY; + for (uint s = 0; s < visible; s++) { + max_score = max(max_score, scores[s]); + } + float sum = 0.0f; + for (uint s = 0; s < visible; s++) { + const float w = exp(scores[s] - max_score); + scores[s] = w; + sum += w; + } + red[0] = max(sum, 1.0e-20f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const float denom = red[0]; + device float *out = (device float *)(heads + + ((uint64_t)token * args.n_head + head) * args.value_dim * sizeof(float)); + for (uint d = tid; d < args.value_dim; d += nth) { + float acc = 0.0f; + for (uint s = 0; s < visible; s++) { + const uint64_t vbase = ((uint64_t)s * args.n_head + head) * args.value_dim; + acc += scores[s] * + glm_dense_cache_load_f32_or_f16(value_cache, + vbase + d, + args.cache_f16); + } + out[d] = acc / denom; + } + return; + } + + if (args.pad0 == 1u) { + if (tid == 0u) { + float max_score = -INFINITY; + for (uint s = 0; s < visible; s++) { + const uint64_t kbase = ((uint64_t)s * args.n_head + head) * args.qk_dim; + float dotv = 0.0f; + for (uint i = 0; i < qk4; i++) { + dotv += dot(q4[i], + glm_dense_cache_load4_f32_or_f16(key_cache, + kbase + 4u * (uint64_t)i, + args.cache_f16)); + } + const float score = dotv * args.scale; + scores[s] = score; + max_score = max(max_score, score); + } + float sum = 0.0f; + for (uint s = 0; s < visible; s++) { + const float w = exp(scores[s] - max_score); + scores[s] = w; + sum += w; + } + red[0] = max(sum, 1.0e-20f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const float denom = red[0]; + device float *out = (device float *)(heads + + ((uint64_t)token * args.n_head + head) * args.value_dim * sizeof(float)); + for (uint d = tid; d < args.value_dim; d += nth) { + float acc = 0.0f; + for (uint s = 0; s < visible; s++) { + const uint64_t vbase = ((uint64_t)s * args.n_head + head) * args.value_dim; + acc += scores[s] * + glm_dense_cache_load_f32_or_f16(value_cache, + vbase + d, + args.cache_f16); + } + out[d] = acc / denom; + } + return; + } + + float local_max = -INFINITY; + for (uint s = tid; s < visible; s += nth) { + const uint64_t kbase = ((uint64_t)s * args.n_head + head) * args.qk_dim; + float dotv = 0.0f; + for (uint i = 0; i < qk4; i++) { + dotv += dot(q4[i], + glm_dense_cache_load4_f32_or_f16(key_cache, + kbase + 4u * (uint64_t)i, + args.cache_f16)); + } + const float score = dotv * args.scale; + scores[s] = score; + local_max = max(local_max, score); + } + red[tid] = local_max; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] = max(red[tid], red[tid + step]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float max_score = red[0]; + + float local_sum = 0.0f; + for (uint s = tid; s < visible; s += nth) { + const float w = exp(scores[s] - max_score); + scores[s] = w; + local_sum += w; + } + red[tid] = local_sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] += red[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float denom = max(red[0], 1.0e-20f); + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *out = (device float *)(heads + + ((uint64_t)token * args.n_head + head) * args.value_dim * sizeof(float)); + for (uint d = tid; d < args.value_dim; d += nth) { + float acc = 0.0f; + for (uint s = 0; s < visible; s++) { + const uint64_t vbase = ((uint64_t)s * args.n_head + head) * args.value_dim; + acc += scores[s] * + glm_dense_cache_load_f32_or_f16(value_cache, + vbase + d, + args.cache_f16); + } + out[d] = acc / denom; + } +} + +kernel void kernel_glm_fill_selected_range( + constant ds4_metal_args_glm_fill_selected_range & args, + device uint32_t *selected, + uint gid [[thread_position_in_grid]]) { + if (gid < args.n_selected) selected[gid] = gid; +} + +kernel void kernel_glm_fill_selected_range_batch( + constant ds4_metal_args_glm_fill_selected_range_batch & args, + device uint32_t *selected, + uint gid [[thread_position_in_grid]]) { + const uint total = args.n_tokens * args.n_selected; + if (gid >= total || args.n_selected == 0u) return; + const uint token = gid / args.n_selected; + const uint slot = gid - token * args.n_selected; + const uint visible = args.pos0 + token + 1u; + selected[gid] = slot < visible ? slot : args.pad_row; +} + +kernel void kernel_glm_indexer_rope_tail_f32( + constant ds4_metal_args_glm_indexer_rope_tail & args, + device char *x, + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint head = tgpig.x; + const uint token = tgpig.y; + if (head >= args.n_head || token >= args.n_tokens) return; + if (args.rot_dim == 0u || args.rot_offset > args.head_dim || + args.rot_dim > args.head_dim - args.rot_offset || (args.rot_dim & 1u) != 0u) return; + + const uint nth = ntg_u.x; + const uint pos = args.pos0 + token; + device float *row = + (device float *)(x + + ((uint64_t)token * args.n_head + head) * args.head_dim * sizeof(float)); + row += args.rot_offset; + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.rot_dim, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + const float theta_base = (float)pos; + const float inv_ndims = -1.0f / (float)args.rot_dim; + for (uint i = tid * 2u; i < args.rot_dim; i += nth * 2u) { + const uint rel_i0 = i; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)rel_i0 * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims * (float)rel_i0); +#endif + float cos_theta; + float sin_theta; + glm_rope_yarn(theta, + args.freq_scale, + corr_dims, + (int)rel_i0, + args.ext_factor, + args.attn_factor, + &cos_theta, + &sin_theta); + const uint j = i + 1u; + const float x0 = row[i]; + const float x1 = row[j]; + row[i] = x0 * cos_theta - x1 * sin_theta; + row[j] = x0 * sin_theta + x1 * cos_theta; + } +} + +static inline float glm_cache_load_f32_or_f16( + device const char *base, + uint64_t index, + uint cache_f16) { + if (cache_f16 != 0u) { + return (float)((device const half *)base)[index]; + } + return ((device const float *)base)[index]; +} + +static inline float glm_cache_load_f16_only( + device const char *base, + uint64_t index) { + return (float)((device const half *)base)[index]; +} + +static inline float2 glm_cache_load_rotated_rope_pair( + device const char *base, + uint64_t rope_base, + uint r, + uint row, + uint qk_rope, + uint cache_f16, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float corr0, + float corr1) { + const float theta_base = (float)row; + const float inv_ndims = -1.0f / (float)qk_rope; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(freq_base)); +#else + const float theta = theta_base * pow(freq_base, inv_ndims * (float)r); +#endif + float corr_dims[2] = {corr0, corr1}; + float cos_theta; + float sin_theta; + glm_rope_yarn(theta, + freq_scale, + corr_dims, + (int)r, + ext_factor, + attn_factor, + &cos_theta, + &sin_theta); + const float x0 = glm_cache_load_f32_or_f16(base, rope_base + r, cache_f16); + const float x1 = glm_cache_load_f32_or_f16(base, rope_base + r + 1u, cache_f16); + return float2(x0 * cos_theta - x1 * sin_theta, + x0 * sin_theta + x1 * cos_theta); +} + +static inline float2 glm_cache_load_rotated_rope_pair_f16_only( + device const char *base, + uint64_t rope_base, + uint r, + uint row, + uint qk_rope, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float corr0, + float corr1) { + const float theta_base = (float)row; + const float inv_ndims = -1.0f / (float)qk_rope; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(freq_base)); +#else + const float theta = theta_base * pow(freq_base, inv_ndims * (float)r); +#endif + float corr_dims[2] = {corr0, corr1}; + float cos_theta; + float sin_theta; + glm_rope_yarn(theta, + freq_scale, + corr_dims, + (int)r, + ext_factor, + attn_factor, + &cos_theta, + &sin_theta); + const float x0 = glm_cache_load_f16_only(base, rope_base + r); + const float x1 = glm_cache_load_f16_only(base, rope_base + r + 1u); + return float2(x0 * cos_theta - x1 * sin_theta, + x0 * sin_theta + x1 * cos_theta); +} + +static inline float glm_q8_0_weight_at( + device const char *row, + uint col) { + const uint block = col >> 5; + const uint qi = col & 31u; + device const char *block_base = row + (uint64_t)block * 34u; + const float d = (float)(*((device const half *)block_base)); + device const int8_t *qs = (device const int8_t *)(block_base + 2u); + return d * (float)qs[qi]; +} + +static inline float glm_q8_0_dot_row_tg_f32( + device const char *row, + threadgroup const float *x, + uint n_cols) { + float acc = 0.0f; + const uint n_blocks = (n_cols + 31u) >> 5; + for (uint block = 0; block < n_blocks; block++) { + device const char *block_base = row + (uint64_t)block * 34u; + const float d = (float)(*((device const half *)block_base)); + device const int8_t *qs = (device const int8_t *)(block_base + 2u); + const uint base = block << 5; + const uint count = min(32u, n_cols - base); + for (uint qi = 0; qi < count; qi++) { + acc += d * (float)qs[qi] * x[base + qi]; + } + } + return acc; +} + +static inline float glm_q8_0_dot_row_tg_f32_512( + device const char *row, + threadgroup const float *x) { + float acc = 0.0f; + for (uint block = 0; block < 16u; block++) { + device const char *block_base = row + (uint64_t)block * 34u; + const float d = (float)(*((device const half *)block_base)); + device const int8_t *qs = (device const int8_t *)(block_base + 2u); + const uint base = block << 5; + FOR_UNROLL (uint qi = 0; qi < 32u; qi++) { + acc += d * (float)qs[qi] * x[base + qi]; + } + } + return acc; +} + +static inline float glm_q8_0_dot_row_tg_f32_fast( + device const char *row, + threadgroup const float *x, + uint n_cols) { + if (n_cols == 512u) { + return glm_q8_0_dot_row_tg_f32_512(row, x); + } + return glm_q8_0_dot_row_tg_f32(row, x, n_cols); +} + +static inline float glm_q8_0_dot_row_dev_f32( + device const char *row, + device const float *x, + uint n_cols) { + float acc = 0.0f; + const uint n_blocks = (n_cols + 31u) >> 5; + for (uint block = 0; block < n_blocks; block++) { + device const char *block_base = row + (uint64_t)block * 34u; + const float d = (float)(*((device const half *)block_base)); + device const int8_t *qs = (device const int8_t *)(block_base + 2u); + const uint base = block << 5; + const uint count = min(32u, n_cols - base); + for (uint qi = 0; qi < count; qi++) { + acc += d * (float)qs[qi] * x[base + qi]; + } + } + return acc; +} + +#define DS4_METAL_GGUF_Q4_0 2u +#define DS4_METAL_GGUF_Q8_0 8u +#define DS4_METAL_GGUF_Q4_K 12u + +static inline uchar2 glm_q4_K_scale_min(int j, int k, device const uchar *q) { + return j < 4 ? uchar2{uchar(q[j + 0 + k] & 63), uchar(q[j + 4 + k] & 63)} + : uchar2{uchar((q[j + 4 + k] & 0x0f) | ((q[j - 4 + k] & 0xc0) >> 2)), + uchar((q[j + 4 + k] >> 4) | ((q[j - 0 + k] & 0xc0) >> 2))}; +} + +static inline float glm_q4_0_weight_at(device const char *row, uint col) { + const uint block = col >> 5; + const uint qi = col & 31u; + device const char *block_base = row + (uint64_t)block * 18u; + const float d = (float)(*((device const half *)block_base)); + device const uchar *qs = (device const uchar *)(block_base + 2u); + /* ggml Q4_0: elems 0..15 = low nibbles of qs[0..15], 16..31 = high. */ + const uchar packed = qs[qi & 15u]; + const uchar q = (qi < 16u) ? (packed & 0x0f) : (packed >> 4); + return d * ((float)q - 8.0f); +} + +static inline float glm_q4_K_weight_at(device const char *row, uint col) { + const uint block = col >> 8u; + const uint idx = col & 255u; + device const char *block_base = row + (uint64_t)block * 144u; + const float d = (float)(*((device const half *)(block_base + 0u))); + const float dmin = (float)(*((device const half *)(block_base + 2u))); + device const uchar *scales = (device const uchar *)(block_base + 4u); + device const uchar *qs = (device const uchar *)(block_base + 16u); + const uint group = idx >> 5u; + const uint l = idx & 31u; + const uchar2 sm = glm_q4_K_scale_min((int)group, 0, scales); + const uint byte_off = (group >> 1u) * 32u + l; + const uint shift = (group & 1u) * 4u; + const uint q = ((uint)qs[byte_off] >> shift) & 0x0fu; + return d * (float)sm.x * (float)q - dmin * (float)sm.y; +} + +static inline float glm_quant_weight_at( + uint weight_type, + device const char *row, + uint col) { + if (weight_type == DS4_METAL_GGUF_Q4_0) return glm_q4_0_weight_at(row, col); + if (weight_type == DS4_METAL_GGUF_Q4_K) return glm_q4_K_weight_at(row, col); + return glm_q8_0_weight_at(row, col); +} + +static inline float glm_q4_0_dot_row_tg_f32( + device const char *row, + threadgroup const float *x, + uint n_cols) { + float acc = 0.0f; + for (uint col = 0; col < n_cols; col++) { + acc += glm_q4_0_weight_at(row, col) * x[col]; + } + return acc; +} + +static inline float glm_q4_K_dot_row_tg_f32( + device const char *row, + threadgroup const float *x, + uint n_cols) { + float acc = 0.0f; + for (uint col = 0; col < n_cols; col++) { + acc += glm_q4_K_weight_at(row, col) * x[col]; + } + return acc; +} + +static inline float glm_quant_dot_row_tg_f32( + uint weight_type, + device const char *row, + threadgroup const float *x, + uint n_cols) { + if (weight_type == DS4_METAL_GGUF_Q4_0) return glm_q4_0_dot_row_tg_f32(row, x, n_cols); + if (weight_type == DS4_METAL_GGUF_Q4_K) return glm_q4_K_dot_row_tg_f32(row, x, n_cols); + return glm_q8_0_dot_row_tg_f32_fast(row, x, n_cols); +} + +/* Per-lane Q4_K row dot: lane l covers elements (g*32 + l) of every + * 32-group so the 144-byte superblocks are read with coalesced per-lane + * bytes; callers simd_sum the result. x lives in threadgroup memory. */ +static inline float glm_q4_K_dot_row_lane_f32( + device const char *row, + threadgroup const float *x, + uint n_cols, + ushort lane) { + float acc = 0.0f; + const uint nblocks = n_cols >> 8u; + for (uint b = 0; b < nblocks; b++) { + device const char *block_base = row + (uint64_t)b * 144u; + const float d = (float)(*((device const half *)(block_base + 0u))); + const float dmin = (float)(*((device const half *)(block_base + 2u))); + device const uchar *scales = (device const uchar *)(block_base + 4u); + device const uchar *qs = (device const uchar *)(block_base + 16u); + threadgroup const float *xb = x + (b << 8u); + FOR_UNROLL (uint g = 0; g < 8u; g++) { + const uchar2 sm = glm_q4_K_scale_min((int)g, 0, scales); + const uint byte_off = (g >> 1u) * 32u + lane; + const uint shift = (g & 1u) * 4u; + const uint q = ((uint)qs[byte_off] >> shift) & 0x0fu; + const float xv = xb[(g << 5u) + lane]; + acc += (d * (float)sm.x * (float)q - dmin * (float)sm.y) * xv; + } + } + return acc; +} + +static inline float glm_quant_dot_row_dev_f32( + uint weight_type, + device const char *row, + device const float *x, + uint n_cols) { + if (weight_type == DS4_METAL_GGUF_Q8_0) return glm_q8_0_dot_row_dev_f32(row, x, n_cols); + float acc = 0.0f; + for (uint col = 0; col < n_cols; col++) { + acc += glm_quant_weight_at(weight_type, row, col) * x[col]; + } + return acc; +} + +kernel void kernel_glm_indexer_score_one( + constant ds4_metal_args_glm_indexer_score_one & args, + device const char *q, + device const float *weights, + device const char *indexer_key_cache, + device float *scores, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint row = tgpig.x; + if (row >= args.n_rows) return; + const uint nth = ntg_u.x; + float score = 0.0f; + for (uint h = 0; h < args.n_head; h++) { + float partial = 0.0f; + device const float *qh = + (device const float *)(q + (uint64_t)h * args.head_dim * sizeof(float)); + for (uint d = tid; d < args.head_dim; d += nth) { + const float k = glm_cache_load_f32_or_f16(indexer_key_cache, + (uint64_t)row * args.head_dim + d, + args.cache_f16); + partial += qh[d] * k; + } + scratch[tid] = partial; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) scratch[tid] += scratch[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0) { + score += max(scratch[0] * args.scale, 0.0f) * weights[h]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0) scores[row] = score; +} + +kernel void kernel_glm_indexer_score_one_direct( + constant ds4_metal_args_glm_indexer_score_one & args, + device const char *q, + device const float *weights, + device const char *indexer_key_cache, + device float *scores, + threadgroup float *shared [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + if (row >= args.n_rows || args.n_head != 32u || args.head_dim != 128u) { + return; + } + + threadgroup float *ktg = shared; + threadgroup float *psum = ktg + 128u; + + if (tid < 128u) { + ktg[tid] = glm_cache_load_f32_or_f16(indexer_key_cache, + (uint64_t)row * 128u + tid, + args.cache_f16); + } + + float acc = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint head0 = 0; head0 < 32u; head0 += 4u) { + const uint head = head0 + (uint)sg; + device const float4 *q4 = (device const float4 *)(q + + (uint64_t)head * 128u * sizeof(float)); + threadgroup const float4 *k4 = (threadgroup const float4 *)ktg; + + float s = dot(q4[lane], k4[lane]); + s = simd_sum(s); + if (lane == 0) { + psum[sg] = max(s * args.scale, 0.0f) * weights[head]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + acc += psum[0]; + acc += psum[1]; + acc += psum[2]; + acc += psum[3]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (tid == 0) { + scores[row] = acc; + } +} + +kernel void kernel_glm_indexer_scores_batch( + constant ds4_metal_args_glm_indexer_scores_batch & args, + device const char *q, + device const char *weights, + device const char *indexer_key_cache, + device char *scores, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint row = tgpig.x; + const uint token = tgpig.y; + if (row >= args.n_rows || token >= args.n_tokens) return; + + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + row; + const uint visible = min(args.pos0 + token + 1u, args.n_rows); + if (row >= visible) { + if (tid == 0) *dst = -INFINITY; + return; + } + + const uint nth = ntg_u.x; + float score = 0.0f; + for (uint h = 0; h < args.n_head; h++) { + float partial = 0.0f; + device const float *qh = (device const float *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)h * args.q_head_stride); + for (uint d = tid; d < args.head_dim; d += nth) { + const float k = glm_cache_load_f32_or_f16(indexer_key_cache, + (uint64_t)row * args.head_dim + d, + args.cache_f16); + partial += qh[d] * k; + } + scratch[tid] = partial; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) scratch[tid] += scratch[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0) { + device const float *w = (device const float *)(weights + + (uint64_t)token * args.weights_token_stride); + score += max(scratch[0] * args.scale, 0.0f) * w[h]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0) *dst = score; +} + +kernel void kernel_glm_indexer_scores_tiled_f32( + constant ds4_metal_args_glm_indexer_scores_batch & args, + device const char *q, + device const char *weights, + device const char *indexer_key_cache, + device char *scores, + threadgroup float *shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint TM = 8; + constexpr uint TN = 32; + constexpr uint TS = 8; + constexpr uint D = 128; + + const uint row_base = tgpig.x * TN; + const uint token_base = tgpig.y * TM; + + threadgroup float *qtg = shared; + threadgroup float *ktg = qtg + TM*D; + threadgroup float *dot = ktg + TN*D; + + const uint last_token = min(token_base + TM, args.n_tokens); + const uint max_visible = last_token > token_base ? + min(args.pos0 + last_token, args.n_rows) : 0u; + + if (row_base >= max_visible) { + for (uint i = tid; i < TM*TN; i += 128) { + const uint tr = i / TN; + const uint rc = i - tr*TN; + const uint token = token_base + tr; + const uint row = row_base + rc; + if (token < args.n_tokens && row < args.n_rows) { + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + row; + *dst = -INFINITY; + } + } + return; + } + + for (uint i = tid; i < TN*D; i += 128) { + const uint rc = i / D; + const uint d = i - rc*D; + const uint row = row_base + rc; + float v = 0.0f; + if (row < args.n_rows) { + v = glm_cache_load_f32_or_f16(indexer_key_cache, + (uint64_t)row * args.head_dim + d, + args.cache_f16); + } + ktg[i] = v; + } + + const uint cell0 = lane; + const uint cell1 = lane + 32u; + const uint token_row0 = cell0 >> 3; + const uint token_row1 = cell1 >> 3; + const uint sub0 = cell0 & 7u; + const uint sub1 = cell1 & 7u; + const uint col0 = (uint)sg * TS + sub0; + const uint col1 = (uint)sg * TS + sub1; + const uint token0 = token_base + token_row0; + const uint token1 = token_base + token_row1; + const uint row0 = row_base + col0; + const uint row1 = row_base + col1; + + float acc0 = 0.0f; + float acc1 = 0.0f; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint head = 0; head < args.n_head; head++) { + for (uint i = tid; i < TM*D; i += 128) { + const uint tr = i / D; + const uint d = i - tr*D; + const uint token = token_base + tr; + float v = 0.0f; + if (token < args.n_tokens) { + device const float *qrow = (device const float *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + v = qrow[d]; + } + qtg[i] = v; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + simdgroup_float8x8 mdot = make_filled_simdgroup_matrix(0.0f); + for (uint db = 0; db < D/TS; db++) { + simdgroup_float8x8 mq; + simdgroup_float8x8 mk; + simdgroup_load(mq, qtg + db*TS, D, 0, false); + simdgroup_load(mk, ktg + ((uint)sg * TS) * D + db*TS, D, 0, true); + simdgroup_multiply_accumulate(mdot, mq, mk, mdot); + } + + simdgroup_store(mdot, dot + (uint)sg * TS, TN, 0, false); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (token0 < args.n_tokens && row0 < args.n_rows) { + device const float *w = (device const float *)(weights + + (uint64_t)token0 * args.weights_token_stride); + const float s = dot[token_row0*TN + col0]; + acc0 += max(s * args.scale, 0.0f) * w[head]; + } + if (token1 < args.n_tokens && row1 < args.n_rows) { + device const float *w = (device const float *)(weights + + (uint64_t)token1 * args.weights_token_stride); + const float s = dot[token_row1*TN + col1]; + acc1 += max(s * args.scale, 0.0f) * w[head]; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (token0 < args.n_tokens && row0 < args.n_rows) { + const uint visible = min(args.pos0 + token0 + 1u, args.n_rows); + device float *dst = (device float *)(scores + + (uint64_t)token0 * args.score_token_stride) + row0; + *dst = row0 < visible ? acc0 : -INFINITY; + } + if (token1 < args.n_tokens && row1 < args.n_rows) { + const uint visible = min(args.pos0 + token1 + 1u, args.n_rows); + device float *dst = (device float *)(scores + + (uint64_t)token1 * args.score_token_stride) + row1; + *dst = row1 < visible ? acc1 : -INFINITY; + } +} + +kernel void kernel_glm_indexer_scores_tiled( + constant ds4_metal_args_glm_indexer_scores_batch & args, + device const char *q, + device const char *weights, + device const char *indexer_key_cache, + device char *scores, + threadgroup float *shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint TM = 8; + constexpr uint TN = 32; + constexpr uint TS = 8; + constexpr uint D = 128; + + const uint row_base = tgpig.x * TN; + const uint token_base = tgpig.y * TM; + + threadgroup half *qtg = (threadgroup half *)shared; + threadgroup half *ktg = qtg + TM*D; + threadgroup float *dot = (threadgroup float *)(ktg + TN*D); + + const uint last_token = min(token_base + TM, args.n_tokens); + const uint max_visible = last_token > token_base ? + min(args.pos0 + last_token, args.n_rows) : 0u; + + if (row_base >= max_visible) { + for (uint i = tid; i < TM*TN; i += 128) { + const uint tr = i / TN; + const uint rc = i - tr*TN; + const uint token = token_base + tr; + const uint row = row_base + rc; + if (token < args.n_tokens && row < args.n_rows) { + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + row; + *dst = -INFINITY; + } + } + return; + } + + for (uint i = tid; i < TN*D; i += 128) { + const uint rc = i / D; + const uint d = i - rc*D; + const uint row = row_base + rc; + half v = half(0.0f); + if (row < args.n_rows) { + v = half(glm_cache_load_f32_or_f16(indexer_key_cache, + (uint64_t)row * args.head_dim + d, + args.cache_f16)); + } + ktg[i] = v; + } + + const uint cell0 = lane; + const uint cell1 = lane + 32u; + const uint token_row0 = cell0 >> 3; + const uint token_row1 = cell1 >> 3; + const uint sub0 = cell0 & 7u; + const uint sub1 = cell1 & 7u; + const uint col0 = (uint)sg * TS + sub0; + const uint col1 = (uint)sg * TS + sub1; + const uint token0 = token_base + token_row0; + const uint token1 = token_base + token_row1; + const uint row0 = row_base + col0; + const uint row1 = row_base + col1; + + float acc0 = 0.0f; + float acc1 = 0.0f; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint head = 0; head < args.n_head; head++) { + for (uint i = tid; i < TM*D; i += 128) { + const uint tr = i / D; + const uint d = i - tr*D; + const uint token = token_base + tr; + half v = half(0.0f); + if (token < args.n_tokens) { + device const float *qrow = (device const float *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + v = half(qrow[d]); + } + qtg[i] = v; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + simdgroup_float8x8 mdot = make_filled_simdgroup_matrix(0.0f); + for (uint db = 0; db < D/TS; db++) { + simdgroup_half8x8 mq; + simdgroup_half8x8 mk; + simdgroup_load(mq, qtg + db*TS, D, 0, false); + simdgroup_load(mk, ktg + ((uint)sg * TS) * D + db*TS, D, 0, true); + simdgroup_multiply_accumulate(mdot, mq, mk, mdot); + } + + simdgroup_store(mdot, dot + (uint)sg * TS, TN, 0, false); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (token0 < args.n_tokens && row0 < args.n_rows) { + device const float *w = (device const float *)(weights + + (uint64_t)token0 * args.weights_token_stride); + const float s = dot[token_row0*TN + col0]; + acc0 += max(s * args.scale, 0.0f) * w[head]; + } + if (token1 < args.n_tokens && row1 < args.n_rows) { + device const float *w = (device const float *)(weights + + (uint64_t)token1 * args.weights_token_stride); + const float s = dot[token_row1*TN + col1]; + acc1 += max(s * args.scale, 0.0f) * w[head]; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (token0 < args.n_tokens && row0 < args.n_rows) { + const uint visible = min(args.pos0 + token0 + 1u, args.n_rows); + device float *dst = (device float *)(scores + + (uint64_t)token0 * args.score_token_stride) + row0; + *dst = row0 < visible ? acc0 : -INFINITY; + } + if (token1 < args.n_tokens && row1 < args.n_rows) { + const uint visible = min(args.pos0 + token1 + 1u, args.n_rows); + device float *dst = (device float *)(scores + + (uint64_t)token1 * args.score_token_stride) + row1; + *dst = row1 < visible ? acc1 : -INFINITY; + } +} + +kernel void kernel_glm_qk_lowrank_q8_0( + constant ds4_metal_args_glm_qk_lowrank & args, + device const char *weight, + device const char *q, + device char *qk_low, + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint head = tgpig.x; + if (head >= args.n_head) return; + const uint nth = ntg_u.x; + device const float *qh = + (device const float *)(q + (uint64_t)head * args.qk_dim * sizeof(float)); + device float *out = + (device float *)(qk_low + (uint64_t)head * args.kv_lora_dim * sizeof(float)); + + for (uint j = tid; j < args.kv_lora_dim; j += nth) { + device const char *row = + weight + ((uint64_t)head * args.kv_lora_dim + j) * args.row_bytes; + out[j] = glm_quant_dot_row_dev_f32(args.weight_type, row, qh, args.qk_nope); + } +} + +kernel void kernel_glm_qk_lowrank_q8_0_glm52( + constant ds4_metal_args_glm_qk_lowrank & args, + device const char *weight, + device const char *q, + device char *qk_low, + threadgroup float *x [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + constexpr uint n_head = 64u; + constexpr uint kv_lora_dim = 512u; + constexpr uint qk_nope = 192u; + constexpr uint qk_dim = 256u; + constexpr uint row_bytes = 204u; + + const uint head = tgpig.x; + if (head >= n_head || + args.n_head != n_head || + args.kv_lora_dim != kv_lora_dim || + args.qk_nope != qk_nope || + args.qk_dim != qk_dim || + args.row_bytes != row_bytes || + args.weight_type != DS4_METAL_GGUF_Q8_0) { + return; + } + const uint nth = ntg_u.x; + device const float *qh = + (device const float *)(q + (uint64_t)head * qk_dim * sizeof(float)); + for (uint d = tid; d < qk_nope; d += nth) { + x[d] = qh[d]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *out = + (device float *)(qk_low + (uint64_t)head * kv_lora_dim * sizeof(float)); + for (uint j = tid; j < kv_lora_dim; j += nth) { + device const char *row = + weight + ((uint64_t)head * kv_lora_dim + j) * row_bytes; + float acc = 0.0f; + for (uint block = 0; block < 6u; block++) { + device const char *block_base = row + (uint64_t)block * 34u; + const float d = (float)(*((device const half *)block_base)); + device const int8_t *qs = (device const int8_t *)(block_base + 2u); + const uint base = block << 5; + FOR_UNROLL (uint qi = 0; qi < 32u; qi++) { + const uint col = base + qi; + acc += d * (float)qs[qi] * x[col]; + } + } + out[j] = acc; + } +} + +// Coalesced GLM 5.2 decode qk-low: one simdgroup per pair of output rows, +// lanes split the 192-wide dot so the 204-byte Q8 rows are read with +// consecutive per-lane bytes. The thread-per-row variant above issues +// strided scalar byte loads from only 64 threadgroups and measures ~7.5x +// off the weight-bandwidth floor. +kernel void kernel_glm_qk_lowrank_q8_0_glm52_sg( + constant ds4_metal_args_glm_qk_lowrank & args, + device const char *weight, + device const char *q, + device char *qk_low, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort3 ntg_u [[threads_per_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + constexpr uint n_head = 64u; + constexpr uint kv_lora_dim = 512u; + constexpr uint qk_nope = 192u; + constexpr uint qk_dim = 256u; + constexpr uint NR = 2u; + + const uint head = tgpig.x; + const uint wt = args.weight_type; + if (head >= n_head || + args.n_head != n_head || + args.kv_lora_dim != kv_lora_dim || + args.qk_nope != qk_nope || + args.qk_dim != qk_dim || + !((wt == DS4_METAL_GGUF_Q8_0 && args.row_bytes == 204u) || + (wt == DS4_METAL_GGUF_Q4_0 && args.row_bytes == 108u))) { + return; + } + const uint row_bytes = args.row_bytes; + + const uint nsg = ntg_u.y; + const uint row0 = (tgpig.y * nsg + (uint)sgitg) * NR; + if (row0 >= kv_lora_dim) return; + + device const float *qh = + (device const float *)(q + (uint64_t)head * qk_dim * sizeof(float)); + float qv[6]; + FOR_UNROLL (uint b = 0; b < 6u; b++) { + qv[b] = qh[(b << 5) + tiisg]; + } + + device float *out = + (device float *)(qk_low + (uint64_t)head * kv_lora_dim * sizeof(float)); + for (uint r = 0; r < NR; r++) { + const uint j = row0 + r; + device const char *row = + weight + ((uint64_t)head * kv_lora_dim + j) * row_bytes; + float acc = 0.0f; + if (wt == DS4_METAL_GGUF_Q8_0) { + FOR_UNROLL (uint b = 0; b < 6u; b++) { + device const char *block_base = row + (uint64_t)b * 34u; + const float d = (float)(*((device const half *)block_base)); + device const int8_t *qs = (device const int8_t *)(block_base + 2u); + acc += d * (float)qs[tiisg] * qv[b]; + } + } else { + /* Q4_0: 18B blocks; elems 0..15 = low nibbles, 16..31 = high. */ + FOR_UNROLL (uint b = 0; b < 6u; b++) { + device const char *block_base = row + (uint64_t)b * 18u; + const float d = (float)(*((device const half *)block_base)); + device const uint8_t *qs = (device const uint8_t *)(block_base + 2u); + const uint byte = qs[tiisg & 15u]; + const float v = (float)((tiisg < 16u) ? (byte & 0xFu) : (byte >> 4)) - 8.0f; + acc += d * v * qv[b]; + } + } + const float sum = simd_sum(acc); + if (tiisg == 0) { + out[j] = sum; + } + } +} + +kernel void kernel_glm_qk_lowrank_q8_0_batch( + constant ds4_metal_args_glm_qk_lowrank_batch & args, + device const char *weight, + device const char *q, + device char *qk_low, + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint head = tgpig.x + args.head_base; + const uint token = tgpig.y; + if (head >= args.n_head || token >= args.n_tokens) return; + const uint nth = ntg_u.x; + const uint qk_dim = args.qk_dim; + const uint64_t q_token_stride = (uint64_t)args.n_head * qk_dim * sizeof(float); + const uint64_t low_token_stride = (uint64_t)args.n_head * args.kv_lora_dim * sizeof(float); + device const float *qh = + (device const float *)(q + + (uint64_t)token * q_token_stride + + (uint64_t)head * qk_dim * sizeof(float)); + device float *out = + (device float *)(qk_low + + (uint64_t)token * low_token_stride + + (uint64_t)head * args.kv_lora_dim * sizeof(float)); + + for (uint j = tid; j < args.kv_lora_dim; j += nth) { + device const char *row = + weight + ((uint64_t)head * args.kv_lora_dim + j) * args.row_bytes; + out[j] = glm_quant_dot_row_dev_f32(args.weight_type, row, qh, args.qk_nope); + } +} + +kernel void kernel_glm_qk_lowrank_q8_0_batch_glm52_t4( + constant ds4_metal_args_glm_qk_lowrank_batch & args, + device const char *weight, + device const char *q, + device char *qk_low, + threadgroup float *x [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + constexpr uint n_head = 64u; + constexpr uint kv_lora_dim = 512u; + constexpr uint qk_nope = 192u; + constexpr uint qk_dim = 256u; + constexpr uint tile_tokens = 4u; + constexpr uint row_bytes = 204u; + + const uint head = tgpig.x + args.head_base; + const uint token0 = tgpig.y * tile_tokens; + const uint nth = ntg_u.x; + const uint64_t q_token_stride = (uint64_t)n_head * qk_dim * sizeof(float); + const uint64_t low_token_stride = (uint64_t)n_head * kv_lora_dim * sizeof(float); + + for (uint t = 0; t < tile_tokens; t++) { + const uint token = token0 + t; + threadgroup float *xt = x + t * qk_nope; + if (token < args.n_tokens) { + device const float *qh = + (device const float *)(q + + (uint64_t)token * q_token_stride + + (uint64_t)head * qk_dim * sizeof(float)); + for (uint d = tid; d < qk_nope; d += nth) { + xt[d] = qh[d]; + } + } else { + for (uint d = tid; d < qk_nope; d += nth) { + xt[d] = 0.0f; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint j = tid; j < kv_lora_dim; j += nth) { + device const char *row = + weight + ((uint64_t)head * kv_lora_dim + j) * row_bytes; + float acc0 = 0.0f; + float acc1 = 0.0f; + float acc2 = 0.0f; + float acc3 = 0.0f; + for (uint block = 0; block < 6u; block++) { + device const char *block_base = row + (uint64_t)block * 34u; + const float d = (float)(*((device const half *)block_base)); + device const int8_t *qs = (device const int8_t *)(block_base + 2u); + const uint base = block << 5; + FOR_UNROLL (uint qi = 0; qi < 32u; qi++) { + const uint col = base + qi; + const float wq = d * (float)qs[qi]; + acc0 += wq * x[col]; + acc1 += wq * x[qk_nope + col]; + acc2 += wq * x[2u * qk_nope + col]; + acc3 += wq * x[3u * qk_nope + col]; + } + } + + if (token0 < args.n_tokens) { + device float *out0 = + (device float *)(qk_low + + (uint64_t)token0 * low_token_stride + + (uint64_t)head * kv_lora_dim * sizeof(float)); + out0[j] = acc0; + } + if (token0 + 1u < args.n_tokens) { + device float *out1 = + (device float *)(qk_low + + (uint64_t)(token0 + 1u) * low_token_stride + + (uint64_t)head * kv_lora_dim * sizeof(float)); + out1[j] = acc1; + } + if (token0 + 2u < args.n_tokens) { + device float *out2 = + (device float *)(qk_low + + (uint64_t)(token0 + 2u) * low_token_stride + + (uint64_t)head * kv_lora_dim * sizeof(float)); + out2[j] = acc2; + } + if (token0 + 3u < args.n_tokens) { + device float *out3 = + (device float *)(qk_low + + (uint64_t)(token0 + 3u) * low_token_stride + + (uint64_t)head * kv_lora_dim * sizeof(float)); + out3[j] = acc3; + } + } +} + +kernel void kernel_glm_value_project_q8_0( + constant ds4_metal_args_glm_qk_lowrank & args, + device const char *weight, + device const char *lora, + device char *heads, + threadgroup float *x [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint head = tgpig.x; + if (head >= args.n_head) return; + const uint nth = ntg_u.x; + device const float *src = + (device const float *)(lora + (uint64_t)head * args.kv_lora_dim * sizeof(float)); + for (uint j = tid; j < args.kv_lora_dim; j += nth) { + x[j] = src[j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *out = + (device float *)(heads + (uint64_t)head * args.qk_dim * sizeof(float)); + for (uint d = tid; d < args.qk_dim; d += nth) { + device const char *row = + weight + ((uint64_t)head * args.qk_dim + d) * args.row_bytes; + out[d] = glm_quant_dot_row_tg_f32(args.weight_type, row, x, args.kv_lora_dim); + } +} + +kernel void kernel_glm_value_project_q8_0_batch_heads( + constant ds4_metal_args_glm_qk_lowrank_batch & args, + device const char *weight, + device const char *lora, + device char *heads, + threadgroup float *x [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint head = tgpig.x + args.head_base; + const uint token = tgpig.y; + if (head >= args.n_head || token >= args.n_tokens) return; + const uint nth = ntg_u.x; + const uint value_dim = args.qk_dim; + const uint64_t lora_token_stride = + (uint64_t)args.n_head * args.kv_lora_dim * sizeof(float); + const uint64_t heads_token_stride = + (uint64_t)args.n_head * value_dim * sizeof(float); + device const float *src = + (device const float *)(lora + + (uint64_t)token * lora_token_stride + + (uint64_t)head * args.kv_lora_dim * sizeof(float)); + device float *out = + (device float *)(heads + + (uint64_t)token * heads_token_stride + + (uint64_t)head * value_dim * sizeof(float)); + + for (uint j = tid; j < args.kv_lora_dim; j += nth) { + x[j] = src[j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint d = tid; d < value_dim; d += nth) { + device const char *row = + weight + ((uint64_t)head * value_dim + d) * args.row_bytes; + out[d] = glm_quant_dot_row_tg_f32(args.weight_type, row, x, args.kv_lora_dim); + } +} + +kernel void kernel_glm_value_project_q8_0_batch_heads_mma( + constant ds4_metal_args_glm_qk_lowrank_batch & args, + device const char *weight, + device const char *lora, + device char *heads, + threadgroup char *shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint NR0 = 64u; + constexpr uint NR1 = 32u; + constexpr uint NK = 32u; + constexpr uint NL0 = 2u; + constexpr uint NL1 = 4u; + + const uint token0 = tgpig.x * NR1; + const uint value0 = tgpig.y * NR0; + const uint head = tgpig.z + args.head_base; + if (head >= args.n_head || token0 >= args.n_tokens || value0 >= args.qk_dim) { + return; + } + + threadgroup half *sa = (threadgroup half *)shmem; + threadgroup half *sb = (threadgroup half *)(shmem + 4096u); + + const uint nr0 = min(NR0, args.qk_dim - value0); + const uint nr1 = min(NR1, args.n_tokens - token0); + + const uint lr0 = min((uint)tid / NL0, nr0 - 1u); + const uint lr1 = min((uint)tid / NL1, nr1 - 1u); + const uint il0 = (uint)tid & 1u; + const uint iy = 8u * ((uint)tid & (NL1 - 1u)); + + const uint64_t lora_token_stride = + (uint64_t)args.n_head * args.kv_lora_dim * sizeof(float); + const uint64_t heads_token_stride = + (uint64_t)args.n_head * args.qk_dim * sizeof(float); + const uint64_t head_lora_base = + (uint64_t)head * args.kv_lora_dim * sizeof(float); + const uint64_t head_out_base = + (uint64_t)head * args.qk_dim * sizeof(float); + + simdgroup_half8x8 ma[4]; + simdgroup_half8x8 mb[2]; + simdgroup_float8x8 mc[8]; + for (uint i = 0; i < 8u; i++) { + mc[i] = make_filled_simdgroup_matrix(0.0f); + } + + for (uint loop_k = 0; loop_k < args.kv_lora_dim; loop_k += NK) { + const uint value = value0 + lr0; + const uint block = loop_k >> 5; + device const char *row = + weight + ((uint64_t)head * args.qk_dim + value) * args.row_bytes; + device const char *block_base = row + (uint64_t)block * 34u; + const float d = (float)(*((device const half *)block_base)); + device const int8_t *qs = (device const int8_t *)(block_base + 2u); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint i = 0; i < 16u; i++) { + const uint k = loop_k + 16u * il0 + i; + const uint sx = 2u * il0 + i / 8u; + const uint sy = ((uint)tid / NL0) / 8u; + const uint lx = ((uint)tid / NL0) & 7u; + const uint ly = i & 7u; + const uint ib = 8u * sx + sy; + const half v = (value < args.qk_dim && k < args.kv_lora_dim) ? + half(d * (float)qs[16u * il0 + i]) : + half(0.0f); + *(sa + 64u * ib + 8u * ly + lx) = v; + } + + const uint token = token0 + lr1; + device const float *y = + (device const float *)(lora + + (uint64_t)token * lora_token_stride + + head_lora_base + + (uint64_t)loop_k * sizeof(float) + + (uint64_t)iy * sizeof(float)); + for (uint i = 0; i < 8u; i++) { + const uint k = loop_k + iy + i; + const uint sx = ((uint)tid) & (NL1 - 1u); + const uint sy = ((uint)tid / NL1) / 8u; + const uint lx = i; + const uint ly = ((uint)tid / NL1) & 7u; + const uint ib = 4u * sx + sy; + const half v = (token < args.n_tokens && k < args.kv_lora_dim) ? + half(y[i]) : + half(0.0f); + *(sb + 64u * ib + 8u * ly + lx) = v; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup const half *lsma = sa + 4u * 64u * ((uint)sg & 1u); + threadgroup const half *lsmb = sb + 2u * 64u * ((uint)sg >> 1); + + for (uint ik = 0; ik < NK / 8u; ik++) { + simdgroup_barrier(mem_flags::mem_none); + + for (uint i = 0; i < 4u; i++) { + simdgroup_load(ma[i], lsma + 64u * i, 8u, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + for (uint i = 0; i < 2u; i++) { + simdgroup_load(mb[i], lsmb + 64u * i, 8u, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + for (uint i = 0; i < 8u; i++) { + simdgroup_multiply_accumulate(mc[i], mb[i / 4u], ma[i & 3u], mc[i]); + } + + lsma += 8u * 64u; + lsmb += 4u * 64u; + } + } + + if (nr0 == NR0 && nr1 == NR1) { + device float *dst = + (device float *)(heads + + (uint64_t)(token0 + 16u * ((uint)sg >> 1)) * heads_token_stride + + head_out_base + + (uint64_t)(value0 + 32u * ((uint)sg & 1u)) * sizeof(float)); + for (uint i = 0; i < 8u; i++) { + simdgroup_store(mc[i], + dst + 8u * (i & 3u) + 8u * (heads_token_stride / sizeof(float)) * (i / 4u), + heads_token_stride / sizeof(float), + 0, + false); + } + } else { + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup float *tmp = (threadgroup float *)shmem; + for (uint i = 0; i < 8u; i++) { + simdgroup_store(mc[i], + tmp + 32u * ((uint)sg & 1u) + + 16u * ((uint)sg >> 1) * NR0 + + 8u * (i & 3u) + 8u * NR0 * (i / 4u), + NR0, + 0, + false); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (sg == 0) { + for (uint t = tid; t < nr1; t += 128u) { + device float *dst = + (device float *)(heads + + (uint64_t)(token0 + t) * heads_token_stride + + head_out_base + + (uint64_t)value0 * sizeof(float)); + threadgroup const float *src = tmp + t * NR0; + for (uint v = 0; v < nr0; v++) { + dst[v] = src[v]; + } + } + } + } +} + +template +kernel void kernel_glm_attention_indexed_decode_split_group8_partial_impl( + constant ds4_metal_args_glm_attention_indexed_decode_split & args, + device const char *q, + device const char *qk_low, + device const char *kv_lora_cache, + device const char *k_rope_cache, + device const uint32_t *selected, + device char *partial_lora, + device char *partial_ms, + threadgroup half4 *scratch [[threadgroup(0)]], + ushort tid_u [[thread_index_in_threadgroup]], + ushort lane_u [[thread_index_in_simdgroup]], + ushort sg_u [[simdgroup_index_in_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + constexpr uint group_heads = 8u; + constexpr uint stage_rows = 16u; + const uint tid = (uint)tid_u; + const uint lane = (uint)lane_u; + const uint head_in_group = (uint)sg_u; + const uint head = tgpig.x * group_heads + head_in_group; + const uint block = tgpig.y; + if (args.n_selected == 0u || + args.cache_f16 == 0u || + args.kv_lora_dim != 512u || + args.qk_rope != 64u || + args.block_rows == 0u || + block >= args.n_blocks) { + return; + } + + const bool valid_head = assume_valid_heads || head < args.n_head; + const uint safe_head = valid_head ? head : 0u; + const uint kv_vecs = args.kv_lora_dim >> 2; + const uint rope_vecs = args.qk_rope >> 2; + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint block_start = block * args.block_rows; + const uint block_end = min(args.n_selected, block_start + args.block_rows); + + threadgroup half4 *kv_shared = scratch; + threadgroup float4 *rope_shared = + (threadgroup float4 *)(kv_shared + stage_rows * kv_vecs); + + device const float *qh = + (device const float *)(q + (uint64_t)safe_head * qk_dim * sizeof(float)); + device const float4 *low4 = + (device const float4 *)(qk_low + + (uint64_t)safe_head * args.kv_lora_dim * sizeof(float)); + + float4 low0 = 0.0f; + float4 low1 = 0.0f; + float4 low2 = 0.0f; + float4 low3 = 0.0f; + float4 qrope = 0.0f; + if (valid_head) { + low0 = low4[lane + 0u]; + low1 = low4[lane + 32u]; + low2 = low4[lane + 64u]; + low3 = low4[lane + 96u]; + if (lane < rope_vecs) { + qrope = *((device const float4 *)(qh + args.qk_nope + lane * 4u)); + } + } + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + + float M = -FLT_MAX / 2.0f; + float S = 0.0f; + float4 o0 = 0.0f; + float4 o1 = 0.0f; + float4 o2 = 0.0f; + float4 o3 = 0.0f; + + for (uint base = block_start; base < block_end; base += stage_rows) { + const uint rows = min(stage_rows, block_end - base); + for (uint off = tid; off < rows * kv_vecs; off += 256u) { + const uint rr = off / kv_vecs; + const uint vv = off - rr * kv_vecs; + const uint row = selected[base + rr]; + const bool valid_row = assume_valid_rows || row < args.cache_cap; + if (valid_row) { + device const half4 *src = + (device const half4 *)((device const half *)kv_lora_cache + + (uint64_t)row * args.kv_lora_dim); + kv_shared[off] = src[vv]; + } else { + kv_shared[off] = half4(half(0.0f)); + } + } + for (uint off = tid; off < rows * rope_vecs; off += 256u) { + const uint rr = off / rope_vecs; + const uint vv = off - rr * rope_vecs; + const uint r = vv * 4u; + const uint row = selected[base + rr]; + const bool valid_row = assume_valid_rows || row < args.cache_cap; + if (valid_row) { + const uint64_t rope_base = (uint64_t)row * args.qk_rope; + const float2 y0 = + glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, + rope_base, + r, + row, + args.qk_rope, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + const float2 y1 = + glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, + rope_base, + r + 2u, + row, + args.qk_rope, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + rope_shared[off] = float4(y0.x, y0.y, y1.x, y1.y); + } else { + rope_shared[off] = float4(0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint rr = 0u; rr < rows; rr++) { + const uint row = selected[base + rr]; + const bool valid_row = assume_valid_rows || row < args.cache_cap; + threadgroup const half4 *kv_row = kv_shared + rr * kv_vecs; + threadgroup const float4 *rope_row = rope_shared + rr * rope_vecs; + float partial = 0.0f; + if (valid_head && valid_row) { + partial += dot(low0, (float4)kv_row[lane + 0u]); + partial += dot(low1, (float4)kv_row[lane + 32u]); + partial += dot(low2, (float4)kv_row[lane + 64u]); + partial += dot(low3, (float4)kv_row[lane + 96u]); + if (lane < rope_vecs) { + partial += dot(qrope, rope_row[lane]); + } + } + const float sum = simd_sum(partial); + const float score = + (valid_head && valid_row) ? sum * args.scale : -FLT_MAX / 2.0f; + if (valid_head && valid_row) { + const float new_m = max(M, score); + const float old_scale = exp(M - new_m); + const float row_scale = exp(score - new_m); + o0 = o0 * old_scale + (float4)kv_row[lane + 0u] * row_scale; + o1 = o1 * old_scale + (float4)kv_row[lane + 32u] * row_scale; + o2 = o2 * old_scale + (float4)kv_row[lane + 64u] * row_scale; + o3 = o3 * old_scale + (float4)kv_row[lane + 96u] * row_scale; + S = S * old_scale + row_scale; + M = new_m; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (valid_head) { + device float4 *out4 = + (device float4 *)(partial_lora + + ((uint64_t)block * args.n_head + head) * + args.kv_lora_dim * sizeof(float)); + out4[lane + 0u] = o0; + out4[lane + 32u] = o1; + out4[lane + 64u] = o2; + out4[lane + 96u] = o3; + if (lane == 0u) { + device float *ms = + (device float *)(partial_ms + + ((uint64_t)block * args.n_head + head) * 2u * sizeof(float)); + ms[0] = M; + ms[1] = S; + } + } +} + +typedef decltype(kernel_glm_attention_indexed_decode_split_group8_partial_impl) + glm_attention_indexed_decode_split_group8_partial_t; + +template [[host_name("kernel_glm_attention_indexed_decode_split_group8_partial")]] +kernel glm_attention_indexed_decode_split_group8_partial_t +kernel_glm_attention_indexed_decode_split_group8_partial_impl; + +template [[host_name("kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads")]] +kernel glm_attention_indexed_decode_split_group8_partial_t +kernel_glm_attention_indexed_decode_split_group8_partial_impl; + +template +static void kernel_glm_attention_indexed_decode_split_group8_reduce_impl( + constant ds4_metal_args_glm_attention_indexed_decode_split & args, + device const char *partial_lora, + device const char *partial_ms, + device const char *value_weight, + device char *heads, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint head = tgpig.x; + const uint n_blocks = FIXED_BLOCKS != 0u ? FIXED_BLOCKS : args.n_blocks; + if (head >= args.n_head || + args.n_selected == 0u || + args.kv_lora_dim != 512u || + n_blocks == 0u || + n_blocks > 64u || + (FIXED_BLOCKS != 0u && args.n_blocks != FIXED_BLOCKS)) { + return; + } + + const uint nth = ntg_u.x; + threadgroup float *red = scratch; + threadgroup float *block_scale = scratch + 256u; + threadgroup float *lora_sum = scratch + 320u; + + float local_m = -FLT_MAX / 2.0f; + if (tid < n_blocks) { + device const float *ms = + (device const float *)(partial_ms + + ((uint64_t)tid * args.n_head + head) * 2u * sizeof(float)); + local_m = ms[1] > 0.0f ? ms[0] : -FLT_MAX / 2.0f; + } + red[tid] = local_m; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] = max(red[tid], red[tid + step]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float max_m = red[0]; + + float local_denom = 0.0f; + if (tid < n_blocks) { + device const float *ms = + (device const float *)(partial_ms + + ((uint64_t)tid * args.n_head + head) * 2u * sizeof(float)); + const float s = ms[1]; + const float e = s > 0.0f ? exp(ms[0] - max_m) : 0.0f; + block_scale[tid] = e; + local_denom = s * e; + } + red[tid] = local_denom; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] += red[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float denom = max(red[0], 1.0e-20f); + + for (uint j = tid; j < args.kv_lora_dim; j += nth) { + float acc = 0.0f; + for (uint b = 0u; b < n_blocks; b++) { + device const float *src = + (device const float *)(partial_lora + + ((uint64_t)b * args.n_head + head) * + args.kv_lora_dim * sizeof(float)); + acc += src[j] * block_scale[b]; + } + lora_sum[j] = acc / denom; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *out = + (device float *)(heads + (uint64_t)head * args.value_dim * sizeof(float)); + if (args.value_type == DS4_METAL_GGUF_Q4_K && + (args.kv_lora_dim & 255u) == 0u) { + /* Lane-split Q4_K value project: one simdgroup per output row with + * coalesced per-lane superblock reads; the per-thread scalar + * fallback below walks the 144-byte rows one element at a time. */ + const uint vp_sg = tid >> 5u; + const uint vp_lane = tid & 31u; + const uint vp_nsg = nth >> 5u; + for (uint d = vp_sg; d < args.value_dim; d += vp_nsg) { + device const char *row = + value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes; + const float part = glm_q4_K_dot_row_lane_f32(row, lora_sum, + args.kv_lora_dim, + (ushort)vp_lane); + const float sum = simd_sum(part); + if (vp_lane == 0u) { + out[d] = sum; + } + } + } else { + for (uint d = tid; d < args.value_dim; d += nth) { + device const char *row = + value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes; + out[d] = glm_quant_dot_row_tg_f32(args.value_type, row, lora_sum, args.kv_lora_dim); + } + } +} + +kernel void kernel_glm_attention_indexed_decode_split_group8_reduce( + constant ds4_metal_args_glm_attention_indexed_decode_split & args, + device const char *partial_lora, + device const char *partial_ms, + device const char *value_weight, + device char *heads, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + kernel_glm_attention_indexed_decode_split_group8_reduce_impl<0>( + args, partial_lora, partial_ms, value_weight, heads, scratch, + tid, ntg_u, tgpig); +} + +kernel void kernel_glm_attention_indexed_decode_split_group8_reduce16( + constant ds4_metal_args_glm_attention_indexed_decode_split & args, + device const char *partial_lora, + device const char *partial_ms, + device const char *value_weight, + device char *heads, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + kernel_glm_attention_indexed_decode_split_group8_reduce_impl<16>( + args, partial_lora, partial_ms, value_weight, heads, scratch, + tid, ntg_u, tgpig); +} + +kernel void kernel_glm_attention_indexed_decode( + constant ds4_metal_args_glm_attention_indexed_decode & args, + device const char *q, + device const char *qk_low, + device const char *kv_lora_cache, + device const char *k_rope_cache, + device const char *value_weight, + device const uint32_t *selected, + device char *heads, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint head = tgpig.x; + if (head >= args.n_head || args.n_selected == 0u) return; + const uint nth = ntg_u.x; + const uint qk_dim = args.qk_nope + args.qk_rope; + threadgroup float *red = scratch; + threadgroup float *scores = scratch + 256u; + threadgroup float *lora_sum = scores + args.n_selected; + + device const float *qh = + (device const float *)(q + (uint64_t)head * qk_dim * sizeof(float)); + device const float *low = + (device const float *)(qk_low + (uint64_t)head * args.kv_lora_dim * sizeof(float)); + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + + if (args.cache_f16 != 0u) { + float local_max = -INFINITY; + for (uint s = tid; s < args.n_selected; s += nth) { + const uint row = selected[s]; + float score = -INFINITY; + if (row < args.cache_cap) { + float dotv = 0.0f; + const uint64_t lora_base = (uint64_t)row * args.kv_lora_dim; + uint j = 0; + for (; j + 3u < args.kv_lora_dim; j += 4u) { + device const half4 *kv4 = + (device const half4 *)((device const half *)kv_lora_cache + lora_base + j); + device const float4 *low4 = + (device const float4 *)(low + j); + const float4 kv = (float4)(*kv4); + const float4 qv = *low4; + dotv += qv.x * kv.x + qv.y * kv.y + + qv.z * kv.z + qv.w * kv.w; + } + if (j < args.kv_lora_dim) { + for (; j < args.kv_lora_dim; j++) { + const float kv = glm_cache_load_f16_only(kv_lora_cache, + lora_base + j); + dotv += low[j] * kv; + } + } + const uint64_t rope_base = (uint64_t)row * args.qk_rope; + for (uint r = 0; r < args.qk_rope; r += 2u) { + const float2 y = glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, + rope_base, + r, + row, + args.qk_rope, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + dotv += qh[args.qk_nope + r] * y.x + + qh[args.qk_nope + r + 1u] * y.y; + } + score = dotv * args.scale; + } + scores[s] = score; + local_max = max(local_max, score); + } + red[tid] = local_max; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] = max(red[tid], red[tid + step]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float max_score = red[0]; + + float local_sum = 0.0f; + for (uint s = tid; s < args.n_selected; s += nth) { + const float w = exp(scores[s] - max_score); + scores[s] = w; + local_sum += w; + } + red[tid] = local_sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] += red[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float denom = max(red[0], 1.0e-20f); + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint j0 = tid * 2u; j0 < args.kv_lora_dim; j0 += nth * 2u) { + const uint j1 = j0 + 1u; + const bool use_j1 = j1 < args.kv_lora_dim; + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint s = 0; s < args.n_selected; s++) { + const uint row = selected[s]; + if (row < args.cache_cap) { + const uint64_t row_base = (uint64_t)row * args.kv_lora_dim; + const float w = scores[s]; + if (use_j1) { + device const half2 *kv2 = + (device const half2 *)((device const half *)kv_lora_cache + row_base + j0); + const float2 kv = (float2)(*kv2); + acc0 += w * kv.x; + acc1 += w * kv.y; + } else { + const float kv0 = glm_cache_load_f16_only(kv_lora_cache, + row_base + j0); + acc0 += w * kv0; + } + } + } + lora_sum[j0] = acc0 / denom; + if (use_j1) lora_sum[j1] = acc1 / denom; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *out = + (device float *)(heads + (uint64_t)head * args.value_dim * sizeof(float)); + for (uint d = tid; d < args.value_dim; d += nth) { + device const char *row = + value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes; + out[d] = glm_quant_dot_row_tg_f32(args.value_type, row, lora_sum, args.kv_lora_dim); + } + return; + } + + float local_max = -INFINITY; + for (uint s = tid; s < args.n_selected; s += nth) { + const uint row = selected[s]; + float score = -INFINITY; + if (row < args.cache_cap) { + float dotv = 0.0f; + const uint64_t lora_base = (uint64_t)row * args.kv_lora_dim; + for (uint j = 0; j < args.kv_lora_dim; j++) { + const float kv = glm_cache_load_f32_or_f16(kv_lora_cache, + lora_base + j, + args.cache_f16); + dotv += low[j] * kv; + } + const uint64_t rope_base = (uint64_t)row * args.qk_rope; + for (uint r = 0; r < args.qk_rope; r += 2u) { + const float2 y = glm_cache_load_rotated_rope_pair(k_rope_cache, + rope_base, + r, + row, + args.qk_rope, + args.cache_f16, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + dotv += qh[args.qk_nope + r] * y.x + + qh[args.qk_nope + r + 1u] * y.y; + } + score = dotv * args.scale; + } + scores[s] = score; + local_max = max(local_max, score); + } + red[tid] = local_max; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] = max(red[tid], red[tid + step]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float max_score = red[0]; + + float local_sum = 0.0f; + for (uint s = tid; s < args.n_selected; s += nth) { + const float w = exp(scores[s] - max_score); + scores[s] = w; + local_sum += w; + } + red[tid] = local_sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] += red[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float denom = max(red[0], 1.0e-20f); + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint j0 = tid; j0 < args.kv_lora_dim; j0 += nth * 2u) { + const uint j1 = j0 + nth; + const bool use_j1 = j1 < args.kv_lora_dim; + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint s = 0; s < args.n_selected; s++) { + const uint row = selected[s]; + if (row < args.cache_cap) { + const uint64_t row_base = (uint64_t)row * args.kv_lora_dim; + const float w = scores[s]; + const float kv0 = glm_cache_load_f32_or_f16(kv_lora_cache, + row_base + j0, + args.cache_f16); + acc0 += w * kv0; + if (use_j1) { + const float kv1 = glm_cache_load_f32_or_f16(kv_lora_cache, + row_base + j1, + args.cache_f16); + acc1 += w * kv1; + } + } + } + lora_sum[j0] = acc0 / denom; + if (use_j1) lora_sum[j1] = acc1 / denom; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *out = + (device float *)(heads + (uint64_t)head * args.value_dim * sizeof(float)); + if (args.value_type == DS4_METAL_GGUF_Q4_K && + (args.kv_lora_dim & 255u) == 0u) { + /* Lane-split Q4_K value project: one simdgroup per output row with + * coalesced per-lane superblock reads; the per-thread scalar + * fallback below walks the 144-byte rows one element at a time. */ + const uint vp_sg = tid >> 5u; + const uint vp_lane = tid & 31u; + const uint vp_nsg = nth >> 5u; + for (uint d = vp_sg; d < args.value_dim; d += vp_nsg) { + device const char *row = + value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes; + const float part = glm_q4_K_dot_row_lane_f32(row, lora_sum, + args.kv_lora_dim, + (ushort)vp_lane); + const float sum = simd_sum(part); + if (vp_lane == 0u) { + out[d] = sum; + } + } + } else { + for (uint d = tid; d < args.value_dim; d += nth) { + device const char *row = + value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes; + out[d] = glm_quant_dot_row_tg_f32(args.value_type, row, lora_sum, args.kv_lora_dim); + } + } +} + +kernel void kernel_glm_attention_indexed_batch( + constant ds4_metal_args_glm_attention_indexed_batch & args, + device const char *q, + device const char *qk_low, + device const char *kv_lora_cache, + device const char *k_rope_cache, + device const char *value_weight, + device const uint32_t *selected, + device char *heads, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint head = tgpig.x; + const uint token = tgpig.y; + if (head >= args.n_head || token >= args.n_tokens || args.n_selected == 0u) return; + const uint nth = ntg_u.x; + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint64_t q_token_stride = (uint64_t)args.n_head * qk_dim * sizeof(float); + const uint64_t low_token_stride = (uint64_t)args.n_head * args.kv_lora_dim * sizeof(float); + const uint64_t heads_token_stride = (uint64_t)args.n_head * args.value_dim * sizeof(float); + threadgroup float *red = scratch; + threadgroup float *scores = scratch + 256u; + threadgroup float *lora_sum = scores + args.n_selected; + + device const float *qh = + (device const float *)(q + + (uint64_t)token * q_token_stride + + (uint64_t)head * qk_dim * sizeof(float)); + device const float *low = + (device const float *)(qk_low + + (uint64_t)token * low_token_stride + + (uint64_t)head * args.kv_lora_dim * sizeof(float)); + device const uint32_t *token_selected = + selected + (uint64_t)token * args.n_selected; + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + + float local_max = -INFINITY; + for (uint s = tid; s < args.n_selected; s += nth) { + const uint row = token_selected[s]; + float score = -INFINITY; + if (row < args.cache_cap) { + float dotv = 0.0f; + const uint64_t lora_base = (uint64_t)row * args.kv_lora_dim; + for (uint j = 0; j < args.kv_lora_dim; j++) { + const float kv = glm_cache_load_f32_or_f16(kv_lora_cache, + lora_base + j, + args.cache_f16); + dotv += low[j] * kv; + } + const uint64_t rope_base = (uint64_t)row * args.qk_rope; + for (uint r = 0; r < args.qk_rope; r += 2u) { + const float2 y = glm_cache_load_rotated_rope_pair(k_rope_cache, + rope_base, + r, + row, + args.qk_rope, + args.cache_f16, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + dotv += qh[args.qk_nope + r] * y.x + + qh[args.qk_nope + r + 1u] * y.y; + } + score = dotv * args.scale; + } + scores[s] = score; + local_max = max(local_max, score); + } + red[tid] = local_max; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] = max(red[tid], red[tid + step]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float max_score = red[0]; + + float local_sum = 0.0f; + for (uint s = tid; s < args.n_selected; s += nth) { + const float w = exp(scores[s] - max_score); + scores[s] = w; + local_sum += w; + } + red[tid] = local_sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) red[tid] += red[tid + step]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float denom = max(red[0], 1.0e-20f); + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint j = tid; j < args.kv_lora_dim; j += nth) { + float acc = 0.0f; + for (uint s = 0; s < args.n_selected; s++) { + const uint row = token_selected[s]; + if (row < args.cache_cap) { + const float kv = glm_cache_load_f32_or_f16(kv_lora_cache, + (uint64_t)row * args.kv_lora_dim + j, + args.cache_f16); + acc += scores[s] * kv; + } + } + lora_sum[j] = acc / denom; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *out = + (device float *)(heads + + (uint64_t)token * heads_token_stride + + (uint64_t)head * args.value_dim * sizeof(float)); + if (args.value_type == DS4_METAL_GGUF_Q4_K && + (args.kv_lora_dim & 255u) == 0u) { + /* Lane-split Q4_K value project: one simdgroup per output row with + * coalesced per-lane superblock reads; the per-thread scalar + * fallback below walks the 144-byte rows one element at a time. */ + const uint vp_sg = tid >> 5u; + const uint vp_lane = tid & 31u; + const uint vp_nsg = nth >> 5u; + for (uint d = vp_sg; d < args.value_dim; d += vp_nsg) { + device const char *row = + value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes; + const float part = glm_q4_K_dot_row_lane_f32(row, lora_sum, + args.kv_lora_dim, + (ushort)vp_lane); + const float sum = simd_sum(part); + if (vp_lane == 0u) { + out[d] = sum; + } + } + } else { + for (uint d = tid; d < args.value_dim; d += nth) { + device const char *row = + value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes; + out[d] = glm_quant_dot_row_tg_f32(args.value_type, row, lora_sum, args.kv_lora_dim); + } + } +} + +kernel void kernel_glm_attention_indexed_batch_group2( + constant ds4_metal_args_glm_attention_indexed_batch & args, + device const char *q, + device const char *qk_low, + device const char *kv_lora_cache, + device const char *k_rope_cache, + device const char *value_weight, + device const uint32_t *selected, + device char *heads, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg_u [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const uint token = tgpig.y; + if (token >= args.n_tokens || args.n_selected == 0u) return; + const uint nth = ntg_u.x; + const uint head0 = tgpig.x * 2u; + const uint head1 = head0 + 1u; + const bool valid0 = head0 < args.n_head; + const bool valid1 = head1 < args.n_head; + if (!valid0 && !valid1) return; + + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint64_t q_token_stride = (uint64_t)args.n_head * qk_dim * sizeof(float); + const uint64_t low_token_stride = (uint64_t)args.n_head * args.kv_lora_dim * sizeof(float); + const uint64_t heads_token_stride = (uint64_t)args.n_head * args.value_dim * sizeof(float); + + threadgroup float *red0 = scratch; + threadgroup float *red1 = red0 + 256u; + threadgroup float *scores0 = red1 + 256u; + threadgroup float *scores1 = scores0 + args.n_selected; + threadgroup float *lora0 = scores1 + args.n_selected; + threadgroup float *lora1 = lora0 + args.kv_lora_dim; + + device const float *qh0 = + (device const float *)(q + + (uint64_t)token * q_token_stride + + (uint64_t)head0 * qk_dim * sizeof(float)); + device const float *qh1 = + (device const float *)(q + + (uint64_t)token * q_token_stride + + (uint64_t)(valid1 ? head1 : head0) * qk_dim * sizeof(float)); + device const float *low0 = + (device const float *)(qk_low + + (uint64_t)token * low_token_stride + + (uint64_t)head0 * args.kv_lora_dim * sizeof(float)); + device const float *low1 = + (device const float *)(qk_low + + (uint64_t)token * low_token_stride + + (uint64_t)(valid1 ? head1 : head0) * args.kv_lora_dim * sizeof(float)); + device const uint32_t *token_selected = + selected + (uint64_t)token * args.n_selected; + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + + float local_max0 = -INFINITY; + float local_max1 = -INFINITY; + for (uint s = tid; s < args.n_selected; s += nth) { + const uint row = token_selected[s]; + float score0 = -INFINITY; + float score1 = -INFINITY; + if (row < args.cache_cap) { + float dot0 = 0.0f; + float dot1 = 0.0f; + const uint64_t lora_base = (uint64_t)row * args.kv_lora_dim; + for (uint j = 0; j < args.kv_lora_dim; j++) { + const float kv = glm_cache_load_f32_or_f16(kv_lora_cache, + lora_base + j, + args.cache_f16); + dot0 += low0[j] * kv; + if (valid1) dot1 += low1[j] * kv; + } + const uint64_t rope_base = (uint64_t)row * args.qk_rope; + for (uint r = 0; r < args.qk_rope; r += 2u) { + const float2 y = glm_cache_load_rotated_rope_pair(k_rope_cache, + rope_base, + r, + row, + args.qk_rope, + args.cache_f16, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + dot0 += qh0[args.qk_nope + r] * y.x + + qh0[args.qk_nope + r + 1u] * y.y; + if (valid1) { + dot1 += qh1[args.qk_nope + r] * y.x + + qh1[args.qk_nope + r + 1u] * y.y; + } + } + score0 = dot0 * args.scale; + if (valid1) score1 = dot1 * args.scale; + } + scores0[s] = score0; + scores1[s] = score1; + local_max0 = max(local_max0, score0); + local_max1 = max(local_max1, score1); + } + red0[tid] = local_max0; + red1[tid] = local_max1; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) { + red0[tid] = max(red0[tid], red0[tid + step]); + red1[tid] = max(red1[tid], red1[tid + step]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float max_score0 = red0[0]; + const float max_score1 = red1[0]; + + float local_sum0 = 0.0f; + float local_sum1 = 0.0f; + for (uint s = tid; s < args.n_selected; s += nth) { + const float w0 = (max_score0 > -INFINITY) ? exp(scores0[s] - max_score0) : 0.0f; + const float w1 = (valid1 && max_score1 > -INFINITY) ? exp(scores1[s] - max_score1) : 0.0f; + scores0[s] = w0; + scores1[s] = w1; + local_sum0 += w0; + local_sum1 += w1; + } + red0[tid] = local_sum0; + red1[tid] = local_sum1; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint step = nth >> 1; step > 0; step >>= 1) { + if (tid < step) { + red0[tid] += red0[tid + step]; + red1[tid] += red1[tid + step]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float denom0 = max(red0[0], 1.0e-20f); + const float denom1 = max(red1[0], 1.0e-20f); + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint j = tid; j < args.kv_lora_dim; j += nth) { + float acc0 = 0.0f; + float acc1 = 0.0f; + for (uint s = 0; s < args.n_selected; s++) { + const uint row = token_selected[s]; + if (row < args.cache_cap) { + const float kv = glm_cache_load_f32_or_f16(kv_lora_cache, + (uint64_t)row * args.kv_lora_dim + j, + args.cache_f16); + acc0 += scores0[s] * kv; + if (valid1) acc1 += scores1[s] * kv; + } + } + lora0[j] = acc0 / denom0; + if (valid1) lora1[j] = acc1 / denom1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint d = tid; d < args.value_dim; d += nth) { + device float *out0 = + (device float *)(heads + + (uint64_t)token * heads_token_stride + + (uint64_t)head0 * args.value_dim * sizeof(float)); + device const char *row0 = + value_weight + ((uint64_t)head0 * args.value_dim + d) * args.value_row_bytes; + out0[d] = glm_quant_dot_row_tg_f32(args.value_type, row0, lora0, args.kv_lora_dim); + + if (valid1) { + device float *out1 = + (device float *)(heads + + (uint64_t)token * heads_token_stride + + (uint64_t)head1 * args.value_dim * sizeof(float)); + device const char *row1 = + value_weight + ((uint64_t)head1 * args.value_dim + d) * args.value_row_bytes; + out1[d] = glm_quant_dot_row_tg_f32(args.value_type, row1, lora1, args.kv_lora_dim); + } + } +} + +template +kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_impl( + constant ds4_metal_args_glm_attention_indexed_batch & args, + device const char *q, + device const char *qk_low, + device const char *kv_lora_cache, + device const char *k_rope_cache, + device const uint32_t *selected, + device char *lora_out, + threadgroup half4 *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid_u [[thread_index_in_threadgroup]], + ushort lane_u [[thread_index_in_simdgroup]], + ushort sg_u [[simdgroup_index_in_threadgroup]]) { + constexpr uint group_heads = 8u; + constexpr uint stage_rows = 16u; + const uint token = tgpig.y; + const uint tid = (uint)tid_u; + const uint lane = (uint)lane_u; + const uint head_in_group = (uint)sg_u; + const uint head = tgpig.x * group_heads + head_in_group + args.head_base; + if (token >= args.n_tokens || + args.n_selected == 0u || + args.cache_f16 == 0u || + args.kv_lora_dim != 512u || + args.qk_rope != 64u) { + return; + } + + const bool valid_head = assume_valid_heads || head < args.n_head; + const uint safe_head = valid_head ? head : 0u; + const uint kv_vecs = args.kv_lora_dim >> 2; + const uint rope_vecs = args.qk_rope >> 2; + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint64_t q_token_stride = (uint64_t)args.n_head * qk_dim * sizeof(float); + const uint64_t low_token_stride = + (uint64_t)args.n_head * args.kv_lora_dim * sizeof(float); + + threadgroup half4 *kv_shared = scratch; + threadgroup float4 *rope_shared = + (threadgroup float4 *)(kv_shared + stage_rows * kv_vecs); + + device const float *qh = + (device const float *)(q + + (uint64_t)token * q_token_stride + + (uint64_t)safe_head * qk_dim * sizeof(float)); + device const float4 *low4 = + (device const float4 *)(qk_low + + (uint64_t)token * low_token_stride + + (uint64_t)safe_head * args.kv_lora_dim * sizeof(float)); + device const uint32_t *token_selected = + selected + (uint64_t)token * args.n_selected; + + float4 low0 = 0.0f; + float4 low1 = 0.0f; + float4 low2 = 0.0f; + float4 low3 = 0.0f; + float4 qrope = 0.0f; + if (valid_head) { + low0 = low4[lane + 0u]; + low1 = low4[lane + 32u]; + low2 = low4[lane + 64u]; + low3 = low4[lane + 96u]; + if (lane < rope_vecs) { + qrope = *((device const float4 *)(qh + args.qk_nope + lane * 4u)); + } + } + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + + float M = -FLT_MAX / 2.0f; + float S = 0.0f; + float4 o0 = 0.0f; + float4 o1 = 0.0f; + float4 o2 = 0.0f; + float4 o3 = 0.0f; + + for (uint base = 0u; base < args.n_selected; base += stage_rows) { + const uint rows = min(stage_rows, args.n_selected - base); + for (uint off = tid; off < rows * kv_vecs; off += 256u) { + const uint rr = off / kv_vecs; + const uint vv = off - rr * kv_vecs; + const uint row = token_selected[base + rr]; + const bool valid_row = assume_valid_rows || row < args.cache_cap; + if (valid_row) { + device const half4 *src = + (device const half4 *)((device const half *)kv_lora_cache + + (uint64_t)row * args.kv_lora_dim); + kv_shared[off] = src[vv]; + } else { + kv_shared[off] = half4(half(0.0f)); + } + } + for (uint off = tid; off < rows * rope_vecs; off += 256u) { + const uint rr = off / rope_vecs; + const uint vv = off - rr * rope_vecs; + const uint r = vv * 4u; + const uint row = token_selected[base + rr]; + const bool valid_row = assume_valid_rows || row < args.cache_cap; + if (valid_row) { + const uint64_t rope_base = (uint64_t)row * args.qk_rope; + const float2 y0 = + glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, + rope_base, + r, + row, + args.qk_rope, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + const float2 y1 = + glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, + rope_base, + r + 2u, + row, + args.qk_rope, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + rope_shared[off] = float4(y0.x, y0.y, y1.x, y1.y); + } else { + rope_shared[off] = float4(0.0f); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint rr = 0u; rr < rows; rr++) { + const uint row = token_selected[base + rr]; + const bool valid_row = assume_valid_rows || row < args.cache_cap; + threadgroup const half4 *kv_row = kv_shared + rr * kv_vecs; + threadgroup const float4 *rope_row = rope_shared + rr * rope_vecs; + float partial = 0.0f; + if (valid_head && valid_row) { + partial += dot(low0, (float4)kv_row[lane + 0u]); + partial += dot(low1, (float4)kv_row[lane + 32u]); + partial += dot(low2, (float4)kv_row[lane + 64u]); + partial += dot(low3, (float4)kv_row[lane + 96u]); + if (lane < rope_vecs) { + partial += dot(qrope, rope_row[lane]); + } + } + const float sum = simd_sum(partial); + const float score = + (valid_head && valid_row) ? sum * args.scale : -FLT_MAX / 2.0f; + if (valid_head && valid_row) { + const float new_m = max(M, score); + const float old_scale = exp(M - new_m); + const float row_scale = exp(score - new_m); + o0 = o0 * old_scale + (float4)kv_row[lane + 0u] * row_scale; + o1 = o1 * old_scale + (float4)kv_row[lane + 32u] * row_scale; + o2 = o2 * old_scale + (float4)kv_row[lane + 64u] * row_scale; + o3 = o3 * old_scale + (float4)kv_row[lane + 96u] * row_scale; + S = S * old_scale + row_scale; + M = new_m; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (valid_head) { + const float inv_s = S > 0.0f ? 1.0f / S : 0.0f; + device float4 *out4 = + (device float4 *)(lora_out + + ((uint64_t)token * args.n_head + head) * + args.kv_lora_dim * sizeof(float)); + out4[lane + 0u] = o0 * inv_s; + out4[lane + 32u] = o1 * inv_s; + out4[lane + 64u] = o2 * inv_s; + out4[lane + 96u] = o3 * inv_s; + } +} + +typedef decltype(kernel_glm_attention_indexed_batch_lora_group8_vec_impl) + glm_attention_indexed_batch_lora_group8_vec_t; + +template [[host_name("kernel_glm_attention_indexed_batch_lora_group8_vec")]] +kernel glm_attention_indexed_batch_lora_group8_vec_t +kernel_glm_attention_indexed_batch_lora_group8_vec_impl; + +template [[host_name("kernel_glm_attention_indexed_batch_lora_group8_vec_valid")]] +kernel glm_attention_indexed_batch_lora_group8_vec_t +kernel_glm_attention_indexed_batch_lora_group8_vec_impl; + +template [[host_name("kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads")]] +kernel glm_attention_indexed_batch_lora_group8_vec_t +kernel_glm_attention_indexed_batch_lora_group8_vec_impl; + +template +kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl( + constant ds4_metal_args_glm_attention_indexed_batch & args, + device const char *q, + device const char *qk_low, + device const char *kv_lora_cache, + device const char *k_rope_cache, + device char *lora_out, + threadgroup half4 *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid_u [[thread_index_in_threadgroup]], + ushort lane_u [[thread_index_in_simdgroup]], + ushort sg_u [[simdgroup_index_in_threadgroup]]) { + constexpr uint group_heads = 8u; + constexpr uint stage_rows = 16u; + const uint token = tgpig.y; + const uint tid = (uint)tid_u; + const uint lane = (uint)lane_u; + const uint head_in_group = (uint)sg_u; + const uint head = tgpig.x * group_heads + head_in_group + args.head_base; + if (token >= args.n_tokens || + args.n_selected == 0u || + args.cache_f16 == 0u || + args.kv_lora_dim != 512u || + args.qk_rope != 64u) { + return; + } + + const uint visible = min(args.n_selected, args.pos0 + token + 1u); + if (visible == 0u) return; + + const bool valid_head = assume_valid_heads || head < args.n_head; + const uint safe_head = valid_head ? head : 0u; + const uint kv_vecs = args.kv_lora_dim >> 2; + const uint rope_vecs = args.qk_rope >> 2; + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint64_t q_token_stride = (uint64_t)args.n_head * qk_dim * sizeof(float); + const uint64_t low_token_stride = + (uint64_t)args.n_head * args.kv_lora_dim * sizeof(float); + + threadgroup half4 *kv_shared = scratch; + threadgroup float4 *rope_shared = + (threadgroup float4 *)(kv_shared + stage_rows * kv_vecs); + + device const float *qh = + (device const float *)(q + + (uint64_t)token * q_token_stride + + (uint64_t)safe_head * qk_dim * sizeof(float)); + device const float4 *low4 = + (device const float4 *)(qk_low + + (uint64_t)token * low_token_stride + + (uint64_t)safe_head * args.kv_lora_dim * sizeof(float)); + + float4 low0 = 0.0f; + float4 low1 = 0.0f; + float4 low2 = 0.0f; + float4 low3 = 0.0f; + float4 qrope = 0.0f; + if (valid_head) { + low0 = low4[lane + 0u]; + low1 = low4[lane + 32u]; + low2 = low4[lane + 64u]; + low3 = low4[lane + 96u]; + if (lane < rope_vecs) { + qrope = *((device const float4 *)(qh + args.qk_nope + lane * 4u)); + } + } + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + + float M = -FLT_MAX / 2.0f; + float S = 0.0f; + float4 o0 = 0.0f; + float4 o1 = 0.0f; + float4 o2 = 0.0f; + float4 o3 = 0.0f; + + for (uint base = 0u; base < visible; base += stage_rows) { + const uint rows = min(stage_rows, visible - base); + for (uint off = tid; off < rows * kv_vecs; off += 256u) { + const uint rr = off / kv_vecs; + const uint vv = off - rr * kv_vecs; + const uint row = base + rr; + device const half4 *src = + (device const half4 *)((device const half *)kv_lora_cache + + (uint64_t)row * args.kv_lora_dim); + kv_shared[off] = src[vv]; + } + for (uint off = tid; off < rows * rope_vecs; off += 256u) { + const uint rr = off / rope_vecs; + const uint vv = off - rr * rope_vecs; + const uint r = vv * 4u; + const uint row = base + rr; + const uint64_t rope_base = (uint64_t)row * args.qk_rope; + const float2 y0 = + glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, + rope_base, + r, + row, + args.qk_rope, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + const float2 y1 = + glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache, + rope_base, + r + 2u, + row, + args.qk_rope, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + rope_shared[off] = float4(y0.x, y0.y, y1.x, y1.y); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint rr = 0u; rr < rows; rr++) { + threadgroup const half4 *kv_row = kv_shared + rr * kv_vecs; + threadgroup const float4 *rope_row = rope_shared + rr * rope_vecs; + float partial = 0.0f; + if (valid_head) { + partial += dot(low0, (float4)kv_row[lane + 0u]); + partial += dot(low1, (float4)kv_row[lane + 32u]); + partial += dot(low2, (float4)kv_row[lane + 64u]); + partial += dot(low3, (float4)kv_row[lane + 96u]); + if (lane < rope_vecs) { + partial += dot(qrope, rope_row[lane]); + } + } + const float sum = simd_sum(partial); + const float score = valid_head ? sum * args.scale : -FLT_MAX / 2.0f; + if (valid_head) { + const float new_m = max(M, score); + const float old_scale = exp(M - new_m); + const float row_scale = exp(score - new_m); + o0 = o0 * old_scale + (float4)kv_row[lane + 0u] * row_scale; + o1 = o1 * old_scale + (float4)kv_row[lane + 32u] * row_scale; + o2 = o2 * old_scale + (float4)kv_row[lane + 64u] * row_scale; + o3 = o3 * old_scale + (float4)kv_row[lane + 96u] * row_scale; + S = S * old_scale + row_scale; + M = new_m; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (valid_head) { + const float inv_s = S > 0.0f ? 1.0f / S : 0.0f; + device float4 *out4 = + (device float4 *)(lora_out + + ((uint64_t)token * args.n_head + head) * + args.kv_lora_dim * sizeof(float)); + out4[lane + 0u] = o0 * inv_s; + out4[lane + 32u] = o1 * inv_s; + out4[lane + 64u] = o2 * inv_s; + out4[lane + 96u] = o3 * inv_s; + } +} + +typedef decltype(kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl) + glm_attention_indexed_batch_lora_group8_vec_causal_t; + +template [[host_name("kernel_glm_attention_indexed_batch_lora_group8_vec_causal")]] +kernel glm_attention_indexed_batch_lora_group8_vec_causal_t +kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl; + +template [[host_name("kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads")]] +kernel glm_attention_indexed_batch_lora_group8_vec_causal_t +kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl; + +kernel void kernel_glm_attention_indexed_batch_group8( + constant ds4_metal_args_glm_attention_indexed_batch & args, + device const char *q, + device const char *qk_low, + device const char *kv_lora_cache, + device const char *k_rope_cache, + device const char *value_weight, + device const uint32_t *selected, + device char *heads, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid_u [[thread_index_in_threadgroup]], + ushort lane_u [[thread_index_in_simdgroup]], + ushort sg_u [[simdgroup_index_in_threadgroup]]) { + const uint token = tgpig.y; + if (token >= args.n_tokens || args.n_selected == 0u) return; + + constexpr uint group_heads = 8u; + constexpr uint stage_rows = 8u; + const uint tid = (uint)tid_u; + const uint lane = (uint)lane_u; + const uint head_in_group = (uint)sg_u; + const uint head = tgpig.x * group_heads + head_in_group + args.head_base; + const bool valid_head = head < args.n_head; + const uint safe_head = valid_head ? head : 0u; + + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint64_t q_token_stride = (uint64_t)args.n_head * qk_dim * sizeof(float); + const uint64_t low_token_stride = (uint64_t)args.n_head * args.kv_lora_dim * sizeof(float); + const uint64_t heads_token_stride = (uint64_t)args.n_head * args.value_dim * sizeof(float); + + threadgroup half *kv_shared = (threadgroup half *)scratch; + threadgroup half *rope_shared = kv_shared + stage_rows * args.kv_lora_dim; + threadgroup float *lora_sums = + (threadgroup float *)(rope_shared + stage_rows * args.qk_rope); + threadgroup float *head_lora = lora_sums + head_in_group * args.kv_lora_dim; + + device const float *qh = + (device const float *)(q + + (uint64_t)token * q_token_stride + + (uint64_t)safe_head * qk_dim * sizeof(float)); + device const float *low = + (device const float *)(qk_low + + (uint64_t)token * low_token_stride + + (uint64_t)safe_head * args.kv_lora_dim * sizeof(float)); + device const uint32_t *token_selected = + selected + (uint64_t)token * args.n_selected; + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + + constexpr uint max_low_cache = 16u; + constexpr uint max_qrope_cache = 4u; + const bool use_low_cache = args.kv_lora_dim <= max_low_cache * 32u; + const bool use_qrope_cache = args.qk_rope <= max_qrope_cache * 32u; + half low_cache[max_low_cache]; + half qrope_cache[max_qrope_cache]; + for (uint k = 0u; k < max_low_cache; k++) { + const uint j = lane + k * 32u; + low_cache[k] = (valid_head && use_low_cache && j < args.kv_lora_dim) ? + (half)low[j] : (half)0.0f; + } + for (uint k = 0u; k < max_qrope_cache; k++) { + const uint r = lane + k * 32u; + qrope_cache[k] = (valid_head && use_qrope_cache && r < args.qk_rope) ? + (half)qh[args.qk_nope + r] : (half)0.0f; + } + + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + head_lora[j] = 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float M = -INFINITY; + float S = 0.0f; + for (uint base = 0u; base < args.n_selected; base += stage_rows) { + const uint rows = min(stage_rows, args.n_selected - base); + const uint kv_count = rows * args.kv_lora_dim; + const uint rope_pairs = args.qk_rope >> 1; + const uint rope_count = rows * rope_pairs; + + for (uint idx = tid; idx < kv_count; idx += 256u) { + const uint rr = idx / args.kv_lora_dim; + const uint j = idx - rr * args.kv_lora_dim; + const uint row = token_selected[base + rr]; + kv_shared[idx] = (row < args.cache_cap) + ? (half)glm_cache_load_f32_or_f16(kv_lora_cache, + (uint64_t)row * args.kv_lora_dim + j, + args.cache_f16) + : (half)0.0f; + } + for (uint idx = tid; idx < rope_count; idx += 256u) { + const uint rr = idx / rope_pairs; + const uint pair = idx - rr * rope_pairs; + const uint r = pair * 2u; + const uint row = token_selected[base + rr]; + threadgroup half *rope_row = rope_shared + rr * args.qk_rope; + if (row < args.cache_cap) { + const float2 y = glm_cache_load_rotated_rope_pair(k_rope_cache, + (uint64_t)row * args.qk_rope, + r, + row, + args.qk_rope, + args.cache_f16, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + rope_row[r] = (half)y.x; + rope_row[r + 1u] = (half)y.y; + } else { + rope_row[r] = (half)0.0f; + rope_row[r + 1u] = (half)0.0f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint rr = 0u; rr < rows; rr++) { + const uint row = token_selected[base + rr]; + const bool valid_row = row < args.cache_cap; + float partial = 0.0f; + if (valid_head && valid_row) { + threadgroup const half *kv_row = kv_shared + rr * args.kv_lora_dim; + threadgroup const half *rope_row = rope_shared + rr * args.qk_rope; + if (use_low_cache) { + for (uint k = 0u; k < max_low_cache; k++) { + const uint j = lane + k * 32u; + if (j < args.kv_lora_dim) { + partial += (float)(low_cache[k] * kv_row[j]); + } + } + } else { + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + partial += low[j] * (float)kv_row[j]; + } + } + if (use_qrope_cache) { + for (uint k = 0u; k < max_qrope_cache; k++) { + const uint r = lane + k * 32u; + if (r < args.qk_rope) { + partial += (float)(qrope_cache[k] * rope_row[r]); + } + } + } else { + for (uint r = lane; r < args.qk_rope; r += 32u) { + partial += qh[args.qk_nope + r] * (float)rope_row[r]; + } + } + } + + const float sum = simd_sum(partial); + const float score = (valid_head && valid_row) ? sum * args.scale : -INFINITY; + if (valid_head && valid_row) { + threadgroup const half *kv_row = kv_shared + rr * args.kv_lora_dim; + const float old_m = M; + const float new_m = max(M, score); + const float old_scale = (old_m == -INFINITY) ? 0.0f : exp(old_m - new_m); + const float row_scale = exp(score - new_m); + S = S * old_scale + row_scale; + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + head_lora[j] = head_lora[j] * old_scale + row_scale * (float)kv_row[j]; + } + M = new_m; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float inv_s = (valid_head && S > 0.0f) ? 1.0f / S : 0.0f; + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + head_lora[j] *= inv_s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (valid_head) { + if (args.value_type == 1u) { + const uint64_t offset = + (uint64_t)token * + ((uint64_t)args.n_head * args.kv_lora_dim * sizeof(float)) + + (uint64_t)head * args.kv_lora_dim * sizeof(float); + device float *out = + (device float *)(heads + offset); + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + out[j] = head_lora[j]; + } + return; + } + device float *out = + (device float *)(heads + + (uint64_t)token * heads_token_stride + + (uint64_t)head * args.value_dim * sizeof(float)); + for (uint d = lane; d < args.value_dim; d += 32u) { + device const char *row = + value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes; + out[d] = glm_quant_dot_row_tg_f32(args.value_type, row, head_lora, args.kv_lora_dim); + } + } +} + +kernel void kernel_glm_attention_indexed_batch_q2_group4( + constant ds4_metal_args_glm_attention_indexed_batch & args, + device const char *q, + device const char *qk_low, + device const char *kv_lora_cache, + device const char *k_rope_cache, + device const char *value_weight, + device const uint32_t *selected, + device char *heads, + threadgroup uint *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid_u [[thread_index_in_threadgroup]], + ushort lane_u [[thread_index_in_simdgroup]], + ushort sg_u [[simdgroup_index_in_threadgroup]]) { + const uint token0 = tgpig.y * 2u; + if (token0 >= args.n_tokens || args.n_selected == 0u) return; + + constexpr uint group_heads = 4u; + constexpr uint stage_rows = 4u; + constexpr uint group_threads = 128u; + const uint token1 = token0 + 1u; + const bool valid1 = token1 < args.n_tokens; + const uint tid = (uint)tid_u; + const uint lane = (uint)lane_u; + const uint head_in_group = (uint)sg_u; + const uint head = tgpig.x * group_heads + head_in_group + args.head_base; + const bool valid_head = head < args.n_head; + const uint safe_head = valid_head ? head : 0u; + + const uint qk_dim = args.qk_nope + args.qk_rope; + const uint64_t q_token_stride = (uint64_t)args.n_head * qk_dim * sizeof(float); + const uint64_t low_token_stride = (uint64_t)args.n_head * args.kv_lora_dim * sizeof(float); + const uint64_t heads_token_stride = (uint64_t)args.n_head * args.value_dim * sizeof(float); + + const uint bit_words = (args.cache_cap + 31u) >> 5; + threadgroup atomic_uint *member_bits = (threadgroup atomic_uint *)scratch; + threadgroup half *kv_shared = (threadgroup half *)(scratch + bit_words); + threadgroup half *rope_shared = kv_shared + stage_rows * args.kv_lora_dim; + threadgroup float *lora_sums = + (threadgroup float *)(rope_shared + stage_rows * args.qk_rope); + threadgroup float *head_lora0 = lora_sums + head_in_group * args.kv_lora_dim; + threadgroup float *head_lora1 = + lora_sums + (group_heads + head_in_group) * args.kv_lora_dim; + + const uint safe_token1 = valid1 ? token1 : token0; + device const float *qh0 = + (device const float *)(q + + (uint64_t)token0 * q_token_stride + + (uint64_t)safe_head * qk_dim * sizeof(float)); + device const float *qh1 = + (device const float *)(q + + (uint64_t)safe_token1 * q_token_stride + + (uint64_t)safe_head * qk_dim * sizeof(float)); + device const float *low0 = + (device const float *)(qk_low + + (uint64_t)token0 * low_token_stride + + (uint64_t)safe_head * args.kv_lora_dim * sizeof(float)); + device const float *low1 = + (device const float *)(qk_low + + (uint64_t)safe_token1 * low_token_stride + + (uint64_t)safe_head * args.kv_lora_dim * sizeof(float)); + device const uint32_t *selected0 = selected + (uint64_t)token0 * args.n_selected; + device const uint32_t *selected1 = selected + (uint64_t)safe_token1 * args.n_selected; + + float corr_dims[2] = {0.0f, 0.0f}; + if (args.ext_factor != 0.0f) { + glm_rope_yarn_corr_dims((int)args.qk_rope, + (int)args.n_ctx_orig, + args.freq_base, + args.beta_fast, + args.beta_slow, + corr_dims); + } + + constexpr uint max_low_cache = 16u; + constexpr uint max_qrope_cache = 4u; + const bool use_low_cache = args.kv_lora_dim <= max_low_cache * 32u; + const bool use_qrope_cache = args.qk_rope <= max_qrope_cache * 32u; + half low_cache0[max_low_cache]; + half low_cache1[max_low_cache]; + half qrope_cache0[max_qrope_cache]; + half qrope_cache1[max_qrope_cache]; + for (uint k = 0u; k < max_low_cache; k++) { + const uint j = lane + k * 32u; + low_cache0[k] = (valid_head && use_low_cache && j < args.kv_lora_dim) ? + (half)low0[j] : (half)0.0f; + low_cache1[k] = (valid_head && valid1 && use_low_cache && j < args.kv_lora_dim) ? + (half)low1[j] : (half)0.0f; + } + for (uint k = 0u; k < max_qrope_cache; k++) { + const uint r = lane + k * 32u; + qrope_cache0[k] = (valid_head && use_qrope_cache && r < args.qk_rope) ? + (half)qh0[args.qk_nope + r] : (half)0.0f; + qrope_cache1[k] = (valid_head && valid1 && use_qrope_cache && r < args.qk_rope) ? + (half)qh1[args.qk_nope + r] : (half)0.0f; + } + + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + head_lora0[j] = 0.0f; + if (valid1) head_lora1[j] = 0.0f; + } + for (uint i = tid; i < bit_words; i += group_threads) { + atomic_store_explicit(member_bits + i, 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint s = tid; s < args.n_selected; s += group_threads) { + const uint row = selected0[s]; + if (row < args.cache_cap) { + const uint mask = 1u << (row & 31u); + atomic_fetch_or_explicit(member_bits + (row >> 5), mask, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float M0 = -INFINITY; + float S0 = 0.0f; + float M1 = -INFINITY; + float S1 = 0.0f; + + if (valid1) { + for (uint base = 0u; base < args.n_selected; base += stage_rows) { + const uint rows = min(stage_rows, args.n_selected - base); + const uint kv_count = rows * args.kv_lora_dim; + const uint rope_pairs = args.qk_rope >> 1; + const uint rope_count = rows * rope_pairs; + + for (uint idx = tid; idx < kv_count; idx += 256u) { + const uint rr = idx / args.kv_lora_dim; + const uint j = idx - rr * args.kv_lora_dim; + const uint row = selected1[base + rr]; + kv_shared[idx] = (row < args.cache_cap) + ? (half)glm_cache_load_f32_or_f16(kv_lora_cache, + (uint64_t)row * args.kv_lora_dim + j, + args.cache_f16) + : (half)0.0f; + } + for (uint idx = tid; idx < rope_count; idx += 256u) { + const uint rr = idx / rope_pairs; + const uint pair = idx - rr * rope_pairs; + const uint r = pair * 2u; + const uint row = selected1[base + rr]; + threadgroup half *rope_row = rope_shared + rr * args.qk_rope; + if (row < args.cache_cap) { + const float2 y = glm_cache_load_rotated_rope_pair(k_rope_cache, + (uint64_t)row * args.qk_rope, + r, + row, + args.qk_rope, + args.cache_f16, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + rope_row[r] = (half)y.x; + rope_row[r + 1u] = (half)y.y; + } else { + rope_row[r] = (half)0.0f; + rope_row[r + 1u] = (half)0.0f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint rr = 0u; rr < rows; rr++) { + const uint row = selected1[base + rr]; + const bool valid_row = row < args.cache_cap; + const bool in_token0 = valid_row && + ((atomic_load_explicit(member_bits + (row >> 5), + memory_order_relaxed) & + (1u << (row & 31u))) != 0u); + threadgroup const half *kv_row = kv_shared + rr * args.kv_lora_dim; + threadgroup const half *rope_row = rope_shared + rr * args.qk_rope; + + float partial0 = 0.0f; + float partial1 = 0.0f; + if (valid_head && valid_row) { + if (use_low_cache) { + for (uint k = 0u; k < max_low_cache; k++) { + const uint j = lane + k * 32u; + if (j < args.kv_lora_dim) { + const half kv = kv_row[j]; + if (in_token0) partial0 += (float)(low_cache0[k] * kv); + partial1 += (float)(low_cache1[k] * kv); + } + } + } else { + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + const float kv = (float)kv_row[j]; + if (in_token0) partial0 += low0[j] * kv; + partial1 += low1[j] * kv; + } + } + if (use_qrope_cache) { + for (uint k = 0u; k < max_qrope_cache; k++) { + const uint r = lane + k * 32u; + if (r < args.qk_rope) { + const half kv = rope_row[r]; + if (in_token0) partial0 += (float)(qrope_cache0[k] * kv); + partial1 += (float)(qrope_cache1[k] * kv); + } + } + } else { + for (uint r = lane; r < args.qk_rope; r += 32u) { + const float kv = (float)rope_row[r]; + if (in_token0) partial0 += qh0[args.qk_nope + r] * kv; + partial1 += qh1[args.qk_nope + r] * kv; + } + } + } + + const float sum0 = simd_sum(partial0); + const float sum1 = simd_sum(partial1); + const float score0 = (valid_head && in_token0) ? sum0 * args.scale : -INFINITY; + const float score1 = (valid_head && valid_row) ? sum1 * args.scale : -INFINITY; + if (valid_head && in_token0) { + const float new_m = max(M0, score0); + const float old_scale = (M0 == -INFINITY) ? 0.0f : exp(M0 - new_m); + const float row_scale = exp(score0 - new_m); + S0 = S0 * old_scale + row_scale; + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + head_lora0[j] = head_lora0[j] * old_scale + row_scale * (float)kv_row[j]; + } + M0 = new_m; + } + if (valid_head && valid_row) { + const float new_m = max(M1, score1); + const float old_scale = (M1 == -INFINITY) ? 0.0f : exp(M1 - new_m); + const float row_scale = exp(score1 - new_m); + S1 = S1 * old_scale + row_scale; + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + head_lora1[j] = head_lora1[j] * old_scale + row_scale * (float)kv_row[j]; + } + M1 = new_m; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (uint i = tid; i < bit_words; i += group_threads) { + atomic_store_explicit(member_bits + i, 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint s = tid; s < args.n_selected; s += group_threads) { + const uint row = selected1[s]; + if (row < args.cache_cap) { + const uint mask = 1u << (row & 31u); + atomic_fetch_or_explicit(member_bits + (row >> 5), mask, memory_order_relaxed); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (uint base = 0u; base < args.n_selected; base += stage_rows) { + const uint rows = min(stage_rows, args.n_selected - base); + const uint kv_count = rows * args.kv_lora_dim; + const uint rope_pairs = args.qk_rope >> 1; + const uint rope_count = rows * rope_pairs; + + for (uint idx = tid; idx < kv_count; idx += 256u) { + const uint rr = idx / args.kv_lora_dim; + const uint j = idx - rr * args.kv_lora_dim; + const uint row = selected0[base + rr]; + kv_shared[idx] = (row < args.cache_cap) + ? (half)glm_cache_load_f32_or_f16(kv_lora_cache, + (uint64_t)row * args.kv_lora_dim + j, + args.cache_f16) + : (half)0.0f; + } + for (uint idx = tid; idx < rope_count; idx += 256u) { + const uint rr = idx / rope_pairs; + const uint pair = idx - rr * rope_pairs; + const uint r = pair * 2u; + const uint row = selected0[base + rr]; + threadgroup half *rope_row = rope_shared + rr * args.qk_rope; + if (row < args.cache_cap) { + const float2 y = glm_cache_load_rotated_rope_pair(k_rope_cache, + (uint64_t)row * args.qk_rope, + r, + row, + args.qk_rope, + args.cache_f16, + args.freq_base, + args.freq_scale, + args.ext_factor, + args.attn_factor, + corr_dims[0], + corr_dims[1]); + rope_row[r] = (half)y.x; + rope_row[r + 1u] = (half)y.y; + } else { + rope_row[r] = (half)0.0f; + rope_row[r + 1u] = (half)0.0f; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint rr = 0u; rr < rows; rr++) { + const uint row = selected0[base + rr]; + const bool valid_row = row < args.cache_cap; + const bool in_token1 = valid1 && valid_row && + ((atomic_load_explicit(member_bits + (row >> 5), + memory_order_relaxed) & + (1u << (row & 31u))) != 0u); + const bool take0 = valid_row && !in_token1; + threadgroup const half *kv_row = kv_shared + rr * args.kv_lora_dim; + threadgroup const half *rope_row = rope_shared + rr * args.qk_rope; + + float partial0 = 0.0f; + if (valid_head && take0) { + if (use_low_cache) { + for (uint k = 0u; k < max_low_cache; k++) { + const uint j = lane + k * 32u; + if (j < args.kv_lora_dim) { + partial0 += (float)(low_cache0[k] * kv_row[j]); + } + } + } else { + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + partial0 += low0[j] * (float)kv_row[j]; + } + } + if (use_qrope_cache) { + for (uint k = 0u; k < max_qrope_cache; k++) { + const uint r = lane + k * 32u; + if (r < args.qk_rope) { + partial0 += (float)(qrope_cache0[k] * rope_row[r]); + } + } + } else { + for (uint r = lane; r < args.qk_rope; r += 32u) { + partial0 += qh0[args.qk_nope + r] * (float)rope_row[r]; + } + } + } + + const float sum0 = simd_sum(partial0); + const float score0 = (valid_head && take0) ? sum0 * args.scale : -INFINITY; + if (valid_head && take0) { + const float new_m = max(M0, score0); + const float old_scale = (M0 == -INFINITY) ? 0.0f : exp(M0 - new_m); + const float row_scale = exp(score0 - new_m); + S0 = S0 * old_scale + row_scale; + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + head_lora0[j] = head_lora0[j] * old_scale + row_scale * (float)kv_row[j]; + } + M0 = new_m; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const float inv_s0 = (valid_head && S0 > 0.0f) ? 1.0f / S0 : 0.0f; + const float inv_s1 = (valid_head && valid1 && S1 > 0.0f) ? 1.0f / S1 : 0.0f; + for (uint j = lane; j < args.kv_lora_dim; j += 32u) { + head_lora0[j] *= inv_s0; + if (valid1) head_lora1[j] *= inv_s1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (valid_head) { + device float *out0 = + (device float *)(heads + + (uint64_t)token0 * heads_token_stride + + (uint64_t)head * args.value_dim * sizeof(float)); + device float *out1 = + (device float *)(heads + + (uint64_t)safe_token1 * heads_token_stride + + (uint64_t)head * args.value_dim * sizeof(float)); + for (uint d = lane; d < args.value_dim; d += 32u) { + device const char *row = + value_weight + ((uint64_t)head * args.value_dim + d) * args.value_row_bytes; + out0[d] = glm_quant_dot_row_tg_f32(args.value_type, row, head_lora0, args.kv_lora_dim); + if (valid1) { + out1[d] = glm_quant_dot_row_tg_f32(args.value_type, row, head_lora1, args.kv_lora_dim); + } + } + } +} + +// GLM-5.2 decode router for one token. Selection uses sigmoid(logit)+bias, +// while route weights are normalized from the unbiased sigmoid probabilities. +kernel void kernel_glm_router_select_one( + constant ds4_metal_args_glm_router_select_one & args, + device const float *logits, + device const float *bias, + device int32_t *selected, + device float *weights, + device float *probs, + threadgroup float *scratch [[threadgroup(0)]], + uint token [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + threadgroup float *sel_scores = scratch; + threadgroup int32_t *idx = (threadgroup int32_t *)(scratch + 256); + device const float *token_logits = logits + (uint64_t)token * args.n_expert; + device int32_t *token_selected = selected + (uint64_t)token * args.n_expert_used; + device float *token_weights = weights + (uint64_t)token * args.n_expert_used; + device float *token_probs = probs + (uint64_t)token * args.n_expert; + + const uint n_expert = min(args.n_expert, 256u); + const bool active = tid < n_expert; + const float p = active ? ds4_glm_router_sigmoid(token_logits[tid]) : 0.0f; + if (active) token_probs[tid] = p; + sel_scores[tid] = active ? p + bias[tid] : -INFINITY; + idx[tid] = (int32_t)tid; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + const uint other = tid ^ j; + if (other > tid) { + const int32_t a = idx[tid]; + const int32_t b = idx[other]; + const bool descending = (tid & k) == 0; + const bool swap = descending + ? ds4_glm_router_better(sel_scores, b, a) + : ds4_glm_router_better(sel_scores, a, b); + if (swap) { + idx[tid] = b; + idx[other] = a; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + const uint k_used = min(args.n_expert_used, n_expert); + if (tid < k_used) { + token_selected[tid] = idx[tid]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tid < k_used) { + float sum = 0.0f; + for (uint i = 0; i < k_used; i++) { + sum += token_probs[(uint)token_selected[i]]; + } + sum = max(sum, 6.103515625e-5f); + token_weights[tid] = token_probs[(uint)token_selected[tid]] / sum * args.expert_weight_scale; + } +} + +// Batched Flash-router weight finalization after selection is already known. +// Six active lanes deliberately match kernel_sum_rows_f32_f32's reduction +// topology. The denominator and divided weights cross threadgroup storage +// boundaries so division cannot be reassociated with the final scale. +kernel void kernel_dsv4_router_weights_batch( + constant float &scale, + device const float *probs, + device const int32_t *selected, + device float *weights, + threadgroup volatile float *scratch [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + ushort tid [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]]) { + if (tid >= 6) return; + + threadgroup volatile float *sum_scratch = scratch; + threadgroup volatile float *denom_scratch = scratch + 32; + threadgroup volatile float *div_scratch = scratch + 33; + const uint out_index = row * 6u + (uint)tid; + const int32_t expert = selected[out_index]; + const float p = probs[row * 256u + (uint)expert]; + + // Keep this sequence identical to kernel_sum_rows_f32_f32 for width 6. + if (sgitg == 0) { + sum_scratch[tiisg] = 0.0f; + } + float sumf = 0.0f; + sumf += p; + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) { + sum_scratch[sgitg] = sumf; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + sumf = sum_scratch[tiisg]; + sumf = simd_sum(sumf); + + if (tid == 0) { + denom_scratch[0] = clamp(sumf, 6.103515625e-5f, INFINITY); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + div_scratch[tid] = p / denom_scratch[0]; + threadgroup_barrier(mem_flags::mem_threadgroup); + weights[out_index] = div_scratch[tid] * scale; +} + +// Decode router selection for one token after the existing +// sqrt(softplus(logit)) probability kernel has run. Bias affects only top-k +// selection. Route-weight normalization deliberately stays in the old one-token +// kernel: even tiny denominator-order changes here are amplified by 43 MoE +// layers, so this kernel only replaces the selection work. +kernel void kernel_dsv4_router_finalize_one( + constant ds4_metal_args_dsv4_router_select_one & args, + device const float *probs, + device const float *bias, + device const int32_t *hash, + device const int32_t *tokens, + device int32_t *selected, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]]) { + if (tid >= 256) return; + + threadgroup float *sel_scores = scratch; + threadgroup int32_t *idx = (threadgroup int32_t *)(scratch + 256); + const float p = probs[tid]; + sel_scores[tid] = args.has_bias ? p + bias[tid] : p; + idx[tid] = (int32_t)tid; + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (args.hash_mode) { + if (tid == 0) { + const uint token = args.use_token_buffer ? (uint)tokens[0] : args.token; + const uint row = min(token, args.hash_rows - 1u); + device const int32_t *src = hash + row * 6u; + for (uint i = 0; i < 6; i++) { + selected[i] = src[i]; + } + } + } else { + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + const uint other = tid ^ j; + if (other > tid) { + if ((tid & k) == 0) { + if (sel_scores[(uint)idx[tid]] < sel_scores[(uint)idx[other]]) { + const int32_t tmp = idx[tid]; + idx[tid] = idx[other]; + idx[other] = tmp; + } + } else { + if (sel_scores[(uint)idx[tid]] > sel_scores[(uint)idx[other]]) { + const int32_t tmp = idx[tid]; + idx[tid] = idx[other]; + idx[other] = tmp; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + if (tid < 6) { + selected[tid] = idx[tid]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +// M3 decode specialization for the non-hash one-token router. Scores and ids +// stay in registers. Intra-SIMD bitonic stages use shuffle-xor; the six stages +// that cross 32-lane SIMD groups exchange through alternating threadgroup +// banks. The next bank's publish barrier proves every prior-bank read finished; +// by the time a bank is reused two cross stages later, no reader can remain. +kernel void kernel_dsv4_router_finalize_one_simd( + constant ds4_metal_args_dsv4_router_select_one & args, + device const float *probs, + device const float *bias, + device const int32_t *hash, + device const int32_t *tokens, + device int32_t *selected, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]]) { + if (tid >= 256 || args.hash_mode) return; + + (void)hash; + (void)tokens; + threadgroup float *score0_tg = scratch; + threadgroup int32_t *idx0_tg = + (threadgroup int32_t *)(scratch + 256); + threadgroup float *score1_tg = scratch + 512; + threadgroup int32_t *idx1_tg = + (threadgroup int32_t *)(scratch + 768); + const float p = probs[tid]; + float score = args.has_bias ? p + bias[tid] : p; + int32_t idx = (int32_t)tid; + uint cross_stage = 0; + + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + float peer_score; + int32_t peer_idx; + bool take_peer; + const bool lower = (tid & j) == 0; + const bool descending = (tid & k) == 0; + + if (j < 32) { + peer_score = simd_shuffle_xor(score, (ushort)j); + peer_idx = simd_shuffle_xor(idx, (ushort)j); + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + } else { + threadgroup float *score_tg = + (cross_stage & 1u) != 0u ? score1_tg : score0_tg; + threadgroup int32_t *idx_tg = + (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg; + score_tg[tid] = score; + idx_tg[tid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint other = tid ^ j; + peer_score = score_tg[other]; + peer_idx = idx_tg[other]; + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + cross_stage++; + } + } + } + + if (tid < 6) { + selected[tid] = idx; + } +} + +// M3 decode specialization that extends the register/TG SIMD selection above +// through the existing six-value serial weight normalization. The selected ids +// cross the same device-memory boundary as the standalone weight kernel; +// volatile TG stores pin its left-fold and scaled-reciprocal rounding points. +kernel void kernel_dsv4_router_finalize_weights_one_simd( + constant ds4_metal_args_dsv4_router_select_one & args, + device const float *probs, + device const float *bias, + device const int32_t *hash, + device const int32_t *tokens, + device int32_t *selected, + device float *weights, + threadgroup float *scratch [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]]) { + if (tid >= 256 || args.hash_mode) return; + + (void)hash; + (void)tokens; + threadgroup float *score0_tg = scratch; + threadgroup int32_t *idx0_tg = + (threadgroup int32_t *)(scratch + 256); + threadgroup float *score1_tg = scratch + 512; + threadgroup int32_t *idx1_tg = + (threadgroup int32_t *)(scratch + 768); + const float p = probs[tid]; + float score = args.has_bias ? p + bias[tid] : p; + int32_t idx = (int32_t)tid; + uint cross_stage = 0; + + for (uint k = 2; k <= 256; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + float peer_score; + int32_t peer_idx; + bool take_peer; + const bool lower = (tid & j) == 0; + const bool descending = (tid & k) == 0; + + if (j < 32) { + peer_score = simd_shuffle_xor(score, (ushort)j); + peer_idx = simd_shuffle_xor(idx, (ushort)j); + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + } else { + threadgroup float *score_tg = + (cross_stage & 1u) != 0u ? score1_tg : score0_tg; + threadgroup int32_t *idx_tg = + (cross_stage & 1u) != 0u ? idx1_tg : idx0_tg; + score_tg[tid] = score; + idx_tg[tid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + + const uint other = tid ^ j; + peer_score = score_tg[other]; + peer_idx = idx_tg[other]; + take_peer = descending + ? (lower ? score < peer_score : score > peer_score) + : (lower ? score > peer_score : score < peer_score); + if (take_peer) { + score = peer_score; + idx = peer_idx; + } + cross_stage++; + } + } + } + + if (tid < 6) { + selected[tid] = idx; + } + threadgroup_barrier(mem_flags::mem_device); + + threadgroup volatile float *norm_scratch = + (threadgroup volatile float *)scratch; + if (tid == 0) { + device const int32_t *s = selected; + norm_scratch[0] = 0.0f; + for (uint i = 0; i < 6; i++) { + norm_scratch[0] = norm_scratch[0] + probs[s[i]]; + } + norm_scratch[0] = max(norm_scratch[0], 6.103515625e-5f); + norm_scratch[1] = 1.5f / norm_scratch[0]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < 6) { + device const int32_t *s = selected; + weights[tid] = probs[s[tid]] * norm_scratch[1]; + } +} + +// Fills the dense compressed-attention mask with -inf. The selected top-k rows +// are enabled by kernel_dsv4_topk_mask_scatter in a second ordered dispatch. +kernel void kernel_dsv4_topk_mask( + constant ds4_metal_args_dsv4_topk_mask & args, + device const char * topk, + device char * dst, + uint gid [[thread_position_in_grid]]) { + const int64_t n = args.ne0 * args.ne1; + if ((int64_t) gid >= n) { + return; + } + + const int64_t ic = gid % args.ne0; + const int64_t it = gid / args.ne0; + + (void)topk; + *((device float *) (dst + ic*args.nb0 + it*args.nb1)) = -INFINITY; +} + +// Enables the selected compressed rows in the dense mask. This replaces the +// old O(n_comp * n_tokens * top_k) membership test with O(top_k * n_tokens) +// writes while preserving exactly the same 0/-inf mask consumed by attention. +kernel void kernel_dsv4_topk_mask_scatter( + constant ds4_metal_args_dsv4_topk_mask & args, + device const char * topk, + device char * dst, + uint gid [[thread_position_in_grid]]) { + const int64_t n = args.ne00 * args.ne01; + if ((int64_t) gid >= n) { + return; + } + + const int64_t ik = gid % args.ne00; + const int64_t it = gid / args.ne00; + const int32_t idx = *((device const int32_t *) (topk + ik*args.nb00 + it*args.nb01)); + if (idx >= 0 && (int64_t)idx < args.ne0) { + *((device float *) (dst + (int64_t)idx*args.nb0 + it*args.nb1)) = 0.0f; + } +} + +// Sorts each token's selected compressed rows by row id. The indexer selects by +// score, but attention scans compressed K/V in cache order in the dense graph. +// Sorting preserves that order while still letting the indexed attention kernel +// touch only the selected rows. +kernel void kernel_dsv4_sort_i32_rows_asc( + constant ds4_metal_args_dsv4_topk_mask & args, + device const char * src, + device char * dst, + threadgroup int32_t * row_tmp [[threadgroup(0)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint n_threads [[threads_per_threadgroup]]) { + const uint top_k = (uint)args.ne00; + if (row >= (uint)args.ne01 || tid >= n_threads) { + return; + } + + for (uint i = tid; i < top_k; i += n_threads) { + row_tmp[i] = *((device const int32_t *) (src + (uint64_t)i*args.nb00 + (uint64_t)row*args.nb01)); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k = 2; k <= top_k; k <<= 1) { + for (uint j = k >> 1; j > 0; j >>= 1) { + for (uint i = tid; i < top_k; i += n_threads) { + const uint other = i ^ j; + if (other > i && other < top_k) { + const int32_t a = row_tmp[i]; + const int32_t b = row_tmp[other]; + const bool up = (i & k) == 0; + if ((up && a > b) || (!up && a < b)) { + row_tmp[i] = b; + row_tmp[other] = a; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + for (uint i = tid; i < top_k; i += n_threads) { + *((device int32_t *) (dst + (uint64_t)i*args.nb00 + (uint64_t)row*args.nb01)) = row_tmp[i]; + } +} + +static inline void dsv4_attend_f32_row_as_f16( + device const char *kv, + uint64_t row_stride, + uint row, + half4 q0, + half4 q1, + half4 q2, + half4 q3, + float scale, + ushort lane, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + device const float4 *kv4 = (device const float4 *)(kv + (uint64_t)row * row_stride); + const half4 k0 = (half4)kv4[lane + 0]; + const half4 k1 = (half4)kv4[lane + 32]; + const half4 k2 = (half4)kv4[lane + 64]; + const half4 k3 = (half4)kv4[lane + 96]; + + float score = dot((float4)q0, (float4)k0) + + dot((float4)q1, (float4)k1) + + dot((float4)q2, (float4)k2) + + dot((float4)q3, (float4)k3); + score = simd_sum(score) * scale; + + const float old_m = M; + const float new_m = max(M, score); + const float old_scale = exp(old_m - new_m); + const float row_scale = exp(score - new_m); + + S = S * old_scale + row_scale; + o0 *= old_scale; + o1 *= old_scale; + o2 *= old_scale; + o3 *= old_scale; + + o0 += (float4)k0 * row_scale; + o1 += (float4)k1 * row_scale; + o2 += (float4)k2 * row_scale; + o3 += (float4)k3 * row_scale; + M = new_m; +} + +static inline void dsv4_attend_shared_f32_row_as_f16( + threadgroup const float4 *kv4, + half4 q0, + half4 q1, + half4 q2, + half4 q3, + float scale, + ushort lane, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + const half4 k0 = (half4)kv4[lane + 0]; + const half4 k1 = (half4)kv4[lane + 32]; + const half4 k2 = (half4)kv4[lane + 64]; + const half4 k3 = (half4)kv4[lane + 96]; + + float score = dot((float4)q0, (float4)k0) + + dot((float4)q1, (float4)k1) + + dot((float4)q2, (float4)k2) + + dot((float4)q3, (float4)k3); + score = simd_sum(score) * scale; + + const float old_m = M; + const float new_m = max(M, score); + const float old_scale = exp(old_m - new_m); + const float row_scale = exp(score - new_m); + + S = S * old_scale + row_scale; + o0 *= old_scale; + o1 *= old_scale; + o2 *= old_scale; + o3 *= old_scale; + + o0 += (float4)k0 * row_scale; + o1 += (float4)k1 * row_scale; + o2 += (float4)k2 * row_scale; + o3 += (float4)k3 * row_scale; + M = new_m; +} + +static inline void dsv4_attend_shared_f32_row_as_f16_at( + threadgroup const float4 *kv4, + uint row_in_tg, + half4 q0, + half4 q1, + half4 q2, + half4 q3, + float scale, + ushort lane, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + dsv4_attend_shared_f32_row_as_f16(kv4 + row_in_tg * 128u, + q0, q1, q2, q3, + scale, + lane, + M, S, + o0, o1, o2, o3); +} + +static inline void dsv4_attend_shared_h4_row( + threadgroup const half4 *kv4, + half4 q0, + half4 q1, + half4 q2, + half4 q3, + float scale, + ushort lane, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + const half4 k0 = kv4[lane + 0]; + const half4 k1 = kv4[lane + 32]; + const half4 k2 = kv4[lane + 64]; + const half4 k3 = kv4[lane + 96]; + + float score = dot((float4)q0, (float4)k0) + + dot((float4)q1, (float4)k1) + + dot((float4)q2, (float4)k2) + + dot((float4)q3, (float4)k3); + score = simd_sum(score) * scale; + + const float old_m = M; + const float new_m = max(M, score); + const float old_scale = exp(old_m - new_m); + const float row_scale = exp(score - new_m); + + S = S * old_scale + row_scale; + o0 *= old_scale; + o1 *= old_scale; + o2 *= old_scale; + o3 *= old_scale; + + o0 += (float4)k0 * row_scale; + o1 += (float4)k1 * row_scale; + o2 += (float4)k2 * row_scale; + o3 += (float4)k3 * row_scale; + M = new_m; +} + +static inline void dsv4_attend_shared_h4_row_at( + threadgroup const half4 *kv4, + uint row_in_tg, + half4 q0, + half4 q1, + half4 q2, + half4 q3, + float scale, + ushort lane, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + dsv4_attend_shared_h4_row(kv4 + row_in_tg * 128u, + q0, q1, q2, q3, + scale, + lane, + M, S, + o0, o1, o2, o3); +} + +static inline half4 dsv4_load_cache_h4( + device const char *kv, + uint64_t row_stride, + uint row, + uint col, + bool f16_rows) { + device const char *base = kv + (uint64_t)row * row_stride; + if (f16_rows) { + return ((device const half4 *)base)[col]; + } + return (half4)((device const float4 *)base)[col]; +} + +static inline void dsv4_attend_sink( + float score, + thread float &M, + thread float &S, + thread float4 &o0, + thread float4 &o1, + thread float4 &o2, + thread float4 &o3) { + const float old_m = M; + const float new_m = max(M, score); + const float old_scale = exp(old_m - new_m); + const float row_scale = exp(score - new_m); + + S = S * old_scale + row_scale; + o0 *= old_scale; + o1 *= old_scale; + o2 *= old_scale; + o3 *= old_scale; + M = new_m; +} + +// DS4 ratio-4 indexed mixed attention. It replaces the dense top-k mask path: +// the threadgroup covers one token and eight heads. Top-k rows and local raw +// rows are the same for all heads of a token, so K/V is staged once in +// threadgroup memory and reused by the eight simdgroups. It keeps the DS4 F16 +// attention rounding by casting Q/K/V to half before the dot/value update. +kernel void kernel_dsv4_indexed_mixed_attention_heads8( + constant ds4_metal_args_dsv4_indexed_attention & args, + device const char *q, + device const char *raw_kv, + device const char *comp_kv, + device const char *topk, + device const char *sinks, + device char *dst, + threadgroup half4 *kv_shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + const uint token = tgpig.x; + const uint head = tgpig.y * 8u + (uint)sg; + if (token >= args.n_tokens || head >= args.n_head) { + return; + } + + device const float4 *q4 = (device const float4 *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + const half4 q0 = (half4)q4[lane + 0]; + const half4 q1 = (half4)q4[lane + 32]; + const half4 q2 = (half4)q4[lane + 64]; + const half4 q3 = (half4)q4[lane + 96]; + + float M = -FLT_MAX/2.0f; + float S = 0.0f; + float4 o0 = 0.0f; + float4 o1 = 0.0f; + float4 o2 = 0.0f; + float4 o3 = 0.0f; + + const uint qpos = args.pos0 + token; + const uint last_pos = args.pos0 + args.n_tokens - 1u; + const uint first_raw_pos = last_pos + 1u - args.n_raw; + const uint raw_last_pos = first_raw_pos + args.n_raw - 1u; + const uint window_first = (args.window != 0u && qpos + 1u > args.window) ? + qpos + 1u - args.window : 0u; + uint first = max(first_raw_pos, window_first); + uint last = min(qpos, raw_last_pos); + + if (first <= last) { + for (uint pos = first; pos <= last; pos++) { + const uint logical = pos - first_raw_pos; + const uint row = (args.raw_start + logical) % args.raw_cap; + device const float4 *src = (device const float4 *)(raw_kv + + (uint64_t)row * args.raw_row_stride); + if (tid < 128) kv_shared[tid] = (half4)src[tid]; + threadgroup_barrier(mem_flags::mem_threadgroup); + dsv4_attend_shared_h4_row(kv_shared, + q0, q1, q2, q3, + args.scale, + lane, + M, S, + o0, o1, o2, o3); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + uint visible = (qpos + 1u) / args.ratio; + visible = min(visible, args.n_comp); + device const int32_t *row_topk = (device const int32_t *)(topk + + (uint64_t)token * args.topk_token_stride); + for (uint i = 0; i < args.top_k; i++) { + const int32_t idx = row_topk[i]; + if (idx < 0) { + continue; + } + if ((uint)idx >= visible) { + break; + } + if (tid < 128) { + kv_shared[tid] = dsv4_load_cache_h4(comp_kv, + args.comp_row_stride, + (uint)idx, + tid, + args.comp_kv_f16 != 0u); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + dsv4_attend_shared_h4_row(kv_shared, + q0, q1, q2, q3, + args.scale, + lane, + M, S, + o0, o1, o2, o3); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + dsv4_attend_sink(((device const float *)sinks)[head], M, S, o0, o1, o2, o3); + + const float inv_s = S == 0.0f ? 0.0f : 1.0f/S; + device float4 *dst4 = (device float4 *)(dst + + (uint64_t)token * args.dst_token_stride + + (uint64_t)head * args.dst_head_stride); + dst4[lane + 0] = o0 * inv_s; + dst4[lane + 32] = o1 * inv_s; + dst4[lane + 64] = o2 * inv_s; + dst4[lane + 96] = o3 * inv_s; +} + +// Decode specialization of kernel_dsv4_indexed_mixed_attention_heads8. +// Generation attends one token at a time, so the ratio-4 indexed path spends a +// visible amount of time repeatedly staging the same K/V row for the eight +// heads in a group. This variant stages sixteen selected rows at once and then +// consumes them sequentially, preserving the row order and online softmax math +// while cutting threadgroup barriers in the long top-k scan. +kernel void kernel_dsv4_indexed_mixed_attention_heads8_rb16( + constant ds4_metal_args_dsv4_indexed_attention & args, + device const char *q, + device const char *raw_kv, + device const char *comp_kv, + device const char *topk, + device const char *sinks, + device char *dst, + threadgroup half4 *kv_shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + const uint token = tgpig.x; + const uint head = tgpig.y * 8u + (uint)sg; + if (token >= args.n_tokens || head >= args.n_head) { + return; + } + + device const float4 *q4 = (device const float4 *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + const half4 q0 = (half4)q4[lane + 0]; + const half4 q1 = (half4)q4[lane + 32]; + const half4 q2 = (half4)q4[lane + 64]; + const half4 q3 = (half4)q4[lane + 96]; + + float M = -FLT_MAX/2.0f; + float S = 0.0f; + float4 o0 = 0.0f; + float4 o1 = 0.0f; + float4 o2 = 0.0f; + float4 o3 = 0.0f; + + const uint qpos = args.pos0 + token; + const uint last_pos = args.pos0 + args.n_tokens - 1u; + const uint first_raw_pos = last_pos + 1u - args.n_raw; + const uint raw_last_pos = first_raw_pos + args.n_raw - 1u; + const uint window_first = (args.window != 0u && qpos + 1u > args.window) ? + qpos + 1u - args.window : 0u; + uint first = max(first_raw_pos, window_first); + uint last = min(qpos, raw_last_pos); + + if (first <= last) { + for (uint pos0 = first; pos0 <= last; pos0 += 16u) { + const uint n_rows = min(16u, last - pos0 + 1u); + for (uint off = (uint)tid; off < n_rows * 128u; off += 256u) { + const uint r = off >> 7; + const uint c = off & 127u; + const uint logical = pos0 + r - first_raw_pos; + const uint row = (args.raw_start + logical) % args.raw_cap; + device const float4 *src = (device const float4 *)(raw_kv + + (uint64_t)row * args.raw_row_stride); + kv_shared[off] = (half4)src[c]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint r = 0; r < n_rows; r++) { + dsv4_attend_shared_h4_row_at(kv_shared, + r, + q0, q1, q2, q3, + args.scale, + lane, + M, S, + o0, o1, o2, o3); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + uint visible = (qpos + 1u) / args.ratio; + visible = min(visible, args.n_comp); + device const int32_t *row_topk = (device const int32_t *)(topk + + (uint64_t)token * args.topk_token_stride); + bool stop = false; + for (uint i = 0; i < args.top_k && !stop; i += 16u) { + uint rows[16]; + uint n_rows = 0; + for (uint j = 0; j < 16u && i + j < args.top_k; j++) { + const int32_t idx = row_topk[i + j]; + if (idx < 0) { + continue; + } + if ((uint)idx >= visible) { + stop = true; + break; + } + rows[n_rows++] = (uint)idx; + } + if (n_rows == 0) { + continue; + } + for (uint off = (uint)tid; off < n_rows * 128u; off += 256u) { + const uint r = off >> 7; + const uint c = off & 127u; + kv_shared[off] = dsv4_load_cache_h4(comp_kv, + args.comp_row_stride, + rows[r], + c, + args.comp_kv_f16 != 0u); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint r = 0; r < n_rows; r++) { + dsv4_attend_shared_h4_row_at(kv_shared, + r, + q0, q1, q2, q3, + args.scale, + lane, + M, S, + o0, o1, o2, o3); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + dsv4_attend_sink(((device const float *)sinks)[head], M, S, o0, o1, o2, o3); + + const float inv_s = S == 0.0f ? 0.0f : 1.0f/S; + device float4 *dst4 = (device float4 *)(dst + + (uint64_t)token * args.dst_token_stride + + (uint64_t)head * args.dst_head_stride); + dst4[lane + 0] = o0 * inv_s; + dst4[lane + 32] = o1 * inv_s; + dst4[lane + 64] = o2 * inv_s; + dst4[lane + 96] = o3 * inv_s; +} + +static inline float dsv4_indexer_dot128_shared_q( + float4 c0, + float4 c1, + float4 c2, + float4 c3, + threadgroup const float4 *q4, + ushort lane) { + float sum = 0.0f; + if (lane < 8) { + const ushort ib = lane >> 1; + const ushort il = lane & 1; + const ushort base = ib*8 + il*4; + sum += dot(c0, q4[base + 0]); + sum += dot(c1, q4[base + 1]); + sum += dot(c2, q4[base + 2]); + sum += dot(c3, q4[base + 3]); + } + return simd_sum(sum); +} + +// Tiled prefill score builder for the sparse-compressed attention indexer. +// +// The kernel covers an 8-token by 32-compressed-row rectangle: K is copied into +// threadgroup memory once, then reused for all 64 indexer heads, while simdgroup +// matrix multiply computes each 8x8 score subtile. +// +// It still writes the exact score matrix consumed by top-k: +// +// score[t,c] = sum_h relu(dot(Q[t,h], K[c])) * W[t,h] * scale +// +// Causal masking is applied on store so invisible compressed rows become -inf. +kernel void kernel_dsv4_indexer_scores_tiled_f32( + constant ds4_metal_args_dsv4_indexer_scores_fused & args, + device const char *q, + device const char *weights, + device const char *index_comp, + device char *scores, + threadgroup float *shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint TM = 8; + constexpr uint TN = 32; + constexpr uint TS = 8; + constexpr uint D = 128; + + const uint c0 = tgpig.x * TN; + const uint t0 = tgpig.y * TM; + + threadgroup float *qtg = shared; // [8][128] + threadgroup float *ktg = qtg + TM*D; // [32][128] + threadgroup float *dot = ktg + TN*D; // [8][32] + + const uint last_token = min(t0 + TM, args.n_tokens); + const uint max_visible = last_token > t0 ? + min((args.pos0 + last_token) / args.ratio, args.n_comp) : 0u; + + if (c0 >= max_visible) { + for (uint i = tid; i < TM*TN; i += 128) { + const uint r = i / TN; + const uint cc = i - r*TN; + const uint token = t0 + r; + const uint comp = c0 + cc; + if (token < args.n_tokens && comp < args.n_comp) { + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + comp; + *dst = -INFINITY; + } + } + return; + } + + for (uint i = tid; i < TN*D; i += 128) { + const uint cc = i / D; + const uint d = i - cc*D; + const uint comp = c0 + cc; + float v = 0.0f; + if (comp < args.n_comp) { + device const float *row = (device const float *)(index_comp + + (uint64_t)comp * args.index_row_stride); + v = row[d]; + } + ktg[i] = v; + } + + const uint cell0 = lane; + const uint cell1 = lane + 32u; + const uint row0 = cell0 >> 3; + const uint row1 = cell1 >> 3; + const uint sub0 = cell0 & 7u; + const uint sub1 = cell1 & 7u; + const uint col0 = (uint)sg * TS + sub0; + const uint col1 = (uint)sg * TS + sub1; + const uint token0 = t0 + row0; + const uint token1 = t0 + row1; + const uint comp0 = c0 + col0; + const uint comp1 = c0 + col1; + + float acc0 = 0.0f; + float acc1 = 0.0f; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint head = 0; head < args.n_head; head++) { + for (uint i = tid; i < TM*D; i += 128) { + const uint r = i / D; + const uint d = i - r*D; + const uint token = t0 + r; + float v = 0.0f; + if (token < args.n_tokens) { + device const float *qrow = (device const float *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + v = qrow[d]; + } + qtg[i] = v; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + simdgroup_float8x8 mdot = make_filled_simdgroup_matrix(0.0f); + for (uint db = 0; db < D/TS; db++) { + simdgroup_float8x8 mq; + simdgroup_float8x8 mk; + simdgroup_load(mq, qtg + db*TS, D, 0, false); + simdgroup_load(mk, ktg + ((uint)sg * TS) * D + db*TS, D, 0, true); + simdgroup_multiply_accumulate(mdot, mq, mk, mdot); + } + + simdgroup_store(mdot, dot + (uint)sg * TS, TN, 0, false); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (token0 < args.n_tokens && comp0 < args.n_comp) { + device const float *w = (device const float *)(weights + + (uint64_t)token0 * args.weights_token_stride); + const float s = dot[row0*TN + col0]; + acc0 += max(s, 0.0f) * (w[head] * args.scale); + } + if (token1 < args.n_tokens && comp1 < args.n_comp) { + device const float *w = (device const float *)(weights + + (uint64_t)token1 * args.weights_token_stride); + const float s = dot[row1*TN + col1]; + acc1 += max(s, 0.0f) * (w[head] * args.scale); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (token0 < args.n_tokens && comp0 < args.n_comp) { + const uint visible = min((args.pos0 + token0 + 1u) / args.ratio, args.n_comp); + device float *dst = (device float *)(scores + + (uint64_t)token0 * args.score_token_stride) + comp0; + *dst = comp0 < visible ? acc0 : -INFINITY; + } + if (token1 < args.n_tokens && comp1 < args.n_comp) { + const uint visible = min((args.pos0 + token1 + 1u) / args.ratio, args.n_comp); + device float *dst = (device float *)(scores + + (uint64_t)token1 * args.score_token_stride) + comp1; + *dst = comp1 < visible ? acc1 : -INFINITY; + } +} + +kernel void kernel_dsv4_indexer_scores_tiled( + constant ds4_metal_args_dsv4_indexer_scores_fused & args, + device const char *q, + device const char *weights, + device const char *index_comp, + device char *scores, + threadgroup float *shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort lane [[thread_index_in_simdgroup]], + ushort sg [[simdgroup_index_in_threadgroup]]) { + constexpr uint TM = 8; + constexpr uint TN = 32; + constexpr uint TS = 8; + constexpr uint D = 128; + + const uint c0 = tgpig.x * TN; + const uint t0 = tgpig.y * TM; + + // Q/K are staged as half but the dot accumulator and final score remain + // float. This is the one intentional precision tradeoff in the indexer: + // the indexer only ranks compressed rows for top-k selection, and long + // context profiling shows this score matrix dominates the prefill slope. + threadgroup half *qtg = (threadgroup half *)shared; // [8][128] + threadgroup half *ktg = qtg + TM*D; // [32][128] + threadgroup float *dot = (threadgroup float *)(ktg + TN*D); // [8][32] + + const uint last_token = min(t0 + TM, args.n_tokens); + const uint max_visible = last_token > t0 ? + min((args.pos0 + last_token) / args.ratio, args.n_comp) : 0u; + + if (c0 >= max_visible) { + for (uint i = tid; i < TM*TN; i += 128) { + const uint r = i / TN; + const uint cc = i - r*TN; + const uint token = t0 + r; + const uint comp = c0 + cc; + if (token < args.n_tokens && comp < args.n_comp) { + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + comp; + *dst = -INFINITY; + } + } + return; + } + + // Stage compressed index rows once. Edge columns are zeroed so the matrix + // loads below can stay regular; guarded stores discard them. + for (uint i = tid; i < TN*D; i += 128) { + const uint cc = i / D; + const uint d = i - cc*D; + const uint comp = c0 + cc; + half v = half(0.0f); + if (comp < args.n_comp) { + device const float *row = (device const float *)(index_comp + + (uint64_t)comp * args.index_row_stride); + v = half(row[d]); + } + ktg[i] = v; + } + + const uint cell0 = lane; + const uint cell1 = lane + 32u; + const uint row0 = cell0 >> 3; + const uint row1 = cell1 >> 3; + const uint sub0 = cell0 & 7u; + const uint sub1 = cell1 & 7u; + const uint col0 = (uint)sg * TS + sub0; + const uint col1 = (uint)sg * TS + sub1; + const uint token0 = t0 + row0; + const uint token1 = t0 + row1; + const uint comp0 = c0 + col0; + const uint comp1 = c0 + col1; + + float acc0 = 0.0f; + float acc1 = 0.0f; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint head = 0; head < args.n_head; head++) { + // Stage Q for the eight-token tile. Each 8x8 matrix load below reads a + // contiguous depth block from this layout. + for (uint i = tid; i < TM*D; i += 128) { + const uint r = i / D; + const uint d = i - r*D; + const uint token = t0 + r; + half v = half(0.0f); + if (token < args.n_tokens) { + device const float *qrow = (device const float *)(q + + (uint64_t)token * args.q_token_stride + + (uint64_t)head * args.q_head_stride); + v = half(qrow[d]); + } + qtg[i] = v; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + simdgroup_float8x8 mdot = make_filled_simdgroup_matrix(0.0f); + for (uint db = 0; db < D/TS; db++) { + simdgroup_half8x8 mq; + simdgroup_half8x8 mk; + simdgroup_load(mq, qtg + db*TS, D, 0, false); + simdgroup_load(mk, ktg + ((uint)sg * TS) * D + db*TS, D, 0, true); + simdgroup_multiply_accumulate(mdot, mq, mk, mdot); + } + + simdgroup_store(mdot, dot + (uint)sg * TS, TN, 0, false); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (token0 < args.n_tokens && comp0 < args.n_comp) { + device const float *w = (device const float *)(weights + + (uint64_t)token0 * args.weights_token_stride); + const float s = dot[row0*TN + col0]; + acc0 += max(s, 0.0f) * (w[head] * args.scale); + } + if (token1 < args.n_tokens && comp1 < args.n_comp) { + device const float *w = (device const float *)(weights + + (uint64_t)token1 * args.weights_token_stride); + const float s = dot[row1*TN + col1]; + acc1 += max(s, 0.0f) * (w[head] * args.scale); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (token0 < args.n_tokens && comp0 < args.n_comp) { + const uint visible = min((args.pos0 + token0 + 1u) / args.ratio, args.n_comp); + device float *dst = (device float *)(scores + + (uint64_t)token0 * args.score_token_stride) + comp0; + *dst = comp0 < visible ? acc0 : -INFINITY; + } + if (token1 < args.n_tokens && comp1 < args.n_comp) { + const uint visible = min((args.pos0 + token1 + 1u) / args.ratio, args.n_comp); + device float *dst = (device float *)(scores + + (uint64_t)token1 * args.score_token_stride) + comp1; + *dst = comp1 < visible ? acc1 : -INFINITY; + } +} + +#ifdef DS4_METAL_HAS_TENSOR +// Retained full-512 prefill indexer score path. This is the part of sparse +// compressed attention that maps cleanly to TensorOps: a regular token by +// compressed-row dot tile. The kernel intentionally leaves top-k selection and +// indexed attention semantics unchanged; all 512 selected rows remain available +// to the later attention kernel. +// +// Each matmul processes a pair of heads (TQ = 2 x TM q rows): the per-element +// dot is still a 128-deep reduction in 32-wide k-steps, so scores are +// bit-identical to single-head tiles while the run count halves. The q tile +// is double-buffered, so the next k-step's stage overlaps the current +// cooperative matmul and each pair needs 5 barriers instead of 10. q and k +// staging use one float4/half4 per lane (each thread covers one row of 8/32 +// consecutive elements), which is the same half(float) conversion per element +// as the scalar form. +kernel void kernel_dsv4_indexer_scores_nax( + constant ds4_metal_args_dsv4_indexer_scores_fused & args, + device const char *q, + device const char *weights, + device const char *index_comp, + device char *scores, + threadgroup half *shared [[threadgroup(0)]], + uint2 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]]) { + constexpr int TM = 16; + constexpr int TQ = 32; + constexpr int TN = 32; + constexpr int NK = 32; + constexpr int D = 128; + constexpr int NUM_THREADS = 128; + + // The 16-token x 32-row tile was the winning NAX shape in local sweeps. A + // wider 64-row compressed tile increased setup/cache pressure and was + // slower despite doing more work per dispatch. + const uint c0 = tgpig.x * TN; + const uint t0 = tgpig.y * TM; + + threadgroup half *qtg = shared; // 2 x [TQ][NK] + threadgroup half *ktg = qtg + 2*TQ*NK; // [32][128] + threadgroup float *dot = (threadgroup float *)(ktg + TN*D); // [TQ][TN], column-major + + const uint last_token = min(t0 + (uint)TM, args.n_tokens); + const uint max_visible = last_token > t0 ? + min((args.pos0 + last_token) / args.ratio, args.n_comp) : 0u; + + if (c0 >= max_visible) { + for (uint i = tid; i < TM*TN; i += NUM_THREADS) { + const uint r = i / TN; + const uint cc = i - r*TN; + const uint token = t0 + r; + const uint comp = c0 + cc; + if (token < args.n_tokens && comp < args.n_comp) { + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + comp; + *dst = -INFINITY; + } + } + return; + } + + { + // One compressed row per 4 threads, 32 consecutive floats per thread. + const uint cc = tid / 4; + const uint comp = c0 + cc; + device const float *krow = nullptr; + if (comp < args.n_comp) { + krow = (device const float *)(index_comp + + (uint64_t)comp * args.index_row_stride); + } + const uint d0 = (tid % 4) * 32; + FOR_UNROLL (uint j = 0; j < 8; j++) { + const float4 kv = krow ? *(device const float4 *)(krow + d0 + 4*j) + : float4(0.0f); + *(threadgroup half4 *)(ktg + cc*D + d0 + 4*j) = half4(kv); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float acc[4]; + #pragma unroll + for (uint j = 0; j < 4; j++) { + acc[j] = 0.0f; + } + + auto tq0 = tensor(qtg, dextents(NK, TQ)); + auto tq1 = tensor(qtg + TQ*NK, dextents(NK, TQ)); + auto tk = tensor(ktg, dextents(D, TN)); + auto td = tensor(dot, dextents(TQ, TN), array({1, TQ})); + + matmul2d< + matmul2d_descriptor(TN, TQ, NK, false, true, false, + matmul2d_descriptor::mode::multiply_accumulate), + execution_simdgroups<4>> mm; + + // One q row per 4 threads, 8 consecutive floats per thread. Row r covers + // head (r / TM) of the pair and token row (r % TM). + const uint q_r = tid / 4; + const uint q_k4 = (tid % 4) * 8; + const uint q_hl = q_r / TM; + const uint q_tr = q_r % TM; + const uint q_token = t0 + q_tr; + device const char *q_row_base = nullptr; + if (q_token < args.n_tokens) { + q_row_base = q + (uint64_t)q_token * args.q_token_stride; + } + + auto stage_q = [&](const uint head0, const uint loop_k, threadgroup half *buf) { + const uint head = head0 + q_hl; + half4 v0 = half4(0.0f); + half4 v1 = half4(0.0f); + if (q_row_base && head < args.n_head) { + device const float4 *src4 = (device const float4 *) + (q_row_base + (uint64_t)head * args.q_head_stride + + (uint64_t)(loop_k + q_k4) * sizeof(float)); + v0 = half4(src4[0]); + v1 = half4(src4[1]); + } + *(threadgroup half4 *)(buf + q_r*NK + q_k4) = v0; + *(threadgroup half4 *)(buf + q_r*NK + q_k4 + 4) = v1; + }; + + for (uint head0 = 0; head0 < args.n_head; head0 += 2) { + auto ct = mm.template get_destination_cooperative_tensor(); + #pragma unroll + for (uint16_t i = 0; i < ct.get_capacity(); i++) { + if (ct.is_valid_element(i)) { + ct[i] = 0.0f; + } + } + + stage_q(head0, 0, qtg); + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint qsel = 0; + FOR_UNROLL (uint i = 0; i < 4; i++) { + auto mk = tk.slice(i*NK, 0); + auto mq = (qsel ? tq1 : tq0).slice(0, 0); + mm.run(mk, mq, ct); + if (i < 3) { + qsel ^= 1u; + stage_q(head0, (i + 1)*NK, qsel ? qtg + TQ*NK : qtg); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + ct.store(td); + threadgroup_barrier(mem_flags::mem_threadgroup); + + #pragma unroll + for (uint j = 0; j < 4; j++) { + const uint linear = (uint)tid + j*NUM_THREADS; + if (linear < TM*TN) { + const uint r = linear / TN; + const uint cc = linear - r*TN; + const uint token = t0 + r; + if (token < args.n_tokens) { + device const float *w = (device const float *)(weights + + (uint64_t)token * args.weights_token_stride); + acc[j] += max(dot[cc*TQ + r], 0.0f) * (w[head0] * args.scale); + if (head0 + 1 < args.n_head) { + acc[j] += max(dot[cc*TQ + TM + r], 0.0f) * (w[head0 + 1] * args.scale); + } + } + } + } + // No barrier here: the next pair's q stage and these dot reads touch + // different buffers, and the next q-stage barrier separates the next + // ct.store from these reads. + } + + #pragma unroll + for (uint j = 0; j < 4; j++) { + const uint linear = (uint)tid + j*NUM_THREADS; + if (linear >= TM*TN) { + continue; + } + const uint r = linear / TN; + const uint cc = linear - r*TN; + const uint token = t0 + r; + const uint comp = c0 + cc; + if (token < args.n_tokens && comp < args.n_comp) { + const uint visible = min((args.pos0 + token + 1u) / args.ratio, args.n_comp); + device float *dst = (device float *)(scores + + (uint64_t)token * args.score_token_stride) + comp; + *dst = comp < visible ? acc[j] : -INFINITY; + } + } +} +#endif + +// Collapses per-head indexer scores into one score per compressed row using the +// learned head weights. Negative head scores are clipped exactly as DS4 expects. +kernel void kernel_dsv4_indexer_weighted_sum( + constant ds4_metal_args_dsv4_indexer_weighted_sum & args, + device const char * scores, + device const char * weights, + device char * dst, + uint gid [[thread_position_in_grid]]) { + const int64_t n = args.ne0 * args.ne1; + if ((int64_t) gid >= n) { + return; + } + + const int64_t ic = gid % args.ne0; + const int64_t it = gid / args.ne0; + + float acc = 0.0f; + for (int64_t ih = 0; ih < args.ne02; ++ih) { + const float s = *((device const float *) (scores + ic*args.nb00 + it*args.nb01 + ih*args.nb02)); + const float w = *((device const float *) (weights + ih*args.nb10 + it*args.nb11)); + acc += max(s, 0.0f) * (w * args.scale); + } + + *((device float *) (dst + ic*args.nb0 + it*args.nb1)) = acc; +} + +// Adds the periodic compressor APE directly to projected scores. The legacy +// path materializes one repeated APE segment per period and then performs this +// same single F32 add; these kernels remove only that intermediate copy graph. +kernel void kernel_dsv4_compressor_score_ape_f32( + constant ds4_metal_args_dsv4_compressor_score_ape & args, + device const float *score, + device const float *ape, + device float *dst, + uint gid [[thread_position_in_grid]]) { + const uint64_t total = (uint64_t)args.n_tokens * args.width; + if ((uint64_t)gid >= total) return; + + const uint token = gid / args.width; + const uint col = gid - token*args.width; + const uint ape_row = (uint)(((uint64_t)args.pos0 + token) % args.ratio); + dst[gid] = score[gid] + ape[(uint64_t)ape_row*args.width + col]; +} + +kernel void kernel_dsv4_compressor_score_ape_f16( + constant ds4_metal_args_dsv4_compressor_score_ape & args, + device const float *score, + device const half *ape, + device float *dst, + uint gid [[thread_position_in_grid]]) { + const uint64_t total = (uint64_t)args.n_tokens * args.width; + if ((uint64_t)gid >= total) return; + + const uint token = gid / args.width; + const uint col = gid - token*args.width; + const uint ape_row = (uint)(((uint64_t)args.pos0 + token) % args.ratio); + dst[gid] = score[gid] + float(ape[(uint64_t)ape_row*args.width + col]); +} + +// Fused softmax-weighted pooling of compressed KV rows. It is used when several +// compressor rows are present; the one-row case deliberately follows the +// unfused softmax/mul/sum graph in Objective-C to keep identical reductions. +kernel void kernel_dsv4_softmax_pool( + constant ds4_metal_args_dsv4_softmax_pool & args, + device const char * kv, + device const char * score, + device char * dst, + uint gid [[thread_position_in_grid]]) { + const int64_t n = args.ne0 * args.ne1; + if ((int64_t) gid >= n) { + return; + } + + const int64_t id = gid % args.ne0; + const int64_t ic = gid / args.ne0; + + float max_s = -INFINITY; + for (int64_t ir = 0; ir < args.ne00; ++ir) { + const float s = *((device const float *) (score + ir*args.nb10 + id*args.nb11 + ic*args.nb12)); + max_s = max(max_s, s); + } + + float sum = 0.0f; + float acc = 0.0f; + for (int64_t ir = 0; ir < args.ne00; ++ir) { + const float s = *((device const float *) (score + ir*args.nb10 + id*args.nb11 + ic*args.nb12)); + const float w = exp(s - max_s); + const float v = *((device const float *) (kv + ir*args.nb00 + id*args.nb01 + ic*args.nb02)); + sum += w; + acc += v*w; + } + + *((device float *) (dst + id*args.nb0 + ic*args.nb1)) = acc/sum; +} + + + +// Tensor-parallel keep-alive: a few threadgroups of FMAs dispatched +// back-to-back on a side queue while TP decode runs. The per-layer gate +// stalls make the real workload look idle to the GPU power manager, which +// otherwise halves the clocks within a second (~2x decode regression); +// this holds them up for negligible bandwidth and a few watts. +kernel void kernel_dsv4_tp_keepalive( + device float * out, + constant uint & iters, + uint tid [[thread_position_in_grid]]) { + float a = out[tid]; + const float b = 1.000001f; + for (uint i = 0; i < iters; i++) { + a = fma(a, b, 0.000001f); + a = fma(a, b, -0.000001f); + } + out[tid] = a; +} + +// Tensor-parallel gate flag: publishes a sequence number to a slab slot the +// CPU service thread spin-reads, replacing the much slower shared-event +// signal for the GPU->CPU direction. Ordering against the partial-output +// kernels comes from the buffer hazard on the shared slab. +kernel void kernel_dsv4_tp_flag_set( + device atomic_uint & flag, + constant uint & value, + uint tid [[thread_position_in_grid]]) { + if (tid == 0) { + atomic_store_explicit(&flag, value, memory_order_relaxed); + } +} + +// Ratio-4 compressor pooling without materializing the [n_comp, 8, head_dim] +// KV and score packs. The row mapping and both reduction loops deliberately +// match kernel_dsv4_softmax_pool so the arithmetic order is unchanged. +kernel void kernel_dsv4_softmax_pool_ratio4_direct( + constant ds4_metal_args_dsv4_softmax_pool_ratio4_direct & args, + device const float * kv, + device const float * score, + device const float * state_kv, + device const float * state_score, + device float * dst, + uint gid [[thread_position_in_grid]]) { + const uint64_t n = (uint64_t)args.head_dim * args.n_comp; + if ((uint64_t)gid >= n || args.head_dim == 0u) { + return; + } + + const uint64_t id = gid % args.head_dim; + const uint64_t ic = gid / args.head_dim; + const uint64_t input_row_stride = 2ull * args.head_dim; + + float max_s = -INFINITY; + float sum = 0.0f; + float acc = 0.0f; + if (ic != 0u) { + const int64_t token_base = (int64_t)ic * 4 - 4; + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + const uint64_t token = (uint64_t)(token_base + ir); + const uint64_t src = token * input_row_stride + + ((uint64_t)ir >> 2u) * args.head_dim + id; + const float s = score[src]; + max_s = max(max_s, s); + } + + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + const uint64_t token = (uint64_t)(token_base + ir); + const uint64_t src = token * input_row_stride + + ((uint64_t)ir >> 2u) * args.head_dim + id; + const float s = score[src]; + const float w = exp(s - max_s); + const float v = kv[src]; + sum += w; + acc += v*w; + } + } else { + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + float s; + if (ir >= 4) { + const uint64_t src = (uint64_t)(ir - 4) * input_row_stride + + args.head_dim + id; + s = score[src]; + } else if (args.replay != 0u) { + s = state_score[(uint64_t)ir * input_row_stride + id]; + } else { + s = -INFINITY; + } + max_s = max(max_s, s); + } + + for (int64_t ir = 0; ir < args.n_rows; ++ir) { + float s; + float v; + if (ir >= 4) { + const uint64_t src = (uint64_t)(ir - 4) * input_row_stride + + args.head_dim + id; + s = score[src]; + v = kv[src]; + } else if (args.replay != 0u) { + const uint64_t src = (uint64_t)ir * input_row_stride + id; + s = state_score[src]; + v = state_kv[src]; + } else { + s = -INFINITY; + v = 0.0f; + } + const float w = exp(s - max_s); + sum += w; + acc += v*w; + } + } + + dst[ic * args.head_dim + id] = acc/sum; +} diff --git a/metal/dsv4_rope.metal b/metal/dsv4_rope.metal new file mode 100644 index 0000000..94a83c7 --- /dev/null +++ b/metal/dsv4_rope.metal @@ -0,0 +1,385 @@ +struct ds4_metal_args_dsv4_rope_tail { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + int32_t n_dims; + int32_t mode; + int32_t n_ctx_orig; + int32_t inverse; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + bool src2; +}; + +struct ds4_metal_args_dsv4_rope_affine_pair { + uint64_t row_bytes; + uint64_t token_bytes; + int32_t head_dim; + int32_t n_dims; + int32_t n_ctx_orig; + int32_t inverse; + uint32_t pos0; + uint32_t pos_step; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +}; + +static float rope_yarn_ramp(const float low, const float high, const int i0) { + const float y = (i0 / 2 - low) / max(0.001f, high - low); + return 1.0f - min(1.0f, max(0.0f, y)); +} + +// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn +// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng. +static void rope_yarn( + float theta_extrap, float freq_scale, float corr_dims[2], int i0, float ext_factor, float mscale, + thread float * cos_theta, thread float * sin_theta) { + // Get n-d rotational scaling corrected for extrapolation + float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + if (ext_factor != 0.0f) { + float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * ext_factor; + theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix; + + // Get n-d magnitude scaling corrected for interpolation + mscale *= 1.0f + 0.1f * log(1.0f / freq_scale); + } + *cos_theta = cos(theta) * mscale; + *sin_theta = sin(theta) * mscale; +} + +// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get +// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))` +static float rope_yarn_corr_factor(int n_dims, int n_ctx_orig, float n_rot, float base) { + return n_dims * log(n_ctx_orig / (n_rot * 2 * M_PI_F)) / (2 * log(base)); +} + +static void rope_yarn_corr_dims( + int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2] +) { + // start and end correction dims + dims[0] = max(0.0f, floor(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_fast, freq_base))); + dims[1] = min(n_dims - 1.0f, ceil(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_slow, freq_base))); +} + +// Applies DeepSeek V4's partial RoPE: the no-position prefix is copied and only +// the rotated tail is transformed. This is used for Q/K after their projections +// and before writing/reading the attention KV state. +kernel void kernel_dsv4_rope_tail_f32( + constant ds4_metal_args_dsv4_rope_tail & args, + device const char * src0, + device const char * src1, + device const char * src2, + device char * dst, + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const int i1 = tgpig[0]; + const int i2 = tgpig[1]; + const int i3 = tgpig[2]; + + const int n_nope = args.ne00 - args.n_dims; + if (n_nope < 0) { + return; + } + + device const int32_t * pos = (device const int32_t *) src1; + + float corr_dims[2]; + rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims); + + const float theta_base = (float) pos[i2]; + const float inv_ndims = -1.f/args.n_dims; + const bool is_neox = args.mode == 2; + + for (int i0 = tid; i0 < args.ne00; i0 += ntg.x) { + device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01; + device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1; + + if (i0 < n_nope) { + *((device float *) (dst_base + i0*args.nb0)) = *((device const float *) (src_base + i0*args.nb00)); + continue; + } + + const int r = i0 - n_nope; + if (is_neox) { + const int n_half = args.n_dims/2; + if (r >= n_half) { + continue; + } + + const int ic = r; + const int rel_i0 = 2*ic; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + // Equivalent to pow(freq_base, k) but expressed through IEEE-754 + // primitives that have tighter precision guarantees than Metal's pow(). + const float theta = theta_base * exp2(inv_ndims * (float)rel_i0 * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims*rel_i0); +#endif + const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f; + + float cos_theta; + float sin_theta; + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, rel_i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + if (args.inverse) { + sin_theta = -sin_theta; + } + + const int j0 = n_nope + ic; + const int j1 = n_nope + ic + n_half; + const float x0 = *((device const float *) (src_base + j0*args.nb00)); + const float x1 = *((device const float *) (src_base + j1*args.nb00)); + + *((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta; + *((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta; + } else { + if ((r & 1) != 0) { + continue; + } + + const int ic = r/2; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims*r); +#endif + const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f; + + float cos_theta; + float sin_theta; + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + if (args.inverse) { + sin_theta = -sin_theta; + } + + const int j0 = n_nope + r; + const int j1 = j0 + 1; + const float x0 = *((device const float *) (src_base + j0*args.nb00)); + const float x1 = *((device const float *) (src_base + j1*args.nb00)); + + *((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta; + *((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta; + } + } +} + +// DS4 only calls the Metal RoPE helper in-place and uses the adjacent-pair +// layout (mode 0). The generic kernel above still copies the no-position +// prefix, even though source and destination alias, and dispatches enough +// lanes for the full head. This specialization maps lanes directly to the +// rotated pairs and deliberately never reads or writes the unchanged prefix. +// Keep the arithmetic below in the same order as the mode-0 branch above so +// the optimized and reference paths remain bit-identical. +kernel void kernel_dsv4_rope_tail_f32_inplace_pair( + constant ds4_metal_args_dsv4_rope_tail & args, + device const char * src0, + device const char * src1, + device const char * src2, + device char * dst, + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + if (args.mode != 0) { + return; + } + + const int i1 = tgpig[0]; + const int i2 = tgpig[1]; + const int i3 = tgpig[2]; + const int n_nope = args.ne00 - args.n_dims; + if (n_nope < 0) { + return; + } + + device const int32_t * pos = (device const int32_t *) src1; + + float corr_dims[2]; + rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims); + + const float theta_base = (float) pos[i2]; + const float inv_ndims = -1.f/args.n_dims; + device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01; + device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1; + + // Keep each pair on the same SIMD lane as the generic in-place dispatch. + // n_nope is 32-aligned for the supported DS4 shapes, so logical tail index + // r ran on lane r%32 in the reference kernel. Compacting pairs would move + // them to lane (r/2)%32 and can perturb fast-math code generation. + for (int r = tid; r < args.n_dims; r += ntg.x) { + if ((r & 1) != 0) { + continue; + } + const int ic = r/2; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims*r); +#endif + const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f; + + float cos_theta; + float sin_theta; + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + if (args.inverse) { + sin_theta = -sin_theta; + } + + const int j0 = n_nope + r; + const int j1 = j0 + 1; + const float x0 = *((device const float *) (src_base + j0*args.nb00)); + const float x1 = *((device const float *) (src_base + j1*args.nb00)); + + *((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta; + *((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta; + } +} + +// The Q/K RoPE calls use the same position and scaling parameters for every +// head. Group four heads into one threadgroup so one 64-thread cohort computes +// the 32 adjacent-pair coefficients and the other cohorts reuse them. Keeping +// r on the same r%32 SIMD lane as the per-head specialization preserves the +// fast-math instruction mapping; only the redundant coefficient work changes. +kernel void kernel_dsv4_rope_tail_f32_inplace_pair_shared4( + constant ds4_metal_args_dsv4_rope_tail & args, + device const char * src0, + device const char * src1, + device const char * src2, + device char * dst, + uint tid [[thread_index_in_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + if (args.mode != 0 || args.n_dims != 64) { + return; + } + + const int n_nope = args.ne00 - args.n_dims; + if (n_nope < 0) { + return; + } + + const uint cohort = tid >> 6; + const int r = (int)(tid & 63u); + threadgroup float cos_shared[32]; + threadgroup float sin_shared[32]; + + device const int32_t * pos = (device const int32_t *) src1; + float corr_dims[2]; + rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims); + + const int i2 = tgpig[1]; + const float theta_base = (float) pos[i2]; + const float inv_ndims = -1.f/args.n_dims; + + if (cohort == 0 && (r & 1) == 0) { + const int ic = r/2; +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims*r); +#endif + const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f; + + float cos_theta; + float sin_theta; + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + if (args.inverse) { + sin_theta = -sin_theta; + } + cos_shared[ic] = cos_theta; + sin_shared[ic] = sin_theta; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + const int i1 = (int)tgpig[0]*4 + (int)cohort; + if (i1 >= args.ne01 || (r & 1) != 0) { + return; + } + + const int i3 = tgpig[2]; + device const char * src_base = src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01; + device char * dst_base = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1; + const int j0 = n_nope + r; + const int j1 = j0 + 1; + const float x0 = *((device const float *) (src_base + j0*args.nb00)); + const float x1 = *((device const float *) (src_base + j1*args.nb00)); + const float cos_theta = cos_shared[r/2]; + const float sin_theta = sin_shared[r/2]; + + *((device float *) (dst_base + j0*args.nb0)) = x0*cos_theta - x1*sin_theta; + *((device float *) (dst_base + j1*args.nb0)) = x0*sin_theta + x1*cos_theta; +} + +// DS4 positions are always affine within one RoPE dispatch. This variant +// reconstructs the same wrapped int32 position in-kernel, avoiding the host +// position array and its buffer binding while preserving the pair lane mapping +// and all floating-point operations of the specialization above. +kernel void kernel_dsv4_rope_tail_f32_inplace_pair_affine( + constant ds4_metal_args_dsv4_rope_affine_pair & args [[buffer(0)]], + device const char * src0 [[buffer(1)]], + device char * dst [[buffer(4)]], + uint tid [[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]], + uint3 tgpig [[threadgroup_position_in_grid]]) { + const int i1 = tgpig[0]; + const int i2 = tgpig[1]; + const int n_nope = args.head_dim - args.n_dims; + if (n_nope < 0) { + return; + } + + float corr_dims[2]; + rope_yarn_corr_dims(args.n_dims, args.n_ctx_orig, args.freq_base, args.beta_fast, args.beta_slow, corr_dims); + + const uint raw_pos = args.pos0 + (uint)i2 * args.pos_step; + const float theta_base = (float)as_type(raw_pos); + const float inv_ndims = -1.f/args.n_dims; + device const char * src_base = + src0 + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes; + device char * dst_base = + dst + (uint64_t)i2*args.token_bytes + (uint64_t)i1*args.row_bytes; + + for (int r = tid; r < args.n_dims; r += ntg.x) { + if ((r & 1) != 0) { + continue; + } +#ifdef DS4_METAL_ROPE_EXP2_LOG2 + const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base)); +#else + const float theta = theta_base * pow(args.freq_base, inv_ndims*r); +#endif + const float freq_factor = 1.0f; + + float cos_theta; + float sin_theta; + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, r, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + if (args.inverse) { + sin_theta = -sin_theta; + } + + const int j0 = n_nope + r; + const int j1 = j0 + 1; + const float x0 = *((device const float *) (src_base + j0*sizeof(float))); + const float x1 = *((device const float *) (src_base + j1*sizeof(float))); + + *((device float *) (dst_base + j0*sizeof(float))) = x0*cos_theta - x1*sin_theta; + *((device float *) (dst_base + j1*sizeof(float))) = x0*sin_theta + x1*cos_theta; + } +} diff --git a/metal/flash_attn.metal b/metal/flash_attn.metal new file mode 100644 index 0000000..900d539 --- /dev/null +++ b/metal/flash_attn.metal @@ -0,0 +1,1441 @@ +#define FC_FLASH_ATTN_EXT_PAD 100 +#define FC_FLASH_ATTN_EXT_BLK 200 +#define FC_FLASH_ATTN_EXT 300 +#define FC_FLASH_ATTN_EXT_VEC 400 +#define FC_FLASH_ATTN_EXT_VEC_REDUCE 500 +#define OP_FLASH_ATTN_EXT_NQPSG 8 +#define OP_FLASH_ATTN_EXT_NCPSG 64 +#define OP_FLASH_ATTN_EXT_VEC_NQPSG 1 +#define OP_FLASH_ATTN_EXT_VEC_NCPSG 32 + +#ifndef PAD2 +#define PAD2(x, n) (((x) + (n) - 1) & ~((n) - 1)) +#endif + +template +void dequantize_f32_t4(device const float4 * src, short il, thread type4 & reg) { + reg = (type4)(*src); +} + +template +void dequantize_f16_t4(device const half4 * src, short il, thread type4 & reg) { + reg = (type4)(*(src)); +} + +template +void dequantize_f32(device const float4x4 * src, short il, thread type4x4 & reg); + +template +void dequantize_f16(device const half4x4 * src, short il, thread type4x4 & reg); + +struct ds4_metal_args_flash_attn_ext_pad { + int32_t ne11; + int32_t ne_12_2; + int32_t ne_12_3; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + uint64_t nb21; + uint64_t nb22; + uint64_t nb23; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; +}; + +struct ds4_metal_args_flash_attn_ext_blk { + int32_t ne01; + int32_t ne30; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; +}; + +struct ds4_metal_args_flash_attn_ext { + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + int32_t ne_12_2; + int32_t ne_12_3; + int32_t ns10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ns20; + uint64_t nb21; + uint64_t nb22; + uint64_t nb23; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; + int32_t ne1; + int32_t ne2; + int32_t ne3; + float scale; + float max_bias; + float m0; + float m1; + int32_t n_head_log2; + float logit_softcap; +}; + +struct ds4_metal_args_flash_attn_ext_vec { + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + int32_t ne_12_2; + int32_t ne_12_3; + int32_t ns10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ns20; + uint64_t nb21; + uint64_t nb22; + uint64_t nb23; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; + int32_t ne1; + int32_t ne2; + int32_t ne3; + float scale; + float max_bias; + float m0; + float m1; + int32_t n_head_log2; + float logit_softcap; +}; + +struct ds4_metal_args_flash_attn_ext_vec_reduce { + int32_t nrows; +}; + +constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]]; +constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]]; + +// DS4 FlashAttention padding: pads the final partial K/V/mask cache block so the +// vector FlashAttention kernel can read full 32-row chunks. +kernel void kernel_flash_attn_ext_pad( + constant ds4_metal_args_flash_attn_ext_pad & args, + device const char * k, + device const char * v, + device const char * mask, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + const int32_t C = FC_flash_attn_ext_pad_ncpsg; + + device char * k_pad = dst; + device char * v_pad = k_pad + args.nb11*C*args.ne_12_2*args.ne_12_3; + device char * mask_pad = v_pad + args.nb21*C*args.ne_12_2*args.ne_12_3; + + const int32_t icp = args.ne11 % C; + const int32_t ic0 = args.ne11 - icp; + + const int32_t i1 = tgpig[0]; + const int32_t i2 = tgpig[1]; + const int32_t i3 = tgpig[2]; + + if (i2 < args.ne_12_2 && i3 < args.ne_12_3) { + device const char * k_src = k + args.nb11*(ic0 + i1) + args.nb12*i2 + args.nb13*i3; + device const char * v_src = v + args.nb21*(ic0 + i1) + args.nb22*i2 + args.nb23*i3; + + device char * k_dst = k_pad + args.nb11*i1 + args.nb11*C*i2 + args.nb11*C*args.ne_12_2*i3; + device char * v_dst = v_pad + args.nb21*i1 + args.nb21*C*i2 + args.nb21*C*args.ne_12_2*i3; + + if (i1 >= icp) { + for (uint64_t i = tiitg; i < args.nb11; i += ntg.x) { + k_dst[i] = 0; + } + for (uint64_t i = tiitg; i < args.nb21; i += ntg.x) { + v_dst[i] = 0; + } + } else { + for (uint64_t i = tiitg; i < args.nb11; i += ntg.x) { + k_dst[i] = k_src[i]; + } + for (uint64_t i = tiitg; i < args.nb21; i += ntg.x) { + v_dst[i] = v_src[i]; + } + } + } + + if (FC_flash_attn_ext_pad_has_mask) { + if (i2 < args.ne32 && i3 < args.ne33) { + for (int ib = i1; ib < args.ne31; ib += C) { + device const half * mask_src = (device const half *)(mask + args.nb31*ib + args.nb32*i2 + args.nb33*i3) + ic0; + device half * mask_dst = (device half *)(mask_pad) + C*ib + C*args.ne31*i2 + C*args.ne31*args.ne32*i3; + + for (int i = tiitg; i < C; i += ntg.x) { + if (i >= icp) { + mask_dst[i] = -MAXHALF; + } else { + mask_dst[i] = mask_src[i]; + } + } + } + } + } +} + +constant int32_t FC_flash_attn_ext_blk_nqptg [[function_constant(FC_FLASH_ATTN_EXT_BLK + 24)]]; +constant int32_t FC_flash_attn_ext_blk_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_BLK + 25)]]; + +// DS4 FlashAttention mask scan: marks blocks so the non-vector kernel can skip +// blocks that are entirely masked or entirely zero. +kernel void kernel_flash_attn_ext_blk( + constant ds4_metal_args_flash_attn_ext_blk & args, + device const char * mask, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]]) { + const int32_t Q = FC_flash_attn_ext_blk_nqptg; + const int32_t C = FC_flash_attn_ext_blk_ncpsg; + + constexpr short NW = N_SIMDWIDTH; + + const int32_t i3 = tgpig[2]/args.ne32; + const int32_t i2 = tgpig[2]%args.ne32; + const int32_t i1 = tgpig[1]; + const int32_t i0 = tgpig[0]; + + char res = i0*C + C > args.ne30 ? 1 : 0; + + if ((C > NW || Q > 1) && res == 0) { + half mmin = MAXHALF; + half mmax = -MAXHALF; + const int32_t q0 = i1*Q; + + FOR_UNROLL (short j = 0; j < Q; ++j) { + if (q0 + j < args.ne31) { + device const half * mask_src = + (device const half *) (mask + (q0 + j)*args.nb31 + i2*args.nb32 + i3*args.nb33) + + i0*C + tiisg; + + FOR_UNROLL (short ii = 0; ii < C/NW; ++ii) { + mmin = min(mmin, mask_src[ii*NW]); + mmax = max(mmax, mask_src[ii*NW]); + } + } + } + + mmin = simd_min(mmin); + mmax = simd_max(mmax); + + if (mmax > -MAXHALF) { + if (mmin == 0.0 && mmax == 0.0) { + res = 2; + } else { + res = 1; + } + } + } + + const int32_t nblk1 = ((args.ne01 + Q - 1)/Q); + const int32_t nblk0 = ((args.ne30 + C - 1)/C); + + if (tiisg == 0) { + dst[((i3*args.ne32 + i2)*nblk1 + i1)*nblk0 + i0] = res; + } +} + +constant bool FC_flash_attn_ext_has_mask [[function_constant(FC_FLASH_ATTN_EXT + 0)]]; +constant bool FC_flash_attn_ext_has_sinks [[function_constant(FC_FLASH_ATTN_EXT + 1)]]; +constant bool FC_flash_attn_ext_has_bias [[function_constant(FC_FLASH_ATTN_EXT + 2)]]; +constant bool FC_flash_attn_ext_has_scap [[function_constant(FC_FLASH_ATTN_EXT + 3)]]; +constant bool FC_flash_attn_ext_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT + 4)]]; + +constant bool FC_flash_attn_ext_bc_mask [[function_constant(FC_FLASH_ATTN_EXT + 10)]]; + +constant int32_t FC_flash_attn_ext_ns10 [[function_constant(FC_FLASH_ATTN_EXT + 20)]]; +constant int32_t FC_flash_attn_ext_ns20 [[function_constant(FC_FLASH_ATTN_EXT + 21)]]; +constant int32_t FC_flash_attn_ext_nsg [[function_constant(FC_FLASH_ATTN_EXT + 22)]]; + +// DS4 non-vector FlashAttention. The only exported instance uses the model's +// 512-wide F16 K/V rows; keeping the template body generic preserves the same +// arithmetic for dense and compressed-attention prefill. +template< + typename q_t, + typename q4_t, + typename q8x8_t, + typename k_t, + typename k4x4_t, + typename k8x8_t, + typename v_t, + typename v4x4_t, + typename v8x8_t, + typename qk_t, + typename qk8x8_t, + typename s_t, + typename s2_t, + typename s8x8_t, + typename o_t, + typename o4_t, + typename o8x8_t, + typename kd4x4_t, + short nl_k, + void (*deq_k)(device const kd4x4_t *, short, thread k4x4_t &), + typename vd4x4_t, + short nl_v, + void (*deq_v)(device const vd4x4_t *, short, thread v4x4_t &), + short DK, + short DV, + short Q, + short C, + short NSG> +void kernel_flash_attn_ext_impl( + constant ds4_metal_args_flash_attn_ext & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * sinks, + device const char * pad, + device const char * blk, + device char * dst, + threadgroup half * shmem_f16, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const ushort iq3 = tgpig[2]; + const ushort iq2 = tgpig[1]; + const ushort iq1 = tgpig[0]*Q; + +#define NS10 (FC_flash_attn_ext_ns10) +#define NS20 (FC_flash_attn_ext_ns20) + + constexpr short KV = 8; + + constexpr short DK4 = DK/4; + constexpr short DK8 = DK/8; + constexpr short DK16 = DK/16; + constexpr short DV4 = DV/4; + constexpr short DV16 = DV/16; + + constexpr short PV = PAD2(DV, 64); + constexpr short PV4 = PV/4; + constexpr short PV8 = PV/8; + + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = Q/NSG; + constexpr short SH = 2*C; + + constexpr short TS = 2*SH; + constexpr short T = DK + 2*PV; + + threadgroup q_t * sq = (threadgroup q_t *) (shmem_f16 + 0*T); + threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*T); + threadgroup o_t * so = (threadgroup o_t *) (shmem_f16 + 0*T + Q*DK); + threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 0*T + Q*DK); + threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + Q*T); + threadgroup s2_t * ss2 = (threadgroup s2_t *) (shmem_f16 + Q*T); + + threadgroup k_t * sk = (threadgroup k_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS); + threadgroup k4x4_t * sk4x4 = (threadgroup k4x4_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS); + + threadgroup v_t * sv = (threadgroup v_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS); + threadgroup v4x4_t * sv4x4 = (threadgroup v4x4_t *) (shmem_f16 + sgitg*(4*16*KV) + Q*T + Q*TS); + + threadgroup half2 * sm2 = (threadgroup half2 *) (shmem_f16 + Q*T + 2*C); + + device const half2 * pm2[NQ]; + + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + pm2[jj] = (device const half2 *) ((device const char *) mask + (iq1 + j)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); + } + + { + const int32_t nblk1 = ((args.ne01 + Q - 1)/Q); + const int32_t nblk0 = ((args.ne11 + C - 1)/C); + + blk += (((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*nblk1 + iq1/Q)*nblk0; + } + + { + q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03; + + const short ikv2 = iq2/(args.ne02/args.ne_12_2); + const short ikv3 = iq3/(args.ne03/args.ne_12_3); + + k += ikv2*args.nb12 + ikv3*args.nb13; + v += ikv2*args.nb22 + ikv3*args.nb23; + } + + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + device const float4 * q4 = (device const float4 *) ((device const char *) q + j*args.nb01); + + for (short i = tiisg; i < DK4; i += NW) { + if (iq1 + j < args.ne01) { + sq4[j*DK4 + i] = (q4_t) q4[i]; + } else { + sq4[j*DK4 + i] = 0; + } + } + } + + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + for (short i = tiisg; i < DV4; i += NW) { + so4[j*PV4 + i] = 0; + } + + for (short i = tiisg; i < SH; i += NW) { + ss[j*SH + i] = 0.0f; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + float S[NQ] = { [0 ... NQ-1] = 0.0f }; + + { + float M[NQ] = { [0 ... NQ-1] = -FLT_MAX/2 }; + + float slope = 1.0f; + + if (FC_flash_attn_ext_has_bias) { + const short h = iq2; + + const float base = h < args.n_head_log2 ? args.m0 : args.m1; + const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1; + + slope = pow(base, exph); + } + + for (int ic0 = 0; ; ++ic0) { + int ic = ic0*C; + if (ic >= args.ne11) { + break; + } + + if (FC_flash_attn_ext_has_kvpad && ic + C > args.ne11) { + k = pad; + v = k + args.nb11*C*args.ne_12_2*args.ne_12_3; + mask = v + args.nb21*C*args.ne_12_2*args.ne_12_3; + + const short ikv2 = iq2/(args.ne02/args.ne_12_2); + const short ikv3 = iq3/(args.ne03/args.ne_12_3); + + k += (ikv2 + ikv3*args.ne_12_2)*args.nb11*C; + v += (ikv2 + ikv3*args.ne_12_2)*args.nb21*C; + + if (!FC_flash_attn_ext_has_mask) { + threadgroup half * sm = (threadgroup half *) (sm2); + + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + for (short i = tiisg; i < C; i += NW) { + if (ic + i >= args.ne11) { + sm[2*j*SH + i] = -MAXHALF; + } + } + } + } else { + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + pm2[jj] = (device const half2 *) ((device const half *) mask + + (iq1 + j)*C + + (iq2%args.ne32)*(C*args.ne31) + + (iq3%args.ne33)*(C*args.ne31*args.ne32)); + } + } + + ic = 0; + } + + char blk_cur = 1; + + if (FC_flash_attn_ext_has_mask) { + blk_cur = blk[ic0]; + + if (blk_cur == 0) { + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + pm2[jj] += NW; + } + + continue; + } + + if (blk_cur == 1) { + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + if (FC_flash_attn_ext_bc_mask) { + sm2[j*SH + tiisg] = (iq1 + j) < args.ne31 ? pm2[jj][tiisg] : half2(-MAXHALF, -MAXHALF); + } else { + sm2[j*SH + tiisg] = pm2[jj][tiisg]; + } + + pm2[jj] += NW; + } + } else if (blk_cur == 2) { + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + pm2[jj] += NW; + } + } + } + + if (is_same::value) { + device const k_t * pk = (device const k_t *) (k + ic*args.nb11); + threadgroup const q_t * pq = sq; + threadgroup s_t * ps = ss; + + pk += sgitg*(8*NS10); + ps += sgitg*(8*1); + + static_assert((C/8) % NSG == 0, ""); + + constexpr short NC = (C/8)/NSG; + + FOR_UNROLL (short cc = 0; cc < NC; ++cc) { + qk8x8_t mqk = make_filled_simdgroup_matrix((qk_t) 0.0f); + + if (DK % 16 != 0) { + k8x8_t mk; + q8x8_t mq; + + FOR_UNROLL (short i = 0; i < DK8; ++i) { + simdgroup_barrier(mem_flags::mem_none); + + simdgroup_load(mk, pk + 8*i, NS10, 0, true); + simdgroup_load(mq, pq + 8*i, DK); + + simdgroup_barrier(mem_flags::mem_none); + + simdgroup_multiply_accumulate(mqk, mq, mk, mqk); + } + } else { + k8x8_t mk[2]; + q8x8_t mq[2]; + + #pragma unroll (MIN(DK8/2, 4*NSG)) + for (short i = 0; i < DK8/2; ++i) { + simdgroup_barrier(mem_flags::mem_none); + + simdgroup_load(mq[0], pq + 0*8 + 16*i, DK); + simdgroup_load(mq[1], pq + 1*8 + 16*i, DK); + + simdgroup_load(mk[0], pk + 0*8 + 16*i, NS10, 0, true); + simdgroup_load(mk[1], pk + 1*8 + 16*i, NS10, 0, true); + + simdgroup_barrier(mem_flags::mem_none); + + simdgroup_multiply_accumulate(mqk, mq[0], mk[0], mqk); + simdgroup_multiply_accumulate(mqk, mq[1], mk[1], mqk); + } + } + + simdgroup_store(mqk, ps, SH, 0, false); + + pk += 8*(NSG*NS10); + ps += 8*(NSG); + } + } else { + for (short ccc = 0; ccc < (C/8)/NSG; ++ccc) { + const short cc = ccc*NSG + sgitg; + + const short tx = tiisg%4; + const short ty = tiisg/4; + + qk8x8_t mqk = make_filled_simdgroup_matrix((qk_t) 0.0f); + + for (short ii = 0; ii < DK16; ii += 4) { + device const kd4x4_t * pk4x4 = (device const kd4x4_t *) (k + ((ic + 8*cc + ty)*args.nb11)); + + if (DK16%4 == 0) { + { + k4x4_t tmp; + deq_k(pk4x4 + (ii + tx)/nl_k, (ii + tx)%nl_k, tmp); + sk4x4[4*ty + tx] = tmp; + } + + simdgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short k = 0; k < 4; ++k) { + k8x8_t mk; + q8x8_t mq; + + simdgroup_load(mk, sk + 16*k + 0*8, 4*16, 0, true); + simdgroup_load(mq, sq + (2*(ii + k) + 0)*8, DK); + simdgroup_multiply_accumulate(mqk, mq, mk, mqk); + + simdgroup_load(mk, sk + 16*k + 1*8, 4*16, 0, true); + simdgroup_load(mq, sq + (2*(ii + k) + 1)*8, DK); + simdgroup_multiply_accumulate(mqk, mq, mk, mqk); + } + } else { + if (ii + tx < DK16) { + k4x4_t tmp; + deq_k(pk4x4 + (ii + tx)/nl_k, (ii + tx)%nl_k, tmp); + sk4x4[4*ty + tx] = tmp; + } + + simdgroup_barrier(mem_flags::mem_threadgroup); + + for (short k = 0; k < 4 && ii + k < DK16; ++k) { + k8x8_t mk; + q8x8_t mq; + + simdgroup_load(mk, sk + 16*k + 0*8, 4*16, 0, true); + simdgroup_load(mq, sq + (2*(ii + k) + 0)*8, DK); + simdgroup_multiply_accumulate(mqk, mq, mk, mqk); + + simdgroup_load(mk, sk + 16*k + 1*8, 4*16, 0, true); + simdgroup_load(mq, sq + (2*(ii + k) + 1)*8, DK); + simdgroup_multiply_accumulate(mqk, mq, mk, mqk); + } + } + } + + simdgroup_store(mqk, ss + 8*cc, SH, 0, false); + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + const float m = M[jj]; + + float2 s2 = ss2[j*SH/2 + tiisg]*args.scale; + + if (FC_flash_attn_ext_has_scap) { + s2 = args.logit_softcap*precise::tanh(s2); + } + + if (blk_cur != 2) { + if (FC_flash_attn_ext_has_bias) { + s2 += s2_t(sm2[j*SH + tiisg])*slope; + } else { + s2 += s2_t(sm2[j*SH + tiisg]); + } + } + + M[jj] = simd_max(max(M[jj], max(s2[0], s2[1]))); + + const float ms = exp(m - M[jj]); + const float2 vs2 = exp(s2 - M[jj]); + + S[jj] = S[jj]*ms + simd_sum(vs2[0] + vs2[1]); + + ss2[j*SH/2 + tiisg] = vs2; + + if (DV4 % NW == 0) { + FOR_UNROLL (short ii = 0; ii < DV4/NW; ++ii) { + const short i = ii*NW + tiisg; + + so4[j*PV4 + i] *= ms; + } + } else { + for (short i = tiisg; i < DV4; i += NW) { + so4[j*PV4 + i] *= ms; + } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + { + if (is_same::value) { + static_assert(PV8 % NSG == 0, ""); + + constexpr short NO = PV8/NSG; + + o8x8_t lo[NO]; + + { + auto sot = so + 8*sgitg; + + FOR_UNROLL (short ii = 0; ii < NO; ++ii) { + simdgroup_load(lo[ii], sot, PV, 0, false); + + sot += 8*NSG; + } + } + + { + device const v_t * pv = (device const v_t *) (v + ic*args.nb21); + + pv += 8*sgitg; + + if (DV <= 64) { + FOR_UNROLL (short cc = 0; cc < C/8; ++cc) { + s8x8_t vs; + simdgroup_load(vs, ss + 8*cc, SH, 0, false); + + FOR_UNROLL (short ii = 0; ii < NO/2; ++ii) { + v8x8_t mv[2]; + + simdgroup_load(mv[0], pv + 0*NSG + 16*ii*NSG, NS20, 0, false); + simdgroup_load(mv[1], pv + 8*NSG + 16*ii*NSG, NS20, 0, false); + + simdgroup_multiply_accumulate(lo[2*ii + 0], vs, mv[0], lo[2*ii + 0]); + simdgroup_multiply_accumulate(lo[2*ii + 1], vs, mv[1], lo[2*ii + 1]); + } + + pv += 8*NS20; + } + } else { + constexpr short NC = (C/8)/2; + + FOR_UNROLL (short cc = 0; cc < NC; ++cc) { + s8x8_t vs[2]; + + simdgroup_load(vs[0], ss + 16*cc + 0, SH, 0, false); + simdgroup_load(vs[1], ss + 16*cc + 8, SH, 0, false); + + FOR_UNROLL (short ii = 0; ii < NO/2; ++ii) { + v8x8_t mv[4]; + + simdgroup_load(mv[0], pv + 0*NSG + 16*ii*NSG + 0*8*NS20, NS20, 0, false); + simdgroup_load(mv[1], pv + 8*NSG + 16*ii*NSG + 0*8*NS20, NS20, 0, false); + simdgroup_load(mv[2], pv + 0*NSG + 16*ii*NSG + 1*8*NS20, NS20, 0, false); + simdgroup_load(mv[3], pv + 8*NSG + 16*ii*NSG + 1*8*NS20, NS20, 0, false); + + simdgroup_multiply_accumulate(lo[2*ii + 0], vs[0], mv[0], lo[2*ii + 0]); + simdgroup_multiply_accumulate(lo[2*ii + 1], vs[0], mv[1], lo[2*ii + 1]); + simdgroup_multiply_accumulate(lo[2*ii + 0], vs[1], mv[2], lo[2*ii + 0]); + simdgroup_multiply_accumulate(lo[2*ii + 1], vs[1], mv[3], lo[2*ii + 1]); + } + + pv += 2*8*NS20; + } + } + } + + { + auto sot = so + 8*sgitg; + + FOR_UNROLL (short ii = 0; ii < NO; ++ii) { + simdgroup_store(lo[ii], sot, PV, 0, false); + + sot += 8*NSG; + } + } + } else { + const short tx = tiisg%4; + const short ty = tiisg/4; + + for (short cc = 0; cc < C/8; ++cc) { + s8x8_t vs; + simdgroup_load(vs, ss + 8*cc, SH, 0, false); + + for (short ii = 4*sgitg; ii < DV16; ii += 4*NSG) { + device const vd4x4_t * pv4x4 = (device const vd4x4_t *) (v + ((ic + 8*cc + ty)*args.nb21)); + + if (DV16%4 == 0) { + { + v4x4_t tmp; + deq_v(pv4x4 + (ii + tx)/nl_v, (ii + tx)%nl_v, tmp); + sv4x4[4*ty + tx] = tmp; + } + + simdgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short k = 0; k < 4; ++k) { + v8x8_t mv[2]; + o8x8_t lo[2]; + + simdgroup_load(mv[0], sv + 16*k + 0*8, 4*16, 0, false); + simdgroup_load(mv[1], sv + 16*k + 1*8, 4*16, 0, false); + simdgroup_load(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false); + simdgroup_load(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false); + + simdgroup_multiply_accumulate(lo[0], vs, mv[0], lo[0]); + simdgroup_multiply_accumulate(lo[1], vs, mv[1], lo[1]); + + simdgroup_store(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false); + simdgroup_store(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false); + } + } else { + if (ii + tx < DV16) { + v4x4_t tmp; + deq_v(pv4x4 + (ii + tx)/nl_v, (ii + tx)%nl_v, tmp); + sv4x4[4*ty + tx] = tmp; + } + + simdgroup_barrier(mem_flags::mem_threadgroup); + + for (short k = 0; k < 4 && ii + k < DV16; ++k) { + v8x8_t mv[2]; + o8x8_t lo[2]; + + simdgroup_load(mv[0], sv + 16*k + 0*8, 4*16, 0, false); + simdgroup_load(mv[1], sv + 16*k + 1*8, 4*16, 0, false); + simdgroup_load(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false); + simdgroup_load(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false); + + simdgroup_multiply_accumulate(lo[0], vs, mv[0], lo[0]); + simdgroup_multiply_accumulate(lo[1], vs, mv[1], lo[1]); + + simdgroup_store(lo[0], so + 8*(2*(ii + k) + 0), PV, 0, false); + simdgroup_store(lo[1], so + 8*(2*(ii + k) + 1), PV, 0, false); + } + } + } + } + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (FC_flash_attn_ext_has_sinks) { + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + const float m = M[jj]; + const float s = tiisg == 0 ? ((device const float *) sinks)[iq2] : -FLT_MAX/2; + + M[jj] = simd_max(max(M[jj], s)); + + const float ms = exp(m - M[jj]); + const float vs = exp(s - M[jj]); + + S[jj] = S[jj]*ms + simd_sum(vs); + + for (short i = tiisg; i < DV4; i += NW) { + so4[j*PV4 + i] *= ms; + } + } + } + } + + for (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + if (iq1 + j >= args.ne01) { + break; + } + + device float4 * dst4 = (device float4 *) dst + ((uint64_t)iq3*args.ne2*args.ne1 + iq2 + (uint64_t)(iq1 + j)*args.ne1)*DV4; + + const float scale = S[jj] == 0.0 ? 0.0f : 1.0f/S[jj]; + + if (DV4 % NW == 0) { + FOR_UNROLL (short ii = 0; ii < DV4/NW; ++ii) { + const short i = ii*NW + tiisg; + + dst4[i] = (float4) so4[j*PV4 + i]*scale; + } + } else { + for (short i = tiisg; i < DV4; i += NW) { + dst4[i] = (float4) so4[j*PV4 + i]*scale; + } + } + } + +#undef NS10 +#undef NS20 +} + +// Batched FlashAttention for prompt/prefill rows. It computes QK, applies mask, +// sinks, ALiBi/softcap options when enabled, and multiplies by V without +// materializing the full attention matrix. +template< + typename q_t, + typename q4_t, + typename q8x8_t, + typename k_t, + typename k4x4_t, + typename k8x8_t, + typename v_t, + typename v4x4_t, + typename v8x8_t, + typename qk_t, + typename qk8x8_t, + typename s_t, + typename s2_t, + typename s8x8_t, + typename o_t, + typename o4_t, + typename o8x8_t, + typename kd4x4_t, + short nl_k, + void (*deq_k)(device const kd4x4_t *, short, thread k4x4_t &), + typename vd4x4_t, + short nl_v, + void (*deq_v)(device const vd4x4_t *, short, thread v4x4_t &), + short DK, + short DV, + short Q = OP_FLASH_ATTN_EXT_NQPSG, + short C = OP_FLASH_ATTN_EXT_NCPSG> +kernel void kernel_flash_attn_ext( + constant ds4_metal_args_flash_attn_ext & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * sinks, + device const char * pad, + device const char * blk, + device char * dst, + threadgroup half * shmem_f16 [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { +#define FWD_TMPL q_t, q4_t, q8x8_t, k_t, k4x4_t, k8x8_t, v_t, v4x4_t, v8x8_t, qk_t, qk8x8_t, s_t, s2_t, s8x8_t, o_t, o4_t, o8x8_t, kd4x4_t, nl_k, deq_k, vd4x4_t, nl_v, deq_v, DK, DV, Q, C +#define FWD_ARGS args, q, k, v, mask, sinks, pad, blk, dst, shmem_f16, tgpig, tiisg, sgitg + switch (FC_flash_attn_ext_nsg) { + case 4: kernel_flash_attn_ext_impl(FWD_ARGS); break; + case 8: kernel_flash_attn_ext_impl(FWD_ARGS); break; + } +#undef FWD_TMPL +#undef FWD_ARGS +} + +#define FA_NONVEC_TYPES \ + half, half4, simdgroup_half8x8, \ + half, half4x4, simdgroup_half8x8, \ + half, half4x4, simdgroup_half8x8, \ + float, simdgroup_float8x8, \ + float, float2, simdgroup_float8x8, \ + float, float4, simdgroup_float8x8 + +typedef decltype(kernel_flash_attn_ext) flash_attn_ext_dk512_t; +typedef decltype(kernel_flash_attn_ext) flash_attn_ext_dk256_t; + +// Host-visible prefill FlashAttention variant for DS4's 512-wide F16 K/V rows. +template [[host_name("kernel_flash_attn_ext_f16_dk512_dv512")]] +kernel flash_attn_ext_dk512_t kernel_flash_attn_ext; + +// Host-visible prefill FlashAttention variant for GLM's 256-wide F16 K/V rows. +template [[host_name("kernel_flash_attn_ext_f16_dk256_dv256")]] +kernel flash_attn_ext_dk256_t kernel_flash_attn_ext; + +#undef FA_NONVEC_TYPES + +constant bool FC_flash_attn_ext_vec_has_mask [[function_constant(FC_FLASH_ATTN_EXT_VEC + 0)]]; +constant bool FC_flash_attn_ext_vec_has_sinks [[function_constant(FC_FLASH_ATTN_EXT_VEC + 1)]]; +constant bool FC_flash_attn_ext_vec_has_bias [[function_constant(FC_FLASH_ATTN_EXT_VEC + 2)]]; +constant bool FC_flash_attn_ext_vec_has_scap [[function_constant(FC_FLASH_ATTN_EXT_VEC + 3)]]; +constant bool FC_flash_attn_ext_vec_has_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 4)]]; +constant bool FC_flash_attn_ext_vec_shared_kvpad [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]]; +constant int32_t FC_flash_attn_ext_vec_ns10 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 20)]]; +constant int32_t FC_flash_attn_ext_vec_ns20 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 21)]]; +constant int32_t FC_flash_attn_ext_vec_nsg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 22)]]; +constant int32_t FC_flash_attn_ext_vec_nwg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 23)]]; + +// Decode FlashAttention for one query row. DS4 uses this in generation to scan +// raw and compressed KV cache chunks, optionally splitting long contexts across +// workgroups and writing partial softmax state for a later reduction. +template< + typename q4_t, + typename k4_t, + typename v4_t, + typename qk_t, + typename s_t, + typename s4_t, + typename o4_t, + typename kd4_t, + short nl_k, + void (*deq_k_t4)(device const kd4_t *, short, thread k4_t &), + typename vd4_t, + short nl_v, + void (*deq_v_t4)(device const vd4_t *, short, thread v4_t &), + short DK, + short DV, + short NE = 4, + short Q = OP_FLASH_ATTN_EXT_VEC_NQPSG, + short C = OP_FLASH_ATTN_EXT_VEC_NCPSG> +kernel void kernel_flash_attn_ext_vec( + constant ds4_metal_args_flash_attn_ext_vec & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * sinks, + device const char * pad, + device char * dst, + threadgroup half * shmem_f16 [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + static_assert(DK % 32 == 0, "DK must be divisible by 32"); + static_assert(DV % 32 == 0, "DV must be divisible by 32"); + +#define NWG (FC_flash_attn_ext_vec_nwg) +#define NSG (FC_flash_attn_ext_vec_nsg) +#define NS10 (FC_flash_attn_ext_vec_ns10) +#define NS20 (FC_flash_attn_ext_vec_ns20) + + const short iwg = tgpig[2]%NWG; + + const ushort iq3 = tgpig[2]/NWG; + const ushort iq2 = tgpig[1]; + const ushort iq1 = tgpig[0]; + + constexpr short DK4 = DK/4; + constexpr short DV4 = DV/4; + constexpr short PK = PAD2(DK, 128); + constexpr short PK4 = PK/4; + constexpr short PV = PAD2(DV, 128); + constexpr short PV4 = PV/4; + constexpr short NW = N_SIMDWIDTH; + constexpr short NL = NW/NE; + constexpr short SH = 4*C; + + static_assert(DK4 % NL == 0, "DK4 must be divisible by NL"); + static_assert(DV4 % NL == 0, "DV4 must be divisible by NL"); + + threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*PK); + threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + sgitg*SH + NSG*PK); + threadgroup s4_t * ss4 = (threadgroup s4_t *) (shmem_f16 + sgitg*SH + NSG*PK); + threadgroup half * sm = (threadgroup half *) (shmem_f16 + sgitg*SH + 2*C + NSG*PK); + threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 2*sgitg*PV + NSG*PK + NSG*SH); + + so4 += tiisg; + + { + q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03; + + const short ikv2 = iq2/(args.ne02/args.ne_12_2); + const short ikv3 = iq3/(args.ne03/args.ne_12_3); + + k += ikv2*args.nb12 + ikv3*args.nb13; + v += ikv2*args.nb22 + ikv3*args.nb23; + } + + device const float4 * q4 = (device const float4 *) ((device const char *) q); + + if (iq1 < args.ne01) { + for (short i = tiisg; i < PK4; i += NW) { + if (i < DK4) { + sq4[i] = (q4_t) q4[i]; + } else { + sq4[i] = (q4_t) 0.0f; + } + } + } + + for (short i = 0; i < DV4/NL; ++i) { + so4[i*NL] = (o4_t) 0.0f; + } + + for (short i = tiisg; i < SH/4; i += NW) { + ss4[i] = (s4_t) 0.0f; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + { + float S = 0.0f; + float M = -FLT_MAX/2; + + const short tx = tiisg%NL; + const short ty = tiisg/NL; + + device const half * pm = (device const half *) (mask + iq1*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); + + float slope = 1.0f; + + if (FC_flash_attn_ext_vec_has_bias) { + const short h = iq2; + + const float base = h < args.n_head_log2 ? args.m0 : args.m1; + const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1; + + slope = pow(base, exph); + } + + for (int ic0 = iwg*NSG + sgitg; ; ic0 += NWG*NSG) { + int ic = ic0*C; + if (ic >= args.ne11) { + break; + } + + if (FC_flash_attn_ext_vec_has_kvpad && ic + C > args.ne11) { + k = pad; + if (FC_flash_attn_ext_vec_shared_kvpad) { + v = k; + mask = k + args.nb11*C*args.ne_12_2*args.ne_12_3 + + args.nb21*C*args.ne_12_2*args.ne_12_3; + } else { + v = k + args.nb11*C*args.ne_12_2*args.ne_12_3; + mask = v + args.nb21*C*args.ne_12_2*args.ne_12_3; + } + + const short ikv2 = iq2/(args.ne02/args.ne_12_2); + const short ikv3 = iq3/(args.ne03/args.ne_12_3); + + k += (ikv2 + ikv3*args.ne_12_2)*args.nb11*C; + v += (ikv2 + ikv3*args.ne_12_2)*args.nb21*C; + + if (!FC_flash_attn_ext_vec_has_mask) { + if (ic + tiisg >= args.ne11) { + sm[tiisg] = -MAXHALF; + } + } else { + pm = (device const half *) (mask) + + iq1*C + + (iq2%args.ne32)*(C*args.ne31) + + (iq3%args.ne33)*(C*args.ne31*args.ne32); + } + + ic = 0; + } + + if (FC_flash_attn_ext_vec_has_mask) { + sm[tiisg] = pm[ic + tiisg]; + } + + if (simd_max(sm[tiisg]) <= -MAXHALF) { + continue; + } + + { + device const k4_t * pk4 = (device const k4_t *) (k + ic*args.nb11); + threadgroup const q4_t * pq4 = sq4; + + pk4 += ty*NS10/4 + tx; + pq4 += tx; + + qk_t mqk[C/NE] = { [ 0 ... C/NE - 1] = 0.0f }; + + FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { + if (is_same::value) { + FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { + mqk[cc] += dot((float4) pk4[cc*NE*NS10/4 + ii*NL], (float4) pq4[ii*NL]); + } + } else { + device const kd4_t * pk = (device const kd4_t *) (k + ((ic + NE*cc + ty)*args.nb11)); + + k4_t mk; + + FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { + const short i = ii*NL + tx; + + deq_k_t4(pk + i/nl_k, i%nl_k, mk); + + mqk[cc] += dot((float4) mk, (float4) sq4[i]); + } + } + + if (NE == 1) { + mqk[cc] = simd_sum(mqk[cc]); + } else { + if (NE <= 1) { + mqk[cc] += simd_shuffle_down(mqk[cc], 16); + } + if (NE <= 2) { + mqk[cc] += simd_shuffle_down(mqk[cc], 8); + } + if (NE <= 4) { + mqk[cc] += simd_shuffle_down(mqk[cc], 4); + } + if (NE <= 8) { + mqk[cc] += simd_shuffle_down(mqk[cc], 2); + } + if (NE <= 16) { + mqk[cc] += simd_shuffle_down(mqk[cc], 1); + } + + mqk[cc] = simd_shuffle(mqk[cc], NL*ty); + } + } + + if (FC_flash_attn_ext_vec_has_mask && + !FC_flash_attn_ext_vec_has_scap && + !FC_flash_attn_ext_vec_has_bias) { + ss[NE*tx + ty] = fma(mqk[tx], args.scale, (qk_t) sm[NE*tx + ty]); + } else { + mqk[tx] *= args.scale; + + if (FC_flash_attn_ext_vec_has_scap) { + mqk[tx] = args.logit_softcap*precise::tanh(mqk[tx]); + } + + if (FC_flash_attn_ext_vec_has_bias) { + mqk[tx] += (qk_t) sm[NE*tx + ty]*slope; + } else { + mqk[tx] += (qk_t) sm[NE*tx + ty]; + } + + ss[NE*tx + ty] = mqk[tx]; + } + } + + simdgroup_barrier(mem_flags::mem_threadgroup); + + { + const float m = M; + const float s = ss[tiisg]; + + M = simd_max(max(M, s)); + + const float ms = exp(m - M); + const float vs = exp(s - M); + + S = S*ms + simd_sum(vs); + + ss[tiisg] = vs; + + if ((DV4/NL % NW == 0) || ty == 0) { + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { + so4[ii*NL] *= ms; + } + } + } + + simdgroup_barrier(mem_flags::mem_threadgroup); + + { + o4_t lo[DV4/NL]; + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { + lo[ii] = 0.0f; + } + + if (is_same::value) { + device const v4_t * pv4 = (device const v4_t *) (v + ic*args.nb21); + + pv4 += ty*NS20/4 + tx; + + const auto sst = ss + ty; + + FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { + lo[ii] += o4_t(float4(pv4[cc*NE*NS20/4 + ii*NL])*float4(sst[cc*NE])); + } + } + } else { + FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { + device const vd4_t * pv4 = (device const vd4_t *) (v + ((ic + NE*cc + ty)*args.nb21)); + + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { + const short i = ii*NL + tx; + + v4_t mv; + deq_v_t4(pv4 + i/nl_v, i%nl_v, mv); + + lo[ii] += o4_t(float4(mv)*float4(ss[NE*cc + ty])); + } + } + } + + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { + if (NE > 1) { + lo[ii][0] += simd_shuffle_down(lo[ii][0], 16); + lo[ii][1] += simd_shuffle_down(lo[ii][1], 16); + lo[ii][2] += simd_shuffle_down(lo[ii][2], 16); + lo[ii][3] += simd_shuffle_down(lo[ii][3], 16); + } + + if (NE > 2) { + lo[ii][0] += simd_shuffle_down(lo[ii][0], 8); + lo[ii][1] += simd_shuffle_down(lo[ii][1], 8); + lo[ii][2] += simd_shuffle_down(lo[ii][2], 8); + lo[ii][3] += simd_shuffle_down(lo[ii][3], 8); + } + + if (NE > 4) { + lo[ii][0] += simd_shuffle_down(lo[ii][0], 4); + lo[ii][1] += simd_shuffle_down(lo[ii][1], 4); + lo[ii][2] += simd_shuffle_down(lo[ii][2], 4); + lo[ii][3] += simd_shuffle_down(lo[ii][3], 4); + } + + if (NE > 8) { + lo[ii][0] += simd_shuffle_down(lo[ii][0], 2); + lo[ii][1] += simd_shuffle_down(lo[ii][1], 2); + lo[ii][2] += simd_shuffle_down(lo[ii][2], 2); + lo[ii][3] += simd_shuffle_down(lo[ii][3], 2); + } + + if (NE > 16) { + lo[ii][0] += simd_shuffle_down(lo[ii][0], 1); + lo[ii][1] += simd_shuffle_down(lo[ii][1], 1); + lo[ii][2] += simd_shuffle_down(lo[ii][2], 1); + lo[ii][3] += simd_shuffle_down(lo[ii][3], 1); + } + } + + if ((DV4/NL % NW == 0) || ty == 0) { + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { + so4[ii*NL] += lo[ii]; + } + } + } + } + + if (FC_flash_attn_ext_vec_has_sinks && sgitg == 0 && iwg == 0) { + const float m = M; + const float s = tiisg == 0 ? ((device const float *) sinks)[iq2] : -FLT_MAX/2; + + M = simd_max(max(M, s)); + + const float ms = exp(m - M); + const float vs = exp(s - M); + + S = S*ms + simd_sum(vs); + + if ((DV4/NL % NW == 0) || ty == 0) { + FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { + so4[ii*NL] *= ms; + } + } + } + + if (tiisg == 0) { + ss[0] = (s_t) S; + ss[1] = (s_t) M; + } + } + + so4 -= tiisg; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short r = NSG/2; r > 0; r >>= 1) { + if (sgitg < r) { + const float S0 = ss[ 0]; + const float S1 = ss[r*(SH/2) + 0]; + + const float M0 = ss[ 1]; + const float M1 = ss[r*(SH/2) + 1]; + + const float M = max(M0, M1); + + const float ms0 = exp(M0 - M); + const float ms1 = exp(M1 - M); + + const float S = S0*ms0 + S1*ms1; + + if (tiisg == 0) { + ss[0] = S; + ss[1] = M; + } + + for (short i = tiisg; i < DV4; i += NW) { + so4[i] = so4[i]*ms0 + so4[i + r*PV4]*ms1; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (sgitg == 0) { + const int64_t nrows = args.ne3*args.ne2*args.ne1; + const int64_t rid = iq3*args.ne2*args.ne1 + iq2 + iq1*args.ne1; + + device float4 * dst4 = (device float4 *) dst; + device float * dst1 = (device float *) dst + nrows*DV*NWG; + + const float S = NWG == 1 ? (ss[0] == 0.0f ? 0.0f : 1.0f/ss[0]) : 1.0f; + + for (short i = tiisg; i < DV4; i += NW) { + dst4[rid*DV4*NWG + NWG*i + iwg] = (float4) so4[i]*S; + } + + if (NWG > 1) { + if (tiisg == 0) { + dst1[rid*(2*NWG) + 2*iwg + 0] = ss[0]; + dst1[rid*(2*NWG) + 2*iwg + 1] = ss[1]; + } + } + } + +#undef NWG +#undef NSG +#undef NS10 +#undef NS20 +} + +#define FA_TYPES \ + half4, \ + half4, \ + half4, \ + float, \ + float, float4, \ + float4 + +#define FA_TYPES_F32 \ + half4, \ + float4, \ + float4, \ + float, \ + float, float4, \ + float4 + +typedef decltype(kernel_flash_attn_ext_vec) flash_attn_ext_vec_t; + +// Host-visible decode FlashAttention variant for DS4's 512-wide F16 K/V rows. +template [[host_name("kernel_flash_attn_ext_vec_f16_dk512_dv512")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec; + +#undef FA_TYPES +#undef FA_TYPES_F32 + +constant int32_t FC_flash_attn_ext_vec_reduce_DV [[function_constant(FC_FLASH_ATTN_EXT_VEC_REDUCE + 0)]]; +constant int32_t FC_flash_attn_ext_vec_reduce_NWG [[function_constant(FC_FLASH_ATTN_EXT_VEC_REDUCE + 1)]]; + +// Reduces split-K decode FlashAttention partials. It combines each workgroup's +// output vector and softmax (sum,max) pair into the final attention result. +kernel void kernel_flash_attn_ext_vec_reduce( + constant ds4_metal_args_flash_attn_ext_vec_reduce & args, + device const char * htmp, + device char * dst, + uint tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { +#define NWG (FC_flash_attn_ext_vec_reduce_NWG) +#define DV (FC_flash_attn_ext_vec_reduce_DV) + + const uint64_t rid = tgpig; + + const short iwg = tiisg; + + device const float * ss = (device const float *) htmp + (uint64_t)args.nrows*DV*NWG; + + float S = ss[rid*(2*NWG) + 2*iwg + 0]; + float M = ss[rid*(2*NWG) + 2*iwg + 1]; + + const float m = simd_max(M); + const float ms = exp(M - m); + + S = simd_sum(S*ms); + S = S == 0.0f ? 0.0f : 1.0f/S; + + const short DV4 = DV/4; + + device const float4 * htmp4 = (device const float4 *) htmp + rid*DV4*NWG; + device float4 * dst4 = (device float4 *) dst + rid*DV4; + + for (short i = sgitg; i < DV4; i += NWG) { + const float4 v = simd_sum(htmp4[i*NWG + iwg]*ms); + + if (iwg == 0) { + dst4[i] = v*S; + } + } + +#undef NWG +#undef DV +} diff --git a/metal/get_rows.metal b/metal/get_rows.metal new file mode 100644 index 0000000..836225f --- /dev/null +++ b/metal/get_rows.metal @@ -0,0 +1,186 @@ +// DS4 Metal get-rows kernel. + +struct ds4_metal_args_get_rows { + int32_t ne00t; + int32_t ne00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +}; + +struct ds4_metal_args_get_rows_q8_0 { + int32_t n_embd; + int32_t n_vocab; + int32_t n_tokens; + uint64_t src_row_bytes; + uint64_t dst_row_bytes; + uint64_t token_stride; +}; + +struct ds4_get_rows_block_q4_0 { + half d; + uchar qs[16]; +}; + +struct ds4_get_rows_block_q4_K { + half d; + half dmin; + uchar scales[12]; + uchar qs[128]; +}; + +static inline uchar2 ds4_get_rows_q4_K_scale_min(int j, int k, device const uchar *q) { + return j < 4 ? uchar2{uchar(q[j + 0 + k] & 63), uchar(q[j + 4 + k] & 63)} + : uchar2{uchar((q[j + 4 + k] & 0x0f) | ((q[j - 4 + k] & 0xc0) >> 2)), + uchar((q[j + 4 + k] >> 4) | ((q[j - 0 + k] & 0xc0) >> 2))}; +} + +// Gathers embedding/table rows by integer ids. DS4 uses this for token +// embeddings and small indexed tables such as router/hash lookup outputs. +template +kernel void kernel_get_rows_f( + constant ds4_metal_args_get_rows & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]]) { + const int32_t iw0 = tgpig.x/args.ne10; + const int32_t i10 = tgpig.x%args.ne10; + const int32_t i11 = tgpig.y; + const int32_t i12 = tgpig.z; + + const int32_t r = ((const device int32_t *) (src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0]; + + const int32_t i02 = i11; + const int32_t i03 = i12; + + auto psrc = (const device T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + r*args.nb01); + auto pdst = ( device T *) (dst + i12*args.nb3 + i11*args.nb2 + i10*args.nb1); + + for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) { + pdst[ind] = psrc[ind]; + + break; + } +} + +typedef decltype(kernel_get_rows_f) get_rows_f_t; + +// Host-visible gather variants for F32, F16, and I32 tables. +template [[host_name("kernel_get_rows_f32")]] kernel get_rows_f_t kernel_get_rows_f; +template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f; +template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f; + +kernel void kernel_get_rows_q8_0_f32( + constant ds4_metal_args_get_rows_q8_0 & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]]) { + const int32_t block = (int32_t)tgpig.x; + const int32_t tok_i = (int32_t)tgpig.y; + if (tok_i >= args.n_tokens) return; + + const int32_t token = + ((const device int32_t *)(src1 + (uint64_t)tok_i*args.token_stride))[0]; + if (token < 0 || token >= args.n_vocab) return; + + const device block_q8_0 *row = + (const device block_q8_0 *)(src0 + (uint64_t)token * args.src_row_bytes); + const device block_q8_0 *qb = row + block; + device float *out = + (device float *)(dst + (uint64_t)tok_i * args.dst_row_bytes); + + const int32_t i0 = block * QK8_0; + const float d = (float)qb->d; + for (int32_t i = (int32_t)tiitg; i < QK8_0; i += (int32_t)ntg.x) { + const int32_t idx = i0 + i; + if (idx < args.n_embd) { + out[idx] = d * (float)qb->qs[i]; + } + } +} + +kernel void kernel_get_rows_q4_0_f32( + constant ds4_metal_args_get_rows_q8_0 & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]]) { + const int32_t block = (int32_t)tgpig.x; + const int32_t tok_i = (int32_t)tgpig.y; + if (tok_i >= args.n_tokens) return; + + const int32_t token = + ((const device int32_t *)(src1 + (uint64_t)tok_i * args.token_stride))[0]; + if (token < 0 || token >= args.n_vocab) return; + + const device ds4_get_rows_block_q4_0 *row = + (const device ds4_get_rows_block_q4_0 *)(src0 + (uint64_t)token * args.src_row_bytes); + const device ds4_get_rows_block_q4_0 *qb = row + block; + device float *out = + (device float *)(dst + (uint64_t)tok_i * args.dst_row_bytes); + + const int32_t i0 = block * 32; + const float d = (float)qb->d; + for (int32_t i = (int32_t)tiitg; i < 32; i += (int32_t)ntg.x) { + const int32_t idx = i0 + i; + if (idx < args.n_embd) { + /* ggml Q4_0: elems 0..15 = low nibbles of qs[0..15], 16..31 = high. */ + const uint8_t packed = qb->qs[(uint)(i & 15)]; + const uint8_t q = (i < 16) ? (packed & 0x0f) : (packed >> 4); + out[idx] = d * ((float)q - 8.0f); + } + } +} + +kernel void kernel_get_rows_q4_K_f32( + constant ds4_metal_args_get_rows_q8_0 & args, + device const char * src0, + device const char * src1, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]]) { + const int32_t block = (int32_t)tgpig.x; + const int32_t tok_i = (int32_t)tgpig.y; + if (tok_i >= args.n_tokens) return; + + const int32_t token = + ((const device int32_t *)(src1 + (uint64_t)tok_i * args.token_stride))[0]; + if (token < 0 || token >= args.n_vocab) return; + + const device ds4_get_rows_block_q4_K *row = + (const device ds4_get_rows_block_q4_K *)(src0 + (uint64_t)token * args.src_row_bytes); + const device ds4_get_rows_block_q4_K *qb = row + block; + device float *out = + (device float *)(dst + (uint64_t)tok_i * args.dst_row_bytes); + + const int32_t i0 = block * 256; + for (int32_t i = (int32_t)tiitg; i < 256; i += (int32_t)ntg.x) { + const int32_t idx = i0 + i; + if (idx < args.n_embd) { + const uint group = (uint)i >> 5u; + const uint l = (uint)i & 31u; + const uchar2 sm = ds4_get_rows_q4_K_scale_min((int)group, 0, qb->scales); + const uint byte_off = (group >> 1u) * 32u + l; + const uint shift = (group & 1u) * 4u; + const uint q = ((uint)qb->qs[byte_off] >> shift) & 0x0fu; + out[idx] = (float)qb->d * (float)sm.x * (float)q - + (float)qb->dmin * (float)sm.y; + } + } +} diff --git a/metal/glu.metal b/metal/glu.metal new file mode 100644 index 0000000..fd5c6fd --- /dev/null +++ b/metal/glu.metal @@ -0,0 +1,63 @@ +struct ds4_metal_args_glu { + int32_t ne00; + uint64_t nb01; + int32_t ne10; + uint64_t nb11; + int32_t ne0; + uint64_t nb1; + int32_t i00; + int32_t i10; + float alpha; + float limit; +}; + +// SwiGLU activation for the FFN inner state. DS4 clamps the shared expert with +// the same swiglu_limit used by routed experts. +kernel void kernel_swiglu_f32( + constant ds4_metal_args_glu & args, + device const char * src0, + device const char * src1, + device char * dst, + uint tgpig[[threadgroup_position_in_grid]], + uint tpitg[[thread_position_in_threadgroup]], + uint ntg[[threads_per_threadgroup]]) { + device const float * src0_row = (device const float *) ((device const char *) src0 + tgpig*args.nb01) + args.i00; + device const float * src1_row = (device const float *) ((device const char *) src1 + tgpig*args.nb11) + args.i10; + device float * dst_row = (device float *) ((device char *) dst + tgpig*args.nb1); + + for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) { + float x0 = src0_row[i0]; + float x1 = src1_row[i0]; + if (args.limit > 1.0e-6f) { + x0 = min(x0, args.limit); + x1 = clamp(x1, -args.limit, args.limit); + } + + const float silu = x0 / (1.0f + exp(-x0)); + + dst_row[i0] = silu*x1*args.alpha; + } +} + +kernel void kernel_swiglu_flat_f32( + constant ds4_metal_args_glu & args, + device const char * src0, + device const char * src1, + device char * dst, + uint i [[thread_position_in_grid]]) { + if (i >= (uint)args.ne0) return; + + device const float * src0_f32 = (device const float *) src0 + args.i00; + device const float * src1_f32 = (device const float *) src1 + args.i10; + device float * dst_f32 = (device float *) dst; + + float x0 = src0_f32[i]; + float x1 = src1_f32[i]; + if (args.limit > 1.0e-6f) { + x0 = min(x0, args.limit); + x1 = clamp(x1, -args.limit, args.limit); + } + + const float silu = x0 / (1.0f + exp(-x0)); + dst_f32[i] = silu*x1*args.alpha; +} diff --git a/metal/moe.metal b/metal/moe.metal new file mode 100644 index 0000000..c1a19e8 --- /dev/null +++ b/metal/moe.metal @@ -0,0 +1,7858 @@ +// DS4 Metal routed-MoE matvec kernels. + +#ifndef QK_K +#define QK_K 256 +#endif +#define N_R0_Q2_K 4 +#define N_R0_GLM_Q2_PAIR2_K 1 +#define N_R0_Q4_K 2 +#define N_R0_Q8_K 2 +#define N_R0_GLM_Q4_PAIR2_K 1 +#define N_R0_GLM_Q4_PAIR_K 4 +#define N_R0_Q5_PAIR_K 4 +#define N_R0_Q5_K 4 +#define N_R0_Q6_K 2 +#define N_R0_IQ2_XXS 4 + +static constant uchar ds4_metal_kmask_iq2xs[8] = { + 1, 2, 4, 8, 16, 32, 64, 128 +}; + +static constant uchar ds4_metal_ksigns_iq2xs[128] = { + 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, + 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, + 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, + 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, + 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, + 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, + 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, + 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, +}; + +static constant ulong ds4_metal_iq2xxs_grid[256] = { + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, + 0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819, + 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819, + 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, + 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, + 0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, + 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b, + 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819, + 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, + 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, + 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, + 0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808, + 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808, + 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, + 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, + 0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, + 0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908, + 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819, + 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, + 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, + 0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, + 0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808, + 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08, + 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, + 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, + 0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, + 0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908, + 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19, + 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, + 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, + 0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, + 0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908, + 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08, + 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, + 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, + 0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, + 0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808, + 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808, + 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, + 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, + 0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, + 0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b, + 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08, + 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, + 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, + 0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, + 0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819, + 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08, + 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, + 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, + 0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, + 0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b, + 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908, + 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, + 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, + 0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, + 0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b, + 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808, + 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, + 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, + 0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, + 0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19, + 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908, +}; + +#define kmask_iq2xs ds4_metal_kmask_iq2xs +#define ksigns_iq2xs ds4_metal_ksigns_iq2xs +#define iq2xxs_grid ds4_metal_iq2xxs_grid + +struct block_q2_K { + uchar scales[QK_K/16]; + uchar qs[QK_K/4]; + half d; + half dmin; +}; + +struct block_q4_K { + half d; + half dmin; + uchar scales[12]; + uchar qs[QK_K/2]; +}; + +struct block_q5_K { + half d; + half dmin; + uchar scales[12]; + uchar qh[QK_K/8]; + uchar qs[QK_K/2]; +}; + +struct block_q6_K { + uchar ql[QK_K/2]; + uchar qh[QK_K/4]; + char scales[QK_K/16]; + half d; +}; + +struct block_iq2_xxs { + half d; + ushort qs[QK_K/8]; +}; + +struct ds4_metal_glm_routed_moe_args { + uint32_t in_dim; + uint32_t mid_dim; + uint32_t out_dim; + uint32_t n_total_expert; + uint32_t n_expert_used; + uint32_t n_tokens; + uint32_t mid_token_stride; + uint32_t down_type; + /* Expert ownership under tensor parallelism: tp_world 0/1 = full + * compute; otherwise each rank owns a contiguous expert range. */ + int32_t tp_rank; + int32_t tp_world; + int32_t tp_expert_base; + uint64_t gate_expert_bytes; + uint64_t gate_row_bytes; + uint64_t up_expert_bytes; + uint64_t up_row_bytes; + uint64_t down_expert_bytes; + uint64_t down_row_bytes; +}; + + +static inline bool ds4_tp_owns_expert(int expert, int n_total, + int tp_rank, int tp_world) { + if (tp_world <= 1) return true; + const int first = tp_rank * (n_total / tp_world); + const int last = tp_rank + 1 == tp_world ? n_total : first + n_total / tp_world; + return expert >= first && expert < last; +} + +struct ds4_metal_dsv4_moe_swiglu_weight_args { + uint32_t width; + uint32_t rows; + uint64_t gate_row_stride; + uint64_t up_row_stride; + uint64_t mid_row_stride; + uint64_t weight_stride; + uint32_t write_clamped; + float clamp_value; +}; + +struct ds4_metal_dsv4_moe_sum6_args { + uint32_t width; + uint32_t tokens; + uint64_t src_token_stride; + uint64_t dst_token_stride; +}; + +// Routed-MoE activation for the selected experts: +// clamp(gate), clamp(up), silu(gate) * up * route_weight. Normal inference +// does not consume gate/up after this point, so the fast path avoids writing the +// clamped intermediates back. A diagnostic env switch can restore those writes +// when comparing the old multi-kernel intermediate tensors. +kernel void kernel_dsv4_moe_swiglu_weight( + constant ds4_metal_dsv4_moe_swiglu_weight_args &args, + device char *gate, + device char *up, + device char *mid, + device const char *weights, + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint ntg [[threads_per_threadgroup]]) { + if (row >= args.rows) return; + + device float *gate_row = (device float *)(gate + (uint64_t)row * args.gate_row_stride); + device float *up_row = (device float *)(up + (uint64_t)row * args.up_row_stride); + device float *mid_row = (device float *)(mid + (uint64_t)row * args.mid_row_stride); + device const float *w = (device const float *)(weights + (uint64_t)row * args.weight_stride); + const float route_weight = w[0]; + const float c = args.clamp_value; + + for (uint i = tid; i < args.width; i += ntg) { + float g = gate_row[i]; + float u = up_row[i]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + if (args.write_clamped != 0) { + gate_row[i] = g; + up_row[i] = u; + } + } + const float silu = g / (1.0f + exp(-g)); + mid_row[i] = silu * u * route_weight; + } +} + +// Same routed-MoE activation as above, but stores the down-projection input in +// half precision. The grouped matmul path converts F32 activations to half +// before MMA anyway, so this cuts the large mid write/read traffic without +// changing the effective matmul input precision. +kernel void kernel_dsv4_moe_swiglu_weight_f16( + constant ds4_metal_dsv4_moe_swiglu_weight_args &args, + device char *gate, + device char *up, + device char *mid, + device const char *weights, + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint ntg [[threads_per_threadgroup]]) { + if (row >= args.rows) return; + + device float *gate_row = (device float *)(gate + (uint64_t)row * args.gate_row_stride); + device float *up_row = (device float *)(up + (uint64_t)row * args.up_row_stride); + device half *mid_row = (device half *)(mid + (uint64_t)row * args.mid_row_stride); + device const float *w = (device const float *)(weights + (uint64_t)row * args.weight_stride); + const float route_weight = w[0]; + const float c = args.clamp_value; + + for (uint i = tid; i < args.width; i += ntg) { + float g = gate_row[i]; + float u = up_row[i]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + if (args.write_clamped != 0) { + gate_row[i] = g; + up_row[i] = u; + } + } + const float silu = g / (1.0f + exp(-g)); + mid_row[i] = (half)(silu * u * route_weight); + } +} + +kernel void kernel_dsv4_moe_sum6_f32( + constant ds4_metal_dsv4_moe_sum6_args &args, + device const char *src, + device char *dst, + uint token[[threadgroup_position_in_grid]], + uint tid[[thread_position_in_threadgroup]], + uint ntg[[threads_per_threadgroup]]) { + if (token >= args.tokens) return; + + device const float *s = + (device const float *)(src + (uint64_t)token * args.src_token_stride); + device float *d = + (device float *)(dst + (uint64_t)token * args.dst_token_stride); + + for (uint col = tid; col < args.width; col += ntg) { + float v = s[col]; + v += s[args.width + col]; + v += s[2u * args.width + col]; + v += s[3u * args.width + col]; + v += s[4u * args.width + col]; + v += s[5u * args.width + col]; + d[col] = v; + } +} + +kernel void kernel_dsv4_moe_sum8_f32( + constant ds4_metal_dsv4_moe_sum6_args &args, + device const char *src, + device char *dst, + uint token[[threadgroup_position_in_grid]], + uint tid[[thread_position_in_threadgroup]], + uint ntg[[threads_per_threadgroup]]) { + if (token >= args.tokens) return; + + device const float *s = + (device const float *)(src + (uint64_t)token * args.src_token_stride); + device float *d = + (device float *)(dst + (uint64_t)token * args.dst_token_stride); + + for (uint col = tid; col < args.width; col += ntg) { + float v = s[col]; + v += s[args.width + col]; + v += s[2u * args.width + col]; + v += s[3u * args.width + col]; + v += s[4u * args.width + col]; + v += s[5u * args.width + col]; + v += s[6u * args.width + col]; + v += s[7u * args.width + col]; + d[col] = v; + } +} + +template +void dequantize_q2_K(device const block_q2_K *xb, short il, thread type4x4 & reg) { + const float d = xb->d; + const float min = xb->dmin; + device const uint8_t * q = (device const uint8_t *)xb->qs; + float dl, ml; + uint8_t sc = xb->scales[il]; + + q = q + 32*(il/8) + 16*(il&1); + il = (il/2)%4; + + half coef = il>1 ? (il>2 ? 1/64.h : 1/16.h) : (il>0 ? 1/4.h : 1.h); + uchar mask = il>1 ? (il>2 ? 192 : 48) : (il>0 ? 12 : 3); + dl = d * (sc & 0xF) * coef, ml = min * (sc >> 4); + for (int i = 0; i < 16; ++i) { + reg[i/4][i%4] = dl * (q[i] & mask) - ml; + } +} + +static inline float ds4_glm_q2_K_value(device const block_q2_K *blocks, uint k) { + const uint block = k / QK_K; + const uint idx = k - block * QK_K; + device const block_q2_K *xb = blocks + block; + const uint group = idx / 16u; + const uint l = idx - group * 16u; + const uint q_base = 32u * (group / 8u) + 16u * (group & 1u); + const uint shift = ((group / 2u) & 3u) * 2u; + const uint q = ((uint)xb->qs[q_base + l] >> shift) & 0x03u; + const uint sc = (uint)xb->scales[group]; + return (float)xb->d * (float)(sc & 0x0fu) * (float)q - + (float)xb->dmin * (float)(sc >> 4u); +} + +static inline uchar2 get_scale_min_k4_just2(int j, int k, device const uchar * q) { + return j < 4 ? uchar2{uchar(q[j+0+k] & 63), uchar(q[j+4+k] & 63)} + : uchar2{uchar((q[j+4+k] & 0xF) | ((q[j-4+k] & 0xc0) >> 2)), + uchar((q[j+4+k] >> 4) | ((q[j-0+k] & 0xc0) >> 2))}; +} + +static inline float ds4_glm_q4_K_value(device const block_q4_K *blocks, uint k) { + const uint block = k / QK_K; + const uint idx = k - block * QK_K; + device const block_q4_K *xb = blocks + block; + const uint group = idx / 32u; + const uint l = idx - group * 32u; + const uchar2 sm = get_scale_min_k4_just2((int)group, 0, xb->scales); + const uint byte_off = (group >> 1u) * 32u + l; + const uint shift = (group & 1u) * 4u; + const uint q = (xb->qs[byte_off] >> shift) & 0x0Fu; + return (float)xb->d * (float)sm.x * (float)q - + (float)xb->dmin * (float)sm.y; +} + +static inline float ds4_glm_q5_K_value(device const block_q5_K *blocks, uint k) { + const uint block = k / QK_K; + const uint idx = k - block * QK_K; + device const block_q5_K *xb = blocks + block; + const uint group = idx / 32u; + const uint l = idx - group * 32u; + const uchar2 sm = get_scale_min_k4_just2((int)group, 0, xb->scales); + const uint ql_base = (group >> 1u) * 32u + l; + const uint shift = (group & 1u) * 4u; + uint q = (xb->qs[ql_base] >> shift) & 0x0Fu; + q += (xb->qh[l] & (uchar)(1u << group)) ? 16u : 0u; + return (float)xb->d * (float)sm.x * (float)q - + (float)xb->dmin * (float)sm.y; +} + +static inline float ds4_glm_q6_K_value(device const block_q6_K *blocks, uint k) { + const uint block = k / QK_K; + const uint idx = k - block * QK_K; + device const block_q6_K *xb = blocks + block; + const uint n128 = idx >> 7u; + const uint r = idx & 127u; + const uint l = r & 31u; + const uint quarter = r >> 5u; + const uint ql_base = n128 * 64u; + const uint qh_base = n128 * 32u; + const uint sc_base = n128 * 8u; + uint q; + int sc; + + if (quarter == 0u) { + q = (xb->ql[ql_base + l] & 0x0Fu) | (((xb->qh[qh_base + l] >> 0u) & 3u) << 4u); + sc = (int)xb->scales[sc_base + l / 16u + 0u]; + } else if (quarter == 1u) { + q = (xb->ql[ql_base + 32u + l] & 0x0Fu) | (((xb->qh[qh_base + l] >> 2u) & 3u) << 4u); + sc = (int)xb->scales[sc_base + l / 16u + 2u]; + } else if (quarter == 2u) { + q = (xb->ql[ql_base + l] >> 4u) | (((xb->qh[qh_base + l] >> 4u) & 3u) << 4u); + sc = (int)xb->scales[sc_base + l / 16u + 4u]; + } else { + q = (xb->ql[ql_base + 32u + l] >> 4u) | (((xb->qh[qh_base + l] >> 6u) & 3u) << 4u); + sc = (int)xb->scales[sc_base + l / 16u + 6u]; + } + + return (float)xb->d * (float)sc * (float)((int)q - 32); +} + +kernel void kernel_glm_q4_K_pair_swiglu_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate, + device const char *up, + device const float *x, + device const int32_t *selected, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + const uint ntg = 256u; + const uint row = tgpig.x; + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (row >= args.mid_dim || slot >= args.n_expert_used || token >= args.n_tokens) return; + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + const uint64_t mid_off = (uint64_t)token * args.mid_token_stride + + (uint64_t)slot * args.mid_dim + row; + const int expert = selected[selected_off]; + if (expert < 0 || (uint)expert >= args.n_total_expert) { + if (tid == 0u) mid[mid_off] = 0.0f; + return; + } + + device const block_q4_K *gate_row = + (device const block_q4_K *)(gate + + (uint64_t)(uint)expert * args.gate_expert_bytes + + (uint64_t)row * args.gate_row_bytes); + device const block_q4_K *up_row = + (device const block_q4_K *)(up + + (uint64_t)(uint)expert * args.up_expert_bytes + + (uint64_t)row * args.up_row_bytes); + + float acc_gate = 0.0f; + float acc_up = 0.0f; + device const float *token_x = x + (uint64_t)token * args.in_dim; + for (uint k = tid; k < args.in_dim; k += ntg) { + const float xv = token_x[k]; + acc_gate += ds4_glm_q4_K_value(gate_row, k) * xv; + acc_up += ds4_glm_q4_K_value(up_row, k) * xv; + } + + scratch[tid] = acc_gate; + scratch[ntg + tid] = acc_up; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = ntg >> 1u; stride > 0u; stride >>= 1u) { + if (tid < stride) { + scratch[tid] += scratch[tid + stride]; + scratch[ntg + tid] += scratch[ntg + tid + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + if (tid == 0u) { + const float g = scratch[0]; + const float u = scratch[ntg]; + const float sw = g / (1.0f + exp(-g)); + mid[mid_off] = sw * u * weights[selected_off]; + } +} + +template +static inline void glm_q2_K_pair_swiglu_simd_f32_impl( + ds4_metal_glm_routed_moe_args args, + device const char *gate, + device const char *up, + device const float *x, + device const float *weights, + device float *mid, + threadgroup float *scratch, + uint3 tgpig, + uint slot, + uint token, + uint64_t selected_off, + int expert, + ushort tiisg, + ushort sgitg) { + const short NSG = 2; + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0; + if (row0 >= args.mid_dim || slot >= args.n_expert_used || token >= args.n_tokens) return; + + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride + + (uint64_t)slot * args.mid_dim; + if (expert < 0 || (uint)expert >= args.n_total_expert) { + if (tiisg == 0u) { + for (short row = 0; + row < N_R0 && row0 + (uint)row < args.mid_dim; + row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const short is = (8 * ir) / 16; + const int nb = args.in_dim / QK_K; + const uint64_t expert_gate = (uint64_t)(uint)expert * args.gate_expert_bytes; + const uint64_t expert_up = (uint64_t)(uint)expert * args.up_expert_bytes; + device const block_q2_K *xg = + (device const block_q2_K *)(gate + expert_gate + + (uint64_t)row0 * args.gate_row_bytes); + device const block_q2_K *xu = + (device const block_q2_K *)(up + expert_up + + (uint64_t)row0 * args.up_row_bytes); + device const float *y = x + (uint64_t)token * args.in_dim; + device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir; + + float yl[32]; + float sumg[N_R0] = {0.f}; + float sumu[N_R0] = {0.f}; + + for (int ib = ix; ib < nb; ib += 4) { + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16]; + yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24]; + } + + device const uint8_t *scg = (device const uint8_t *)xg[ib].scales + 8 * iq + is; + device const uint16_t *qsg = (device const uint16_t *)xg[ib].qs + 16 * iq + 4 * ir; + device const half *dhg = &xg[ib].d; + device const uint8_t *scu = (device const uint8_t *)xu[ib].scales + 8 * iq + is; + device const uint16_t *qsu = (device const uint16_t *)xu[ib].qs + 16 * iq + 4 * ir; + device const half *dhu = &xu[ib].d; + + for (short row = 0; + row < N_R0 && row0 + (uint)row < args.mid_dim; + row++) { + float4 acc1g = {0.f, 0.f, 0.f, 0.f}; + float4 acc2g = {0.f, 0.f, 0.f, 0.f}; + float4 acc1u = {0.f, 0.f, 0.f, 0.f}; + float4 acc2u = {0.f, 0.f, 0.f, 0.f}; + + for (int i = 0; i < 8; i += 2) { + acc1g[0] += yl[i + 0] * (qsg[i / 2] & 0x0003); + acc2g[0] += yl[i + 1] * (qsg[i / 2] & 0x0300); + acc1g[1] += yl[i + 8] * (qsg[i / 2] & 0x000c); + acc2g[1] += yl[i + 9] * (qsg[i / 2] & 0x0c00); + acc1g[2] += yl[i + 16] * (qsg[i / 2] & 0x0030); + acc2g[2] += yl[i + 17] * (qsg[i / 2] & 0x3000); + acc1g[3] += yl[i + 24] * (qsg[i / 2] & 0x00c0); + acc2g[3] += yl[i + 25] * (qsg[i / 2] & 0xc000); + + acc1u[0] += yl[i + 0] * (qsu[i / 2] & 0x0003); + acc2u[0] += yl[i + 1] * (qsu[i / 2] & 0x0300); + acc1u[1] += yl[i + 8] * (qsu[i / 2] & 0x000c); + acc2u[1] += yl[i + 9] * (qsu[i / 2] & 0x0c00); + acc1u[2] += yl[i + 16] * (qsu[i / 2] & 0x0030); + acc2u[2] += yl[i + 17] * (qsu[i / 2] & 0x3000); + acc1u[3] += yl[i + 24] * (qsu[i / 2] & 0x00c0); + acc2u[3] += yl[i + 25] * (qsu[i / 2] & 0xc000); + } + + const float dg = dhg[0]; + const float mg = dhg[1] * 1.f / 16.f; + sumg[row] += dg * ((acc1g[0] + 1.f / 256.f * acc2g[0]) * (scg[0] & 0xF) * 1.f / 1.f + + (acc1g[1] + 1.f / 256.f * acc2g[1]) * (scg[2] & 0xF) * 1.f / 4.f + + (acc1g[2] + 1.f / 256.f * acc2g[2]) * (scg[4] & 0xF) * 1.f / 16.f + + (acc1g[3] + 1.f / 256.f * acc2g[3]) * (scg[6] & 0xF) * 1.f / 64.f) - + mg * (sumy[0] * (scg[0] & 0xF0) + sumy[1] * (scg[2] & 0xF0) + + sumy[2] * (scg[4] & 0xF0) + sumy[3] * (scg[6] & 0xF0)); + + const float du = dhu[0]; + const float mu = dhu[1] * 1.f / 16.f; + sumu[row] += du * ((acc1u[0] + 1.f / 256.f * acc2u[0]) * (scu[0] & 0xF) * 1.f / 1.f + + (acc1u[1] + 1.f / 256.f * acc2u[1]) * (scu[2] & 0xF) * 1.f / 4.f + + (acc1u[2] + 1.f / 256.f * acc2u[2]) * (scu[4] & 0xF) * 1.f / 16.f + + (acc1u[3] + 1.f / 256.f * acc2u[3]) * (scu[6] & 0xF) * 1.f / 64.f) - + mu * (sumy[0] * (scu[0] & 0xF0) + sumy[1] * (scu[2] & 0xF0) + + sumy[2] * (scu[4] & 0xF0) + sumy[3] * (scu[6] & 0xF0)); + + qsg += args.gate_row_bytes / 2; + scg += args.gate_row_bytes; + dhg += args.gate_row_bytes / 2; + qsu += args.up_row_bytes / 2; + scu += args.up_row_bytes; + dhu += args.up_row_bytes / 2; + } + + y4 += 4 * QK_K; + } + + for (short row = 0; + row < N_R0 && row0 + (uint)row < args.mid_dim; + row++) { + const float g = simd_sum(sumg[row]); + const float u = simd_sum(sumu[row]); + if (tiisg == 0u) { + const float sw = g / (1.0f + exp(-g)); + mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off]; + } + } + + (void)scratch; +} + +kernel void kernel_glm_q2_K_pair_swiglu_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate, + device const char *up, + device const float *x, + device const int32_t *selected, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + const int expert = selected[selected_off]; + if (!ds4_tp_owns_expert(expert, args.n_total_expert, + args.tp_rank, args.tp_world)) return; + glm_q2_K_pair_swiglu_simd_f32_impl( + args, gate, up, x, weights, mid, scratch, + tgpig, slot, token, selected_off, expert - args.tp_expert_base, + tiisg, sgitg); +} + +kernel void kernel_glm_q2_K_addr_pair_swiglu2_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const uint64_t *gate_addrs, + device const uint64_t *up_addrs, + device const float *x, + device const int32_t *selected, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + const int expert = selected[selected_off]; + const short NSG = 2; + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_GLM_Q2_PAIR2_K; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride + + (uint64_t)slot * args.mid_dim; + if (row0 >= args.mid_dim) return; + + if (expert < 0 || (uint)expert >= args.n_total_expert) { + if (tiisg == 0u) { + for (short row = 0; + row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim; + row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + const uint64_t gate_addr = gate_addrs[(uint)expert]; + const uint64_t up_addr = up_addrs[(uint)expert]; + if (gate_addr == 0 || up_addr == 0) { + if (tiisg == 0u) { + for (short row = 0; + row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim; + row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + + ds4_metal_glm_routed_moe_args local = args; + local.n_total_expert = 1; + local.gate_expert_bytes = 0; + local.up_expert_bytes = 0; + glm_q2_K_pair_swiglu_simd_f32_impl( + local, + reinterpret_cast(gate_addr), + reinterpret_cast(up_addr), + x, weights, mid, scratch, + tgpig, slot, token, selected_off, 0, tiisg, sgitg); +} + +kernel void kernel_glm_q2_K_addr_pair_swiglu2_f32_masked( + constant ds4_metal_glm_routed_moe_args &args, + constant uint32_t &active_mask, + device const uint64_t *gate_addrs, + device const uint64_t *up_addrs, + device const float *x, + device const int32_t *selected, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + if ((active_mask & (1u << slot)) == 0u) return; + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + const int expert = selected[selected_off]; + const short NSG = 2; + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_GLM_Q2_PAIR2_K; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride + + (uint64_t)slot * args.mid_dim; + if (row0 >= args.mid_dim) return; + + if (expert < 0 || (uint)expert >= args.n_total_expert) { + if (tiisg == 0u) { + for (short row = 0; + row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim; + row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + const uint64_t gate_addr = gate_addrs[(uint)expert]; + const uint64_t up_addr = up_addrs[(uint)expert]; + if (gate_addr == 0 || up_addr == 0) { + if (tiisg == 0u) { + for (short row = 0; + row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim; + row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + + ds4_metal_glm_routed_moe_args local = args; + local.n_total_expert = 1; + local.gate_expert_bytes = 0; + local.up_expert_bytes = 0; + glm_q2_K_pair_swiglu_simd_f32_impl( + local, + reinterpret_cast(gate_addr), + reinterpret_cast(up_addr), + x, weights, mid, scratch, + tgpig, slot, token, selected_off, 0, tiisg, sgitg); +} + +template +static inline void glm_q4_K_pair_swiglu_simd_f32_impl( + ds4_metal_glm_routed_moe_args args, + device const char *gate, + device const char *up, + device const float *x, + device const float *weights, + device float *mid, + threadgroup float *scratch, + uint3 tgpig, + uint slot, + uint token, + uint64_t selected_off, + int expert, + ushort tiisg, + ushort sgitg) { + const short NSG = 2; + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0; + if (row0 >= args.mid_dim || slot >= args.n_expert_used || token >= args.n_tokens) return; + + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride + + (uint64_t)slot * args.mid_dim; + if (expert < 0 || (uint)expert >= args.n_total_expert) { + if (tiisg == 0u) { + for (short row = 0; + row < N_R0 && row0 + (uint)row < args.mid_dim; + row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const int nb = args.in_dim / QK_K; + const uint64_t expert_gate = (uint64_t)(uint)expert * args.gate_expert_bytes; + const uint64_t expert_up = (uint64_t)(uint)expert * args.up_expert_bytes; + device const block_q4_K *xg = + (device const block_q4_K *)(gate + expert_gate + + (uint64_t)row0 * args.gate_row_bytes); + device const block_q4_K *xu = + (device const block_q4_K *)(up + expert_up + + (uint64_t)row0 * args.up_row_bytes); + device const float *y = x + (uint64_t)token * args.in_dim; + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + float sumg[N_R0] = {0.f}; + float sumu[N_R0] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *scg = (device const uint16_t *)xg[ib].scales + iq; + device const uint16_t *qg1 = (device const uint16_t *)xg[ib].qs + 16 * iq + 4 * ir; + device const half *dhg = &xg[ib].d; + device const uint16_t *scu = (device const uint16_t *)xu[ib].scales + iq; + device const uint16_t *qu1 = (device const uint16_t *)xu[ib].qs + 16 * iq + 4 * ir; + device const half *dhu = &xu[ib].d; + + for (short row = 0; + row < N_R0 && row0 + (uint)row < args.mid_dim; + row++) { + sc16[0] = scg[0] & kmask1; + sc16[1] = scg[2] & kmask1; + sc16[2] = ((scg[4] >> 0) & kmask2) | ((scg[0] & kmask3) >> 2); + sc16[3] = ((scg[4] >> 4) & kmask2) | ((scg[2] & kmask3) >> 2); + + device const uint16_t *qg2 = qg1 + 32; + float4 acc1g = {0.f, 0.f, 0.f, 0.f}; + float4 acc2g = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1g[0] += yl[2 * i + 0] * (qg1[i] & 0x000F); + acc1g[1] += yl[2 * i + 1] * (qg1[i] & 0x0F00); + acc1g[2] += yl[2 * i + 8] * (qg1[i] & 0x00F0); + acc1g[3] += yl[2 * i + 9] * (qg1[i] & 0xF000); + acc2g[0] += yh[2 * i + 0] * (qg2[i] & 0x000F); + acc2g[1] += yh[2 * i + 1] * (qg2[i] & 0x0F00); + acc2g[2] += yh[2 * i + 8] * (qg2[i] & 0x00F0); + acc2g[3] += yh[2 * i + 9] * (qg2[i] & 0xF000); + } + + sumg[row] += dhg[0] * ((acc1g[0] + 1.f / 256.f * acc1g[1]) * sc8[0] + + (acc1g[2] + 1.f / 256.f * acc1g[3]) * sc8[1] * 1.f / 16.f + + (acc2g[0] + 1.f / 256.f * acc2g[1]) * sc8[4] + + (acc2g[2] + 1.f / 256.f * acc2g[3]) * sc8[5] * 1.f / 16.f) - + dhg[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + sc16[0] = scu[0] & kmask1; + sc16[1] = scu[2] & kmask1; + sc16[2] = ((scu[4] >> 0) & kmask2) | ((scu[0] & kmask3) >> 2); + sc16[3] = ((scu[4] >> 4) & kmask2) | ((scu[2] & kmask3) >> 2); + + device const uint16_t *qu2 = qu1 + 32; + float4 acc1u = {0.f, 0.f, 0.f, 0.f}; + float4 acc2u = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1u[0] += yl[2 * i + 0] * (qu1[i] & 0x000F); + acc1u[1] += yl[2 * i + 1] * (qu1[i] & 0x0F00); + acc1u[2] += yl[2 * i + 8] * (qu1[i] & 0x00F0); + acc1u[3] += yl[2 * i + 9] * (qu1[i] & 0xF000); + acc2u[0] += yh[2 * i + 0] * (qu2[i] & 0x000F); + acc2u[1] += yh[2 * i + 1] * (qu2[i] & 0x0F00); + acc2u[2] += yh[2 * i + 8] * (qu2[i] & 0x00F0); + acc2u[3] += yh[2 * i + 9] * (qu2[i] & 0xF000); + } + + sumu[row] += dhu[0] * ((acc1u[0] + 1.f / 256.f * acc1u[1]) * sc8[0] + + (acc1u[2] + 1.f / 256.f * acc1u[3]) * sc8[1] * 1.f / 16.f + + (acc2u[0] + 1.f / 256.f * acc2u[1]) * sc8[4] + + (acc2u[2] + 1.f / 256.f * acc2u[3]) * sc8[5] * 1.f / 16.f) - + dhu[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + qg1 += args.gate_row_bytes / 2; + scg += args.gate_row_bytes / 2; + dhg += args.gate_row_bytes / 2; + qu1 += args.up_row_bytes / 2; + scu += args.up_row_bytes / 2; + dhu += args.up_row_bytes / 2; + } + + y4 += 4 * QK_K; + } + + for (int row = 0; + row < N_R0 && row0 + (uint)row < args.mid_dim; + ++row) { + const float g = simd_sum(sumg[row]); + const float u = simd_sum(sumu[row]); + if (tiisg == 0u) { + const float sw = g / (1.0f + exp(-g)); + mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off]; + } + } + + (void)scratch; +} + +kernel void kernel_glm_q4_K_pair_swiglu2_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate, + device const char *up, + device const float *x, + device const int32_t *selected, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + const int expert = selected[selected_off]; + glm_q4_K_pair_swiglu_simd_f32_impl( + args, gate, up, x, weights, mid, scratch, + tgpig, slot, token, selected_off, expert, tiisg, sgitg); +} + +kernel void kernel_glm_q4_K_addr_pair_swiglu_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const uint64_t *gate_addrs, + device const uint64_t *up_addrs, + device const float *x, + device const int32_t *selected, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + const int expert = selected[selected_off]; + const short NSG = 2; + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q4_K; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride + + (uint64_t)slot * args.mid_dim; + if (row0 >= args.mid_dim) return; + + if (expert < 0 || (uint)expert >= args.n_total_expert) { + if (tiisg == 0u) { + for (short row = 0; + row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim; + row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + const uint64_t gate_addr = gate_addrs[(uint)expert]; + const uint64_t up_addr = up_addrs[(uint)expert]; + if (gate_addr == 0 || up_addr == 0) { + if (tiisg == 0u) { + for (short row = 0; + row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim; + row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + + ds4_metal_glm_routed_moe_args local = args; + local.n_total_expert = 1; + local.gate_expert_bytes = 0; + local.up_expert_bytes = 0; + glm_q4_K_pair_swiglu_simd_f32_impl( + local, + reinterpret_cast(gate_addr), + reinterpret_cast(up_addr), + x, weights, mid, scratch, + tgpig, slot, token, selected_off, 0, tiisg, sgitg); +} + +kernel void kernel_glm_q4_K_addr_pair_swiglu_f32_masked( + constant ds4_metal_glm_routed_moe_args &args, + constant uint32_t &active_mask, + device const uint64_t *gate_addrs, + device const uint64_t *up_addrs, + device const float *x, + device const int32_t *selected, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + if ((active_mask & (1u << slot)) == 0u) return; + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + const int expert = selected[selected_off]; + const short NSG = 2; + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q4_K; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride + + (uint64_t)slot * args.mid_dim; + if (row0 >= args.mid_dim) return; + + if (expert < 0 || (uint)expert >= args.n_total_expert) { + if (tiisg == 0u) { + for (short row = 0; + row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim; + row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + const uint64_t gate_addr = gate_addrs[(uint)expert]; + const uint64_t up_addr = up_addrs[(uint)expert]; + if (gate_addr == 0 || up_addr == 0) { + if (tiisg == 0u) { + for (short row = 0; + row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim; + row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + + ds4_metal_glm_routed_moe_args local = args; + local.n_total_expert = 1; + local.gate_expert_bytes = 0; + local.up_expert_bytes = 0; + glm_q4_K_pair_swiglu_simd_f32_impl( + local, + reinterpret_cast(gate_addr), + reinterpret_cast(up_addr), + x, weights, mid, scratch, + tgpig, slot, token, selected_off, 0, tiisg, sgitg); +} + +kernel void kernel_glm_q4_K_slots6_pair_swiglu_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate0, + device const char *gate1, + device const char *gate2, + device const char *gate3, + device const char *gate4, + device const char *gate5, + device const char *up0, + device const char *up1, + device const char *up2, + device const char *up3, + device const char *up4, + device const char *up5, + device const float *x, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + + device const char *gate_cur = gate0; + device const char *up_cur = up0; + switch (slot) { + case 1: gate_cur = gate1; up_cur = up1; break; + case 2: gate_cur = gate2; up_cur = up2; break; + case 3: gate_cur = gate3; up_cur = up3; break; + case 4: gate_cur = gate4; up_cur = up4; break; + case 5: gate_cur = gate5; up_cur = up5; break; + default: break; + } + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + glm_q4_K_pair_swiglu_simd_f32_impl( + args, gate_cur, up_cur, x, weights, mid, scratch, + tgpig, slot, token, selected_off, 0, tiisg, sgitg); +} + +kernel void kernel_glm_q4_K_slots8_pair_swiglu_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate0, + device const char *gate1, + device const char *gate2, + device const char *gate3, + device const char *gate4, + device const char *gate5, + device const char *gate6, + device const char *gate7, + device const char *up0, + device const char *up1, + device const char *up2, + device const char *up3, + device const char *up4, + device const char *up5, + device const char *up6, + device const char *up7, + device const float *x, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + + device const char *gate_cur = gate0; + device const char *up_cur = up0; + switch (slot) { + case 1: gate_cur = gate1; up_cur = up1; break; + case 2: gate_cur = gate2; up_cur = up2; break; + case 3: gate_cur = gate3; up_cur = up3; break; + case 4: gate_cur = gate4; up_cur = up4; break; + case 5: gate_cur = gate5; up_cur = up5; break; + case 6: gate_cur = gate6; up_cur = up6; break; + case 7: gate_cur = gate7; up_cur = up7; break; + default: break; + } + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + glm_q4_K_pair_swiglu_simd_f32_impl( + args, gate_cur, up_cur, x, weights, mid, scratch, + tgpig, slot, token, selected_off, 0, tiisg, sgitg); +} + +kernel void kernel_glm_q4_K_slots6_pair_swiglu4_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate0, + device const char *gate1, + device const char *gate2, + device const char *gate3, + device const char *gate4, + device const char *gate5, + device const char *up0, + device const char *up1, + device const char *up2, + device const char *up3, + device const char *up4, + device const char *up5, + device const float *x, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + + device const char *gate_cur = gate0; + device const char *up_cur = up0; + switch (slot) { + case 1: gate_cur = gate1; up_cur = up1; break; + case 2: gate_cur = gate2; up_cur = up2; break; + case 3: gate_cur = gate3; up_cur = up3; break; + case 4: gate_cur = gate4; up_cur = up4; break; + case 5: gate_cur = gate5; up_cur = up5; break; + default: break; + } + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + glm_q4_K_pair_swiglu_simd_f32_impl( + args, gate_cur, up_cur, x, weights, mid, scratch, + tgpig, slot, token, selected_off, 0, tiisg, sgitg); +} + +kernel void kernel_glm_q4_K_slots8_pair_swiglu4_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate0, + device const char *gate1, + device const char *gate2, + device const char *gate3, + device const char *gate4, + device const char *gate5, + device const char *gate6, + device const char *gate7, + device const char *up0, + device const char *up1, + device const char *up2, + device const char *up3, + device const char *up4, + device const char *up5, + device const char *up6, + device const char *up7, + device const float *x, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + + device const char *gate_cur = gate0; + device const char *up_cur = up0; + switch (slot) { + case 1: gate_cur = gate1; up_cur = up1; break; + case 2: gate_cur = gate2; up_cur = up2; break; + case 3: gate_cur = gate3; up_cur = up3; break; + case 4: gate_cur = gate4; up_cur = up4; break; + case 5: gate_cur = gate5; up_cur = up5; break; + case 6: gate_cur = gate6; up_cur = up6; break; + case 7: gate_cur = gate7; up_cur = up7; break; + default: break; + } + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + glm_q4_K_pair_swiglu_simd_f32_impl( + args, gate_cur, up_cur, x, weights, mid, scratch, + tgpig, slot, token, selected_off, 0, tiisg, sgitg); +} + +kernel void kernel_glm_q4_K_pair_swiglu4_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate, + device const char *up, + device const float *x, + device const int32_t *selected, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + const int expert = selected[selected_off]; + glm_q4_K_pair_swiglu_simd_f32_impl( + args, gate, up, x, weights, mid, scratch, + tgpig, slot, token, selected_off, expert, tiisg, sgitg); +} + +kernel void kernel_glm_q4_K_pair_swiglu2_mapped_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate, + device const char *up, + device const float *x, + device const uint32_t *htpe, + device const int32_t *hids, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint expert = tgpig.z; + if (expert >= args.n_total_expert) return; + const uint count = htpe[expert]; + const uint map_base = tgpig.y * 32u; + for (uint i = 0; i < 32u; i++) { + const uint map_row = map_base + i; + if (map_row >= count) break; + const int id = hids[(uint64_t)expert * args.n_tokens + map_row]; + if (id < 0) continue; + const uint token = (uint)id / args.n_expert_used; + const uint slot = (uint)id - token * args.n_expert_used; + if (slot >= args.n_expert_used || token >= args.n_tokens) continue; + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + glm_q4_K_pair_swiglu_simd_f32_impl( + args, gate, up, x, weights, mid, scratch, + tgpig, slot, token, selected_off, (int)expert, tiisg, sgitg); + } +} + +kernel void kernel_glm_q4_K_pair_swiglu2_mapped_row_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate, + device const char *up, + device const float *x, + device const uint32_t *htpe, + device const int32_t *hids, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint expert = tgpig.z; + const uint map_row = tgpig.y; + if (expert >= args.n_total_expert || map_row >= htpe[expert]) return; + const int id = hids[(uint64_t)expert * args.n_tokens + map_row]; + if (id < 0) return; + const uint token = (uint)id / args.n_expert_used; + const uint slot = (uint)id - token * args.n_expert_used; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + glm_q4_K_pair_swiglu_simd_f32_impl( + args, gate, up, x, weights, mid, scratch, + tgpig, slot, token, selected_off, (int)expert, tiisg, sgitg); +} + +static inline void glm_q5_K_pair_swiglu_f32_impl( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate, + device const char *up, + device const float *x, + device const float *weights, + device float *mid, + threadgroup float *scratch, + uint3 tgpig, + uint slot, + uint token, + uint64_t selected_off, + int expert, + ushort tiisg, + ushort sgitg) { + const short NSG = 2; + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q5_PAIR_K; + if (row0 >= args.mid_dim || slot >= args.n_expert_used || token >= args.n_tokens) return; + + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride + + (uint64_t)slot * args.mid_dim; + if (expert < 0 || (uint)expert >= args.n_total_expert) { + if (tiisg == 0u) { + for (short row = 0; row < N_R0_Q5_PAIR_K && row0 + (uint)row < args.mid_dim; row++) { + mid[mid_base + row0 + (uint)row] = 0.0f; + } + } + return; + } + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const uint bit0 = 2u * (uint)iq; + const uint bit1 = bit0 + 1u; + const uint bit2 = bit0 + 4u; + const uint bit3 = bit0 + 5u; + const int nb = args.in_dim / QK_K; + + const uint64_t expert_gate = (uint64_t)(uint)expert * args.gate_expert_bytes; + const uint64_t expert_up = (uint64_t)(uint)expert * args.up_expert_bytes; + device const block_q5_K *gate_rows = + (device const block_q5_K *)(gate + + expert_gate + (uint64_t)row0 * args.gate_row_bytes); + device const block_q5_K *up_rows = + (device const block_q5_K *)(up + + expert_up + (uint64_t)row0 * args.up_row_bytes); + device const float *y = x + (uint64_t)token * args.in_dim; + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + float sumg[N_R0_Q5_PAIR_K] = {0.f}; + float sumu[N_R0_Q5_PAIR_K] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *scg = (device const uint16_t *)gate_rows[ib].scales + iq; + device const uint16_t *qg1 = (device const uint16_t *)gate_rows[ib].qs + 16 * iq + 4 * ir; + device const uint16_t *qhg = (device const uint16_t *)gate_rows[ib].qh + 4 * ir; + device const half *dhg = &gate_rows[ib].d; + device const uint16_t *scu = (device const uint16_t *)up_rows[ib].scales + iq; + device const uint16_t *qu1 = (device const uint16_t *)up_rows[ib].qs + 16 * iq + 4 * ir; + device const uint16_t *qhu = (device const uint16_t *)up_rows[ib].qh + 4 * ir; + device const half *dhu = &up_rows[ib].d; + + for (short row = 0; row < N_R0_Q5_PAIR_K && row0 + (uint)row < args.mid_dim; row++) { + sc16[0] = scg[0] & kmask1; + sc16[1] = scg[2] & kmask1; + sc16[2] = ((scg[4] >> 0) & kmask2) | ((scg[0] & kmask3) >> 2); + sc16[3] = ((scg[4] >> 4) & kmask2) | ((scg[2] & kmask3) >> 2); + + device const uint16_t *qg2 = qg1 + 32; + float4 accg = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + const uint ql1 = (uint)qg1[i]; + const uint ql2 = (uint)qg2[i]; + const uint hb = (uint)qhg[i]; + accg[0] += yl[2 * i + 0] * + (float)((ql1 & 0x000Fu) + (((hb >> bit0) & 1u) << 4u)) + + yl[2 * i + 1] * + (float)(((ql1 >> 8u) & 0x000Fu) + (((hb >> (bit0 + 8u)) & 1u) << 4u)); + accg[1] += yl[2 * i + 8] * + (float)(((ql1 >> 4u) & 0x000Fu) + (((hb >> bit1) & 1u) << 4u)) + + yl[2 * i + 9] * + (float)(((ql1 >> 12u) & 0x000Fu) + (((hb >> (bit1 + 8u)) & 1u) << 4u)); + accg[2] += yh[2 * i + 0] * + (float)((ql2 & 0x000Fu) + (((hb >> bit2) & 1u) << 4u)) + + yh[2 * i + 1] * + (float)(((ql2 >> 8u) & 0x000Fu) + (((hb >> (bit2 + 8u)) & 1u) << 4u)); + accg[3] += yh[2 * i + 8] * + (float)(((ql2 >> 4u) & 0x000Fu) + (((hb >> bit3) & 1u) << 4u)) + + yh[2 * i + 9] * + (float)(((ql2 >> 12u) & 0x000Fu) + (((hb >> (bit3 + 8u)) & 1u) << 4u)); + } + + sumg[row] += dhg[0] * (accg[0] * sc8[0] + accg[1] * sc8[1] + + accg[2] * sc8[4] + accg[3] * sc8[5]) - + dhg[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + sc16[0] = scu[0] & kmask1; + sc16[1] = scu[2] & kmask1; + sc16[2] = ((scu[4] >> 0) & kmask2) | ((scu[0] & kmask3) >> 2); + sc16[3] = ((scu[4] >> 4) & kmask2) | ((scu[2] & kmask3) >> 2); + + device const uint16_t *qu2 = qu1 + 32; + float4 accu = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + const uint ql1 = (uint)qu1[i]; + const uint ql2 = (uint)qu2[i]; + const uint hb = (uint)qhu[i]; + accu[0] += yl[2 * i + 0] * + (float)((ql1 & 0x000Fu) + (((hb >> bit0) & 1u) << 4u)) + + yl[2 * i + 1] * + (float)(((ql1 >> 8u) & 0x000Fu) + (((hb >> (bit0 + 8u)) & 1u) << 4u)); + accu[1] += yl[2 * i + 8] * + (float)(((ql1 >> 4u) & 0x000Fu) + (((hb >> bit1) & 1u) << 4u)) + + yl[2 * i + 9] * + (float)(((ql1 >> 12u) & 0x000Fu) + (((hb >> (bit1 + 8u)) & 1u) << 4u)); + accu[2] += yh[2 * i + 0] * + (float)((ql2 & 0x000Fu) + (((hb >> bit2) & 1u) << 4u)) + + yh[2 * i + 1] * + (float)(((ql2 >> 8u) & 0x000Fu) + (((hb >> (bit2 + 8u)) & 1u) << 4u)); + accu[3] += yh[2 * i + 8] * + (float)(((ql2 >> 4u) & 0x000Fu) + (((hb >> bit3) & 1u) << 4u)) + + yh[2 * i + 9] * + (float)(((ql2 >> 12u) & 0x000Fu) + (((hb >> (bit3 + 8u)) & 1u) << 4u)); + } + + sumu[row] += dhu[0] * (accu[0] * sc8[0] + accu[1] * sc8[1] + + accu[2] * sc8[4] + accu[3] * sc8[5]) - + dhu[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + qg1 += args.gate_row_bytes / 2; + qhg += args.gate_row_bytes / 2; + scg += args.gate_row_bytes / 2; + dhg += args.gate_row_bytes / 2; + qu1 += args.up_row_bytes / 2; + qhu += args.up_row_bytes / 2; + scu += args.up_row_bytes / 2; + dhu += args.up_row_bytes / 2; + } + + y4 += 4 * QK_K; + } + + for (short row = 0; row < N_R0_Q5_PAIR_K && row0 + (uint)row < args.mid_dim; row++) { + const float g = simd_sum(sumg[row]); + const float u = simd_sum(sumu[row]); + if (tiisg == 0u) { + const float sw = g / (1.0f + exp(-g)); + mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off]; + } + } + + (void)scratch; +} + +kernel void kernel_glm_q5_K_pair_swiglu_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate, + device const char *up, + device const float *x, + device const int32_t *selected, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + const int expert = selected[selected_off]; + glm_q5_K_pair_swiglu_f32_impl( + args, gate, up, x, weights, mid, scratch, + tgpig, slot, token, selected_off, expert, tiisg, sgitg); +} + +kernel void kernel_glm_q5_K_slots6_pair_swiglu_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate0, + device const char *gate1, + device const char *gate2, + device const char *gate3, + device const char *gate4, + device const char *gate5, + device const char *up0, + device const char *up1, + device const char *up2, + device const char *up3, + device const char *up4, + device const char *up5, + device const float *x, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + + device const char *gate_cur = gate0; + device const char *up_cur = up0; + switch (slot) { + case 1: gate_cur = gate1; up_cur = up1; break; + case 2: gate_cur = gate2; up_cur = up2; break; + case 3: gate_cur = gate3; up_cur = up3; break; + case 4: gate_cur = gate4; up_cur = up4; break; + case 5: gate_cur = gate5; up_cur = up5; break; + default: break; + } + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + glm_q5_K_pair_swiglu_f32_impl( + args, gate_cur, up_cur, x, weights, mid, scratch, + tgpig, slot, token, selected_off, 0, tiisg, sgitg); +} + +kernel void kernel_glm_q5_K_slots8_pair_swiglu_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate0, + device const char *gate1, + device const char *gate2, + device const char *gate3, + device const char *gate4, + device const char *gate5, + device const char *gate6, + device const char *gate7, + device const char *up0, + device const char *up1, + device const char *up2, + device const char *up3, + device const char *up4, + device const char *up5, + device const char *up6, + device const char *up7, + device const float *x, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint slot = tgpig.y; + const uint token = tgpig.z; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + + device const char *gate_cur = gate0; + device const char *up_cur = up0; + switch (slot) { + case 1: gate_cur = gate1; up_cur = up1; break; + case 2: gate_cur = gate2; up_cur = up2; break; + case 3: gate_cur = gate3; up_cur = up3; break; + case 4: gate_cur = gate4; up_cur = up4; break; + case 5: gate_cur = gate5; up_cur = up5; break; + case 6: gate_cur = gate6; up_cur = up6; break; + case 7: gate_cur = gate7; up_cur = up7; break; + default: break; + } + + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + glm_q5_K_pair_swiglu_f32_impl( + args, gate_cur, up_cur, x, weights, mid, scratch, + tgpig, slot, token, selected_off, 0, tiisg, sgitg); +} + +kernel void kernel_glm_q5_K_pair_swiglu_mapped_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate, + device const char *up, + device const float *x, + device const uint32_t *htpe, + device const int32_t *hids, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint expert = tgpig.z; + if (expert >= args.n_total_expert) return; + const uint count = htpe[expert]; + const uint map_base = tgpig.y * 32u; + for (uint i = 0; i < 32u; i++) { + const uint map_row = map_base + i; + if (map_row >= count) break; + const int id = hids[(uint64_t)expert * args.n_tokens + map_row]; + if (id < 0) continue; + const uint token = (uint)id / args.n_expert_used; + const uint slot = (uint)id - token * args.n_expert_used; + if (slot >= args.n_expert_used || token >= args.n_tokens) continue; + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + glm_q5_K_pair_swiglu_f32_impl( + args, gate, up, x, weights, mid, scratch, + tgpig, slot, token, selected_off, (int)expert, tiisg, sgitg); + } +} + +kernel void kernel_glm_q5_K_pair_swiglu_mapped_row_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *gate, + device const char *up, + device const float *x, + device const uint32_t *htpe, + device const int32_t *hids, + device const float *weights, + device float *mid, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const uint expert = tgpig.z; + const uint map_row = tgpig.y; + if (expert >= args.n_total_expert || map_row >= htpe[expert]) return; + const int id = hids[(uint64_t)expert * args.n_tokens + map_row]; + if (id < 0) return; + const uint token = (uint)id / args.n_expert_used; + const uint slot = (uint)id - token * args.n_expert_used; + if (slot >= args.n_expert_used || token >= args.n_tokens) return; + const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot; + glm_q5_K_pair_swiglu_f32_impl( + args, gate, up, x, weights, mid, scratch, + tgpig, slot, token, selected_off, (int)expert, tiisg, sgitg); +} + +kernel void kernel_glm_q5_K_down_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *down, + device const int32_t *selected, + device const float *mid, + device float *out, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = 2; + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q5_K; + const uint token = tgpig.y; + if (row0 >= args.out_dim || token >= args.n_tokens) return; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const uint bit0 = 2u * (uint)iq; + const uint bit1 = bit0 + 1u; + const uint bit2 = bit0 + 4u; + const uint bit3 = bit0 + 5u; + const int nb = args.mid_dim / QK_K; + + float sumf[N_R0_Q5_K] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + const uint64_t selected_base = (uint64_t)token * args.n_expert_used; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + const int expert = selected[selected_base + slot]; + if (expert < 0 || (uint)expert >= args.n_total_expert) continue; + device const block_q5_K *x = + (device const block_q5_K *)(down + + (uint64_t)(uint)expert * args.down_expert_bytes + + (uint64_t)row0 * args.down_row_bytes); + device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim; + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const uint16_t *qh = (device const uint16_t *)x[ib].qh + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + float4 acc = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + const uint ql1 = (uint)q1[i]; + const uint ql2 = (uint)q2[i]; + const uint hb = (uint)qh[i]; + acc[0] += yl[2 * i + 0] * + (float)((ql1 & 0x000Fu) + (((hb >> bit0) & 1u) << 4u)) + + yl[2 * i + 1] * + (float)(((ql1 >> 8u) & 0x000Fu) + (((hb >> (bit0 + 8u)) & 1u) << 4u)); + acc[1] += yl[2 * i + 8] * + (float)(((ql1 >> 4u) & 0x000Fu) + (((hb >> bit1) & 1u) << 4u)) + + yl[2 * i + 9] * + (float)(((ql1 >> 12u) & 0x000Fu) + (((hb >> (bit1 + 8u)) & 1u) << 4u)); + acc[2] += yh[2 * i + 0] * + (float)((ql2 & 0x000Fu) + (((hb >> bit2) & 1u) << 4u)) + + yh[2 * i + 1] * + (float)(((ql2 >> 8u) & 0x000Fu) + (((hb >> (bit2 + 8u)) & 1u) << 4u)); + acc[3] += yh[2 * i + 8] * + (float)(((ql2 >> 4u) & 0x000Fu) + (((hb >> bit3) & 1u) << 4u)) + + yh[2 * i + 9] * + (float)(((ql2 >> 12u) & 0x000Fu) + (((hb >> (bit3 + 8u)) & 1u) << 4u)); + } + + sumf[row] += dh[0] * (acc[0] * sc8[0] + acc[1] * sc8[1] + + acc[2] * sc8[4] + acc[3] * sc8[5]) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + q1 += args.down_row_bytes / 2; + qh += args.down_row_bytes / 2; + sc += args.down_row_bytes / 2; + dh += args.down_row_bytes / 2; + } + + y4 += 4 * QK_K; + } + } + + for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0u) { + out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; + } + } +} + +kernel void kernel_glm_q2_K_down_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *down, + device const int32_t *selected, + device const float *mid, + device float *out, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = 2; + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q2_K; + const uint token = tgpig.y; + if (row0 >= args.out_dim || token >= args.n_tokens) return; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const short is = (8 * ir) / 16; + const int nb = args.mid_dim / QK_K; + float sumf[N_R0_Q2_K] = {0.f}; + const uint64_t selected_base = (uint64_t)token * args.n_expert_used; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + const int expert = selected[selected_base + slot]; + if (expert < 0 || (uint)expert >= args.n_total_expert) continue; + if (!ds4_tp_owns_expert(expert, args.n_total_expert, + args.tp_rank, args.tp_world)) continue; + device const block_q2_K *x = + (device const block_q2_K *)(down + + (uint64_t)(uint)(expert - args.tp_expert_base) * args.down_expert_bytes + + (uint64_t)row0 * args.down_row_bytes); + device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim; + device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[32]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16]; + yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24]; + } + + device const uint8_t *sc = (device const uint8_t *)x[ib].scales + 8 * iq + is; + device const uint16_t *qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + for (short row = 0; row < N_R0_Q2_K && row0 + (uint)row < args.out_dim; row++) { + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + for (int i = 0; i < 8; i += 2) { + acc1[0] += yl[i + 0] * (qs[i / 2] & 0x0003); + acc2[0] += yl[i + 1] * (qs[i / 2] & 0x0300); + acc1[1] += yl[i + 8] * (qs[i / 2] & 0x000c); + acc2[1] += yl[i + 9] * (qs[i / 2] & 0x0c00); + acc1[2] += yl[i + 16] * (qs[i / 2] & 0x0030); + acc2[2] += yl[i + 17] * (qs[i / 2] & 0x3000); + acc1[3] += yl[i + 24] * (qs[i / 2] & 0x00c0); + acc2[3] += yl[i + 25] * (qs[i / 2] & 0xc000); + } + + const float d = dh[0]; + const float m = dh[1] * 1.f / 16.f; + sumf[row] += d * ((acc1[0] + 1.f / 256.f * acc2[0]) * (sc[0] & 0xF) * 1.f / 1.f + + (acc1[1] + 1.f / 256.f * acc2[1]) * (sc[2] & 0xF) * 1.f / 4.f + + (acc1[2] + 1.f / 256.f * acc2[2]) * (sc[4] & 0xF) * 1.f / 16.f + + (acc1[3] + 1.f / 256.f * acc2[3]) * (sc[6] & 0xF) * 1.f / 64.f) - + m * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + + sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0)); + + qs += args.down_row_bytes / 2; + sc += args.down_row_bytes; + dh += args.down_row_bytes / 2; + } + + y4 += 4 * QK_K; + } + } + + for (short row = 0; row < N_R0_Q2_K && row0 + (uint)row < args.out_dim; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0u) { + out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; + } + } + + (void)scratch; +} + +kernel void kernel_glm_q2_K_addr_down_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const uint64_t *down_addrs, + device const int32_t *selected, + device const float *mid, + device float *out, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = 2; + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q2_K; + const uint token = tgpig.y; + if (row0 >= args.out_dim || token >= args.n_tokens) return; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const short is = (8 * ir) / 16; + const int nb = args.mid_dim / QK_K; + float sumf[N_R0_Q2_K] = {0.f}; + const uint64_t selected_base = (uint64_t)token * args.n_expert_used; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + const int expert = selected[selected_base + slot]; + if (expert < 0 || (uint)expert >= args.n_total_expert) continue; + const uint64_t down_addr = down_addrs[(uint)expert]; + if (down_addr == 0) continue; + device const block_q2_K *x = + (device const block_q2_K *)(reinterpret_cast(down_addr) + + (uint64_t)row0 * args.down_row_bytes); + device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim; + device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[32]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16]; + yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24]; + } + + device const uint8_t *sc = (device const uint8_t *)x[ib].scales + 8 * iq + is; + device const uint16_t *qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + for (short row = 0; row < N_R0_Q2_K && row0 + (uint)row < args.out_dim; row++) { + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + for (int i = 0; i < 8; i += 2) { + acc1[0] += yl[i + 0] * (qs[i / 2] & 0x0003); + acc2[0] += yl[i + 1] * (qs[i / 2] & 0x0300); + acc1[1] += yl[i + 8] * (qs[i / 2] & 0x000c); + acc2[1] += yl[i + 9] * (qs[i / 2] & 0x0c00); + acc1[2] += yl[i + 16] * (qs[i / 2] & 0x0030); + acc2[2] += yl[i + 17] * (qs[i / 2] & 0x3000); + acc1[3] += yl[i + 24] * (qs[i / 2] & 0x00c0); + acc2[3] += yl[i + 25] * (qs[i / 2] & 0xc000); + } + + const float d = dh[0]; + const float m = dh[1] * 1.f / 16.f; + sumf[row] += d * ((acc1[0] + 1.f / 256.f * acc2[0]) * (sc[0] & 0xF) * 1.f / 1.f + + (acc1[1] + 1.f / 256.f * acc2[1]) * (sc[2] & 0xF) * 1.f / 4.f + + (acc1[2] + 1.f / 256.f * acc2[2]) * (sc[4] & 0xF) * 1.f / 16.f + + (acc1[3] + 1.f / 256.f * acc2[3]) * (sc[6] & 0xF) * 1.f / 64.f) - + m * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + + sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0)); + + qs += args.down_row_bytes / 2; + sc += args.down_row_bytes; + dh += args.down_row_bytes / 2; + } + + y4 += 4 * QK_K; + } + } + + for (short row = 0; row < N_R0_Q2_K && row0 + (uint)row < args.out_dim; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0u) { + out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; + } + } + + (void)scratch; +} + +kernel void kernel_glm_q4_K_down_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *down, + device const int32_t *selected, + device const float *mid, + device float *out, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + const uint ntg = 256u; + const uint row = tgpig.x; + const uint token = tgpig.y; + if (row >= args.out_dim || token >= args.n_tokens) return; + + float acc = 0.0f; + const uint64_t selected_base = (uint64_t)token * args.n_expert_used; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + const int expert = selected[selected_base + slot]; + if (expert < 0 || (uint)expert >= args.n_total_expert) continue; + device const block_q4_K *down_row = + (device const block_q4_K *)(down + + (uint64_t)(uint)expert * args.down_expert_bytes + + (uint64_t)row * args.down_row_bytes); + device const float *slot_mid = mid + mid_base + (uint64_t)slot * args.mid_dim; + for (uint k = tid; k < args.mid_dim; k += ntg) { + acc += ds4_glm_q4_K_value(down_row, k) * slot_mid[k]; + } + } + + scratch[tid] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = ntg >> 1u; stride > 0u; stride >>= 1u) { + if (tid < stride) scratch[tid] += scratch[tid + stride]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0u) { + out[(uint64_t)token * args.out_dim + row] = scratch[0]; + } +} + +kernel void kernel_glm_q4_K_addr_down_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const uint64_t *down_addrs, + device const int32_t *selected, + device const float *mid, + device float *out, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + const uint ntg = 256u; + const uint row = tgpig.x; + const uint token = tgpig.y; + if (row >= args.out_dim || token >= args.n_tokens) return; + + float acc = 0.0f; + const uint64_t selected_base = (uint64_t)token * args.n_expert_used; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + const int expert = selected[selected_base + slot]; + if (expert < 0 || (uint)expert >= args.n_total_expert) continue; + const uint64_t down_addr = down_addrs[(uint)expert]; + if (down_addr == 0) continue; + device const block_q4_K *down_row = + (device const block_q4_K *)(reinterpret_cast(down_addr) + + (uint64_t)row * args.down_row_bytes); + device const float *slot_mid = mid + mid_base + (uint64_t)slot * args.mid_dim; + for (uint k = tid; k < args.mid_dim; k += ntg) { + acc += ds4_glm_q4_K_value(down_row, k) * slot_mid[k]; + } + } + + scratch[tid] = acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = ntg >> 1u; stride > 0u; stride >>= 1u) { + if (tid < stride) scratch[tid] += scratch[tid + stride]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (tid == 0u) { + out[(uint64_t)token * args.out_dim + row] = scratch[0]; + } +} + +kernel void kernel_glm_q4_K_down_simd_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *down, + device const int32_t *selected, + device const float *mid, + device float *out, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = 2; + const short nr0 = N_R0_Q4_K; + const int nb = args.mid_dim / QK_K; + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * (uint)nr0; + const uint token = tgpig.y; + if (row0 >= args.out_dim || token >= args.n_tokens) return; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + float sumf[N_R0_Q4_K] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + const uint64_t selected_base = (uint64_t)token * args.n_expert_used; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + const int expert = selected[selected_base + slot]; + if (expert < 0 || (uint)expert >= args.n_total_expert) continue; + + device const block_q4_K *x = + (device const block_q4_K *)(down + + (uint64_t)(uint)expert * args.down_expert_bytes + + (uint64_t)row0 * args.down_row_bytes); + device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim; + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0 && row0 + (uint)row < args.out_dim; row++) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + q1 += args.down_row_bytes / 2; + sc += args.down_row_bytes / 2; + dh += args.down_row_bytes / 2; + } + + y4 += 4 * QK_K; + } + } + + for (short row = 0; row < nr0 && row0 + (uint)row < args.out_dim; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0u) { + out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; + } + } +} + +kernel void kernel_glm_q4_K_addr_down_simd_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const uint64_t *down_addrs, + device const int32_t *selected, + device const float *mid, + device float *out, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = 2; + const short nr0 = N_R0_Q4_K; + const int nb = args.mid_dim / QK_K; + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * (uint)nr0; + const uint token = tgpig.y; + if (row0 >= args.out_dim || token >= args.n_tokens) return; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + float sumf[N_R0_Q4_K] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + const uint64_t selected_base = (uint64_t)token * args.n_expert_used; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + const int expert = selected[selected_base + slot]; + if (expert < 0 || (uint)expert >= args.n_total_expert) continue; + const uint64_t down_addr = down_addrs[(uint)expert]; + if (down_addr == 0) continue; + + device const block_q4_K *x = + (device const block_q4_K *)(reinterpret_cast(down_addr) + + (uint64_t)row0 * args.down_row_bytes); + device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim; + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0 && row0 + (uint)row < args.out_dim; row++) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + q1 += args.down_row_bytes / 2; + sc += args.down_row_bytes / 2; + dh += args.down_row_bytes / 2; + } + + y4 += 4 * QK_K; + } + } + + for (short row = 0; row < nr0 && row0 + (uint)row < args.out_dim; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0u) { + out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; + } + } +} + +kernel void kernel_glm_q6_K_down_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *down, + device const int32_t *selected, + device const float *mid, + device float *out, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = 2; + constexpr uint kmask1 = 0x03u; + constexpr uint kmask2 = 0x0Cu; + constexpr uint kmask3 = 0x30u; + constexpr uint kmask4 = 0xC0u; + + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q6_K; + const uint token = tgpig.y; + if (row0 >= args.out_dim || token >= args.n_tokens) return; + + const int nb = args.mid_dim / QK_K; + float sumf[N_R0_Q6_K] = {0.f}; + float yl[16]; + const short tid = tiisg / 2; + const short ix = tiisg % 2; + const short ip = tid / 8; + const short il = tid % 8; + const short l0 = 4 * il; + const short is = 8 * ip + l0 / 16; + const short y_offset = 128 * ip + l0; + const short q_offset_l = 64 * ip + l0; + const short q_offset_h = 32 * ip + l0; + + const uint64_t selected_base = (uint64_t)token * args.n_expert_used; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + const int expert = selected[selected_base + slot]; + if (expert < 0 || (uint)expert >= args.n_total_expert) continue; + device const block_q6_K *x = + (device const block_q6_K *)(down + + (uint64_t)(uint)expert * args.down_expert_bytes + + (uint64_t)row0 * args.down_row_bytes); + device const float *yy = mid + mid_base + (uint64_t)slot * args.mid_dim; + + for (int ib = ix; ib < nb; ib += 2) { + device const uchar *q1 = x[ib].ql + q_offset_l; + device const uchar *q2 = q1 + 32; + device const uchar *qh = x[ib].qh + q_offset_h; + device const char *sc = x[ib].scales + is; + device const half *dh = &x[ib].d; + device const float *y = yy + (uint64_t)ib * QK_K + y_offset; + + for (short l = 0; l < 4; ++l) { + yl[4 * l + 0] = y[l + 0]; + yl[4 * l + 1] = y[l + 32]; + yl[4 * l + 2] = y[l + 64]; + yl[4 * l + 3] = y[l + 96]; + } + + for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) { + float4 sums = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short l = 0; l < 4; ++l) { + const uint h = (uint)qh[l]; + sums[0] += yl[4 * l + 0] * + (float)((int)((q1[l] & 0x0Fu) | ((h & kmask1) << 4u)) - 32); + sums[1] += yl[4 * l + 1] * + (float)((int)((q2[l] & 0x0Fu) | ((h & kmask2) << 2u)) - 32); + sums[2] += yl[4 * l + 2] * + (float)((int)((q1[l] >> 4u) | (h & kmask3)) - 32); + sums[3] += yl[4 * l + 3] * + (float)((int)((q2[l] >> 4u) | ((h & kmask4) >> 2u)) - 32); + } + + sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] + + sums[2] * sc[4] + sums[3] * sc[6]); + + q1 += args.down_row_bytes; + q2 += args.down_row_bytes; + qh += args.down_row_bytes; + sc += args.down_row_bytes; + dh += args.down_row_bytes / 2; + } + } + } + + for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0u) { + out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; + } + } +} + +kernel void kernel_glm_q5_K_slots6_down_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *down0, + device const char *down1, + device const char *down2, + device const char *down3, + device const char *down4, + device const char *down5, + device const float *mid, + device float *out, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = 2; + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q5_K; + const uint token = tgpig.y; + if (row0 >= args.out_dim || token >= args.n_tokens) return; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const uint bit0 = 2u * (uint)iq; + const uint bit1 = bit0 + 1u; + const uint bit2 = bit0 + 4u; + const uint bit3 = bit0 + 5u; + const int nb = args.mid_dim / QK_K; + + float sumf[N_R0_Q5_K] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + device const char *down_cur = down0; + switch (slot) { + case 1: down_cur = down1; break; + case 2: down_cur = down2; break; + case 3: down_cur = down3; break; + case 4: down_cur = down4; break; + case 5: down_cur = down5; break; + default: break; + } + + device const block_q5_K *x = + (device const block_q5_K *)(down_cur + + (uint64_t)row0 * args.down_row_bytes); + device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim; + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const uint16_t *qh = (device const uint16_t *)x[ib].qh + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + float4 acc = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + const uint ql1 = (uint)q1[i]; + const uint ql2 = (uint)q2[i]; + const uint hb = (uint)qh[i]; + acc[0] += yl[2 * i + 0] * + (float)((ql1 & 0x000Fu) + (((hb >> bit0) & 1u) << 4u)) + + yl[2 * i + 1] * + (float)(((ql1 >> 8u) & 0x000Fu) + (((hb >> (bit0 + 8u)) & 1u) << 4u)); + acc[1] += yl[2 * i + 8] * + (float)(((ql1 >> 4u) & 0x000Fu) + (((hb >> bit1) & 1u) << 4u)) + + yl[2 * i + 9] * + (float)(((ql1 >> 12u) & 0x000Fu) + (((hb >> (bit1 + 8u)) & 1u) << 4u)); + acc[2] += yh[2 * i + 0] * + (float)((ql2 & 0x000Fu) + (((hb >> bit2) & 1u) << 4u)) + + yh[2 * i + 1] * + (float)(((ql2 >> 8u) & 0x000Fu) + (((hb >> (bit2 + 8u)) & 1u) << 4u)); + acc[3] += yh[2 * i + 8] * + (float)(((ql2 >> 4u) & 0x000Fu) + (((hb >> bit3) & 1u) << 4u)) + + yh[2 * i + 9] * + (float)(((ql2 >> 12u) & 0x000Fu) + (((hb >> (bit3 + 8u)) & 1u) << 4u)); + } + + sumf[row] += dh[0] * (acc[0] * sc8[0] + acc[1] * sc8[1] + + acc[2] * sc8[4] + acc[3] * sc8[5]) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + q1 += args.down_row_bytes / 2; + qh += args.down_row_bytes / 2; + sc += args.down_row_bytes / 2; + dh += args.down_row_bytes / 2; + } + + y4 += 4 * QK_K; + } + } + + for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0u) { + out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; + } + } +} + +kernel void kernel_glm_q5_K_slots8_down_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *down0, + device const char *down1, + device const char *down2, + device const char *down3, + device const char *down4, + device const char *down5, + device const char *down6, + device const char *down7, + device const float *mid, + device float *out, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = 2; + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q5_K; + const uint token = tgpig.y; + if (row0 >= args.out_dim || token >= args.n_tokens) return; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const uint bit0 = 2u * (uint)iq; + const uint bit1 = bit0 + 1u; + const uint bit2 = bit0 + 4u; + const uint bit3 = bit0 + 5u; + const int nb = args.mid_dim / QK_K; + + float sumf[N_R0_Q5_K] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + device const char *down_cur = down0; + switch (slot) { + case 1: down_cur = down1; break; + case 2: down_cur = down2; break; + case 3: down_cur = down3; break; + case 4: down_cur = down4; break; + case 5: down_cur = down5; break; + case 6: down_cur = down6; break; + case 7: down_cur = down7; break; + default: break; + } + + device const block_q5_K *x = + (device const block_q5_K *)(down_cur + + (uint64_t)row0 * args.down_row_bytes); + device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim; + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const uint16_t *qh = (device const uint16_t *)x[ib].qh + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + float4 acc = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + const uint ql1 = (uint)q1[i]; + const uint ql2 = (uint)q2[i]; + const uint hb = (uint)qh[i]; + acc[0] += yl[2 * i + 0] * + (float)((ql1 & 0x000Fu) + (((hb >> bit0) & 1u) << 4u)) + + yl[2 * i + 1] * + (float)(((ql1 >> 8u) & 0x000Fu) + (((hb >> (bit0 + 8u)) & 1u) << 4u)); + acc[1] += yl[2 * i + 8] * + (float)(((ql1 >> 4u) & 0x000Fu) + (((hb >> bit1) & 1u) << 4u)) + + yl[2 * i + 9] * + (float)(((ql1 >> 12u) & 0x000Fu) + (((hb >> (bit1 + 8u)) & 1u) << 4u)); + acc[2] += yh[2 * i + 0] * + (float)((ql2 & 0x000Fu) + (((hb >> bit2) & 1u) << 4u)) + + yh[2 * i + 1] * + (float)(((ql2 >> 8u) & 0x000Fu) + (((hb >> (bit2 + 8u)) & 1u) << 4u)); + acc[3] += yh[2 * i + 8] * + (float)(((ql2 >> 4u) & 0x000Fu) + (((hb >> bit3) & 1u) << 4u)) + + yh[2 * i + 9] * + (float)(((ql2 >> 12u) & 0x000Fu) + (((hb >> (bit3 + 8u)) & 1u) << 4u)); + } + + sumf[row] += dh[0] * (acc[0] * sc8[0] + acc[1] * sc8[1] + + acc[2] * sc8[4] + acc[3] * sc8[5]) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + q1 += args.down_row_bytes / 2; + qh += args.down_row_bytes / 2; + sc += args.down_row_bytes / 2; + dh += args.down_row_bytes / 2; + } + + y4 += 4 * QK_K; + } + } + + for (short row = 0; row < N_R0_Q5_K && row0 + (uint)row < args.out_dim; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0u) { + out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; + } + } +} + +kernel void kernel_glm_q6_K_slots6_down_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *down0, + device const char *down1, + device const char *down2, + device const char *down3, + device const char *down4, + device const char *down5, + device const float *mid, + device float *out, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = 2; + constexpr uint kmask1 = 0x03u; + constexpr uint kmask2 = 0x0Cu; + constexpr uint kmask3 = 0x30u; + constexpr uint kmask4 = 0xC0u; + + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q6_K; + const uint token = tgpig.y; + if (row0 >= args.out_dim || token >= args.n_tokens) return; + + const int nb = args.mid_dim / QK_K; + float sumf[N_R0_Q6_K] = {0.f}; + float yl[16]; + const short tid = tiisg / 2; + const short ix = tiisg % 2; + const short ip = tid / 8; + const short il = tid % 8; + const short l0 = 4 * il; + const short is = 8 * ip + l0 / 16; + const short y_offset = 128 * ip + l0; + const short q_offset_l = 64 * ip + l0; + const short q_offset_h = 32 * ip + l0; + + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + device const char *down_cur = down0; + switch (slot) { + case 1: down_cur = down1; break; + case 2: down_cur = down2; break; + case 3: down_cur = down3; break; + case 4: down_cur = down4; break; + case 5: down_cur = down5; break; + default: break; + } + + device const block_q6_K *x = + (device const block_q6_K *)(down_cur + + (uint64_t)row0 * args.down_row_bytes); + device const float *yy = mid + mid_base + (uint64_t)slot * args.mid_dim; + + for (int ib = ix; ib < nb; ib += 2) { + device const uchar *q1 = x[ib].ql + q_offset_l; + device const uchar *q2 = q1 + 32; + device const uchar *qh = x[ib].qh + q_offset_h; + device const char *sc = x[ib].scales + is; + device const half *dh = &x[ib].d; + device const float *y = yy + (uint64_t)ib * QK_K + y_offset; + + for (short l = 0; l < 4; ++l) { + yl[4 * l + 0] = y[l + 0]; + yl[4 * l + 1] = y[l + 32]; + yl[4 * l + 2] = y[l + 64]; + yl[4 * l + 3] = y[l + 96]; + } + + for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) { + float4 sums = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short l = 0; l < 4; ++l) { + const uint h = (uint)qh[l]; + sums[0] += yl[4 * l + 0] * + (float)((int)((q1[l] & 0x0Fu) | ((h & kmask1) << 4u)) - 32); + sums[1] += yl[4 * l + 1] * + (float)((int)((q2[l] & 0x0Fu) | ((h & kmask2) << 2u)) - 32); + sums[2] += yl[4 * l + 2] * + (float)((int)((q1[l] >> 4u) | (h & kmask3)) - 32); + sums[3] += yl[4 * l + 3] * + (float)((int)((q2[l] >> 4u) | ((h & kmask4) >> 2u)) - 32); + } + + sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] + + sums[2] * sc[4] + sums[3] * sc[6]); + + q1 += args.down_row_bytes; + q2 += args.down_row_bytes; + qh += args.down_row_bytes; + sc += args.down_row_bytes; + dh += args.down_row_bytes / 2; + } + } + } + + for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0u) { + out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; + } + } +} + +kernel void kernel_glm_q6_K_slots8_down_f32( + constant ds4_metal_glm_routed_moe_args &args, + device const char *down0, + device const char *down1, + device const char *down2, + device const char *down3, + device const char *down4, + device const char *down5, + device const char *down6, + device const char *down7, + device const float *mid, + device float *out, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + const short NSG = 2; + constexpr uint kmask1 = 0x03u; + constexpr uint kmask2 = 0x0Cu; + constexpr uint kmask3 = 0x30u; + constexpr uint kmask4 = 0xC0u; + + const uint row0 = ((uint)tgpig.x * (uint)NSG + (uint)sgitg) * N_R0_Q6_K; + const uint token = tgpig.y; + if (row0 >= args.out_dim || token >= args.n_tokens) return; + + const int nb = args.mid_dim / QK_K; + float sumf[N_R0_Q6_K] = {0.f}; + float yl[16]; + const short tid = tiisg / 2; + const short ix = tiisg % 2; + const short ip = tid / 8; + const short il = tid % 8; + const short l0 = 4 * il; + const short is = 8 * ip + l0 / 16; + const short y_offset = 128 * ip + l0; + const short q_offset_l = 64 * ip + l0; + const short q_offset_h = 32 * ip + l0; + + const uint64_t mid_base = (uint64_t)token * args.mid_token_stride; + for (uint slot = 0; slot < args.n_expert_used; slot++) { + device const char *down_cur = down0; + switch (slot) { + case 1: down_cur = down1; break; + case 2: down_cur = down2; break; + case 3: down_cur = down3; break; + case 4: down_cur = down4; break; + case 5: down_cur = down5; break; + case 6: down_cur = down6; break; + case 7: down_cur = down7; break; + default: break; + } + + device const block_q6_K *x = + (device const block_q6_K *)(down_cur + + (uint64_t)row0 * args.down_row_bytes); + device const float *yy = mid + mid_base + (uint64_t)slot * args.mid_dim; + + for (int ib = ix; ib < nb; ib += 2) { + device const uchar *q1 = x[ib].ql + q_offset_l; + device const uchar *q2 = q1 + 32; + device const uchar *qh = x[ib].qh + q_offset_h; + device const char *sc = x[ib].scales + is; + device const half *dh = &x[ib].d; + device const float *y = yy + (uint64_t)ib * QK_K + y_offset; + + for (short l = 0; l < 4; ++l) { + yl[4 * l + 0] = y[l + 0]; + yl[4 * l + 1] = y[l + 32]; + yl[4 * l + 2] = y[l + 64]; + yl[4 * l + 3] = y[l + 96]; + } + + for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) { + float4 sums = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short l = 0; l < 4; ++l) { + const uint h = (uint)qh[l]; + sums[0] += yl[4 * l + 0] * + (float)((int)((q1[l] & 0x0Fu) | ((h & kmask1) << 4u)) - 32); + sums[1] += yl[4 * l + 1] * + (float)((int)((q2[l] & 0x0Fu) | ((h & kmask2) << 2u)) - 32); + sums[2] += yl[4 * l + 2] * + (float)((int)((q1[l] >> 4u) | (h & kmask3)) - 32); + sums[3] += yl[4 * l + 3] * + (float)((int)((q2[l] >> 4u) | ((h & kmask4) >> 2u)) - 32); + } + + sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] + + sums[2] * sc[4] + sums[3] * sc[6]); + + q1 += args.down_row_bytes; + q2 += args.down_row_bytes; + qh += args.down_row_bytes; + sc += args.down_row_bytes; + dh += args.down_row_bytes / 2; + } + } + } + + for (short row = 0; row < N_R0_Q6_K && row0 + (uint)row < args.out_dim; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0u) { + out[(uint64_t)token * args.out_dim + row0 + (uint)row] = sum_all; + } + } +} + +template +void dequantize_q4_K(device const block_q4_K *xb, short il, thread type4x4 ®) { + device const uchar *q = xb->qs; + + short is = (il / 4) * 2; + q = q + (il / 4) * 32 + 16 * (il & 1); + il = il & 3; + const uchar2 sc = get_scale_min_k4_just2(is, il / 2, xb->scales); + const float d = il < 2 ? + (float)xb->d : + (float)xb->d * (1.0f / 16.0f); + const float min = (float)xb->dmin; + const float dl = d * sc[0]; + const float ml = min * sc[1]; + + const ushort mask = il < 2 ? 0x0F : 0xF0; + for (int i = 0; i < 16; ++i) { + reg[i / 4][i % 4] = dl * (q[i] & mask) - ml; + } +} + +template +void dequantize_q5_K(device const block_q5_K *xb, short il, thread type4x4 ®) { + const short group = il / 2; + const short l0 = (il & 1) * 16; + const uchar2 sm = get_scale_min_k4_just2(group, 0, xb->scales); + const float dl = (float)xb->d * (float)sm.x; + const float ml = (float)xb->dmin * (float)sm.y; + device const uchar *q = xb->qs + (group / 2) * 32 + l0; + device const uchar *qh = xb->qh + l0; + const uint shift = (uint)(group & 1) * 4u; + const uchar hmask = (uchar)(1u << (uint)group); + + for (int i = 0; i < 16; ++i) { + uint v = ((uint)q[i] >> shift) & 0x0Fu; + v += (qh[i] & hmask) ? 16u : 0u; + reg[i / 4][i % 4] = dl * (float)v - ml; + } +} + +template +void dequantize_q6_K(device const block_q6_K *xb, short il, thread type4x4 ®) { + const short n128 = il / 8; + const short il128 = il & 7; + const short quarter = il128 / 2; + const short l0 = (il128 & 1) * 16; + const uint ql_base = (uint)n128 * 64u; + const uint qh_base = (uint)n128 * 32u; + const uint sc_base = (uint)n128 * 8u + (uint)quarter * 2u + (uint)(il128 & 1); + const float d = (float)xb->d * (float)((int)xb->scales[sc_base]); + + for (int i = 0; i < 16; ++i) { + const uint l = (uint)l0 + (uint)i; + uint v; + if (quarter == 0) { + v = ((uint)xb->ql[ql_base + l] & 0x0Fu) | + ((((uint)xb->qh[qh_base + l] >> 0u) & 3u) << 4u); + } else if (quarter == 1) { + v = ((uint)xb->ql[ql_base + 32u + l] & 0x0Fu) | + ((((uint)xb->qh[qh_base + l] >> 2u) & 3u) << 4u); + } else if (quarter == 2) { + v = ((uint)xb->ql[ql_base + l] >> 4u) | + ((((uint)xb->qh[qh_base + l] >> 4u) & 3u) << 4u); + } else { + v = ((uint)xb->ql[ql_base + 32u + l] >> 4u) | + ((((uint)xb->qh[qh_base + l] >> 6u) & 3u) << 4u); + } + reg[i / 4][i % 4] = d * (float)((int)v - 32); + } +} + +template +void dequantize_iq2_xxs(device const block_iq2_xxs * xb, short il, thread type4x4 & reg) { + const float d = xb->d; + const int ib32 = il/2; + il = il%2; + device const uint16_t * q2 = xb->qs + 4*ib32; + const uint32_t aux32_g = q2[0] | (q2[1] << 16); + const uint32_t aux32_s = q2[2] | (q2[3] << 16); + thread const uint8_t * aux8 = (thread const uint8_t *)&aux32_g; + const float dl = d * (0.5f + (aux32_s >> 28)) * 0.25f; + constant uint8_t * grid = (constant uint8_t *)(iq2xxs_grid + aux8[2*il+0]); + uint8_t signs = ksigns_iq2xs[(aux32_s >> 14*il) & 127]; + for (int i = 0; i < 8; ++i) { + reg[i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f); + } + grid = (constant uint8_t *)(iq2xxs_grid + aux8[2*il+1]); + signs = ksigns_iq2xs[(aux32_s >> (14*il+7)) & 127]; + for (int i = 0; i < 8; ++i) { + reg[2+i/4][i%4] = dl * grid[i] * (signs & kmask_iq2xs[i] ? -1.f : 1.f); + } +} + +struct ds4_metal_args_mul_mv_id { + int32_t nei0; + int32_t nei1; + uint64_t nbi1; + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + int32_t ne10; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + int32_t ne0; + int32_t ne1; + uint64_t nb1; + int32_t nr0; + /* Tensor-parallel expert ownership: when tp_world > 1, routed-MoE + * kernels skip experts outside this rank's contiguous id range so this + * rank never touches their weights. Zero-initialized (tp_world 0) means + * no split; only the routed id kernels read these. */ + int32_t tp_rank; + int32_t tp_world; + /* When set, the sum6 kernels add this per-row vector (the shared-expert + * partial) into their accumulation, folding the TP local-add. */ + int32_t tp_addend; + /* First expert id present at the bound blob base (TP mappings bind only + * the owned range; kernels rebase before pointer math). */ + int32_t tp_expert_base; +}; + +struct ds4_metal_moe_expert_group_args { + uint32_t expert_base; + uint32_t expert_count; + uint32_t accumulate; + uint32_t pad0; +}; + +struct ds4_metal_q4_gather_slots6_args { + uint64_t expert_bytes; + uint32_t group_size; + uint32_t n_slots; +}; + +struct ds4_metal_q4_expert_table { + array experts [[id(0)]]; +}; + +struct ds4_metal_expert_address_table { + device const uint64_t *addrs; +}; + +struct ds4_metal_stream_expert_validate_args { + uint32_t n_total_expert; + uint32_t n_expert; +}; + +struct ds4_metal_stream_expert_split_args { + uint32_t active_mask; + uint32_t accumulate; +}; + +struct ds4_metal_args_mul_mm_id_map0 { + int32_t ne02; + int32_t ne10; + int32_t ne11; + uint64_t nb11; + uint64_t nb12; + int32_t ne21; + int32_t ne20; + uint64_t nb21; +}; + +struct ds4_metal_args_mul_mm_id { + int32_t ne00; + int32_t ne02; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne20; + int32_t ne21; + int32_t ne0; + int32_t ne1; + int16_t r2; + int16_t r3; + /* Tensor-parallel expert ownership; tp_world<=1 disables the split. + * tp_expert_base rebases expert ids when the bind covers only the + * owned contiguous range of the expert blob. */ + int32_t tp_rank; + int32_t tp_world; + int32_t tp_expert_base; +}; + +template +void kernel_mul_mv_q2_K_f32_impl( + args_t args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + + const int nb = args.ne00/QK_K; + + const int r0 = tgpig.x; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const int first_row = (r0 * NSG + sgitg) * nr0; + + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + + const uint64_t offset0 = first_row*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; + + device const block_q2_K * x = (device const block_q2_K *) (src0 + offset0); + device const float * y = (device const float *) (src1 + offset1); + + float yl[32]; + float sumf[nr0]={0.f}; + + const short ix = tiisg/8; // 0...3 + const short it = tiisg%8; // 0...7 + const short iq = it/4; // 0 or 1 + const short ir = it%4; // 0...3 + const short is = (8*ir)/16;// 0 or 1 + + device const float * y4 = y + ix * QK_K + 128 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + for (short i = 0; i < 8; ++i) { + yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0]; + yl[i+ 8] = y4[i+32]; sumy[1] += yl[i+ 8]; + yl[i+16] = y4[i+64]; sumy[2] += yl[i+16]; + yl[i+24] = y4[i+96]; sumy[3] += yl[i+24]; + } + + device const uint8_t * sc = (device const uint8_t *)x[ib].scales + 8*iq + is; + device const uint16_t * qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half * dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + for (int i = 0; i < 8; i += 2) { + acc1[0] += yl[i+ 0] * (qs[i/2] & 0x0003); + acc2[0] += yl[i+ 1] * (qs[i/2] & 0x0300); + acc1[1] += yl[i+ 8] * (qs[i/2] & 0x000c); + acc2[1] += yl[i+ 9] * (qs[i/2] & 0x0c00); + acc1[2] += yl[i+16] * (qs[i/2] & 0x0030); + acc2[2] += yl[i+17] * (qs[i/2] & 0x3000); + acc1[3] += yl[i+24] * (qs[i/2] & 0x00c0); + acc2[3] += yl[i+25] * (qs[i/2] & 0xc000); + } + float dall = dh[0]; + float dmin = dh[1] * 1.f/16.f; + sumf[row] += dall * ((acc1[0] + 1.f/256.f * acc2[0]) * (sc[0] & 0xF) * 1.f/ 1.f + + (acc1[1] + 1.f/256.f * acc2[1]) * (sc[2] & 0xF) * 1.f/ 4.f + + (acc1[2] + 1.f/256.f * acc2[2]) * (sc[4] & 0xF) * 1.f/16.f + + (acc1[3] + 1.f/256.f * acc2[3]) * (sc[6] & 0xF) * 1.f/64.f) - + dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0)); + + qs += args.nb01/2; + sc += args.nb01; + dh += args.nb01/2; + } + + y4 += 4 * QK_K; + } + + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + + for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { + float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) { + dst_f32[first_row + row] = sum_all; + } + } +} + +template +void kernel_mul_mv_q4_K_f32_impl( + args_t args, + device const char *src0, + device const char *src1, + device char *dst, + threadgroup char *shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + const int nb = args.ne00 / QK_K; + + const int r0 = tgpig.x; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const int first_row = (r0 * NSG + sgitg) * nr0; + + const uint i12 = im % args.ne12; + const uint i13 = im / args.ne12; + + const uint64_t offset0 = first_row * args.nb01 + (i12 / args.r2) * args.nb02 + (i13 / args.r3) * args.nb03; + const uint64_t offset1 = r1 * args.nb11 + i12 * args.nb12 + i13 * args.nb13; + + device const block_q4_K *x = (device const block_q4_K *)(src0 + offset0); + device const float *y = (device const float *)(src1 + offset1); + + float yl[16]; + float yh[16]; + float sumf[nr0] = {0.f}; + + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int ib = ix; ib < nb; ib += 4) { + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + q1 += args.nb01 / 2; + sc += args.nb01 / 2; + dh += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + + device float *dst_f32 = (device float *)dst + (uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0; + + for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { + float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) { + dst_f32[first_row + row] = sum_all; + } + } + + (void)shmem; +} + +template +void kernel_mul_mv_q8_K_f32_impl( + args_t args, + device const char *src0, + device const char *src1, + device char *dst, + threadgroup char *shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + + const int nb = args.ne00 / QK_K; + + const int r0 = tgpig.x; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const int first_row = (r0 * NSG + sgitg) * nr0; + + const uint i12 = im % args.ne12; + const uint i13 = im / args.ne12; + + const uint64_t offset0 = first_row * args.nb01 + (i12 / args.r2) * args.nb02 + (i13 / args.r3) * args.nb03; + const uint64_t offset1 = r1 * args.nb11 + i12 * args.nb12 + i13 * args.nb13; + + device const block_q8_K *x = (device const block_q8_K *)(src0 + offset0); + device const float *y = (device const float *)(src1 + offset1); + + float sumf[nr0] = {0.f}; + + const short ix = tiisg / 8; + const short il = tiisg % 8; + const int ib0 = ix; + + device const float *yb = y + ib0 * QK_K + il * 32; + + for (int ib = ib0; ib < nb; ib += 4) { + float yl[32]; + for (short i = 0; i < 32; ++i) { + yl[i] = yb[i]; + } + + for (short row = 0; row < nr0; row++) { + device const block_q8_K *xr = + (device const block_q8_K *)((device const char *)x + (uint64_t)row * args.nb01); + device const int8_t *qs = xr[ib].qs + il * 32; + + float sumq = 0.f; + for (short i = 0; i < 32; ++i) { + sumq += (float)qs[i] * yl[i]; + } + + sumf[row] += sumq * xr[ib].d; + } + + yb += 4 * QK_K; + } + + device float *dst_f32 = + (device float *)dst + (uint64_t)im * args.ne0 * args.ne1 + (uint64_t)r1 * args.ne0; + + for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { + float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) { + dst_f32[first_row + row] = sum_all; + } + } + + (void)shmem; +} + +template +void kernel_mul_mv_iq2_xxs_f32_impl( + args_t args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + + const int nb = args.ne00/QK_K; + + const int r0 = tgpig.x; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const int first_row = (r0 * NSG + sgitg) * nr0; + + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + + const uint64_t offset0 = first_row*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; + + device const block_iq2_xxs * x = (device const block_iq2_xxs *) (src0 + offset0); + device const float * y = (device const float *) (src1 + offset1); + + float yl[32]; + float sumf[nr0]={0.f}; + + const int nb32 = nb * (QK_K / 32); + + threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem); + threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256); + { + int nval = 4; + int pos = (32*sgitg + tiisg)*nval; + for (int i = 0; i < nval; ++i) svalues[pos + i] = ds4_metal_iq2xxs_grid[pos + i]; + nval = 2; + pos = (32*sgitg + tiisg)*nval; + for (int i = 0; i < nval; ++i) ssigns[pos+i] = ds4_metal_ksigns_iq2xs[pos+i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const int ix = tiisg; + + device const float * y4 = y + 32 * ix; + + for (int ib32 = ix; ib32 < nb32; ib32 += 32) { + for (short i = 0; i < 32; ++i) { + yl[i] = y4[i]; + } + + const int ibl = ib32 / (QK_K / 32); + const int ib = ib32 % (QK_K / 32); + + device const block_iq2_xxs * xr = x + ibl; + device const uint16_t * q2 = xr->qs + 4 * ib; + device const half * dh = &xr->d; + + for (short row = 0; row < nr0; row++) { + const float db = dh[0]; + device const uint8_t * aux8 = (device const uint8_t *)q2; + const uint32_t aux32 = q2[2] | (q2[3] << 16); + const float d = db * (0.5f + (aux32 >> 28)); + + float sum = 0; + for (short l = 0; l < 4; ++l) { + const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + aux8[l]); + const uint8_t signs = ssigns[(aux32 >> 7*l) & 127]; + for (short j = 0; j < 8; ++j) { + sum += yl[8*l + j] * grid[j] * (signs & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f); + } + } + sumf[row] += d * sum; + + dh += args.nb01/2; + q2 += args.nb01/2; + } + + y4 += 32 * 32; + } + + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + + for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { + float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) { + dst_f32[first_row + row] = sum_all * 0.25f; + } + } +} + +template +void kernel_mul_mv_iq2_xxs_pair_f32_impl( + ds4_metal_args_mul_mv args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + + const int nb = args.ne00/QK_K; + + const int r0 = tgpig.x; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const int first_row = (r0 * NSG + sgitg) * nr0; + + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + + const uint64_t offset0 = first_row*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; + + device const block_iq2_xxs * xg = (device const block_iq2_xxs *) (src0_gate + offset0); + device const block_iq2_xxs * xu = (device const block_iq2_xxs *) (src0_up + offset0); + device const float * y = (device const float *) (src1 + offset1); + + float yl[32]; + float sumg[nr0]={0.f}; + float sumu[nr0]={0.f}; + + const int nb32 = nb * (QK_K / 32); + + threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem); + threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256); + { + int nval = 4; + int pos = (32*sgitg + tiisg)*nval; + for (int i = 0; i < nval; ++i) svalues[pos + i] = ds4_metal_iq2xxs_grid[pos + i]; + nval = 2; + pos = (32*sgitg + tiisg)*nval; + for (int i = 0; i < nval; ++i) ssigns[pos+i] = ds4_metal_ksigns_iq2xs[pos+i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const int ix = tiisg; + device const float * y4 = y + 32 * ix; + + for (int ib32 = ix; ib32 < nb32; ib32 += 32) { + for (short i = 0; i < 32; ++i) { + yl[i] = y4[i]; + } + + const int ibl = ib32 / (QK_K / 32); + const int ib = ib32 % (QK_K / 32); + + device const block_iq2_xxs * xgr = xg + ibl; + device const block_iq2_xxs * xur = xu + ibl; + device const uint16_t * qg = xgr->qs + 4 * ib; + device const uint16_t * qu = xur->qs + 4 * ib; + device const half * dhg = &xgr->d; + device const half * dhu = &xur->d; + + for (short row = 0; row < nr0; row++) { + device const uint8_t * aux8g = (device const uint8_t *)qg; + device const uint8_t * aux8u = (device const uint8_t *)qu; + const uint32_t aux32g = qg[2] | (qg[3] << 16); + const uint32_t aux32u = qu[2] | (qu[3] << 16); + const float dg = (float)dhg[0] * (0.5f + (aux32g >> 28)); + const float du = (float)dhu[0] * (0.5f + (aux32u >> 28)); + + float sg = 0; + float su = 0; + for (short l = 0; l < 4; ++l) { + const threadgroup uint8_t * gridg = (const threadgroup uint8_t *)(svalues + aux8g[l]); + const threadgroup uint8_t * gridu = (const threadgroup uint8_t *)(svalues + aux8u[l]); + const uint8_t signg = ssigns[(aux32g >> 7*l) & 127]; + const uint8_t signu = ssigns[(aux32u >> 7*l) & 127]; + for (short j = 0; j < 8; ++j) { + const float v = yl[8*l + j]; + sg += v * gridg[j] * (signg & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f); + su += v * gridu[j] * (signu & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f); + } + } + sumg[row] += dg * sg; + sumu[row] += du * su; + + dhg += args.nb01/2; + dhu += args.nb01/2; + qg += args.nb01/2; + qu += args.nb01/2; + } + + y4 += 32 * 32; + } + + device float * dst_gate_f32 = (device float *) dst_gate + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + device float * dst_up_f32 = (device float *) dst_up + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + + for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { + const float sum_gate = simd_sum(sumg[row]); + const float sum_up = simd_sum(sumu[row]); + if (tiisg == 0) { + dst_gate_f32[first_row + row] = sum_gate * 0.25f; + dst_up_f32[first_row + row] = sum_up * 0.25f; + } + } +} + +typedef void (kernel_mul_mv2_disp_t)( + ds4_metal_args_mul_mv args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg); + +template +void mmv_fn( + ds4_metal_args_mul_mv args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiitg, + ushort tiisg, + ushort sgitg) { + disp_fn(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +} + +typedef decltype(mmv_fn>) mul_mv_id_disp_fn_t; + +// Decode-time expert matvec. The ids tensor selects the routed expert for each +// slot, then this wrapper invokes the quantized row kernel for Q8_0, Q2_K, or +// IQ2_XXS weights without materializing per-expert dispatches on the CPU. +template +kernel void kernel_mul_mv_id( + constant ds4_metal_args_mul_mv_id & args, + device const char * src0s, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + (void)tiitg; + + const int iid1 = tgpig.z/args.nei0; + const int idx = tgpig.z%args.nei0; + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *) (ids + iid1*args.nbi1))[idx]; + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + const int64_t i1 = idx; + const int64_t i2 = i12; + + device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float); + + if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) { + /* Unowned expert under the TP split: zero this threadgroup's output + * rows so the downstream expert-sum stages stay unchanged. */ + const short NSG = FC_mul_mv_nsg; + const int row0 = (tgpig.x * NSG + sgitg) * args.nr0; + device float *dst_f32 = (device float *)dst_cur; + for (int r = 0; r < args.nr0 && row0 + r < args.ne0; r++) { + if (tiisg == 0) dst_f32[row0 + r] = 0.0f; + } + return; + } + + device const char * src0_cur = + src0s + (int64_t)(i02 - args.tp_expert_base)*args.nb02; + device const char * src1_cur = src1 + i11*args.nb11 + i12*args.nb12; + + ds4_metal_args_mul_mv args0 = { + /*.ne00 =*/ args.ne00, + /*.ne01 =*/ args.ne01, + /*.ne02 =*/ 1, + /*.nb00 =*/ args.nb00, + /*.nb01 =*/ args.nb01, + /*.nb02 =*/ args.nb02, + /*.nb03 =*/ args.nb02, + /*.ne10 =*/ args.ne10, + /*.ne11 =*/ 1, + /*.ne12 =*/ 1, + /*.nb10 =*/ args.nb10, + /*.nb11 =*/ args.nb11, + /*.nb12 =*/ args.nb12, + /*.nb13 =*/ args.nb12, + /*.ne0 =*/ args.ne0, + /*.ne1 =*/ 1, + /*.nr0 =*/ args.nr0, + /*.r2 =*/ 1, + /*.r3 =*/ 1, + }; + + disp_fn( + args0, + /* src0 */ src0_cur, + /* src1 */ src1_cur, + /* dst */ dst_cur, + shmem, + tgpig, + tiitg, + tiisg, + sgitg); +} + +typedef decltype(kernel_mul_mv_id>>) kernel_mul_mv_id_q_t; +typedef decltype(kernel_mul_mv_id>>) kernel_mul_mv_id_q8_0_t; + +// Host-visible decode MoE matvec variants for the DS4 quant formats. +template [[host_name("kernel_mul_mv_id_q8_0_f32")]] kernel kernel_mul_mv_id_q8_0_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_q2_K_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_q4_K_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_q8_K_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_q_t kernel_mul_mv_id>>; + +// Plain dense Q4_K matvec on the classic impl. The mul_mv_ext family used +// by the generic dense path tops out around 220 GB/s on M5 for the GLM +// DenseQ4 decode shapes; this impl streams the same rows at 530-650 GB/s. +kernel void kernel_mul_mv_q4_K_dense_f32( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + kernel_mul_mv_q4_K_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +} + +// DS4 attention output low projection, specialized for the fixed block +// diagonal mapping used by the model: +// +// low[token, group, rank] = heads[token, group, :] * Woa[group, rank, :] +// +// The generic GGML-style id matvec supports arbitrary routed expert ids. Here +// the id is always equal to the group number, so this wrapper keeps the exact +// Q8_0 dot kernel but removes the id-buffer load and the CPU-side id table. +kernel void kernel_dsv4_attn_out_low_q8_0_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src0s, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z/args.nei0; + const int idx = tgpig.z%args.nei0; + + tgpig.z = 0; + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char * src0_cur = src0s + idx*args.nb02; + device const char * src1_cur = src1 + i11*args.nb11 + i12*args.nb12; + device char * dst_cur = dst + (idx*args.ne0 + i12*args.ne1*args.ne0)*sizeof(float); + + ds4_metal_args_mul_mv args0 = { + /*.ne00 =*/ args.ne00, + /*.ne01 =*/ args.ne01, + /*.ne02 =*/ 1, + /*.nb00 =*/ args.nb00, + /*.nb01 =*/ args.nb01, + /*.nb02 =*/ args.nb02, + /*.nb03 =*/ args.nb02, + /*.ne10 =*/ args.ne10, + /*.ne11 =*/ 1, + /*.ne12 =*/ 1, + /*.nb10 =*/ args.nb10, + /*.nb11 =*/ args.nb11, + /*.nb12 =*/ args.nb12, + /*.nb13 =*/ args.nb12, + /*.ne0 =*/ args.ne0, + /*.ne1 =*/ 1, + /*.nr0 =*/ args.nr0, + /*.r2 =*/ 1, + /*.r3 =*/ 1, + }; + + kernel_mul_mv_q8_0_f32_impl( + args0, + src0_cur, + src1_cur, + dst_cur, + shmem, + tgpig, + tiisg, + sgitg); +} + +kernel void kernel_dsv4_attn_out_low_q4_K_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src0s, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + (void)tiitg; + + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char * src0_cur = src0s + idx * args.nb02; + device const char * src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + device char * dst_cur = dst + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + /*.ne00 =*/ args.ne00, + /*.ne01 =*/ args.ne01, + /*.ne02 =*/ 1, + /*.nb00 =*/ args.nb00, + /*.nb01 =*/ args.nb01, + /*.nb02 =*/ args.nb02, + /*.nb03 =*/ args.nb02, + /*.ne10 =*/ args.ne10, + /*.ne11 =*/ 1, + /*.ne12 =*/ 1, + /*.nb10 =*/ args.nb10, + /*.nb11 =*/ args.nb11, + /*.nb12 =*/ args.nb12, + /*.nb13 =*/ args.nb12, + /*.ne0 =*/ args.ne0, + /*.ne1 =*/ 1, + /*.nr0 =*/ args.nr0, + /*.r2 =*/ 1, + /*.r3 =*/ 1, + }; + + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_cur, + src1_cur, + dst_cur, + shmem, + tgpig, + tiisg, + sgitg); +} + +kernel void kernel_mul_mv_id_iq2_xxs_pair_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z/args.nei0; + const int idx = tgpig.z%args.nei0; + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *) (ids + iid1*args.nbi1))[idx]; + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device char * dst_gate_cur = dst_gate + (idx*args.ne0 + i12*args.ne1*args.ne0)*sizeof(float); + device char * dst_up_cur = dst_up + (idx*args.ne0 + i12*args.ne1*args.ne0)*sizeof(float); + + if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) { + /* Unowned expert: zero this threadgroup's gate/up rows so the + * separate swiglu-weight stage yields zero mid rows. */ + const short NSG_z = FC_mul_mv_nsg; + const int row0_z = (tgpig.x * NSG_z + sgitg) * args.nr0; + device float *zg = (device float *)dst_gate_cur; + device float *zu = (device float *)dst_up_cur; + for (int r = 0; r < args.nr0 && row0_z + r < args.ne0; r++) { + if (tiisg == 0) { + zg[row0_z + r] = 0.0f; + zu[row0_z + r] = 0.0f; + } + } + return; + } + + device const char * src0_gate_cur = src0_gate + (int64_t)(i02 - args.tp_expert_base)*args.nb02; + device const char * src0_up_cur = src0_up + (int64_t)(i02 - args.tp_expert_base)*args.nb02; + device const char * src1_cur = src1 + i11*args.nb11 + i12*args.nb12; + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + (void)tiitg; + kernel_mul_mv_iq2_xxs_pair_f32_impl( + args0, + src0_gate_cur, + src0_up_cur, + src1_cur, + dst_gate_cur, + dst_up_cur, + shmem, + tgpig, + tiisg, + sgitg); +} + +// Decode-only routed expert gate/up projection fused with the DS4 activation: +// +// mid = silu(clamp(gate)) * clamp(up) * route_weight +// +// The quantized dot products are intentionally the same IQ2_XXS paired path as +// `kernel_mul_mv_id_iq2_xxs_pair_f32`. The only extra work is done by lane 0 +// after each exact reduced row has been produced. This removes the separate +// routed activation dispatch and avoids rereading the gate/up rows before the +// down projection. The host uses this only for the normal release path where +// diagnostics do not request clamped gate/up intermediates. +kernel void kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * ids, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *) (ids + iid1 * args.nbi1))[idx]; + if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) return; + const int i02b = i02 - args.tp_expert_base; + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_IQ2_XXS; + const int nb32 = nb * (QK_K / 32); + + device const block_iq2_xxs *xg = + (device const block_iq2_xxs *)(src0_gate + (int64_t)i02b * args.nb02 + (uint64_t)first_row * args.nb01); + device const block_iq2_xxs *xu = + (device const block_iq2_xxs *)(src0_up + (int64_t)i02b * args.nb02 + (uint64_t)first_row * args.nb01); + device const float *y = + (device const float *)(src1 + i11 * args.nb11 + i12 * args.nb12); + + float yl[32]; + float sumg[N_R0_IQ2_XXS] = {0.f}; + float sumu[N_R0_IQ2_XXS] = {0.f}; + + threadgroup uint64_t *svalues = (threadgroup uint64_t *)(shmem); + threadgroup uint8_t *ssigns = (threadgroup uint8_t *)(svalues + 256); + { + int nval = 4; + int pos = (32 * sgitg + tiisg) * nval; + for (int i = 0; i < nval; ++i) svalues[pos + i] = ds4_metal_iq2xxs_grid[pos + i]; + nval = 2; + pos = (32 * sgitg + tiisg) * nval; + for (int i = 0; i < nval; ++i) ssigns[pos + i] = ds4_metal_ksigns_iq2xs[pos + i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const int ix = tiisg; + device const float *y4 = y + 32 * ix; + + for (int ib32 = ix; ib32 < nb32; ib32 += 32) { + for (short i = 0; i < 32; ++i) { + yl[i] = y4[i]; + } + + const int ibl = ib32 / (QK_K / 32); + const int ib = ib32 % (QK_K / 32); + + device const block_iq2_xxs *xgr = xg + ibl; + device const block_iq2_xxs *xur = xu + ibl; + device const uint16_t *qg = xgr->qs + 4 * ib; + device const uint16_t *qu = xur->qs + 4 * ib; + device const half *dhg = &xgr->d; + device const half *dhu = &xur->d; + + for (short row = 0; row < N_R0_IQ2_XXS; row++) { + device const uint8_t *aux8g = (device const uint8_t *)qg; + device const uint8_t *aux8u = (device const uint8_t *)qu; + const uint32_t aux32g = qg[2] | (qg[3] << 16); + const uint32_t aux32u = qu[2] | (qu[3] << 16); + const float dg = (float)dhg[0] * (0.5f + (aux32g >> 28)); + const float du = (float)dhu[0] * (0.5f + (aux32u >> 28)); + + float sg = 0; + float su = 0; + for (short l = 0; l < 4; ++l) { + const threadgroup uint8_t *gridg = (const threadgroup uint8_t *)(svalues + aux8g[l]); + const threadgroup uint8_t *gridu = (const threadgroup uint8_t *)(svalues + aux8u[l]); + const uint8_t signg = ssigns[(aux32g >> 7 * l) & 127]; + const uint8_t signu = ssigns[(aux32u >> 7 * l) & 127]; + for (short j = 0; j < 8; ++j) { + const float v = yl[8 * l + j]; + sg += v * gridg[j] * (signg & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f); + su += v * gridu[j] * (signu & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f); + } + } + sumg[row] += dg * sg; + sumu[row] += du * su; + + dhg += args.nb01 / 2; + dhu += args.nb01 / 2; + qg += args.nb01 / 2; + qu += args.nb01 / 2; + } + + y4 += 32 * 32; + } + + device float *dst_gate_f32 = + (device float *)dst_gate + (uint64_t)i12 * args.ne0 * args.ne1 + (uint64_t)i11 * args.ne0; + device float *dst_up_f32 = + (device float *)dst_up + (uint64_t)i12 * args.ne0 * args.ne1 + (uint64_t)i11 * args.ne0; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *dst_mid_f32 = + (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = + (device const float *)(weights + pair_row * act.weight_stride); + + const float c = act.clamp_value; + const float route_weight = route_w[0]; + for (int row = 0; row < N_R0_IQ2_XXS && first_row + row < args.ne0; ++row) { + const float sum_gate = simd_sum(sumg[row]); + const float sum_up = simd_sum(sumu[row]); + if (tiisg == 0) { + const uint out_row = first_row + row; + const float gate = sum_gate * 0.25f; + const float up = sum_up * 0.25f; + float g = gate; + float u = up; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + dst_gate_f32[out_row] = gate; + dst_up_f32[out_row] = up; + const float silu = g / (1.0f + exp(-g)); + dst_mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + +kernel void kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const char * src0_gate0, + device const char * src0_gate1, + device const char * src0_gate2, + device const char * src0_gate3, + device const char * src0_gate4, + device const char * src0_gate5, + device const char * src0_up0, + device const char * src0_up1, + device const char * src0_up2, + device const char * src0_up3, + device const char * src0_up4, + device const char * src0_up5, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char *src0_gate_cur = src0_gate0; + device const char *src0_up_cur = src0_up0; + switch (idx) { + case 1: src0_gate_cur = src0_gate1; src0_up_cur = src0_up1; break; + case 2: src0_gate_cur = src0_gate2; src0_up_cur = src0_up2; break; + case 3: src0_gate_cur = src0_gate3; src0_up_cur = src0_up3; break; + case 4: src0_gate_cur = src0_gate4; src0_up_cur = src0_up4; break; + case 5: src0_gate_cur = src0_gate5; src0_up_cur = src0_up5; break; + default: break; + } + + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + kernel_mul_mv_iq2_xxs_pair_f32_impl( + args0, + src0_gate_cur, + src0_up_cur, + src1_cur, + dst_gate_cur, + dst_up_cur, + shmem, + tgpig, + tiisg, + sgitg); + + const short NSG = FC_mul_mv_nsg; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_IQ2_XXS; + device float *gate_f32 = (device float *)dst_gate_cur; + device float *up_f32 = (device float *)dst_up_cur; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride); + const float c = act.clamp_value; + const float route_weight = route_w[0]; + + if (tiisg == 0) { + for (int row = 0; row < N_R0_IQ2_XXS && first_row + row < args.ne0; ++row) { + const uint out_row = first_row + row; + float g = gate_f32[out_row]; + float u = up_f32[out_row]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + +kernel void kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const uint64_t * gate_addrs, + device const uint64_t * up_addrs, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * ids, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + if (i02 < 0 || i02 >= args.ne02 || i02 >= 384) { + return; + } + const uint64_t gate_addr = gate_addrs[(uint)i02]; + const uint64_t up_addr = up_addrs[(uint)i02]; + if (gate_addr == 0 || up_addr == 0) { + return; + } + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char *src0_gate_cur = + reinterpret_cast(gate_addr); + device const char *src0_up_cur = + reinterpret_cast(up_addr); + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + kernel_mul_mv_iq2_xxs_pair_f32_impl( + args0, + src0_gate_cur, + src0_up_cur, + src1_cur, + dst_gate_cur, + dst_up_cur, + shmem, + tgpig, + tiisg, + sgitg); + + const short NSG = FC_mul_mv_nsg; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_IQ2_XXS; + device float *gate_f32 = (device float *)dst_gate_cur; + device float *up_f32 = (device float *)dst_up_cur; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride); + const float c = act.clamp_value; + const float route_weight = route_w[0]; + + if (tiisg == 0) { + for (int row = 0; row < N_R0_IQ2_XXS && first_row + row < args.ne0; ++row) { + const uint out_row = first_row + row; + float g = gate_f32[out_row]; + float u = up_f32[out_row]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + +kernel void kernel_mul_mv_addr_iq2_xxs_f32( + constant ds4_metal_args_mul_mv_id & args, + device const uint64_t * addrs, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + (void)tiitg; + + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + if (i02 < 0 || i02 >= args.ne02 || i02 >= 384) { + return; + } + const uint64_t addr = addrs[(uint)i02]; + if (addr == 0) { + return; + } + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char *src0_cur = reinterpret_cast(addr); + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + device char *dst_cur = dst + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + kernel_mul_mv_iq2_xxs_f32_impl( + args0, + src0_cur, + src1_cur, + dst_cur, + shmem, + tgpig, + tiisg, + sgitg); +} + +kernel void kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + constant ds4_metal_stream_expert_split_args & split, + device const uint64_t * gate_addrs, + device const uint64_t * up_addrs, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * ids, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + if ((split.active_mask & (1u << (uint)idx)) == 0) { + return; + } + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + if (i02 < 0 || i02 >= args.ne02 || i02 >= 384) { + return; + } + const uint64_t gate_addr = gate_addrs[(uint)i02]; + const uint64_t up_addr = up_addrs[(uint)i02]; + if (gate_addr == 0 || up_addr == 0) { + return; + } + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char *src0_gate_cur = + reinterpret_cast(gate_addr); + device const char *src0_up_cur = + reinterpret_cast(up_addr); + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + kernel_mul_mv_iq2_xxs_pair_f32_impl( + args0, + src0_gate_cur, + src0_up_cur, + src1_cur, + dst_gate_cur, + dst_up_cur, + shmem, + tgpig, + tiisg, + sgitg); + + const short NSG = FC_mul_mv_nsg; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_IQ2_XXS; + device float *gate_f32 = (device float *)dst_gate_cur; + device float *up_f32 = (device float *)dst_up_cur; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride); + const float c = act.clamp_value; + const float route_weight = route_w[0]; + + if (tiisg == 0) { + for (int row = 0; row < N_R0_IQ2_XXS && first_row + row < args.ne0; ++row) { + const uint out_row = first_row + row; + float g = gate_f32[out_row]; + float u = up_f32[out_row]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + +kernel void kernel_stream_expert_cache_validate( + constant ds4_metal_stream_expert_validate_args & args, + device const char * ids, + device const uint64_t * gate_addrs, + device const uint64_t * up_addrs, + device const uint64_t * down_addrs, + device uint32_t * status, + uint tid [[thread_position_in_grid]]) { + if (tid != 0) return; + + uint32_t miss_mask = 0; + uint32_t invalid_mask = 0; + const uint32_t n_expert = min(args.n_expert, (uint32_t)6); + device const int32_t *selected = (device const int32_t *)ids; + + status[3] = n_expert; + for (uint32_t i = 0; i < 6; i++) { + const int32_t expert = i < n_expert ? selected[i] : -1; + status[4 + i] = as_type(expert); + if (i >= n_expert) continue; + if (expert < 0 || + (uint32_t)expert >= args.n_total_expert || + (uint32_t)expert >= 384) { + invalid_mask |= (1u << i); + continue; + } + const uint32_t e = (uint32_t)expert; + if (gate_addrs[e] == 0 || up_addrs[e] == 0 || down_addrs[e] == 0) { + miss_mask |= (1u << i); + } + } + + status[0] = (miss_mask == 0 && invalid_mask == 0) ? 1u : 0u; + status[1] = miss_mask; + status[2] = invalid_mask; +} + +kernel void kernel_mul_mv_id_q4_K_pair_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) { + const short NSG_z = FC_mul_mv_nsg; + const int row0_z = (tgpig.x * NSG_z + sgitg) * args.nr0; + device float *zg = (device float *)dst_gate_cur; + device float *zu = (device float *)dst_up_cur; + for (int r = 0; r < args.nr0 && row0_z + r < args.ne0; r++) { + if (tiisg == 0) { + zg[row0_z + r] = 0.0f; + zu[row0_z + r] = 0.0f; + } + } + return; + } + + device const char *src0_gate_cur = src0_gate + (int64_t)(i02 - args.tp_expert_base) * args.nb02; + device const char *src0_up_cur = src0_up + (int64_t)(i02 - args.tp_expert_base) * args.nb02; + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + (void)tiitg; + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_gate_cur, + src1_cur, + dst_gate_cur, + shmem, + tgpig, + tiisg, + sgitg); + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_up_cur, + src1_cur, + dst_up_cur, + shmem, + tgpig, + tiisg, + sgitg); +} + +// Same release-path fusion as the IQ2_XXS kernel above for the Q4_K expert +// variant. The Q4 pair path reuses the existing exact matvec implementation +// for gate and up, then the same lane that wrote each row derives the routed +// SwiGLU input. This keeps Q4 behavior aligned with the Q2 optimization while +// preserving the old pair projection arithmetic. +kernel void kernel_mul_mv_id_q4_K_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * ids, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + if (!ds4_tp_owns_expert(i02, args.ne02, args.tp_rank, args.tp_world)) return; + const int i02b = i02 - args.tp_expert_base; + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char *src0_gate_cur = src0_gate + (int64_t)i02b * args.nb02; + device const char *src0_up_cur = src0_up + (int64_t)i02b * args.nb02; + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + const short NSG = FC_mul_mv_nsg; + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K; + device float *gate_f32 = (device float *)dst_gate_cur; + device float *up_f32 = (device float *)dst_up_cur; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride); + const float c = act.clamp_value; + const float route_weight = route_w[0]; + + device const block_q4_K *xg = + (device const block_q4_K *)(src0_gate_cur + (uint64_t)first_row * args.nb01); + device const block_q4_K *xu = + (device const block_q4_K *)(src0_up_cur + (uint64_t)first_row * args.nb01); + device const float *y = (device const float *)src1_cur; + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + float sumg[N_R0_Q4_K] = {0.f}; + float sumu[N_R0_Q4_K] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *scg = (device const uint16_t *)xg[ib].scales + iq; + device const uint16_t *qg1 = (device const uint16_t *)xg[ib].qs + 16 * iq + 4 * ir; + device const half *dhg = &xg[ib].d; + device const uint16_t *scu = (device const uint16_t *)xu[ib].scales + iq; + device const uint16_t *qu1 = (device const uint16_t *)xu[ib].qs + 16 * iq + 4 * ir; + device const half *dhu = &xu[ib].d; + + for (short row = 0; row < N_R0_Q4_K; row++) { + sc16[0] = scg[0] & kmask1; + sc16[1] = scg[2] & kmask1; + sc16[2] = ((scg[4] >> 0) & kmask2) | ((scg[0] & kmask3) >> 2); + sc16[3] = ((scg[4] >> 4) & kmask2) | ((scg[2] & kmask3) >> 2); + + device const uint16_t *qg2 = qg1 + 32; + float4 acc1g = {0.f, 0.f, 0.f, 0.f}; + float4 acc2g = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1g[0] += yl[2 * i + 0] * (qg1[i] & 0x000F); + acc1g[1] += yl[2 * i + 1] * (qg1[i] & 0x0F00); + acc1g[2] += yl[2 * i + 8] * (qg1[i] & 0x00F0); + acc1g[3] += yl[2 * i + 9] * (qg1[i] & 0xF000); + acc2g[0] += yh[2 * i + 0] * (qg2[i] & 0x000F); + acc2g[1] += yh[2 * i + 1] * (qg2[i] & 0x0F00); + acc2g[2] += yh[2 * i + 8] * (qg2[i] & 0x00F0); + acc2g[3] += yh[2 * i + 9] * (qg2[i] & 0xF000); + } + + sumg[row] += dhg[0] * ((acc1g[0] + 1.f / 256.f * acc1g[1]) * sc8[0] + + (acc1g[2] + 1.f / 256.f * acc1g[3]) * sc8[1] * 1.f / 16.f + + (acc2g[0] + 1.f / 256.f * acc2g[1]) * sc8[4] + + (acc2g[2] + 1.f / 256.f * acc2g[3]) * sc8[5] * 1.f / 16.f) - + dhg[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + sc16[0] = scu[0] & kmask1; + sc16[1] = scu[2] & kmask1; + sc16[2] = ((scu[4] >> 0) & kmask2) | ((scu[0] & kmask3) >> 2); + sc16[3] = ((scu[4] >> 4) & kmask2) | ((scu[2] & kmask3) >> 2); + + device const uint16_t *qu2 = qu1 + 32; + float4 acc1u = {0.f, 0.f, 0.f, 0.f}; + float4 acc2u = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1u[0] += yl[2 * i + 0] * (qu1[i] & 0x000F); + acc1u[1] += yl[2 * i + 1] * (qu1[i] & 0x0F00); + acc1u[2] += yl[2 * i + 8] * (qu1[i] & 0x00F0); + acc1u[3] += yl[2 * i + 9] * (qu1[i] & 0xF000); + acc2u[0] += yh[2 * i + 0] * (qu2[i] & 0x000F); + acc2u[1] += yh[2 * i + 1] * (qu2[i] & 0x0F00); + acc2u[2] += yh[2 * i + 8] * (qu2[i] & 0x00F0); + acc2u[3] += yh[2 * i + 9] * (qu2[i] & 0xF000); + } + + sumu[row] += dhu[0] * ((acc1u[0] + 1.f / 256.f * acc1u[1]) * sc8[0] + + (acc1u[2] + 1.f / 256.f * acc1u[3]) * sc8[1] * 1.f / 16.f + + (acc2u[0] + 1.f / 256.f * acc2u[1]) * sc8[4] + + (acc2u[2] + 1.f / 256.f * acc2u[3]) * sc8[5] * 1.f / 16.f) - + dhu[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + qg1 += args.nb01 / 2; + scg += args.nb01 / 2; + dhg += args.nb01 / 2; + qu1 += args.nb01 / 2; + scu += args.nb01 / 2; + dhu += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + + for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) { + const float gate = simd_sum(sumg[row]); + const float up = simd_sum(sumu[row]); + if (tiisg == 0) { + const uint out_row = first_row + row; + float g = gate; + float u = up; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + gate_f32[out_row] = gate; + up_f32[out_row] = up; + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + +kernel void kernel_mul_mv_table_q4_K_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const ds4_metal_q4_expert_table & gate_table, + device const ds4_metal_q4_expert_table & up_table, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * ids, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + if (i02 < 0 || i02 >= args.ne02 || i02 >= 384) { + return; + } + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char *src0_gate_cur = gate_table.experts[(uint)i02]; + device const char *src0_up_cur = up_table.experts[(uint)i02]; + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + const short NSG = FC_mul_mv_nsg; + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K; + device float *gate_f32 = (device float *)dst_gate_cur; + device float *up_f32 = (device float *)dst_up_cur; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride); + const float c = act.clamp_value; + const float route_weight = route_w[0]; + + device const block_q4_K *xg = + (device const block_q4_K *)(src0_gate_cur + (uint64_t)first_row * args.nb01); + device const block_q4_K *xu = + (device const block_q4_K *)(src0_up_cur + (uint64_t)first_row * args.nb01); + device const float *y = (device const float *)src1_cur; + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + float sumg[N_R0_Q4_K] = {0.f}; + float sumu[N_R0_Q4_K] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *scg = (device const uint16_t *)xg[ib].scales + iq; + device const uint16_t *qg1 = (device const uint16_t *)xg[ib].qs + 16 * iq + 4 * ir; + device const half *dhg = &xg[ib].d; + device const uint16_t *scu = (device const uint16_t *)xu[ib].scales + iq; + device const uint16_t *qu1 = (device const uint16_t *)xu[ib].qs + 16 * iq + 4 * ir; + device const half *dhu = &xu[ib].d; + + for (short row = 0; row < N_R0_Q4_K; row++) { + sc16[0] = scg[0] & kmask1; + sc16[1] = scg[2] & kmask1; + sc16[2] = ((scg[4] >> 0) & kmask2) | ((scg[0] & kmask3) >> 2); + sc16[3] = ((scg[4] >> 4) & kmask2) | ((scg[2] & kmask3) >> 2); + + device const uint16_t *qg2 = qg1 + 32; + float4 acc1g = {0.f, 0.f, 0.f, 0.f}; + float4 acc2g = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1g[0] += yl[2 * i + 0] * (qg1[i] & 0x000F); + acc1g[1] += yl[2 * i + 1] * (qg1[i] & 0x0F00); + acc1g[2] += yl[2 * i + 8] * (qg1[i] & 0x00F0); + acc1g[3] += yl[2 * i + 9] * (qg1[i] & 0xF000); + acc2g[0] += yh[2 * i + 0] * (qg2[i] & 0x000F); + acc2g[1] += yh[2 * i + 1] * (qg2[i] & 0x0F00); + acc2g[2] += yh[2 * i + 8] * (qg2[i] & 0x00F0); + acc2g[3] += yh[2 * i + 9] * (qg2[i] & 0xF000); + } + + sumg[row] += dhg[0] * ((acc1g[0] + 1.f / 256.f * acc1g[1]) * sc8[0] + + (acc1g[2] + 1.f / 256.f * acc1g[3]) * sc8[1] * 1.f / 16.f + + (acc2g[0] + 1.f / 256.f * acc2g[1]) * sc8[4] + + (acc2g[2] + 1.f / 256.f * acc2g[3]) * sc8[5] * 1.f / 16.f) - + dhg[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + sc16[0] = scu[0] & kmask1; + sc16[1] = scu[2] & kmask1; + sc16[2] = ((scu[4] >> 0) & kmask2) | ((scu[0] & kmask3) >> 2); + sc16[3] = ((scu[4] >> 4) & kmask2) | ((scu[2] & kmask3) >> 2); + + device const uint16_t *qu2 = qu1 + 32; + float4 acc1u = {0.f, 0.f, 0.f, 0.f}; + float4 acc2u = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1u[0] += yl[2 * i + 0] * (qu1[i] & 0x000F); + acc1u[1] += yl[2 * i + 1] * (qu1[i] & 0x0F00); + acc1u[2] += yl[2 * i + 8] * (qu1[i] & 0x00F0); + acc1u[3] += yl[2 * i + 9] * (qu1[i] & 0xF000); + acc2u[0] += yh[2 * i + 0] * (qu2[i] & 0x000F); + acc2u[1] += yh[2 * i + 1] * (qu2[i] & 0x0F00); + acc2u[2] += yh[2 * i + 8] * (qu2[i] & 0x00F0); + acc2u[3] += yh[2 * i + 9] * (qu2[i] & 0xF000); + } + + sumu[row] += dhu[0] * ((acc1u[0] + 1.f / 256.f * acc1u[1]) * sc8[0] + + (acc1u[2] + 1.f / 256.f * acc1u[3]) * sc8[1] * 1.f / 16.f + + (acc2u[0] + 1.f / 256.f * acc2u[1]) * sc8[4] + + (acc2u[2] + 1.f / 256.f * acc2u[3]) * sc8[5] * 1.f / 16.f) - + dhu[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + + qg1 += args.nb01 / 2; + scg += args.nb01 / 2; + dhg += args.nb01 / 2; + qu1 += args.nb01 / 2; + scu += args.nb01 / 2; + dhu += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + + for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) { + const float gate = simd_sum(sumg[row]); + const float up = simd_sum(sumu[row]); + if (tiisg == 0) { + const uint out_row = first_row + row; + float g = gate; + float u = up; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + gate_f32[out_row] = gate; + up_f32[out_row] = up; + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + +kernel void kernel_mul_mv_addr_q4_K_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const ulong * gate_addrs, + device const ulong * up_addrs, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * ids, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t i02 = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + if (i02 < 0 || i02 >= args.ne02 || i02 >= 384) { + return; + } + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char *src0_gate_cur = + reinterpret_cast(gate_addrs[(uint)i02]); + device const char *src0_up_cur = + reinterpret_cast(up_addrs[(uint)i02]); + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_gate_cur, + src1_cur, + dst_gate_cur, + shmem, + tgpig, + tiisg, + sgitg); + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_up_cur, + src1_cur, + dst_up_cur, + shmem, + tgpig, + tiisg, + sgitg); + + const short NSG = FC_mul_mv_nsg; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K; + device float *gate_f32 = (device float *)dst_gate_cur; + device float *up_f32 = (device float *)dst_up_cur; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride); + const float c = act.clamp_value; + const float route_weight = route_w[0]; + + if (tiisg == 0) { + for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) { + const uint out_row = first_row + row; + float g = gate_f32[out_row]; + float u = up_f32[out_row]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + +kernel void kernel_q4_gather_slots6( + constant ds4_metal_q4_gather_slots6_args &args, + device const char *src_group0, + device const char *src_group1, + device const char *src_group2, + device const char *src_group3, + device const char *src_group4, + device const char *src_group5, + device const int32_t *ids, + device char *dst, + uint3 tgpig [[threadgroup_position_in_grid]], + uint tiitg [[thread_index_in_threadgroup]]) { + const uint slot = tgpig.y; + if (slot >= args.n_slots || args.group_size == 0) return; + + const int32_t expert = ids[slot]; + if (expert < 0) return; + + const uint expert_u = (uint)expert; + const uint group = expert_u / args.group_size; + if (group >= 6) return; + + const uint local_expert = expert_u - group * args.group_size; + device const char *src_group = src_group0; + switch (group) { + case 1: src_group = src_group1; break; + case 2: src_group = src_group2; break; + case 3: src_group = src_group3; break; + case 4: src_group = src_group4; break; + case 5: src_group = src_group5; break; + default: break; + } + + const uint64_t chunk = (uint64_t)tgpig.x * 256ul + (uint64_t)tiitg; + const uint64_t n_chunks = args.expert_bytes >> 4; + if (chunk >= n_chunks) return; + + device const uint4 *src = + (device const uint4 *)(src_group + (uint64_t)local_expert * args.expert_bytes); + device uint4 *out = + (device uint4 *)(dst + (uint64_t)slot * args.expert_bytes); + out[chunk] = src[chunk]; +} + +kernel void kernel_mul_mv_slots6_q4_K_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const char * src0_gate0, + device const char * src0_gate1, + device const char * src0_gate2, + device const char * src0_gate3, + device const char * src0_gate4, + device const char * src0_gate5, + device const char * src0_up0, + device const char * src0_up1, + device const char * src0_up2, + device const char * src0_up3, + device const char * src0_up4, + device const char * src0_up5, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char *src0_gate_cur = src0_gate0; + device const char *src0_up_cur = src0_up0; + switch (idx) { + case 1: src0_gate_cur = src0_gate1; src0_up_cur = src0_up1; break; + case 2: src0_gate_cur = src0_gate2; src0_up_cur = src0_up2; break; + case 3: src0_gate_cur = src0_gate3; src0_up_cur = src0_up3; break; + case 4: src0_gate_cur = src0_gate4; src0_up_cur = src0_up4; break; + case 5: src0_gate_cur = src0_gate5; src0_up_cur = src0_up5; break; + default: break; + } + + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_gate_cur, + src1_cur, + dst_gate_cur, + shmem, + tgpig, + tiisg, + sgitg); + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_up_cur, + src1_cur, + dst_up_cur, + shmem, + tgpig, + tiisg, + sgitg); + + const short NSG = FC_mul_mv_nsg; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K; + device float *gate_f32 = (device float *)dst_gate_cur; + device float *up_f32 = (device float *)dst_up_cur; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride); + const float c = act.clamp_value; + const float route_weight = route_w[0]; + + if (tiisg == 0) { + for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) { + const uint out_row = first_row + row; + float g = gate_f32[out_row]; + float u = up_f32[out_row]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + +static inline device const char *ds4_q4_group24_select( + uint32_t group_id, + device const char *src00, + device const char *src01, + device const char *src02, + device const char *src03, + device const char *src04, + device const char *src05, + device const char *src06, + device const char *src07, + device const char *src08, + device const char *src09, + device const char *src10, + device const char *src11, + device const char *src12, + device const char *src13, + device const char *src14, + device const char *src15, + device const char *src16, + device const char *src17, + device const char *src18, + device const char *src19, + device const char *src20, + device const char *src21, + device const char *src22, + device const char *src23) { + switch (group_id) { + case 1: return src01; + case 2: return src02; + case 3: return src03; + case 4: return src04; + case 5: return src05; + case 6: return src06; + case 7: return src07; + case 8: return src08; + case 9: return src09; + case 10: return src10; + case 11: return src11; + case 12: return src12; + case 13: return src13; + case 14: return src14; + case 15: return src15; + case 16: return src16; + case 17: return src17; + case 18: return src18; + case 19: return src19; + case 20: return src20; + case 21: return src21; + case 22: return src22; + case 23: return src23; + default: return src00; + } +} + +kernel void kernel_mul_mv_group6_q4_K_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const char * src0_gate0, + device const char * src0_gate1, + device const char * src0_gate2, + device const char * src0_gate3, + device const char * src0_gate4, + device const char * src0_gate5, + device const char * src0_up0, + device const char * src0_up1, + device const char * src0_up2, + device const char * src0_up3, + device const char * src0_up4, + device const char * src0_up5, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * ids, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr uint32_t expert_group_size = 64; + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t expert = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + if (expert < 0) { + return; + } + const uint32_t expert_u = (uint32_t)expert; + const uint32_t group_id = expert_u / expert_group_size; + if (group_id >= 6) { + return; + } + const uint32_t expert_local = expert_u - group_id * expert_group_size; + + device const char *src0_gate_cur = src0_gate0; + device const char *src0_up_cur = src0_up0; + switch (group_id) { + case 1: src0_gate_cur = src0_gate1; src0_up_cur = src0_up1; break; + case 2: src0_gate_cur = src0_gate2; src0_up_cur = src0_up2; break; + case 3: src0_gate_cur = src0_gate3; src0_up_cur = src0_up3; break; + case 4: src0_gate_cur = src0_gate4; src0_up_cur = src0_up4; break; + case 5: src0_gate_cur = src0_gate5; src0_up_cur = src0_up5; break; + default: break; + } + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + src0_gate_cur += (uint64_t)expert_local * args.nb02; + src0_up_cur += (uint64_t)expert_local * args.nb02; + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_gate_cur, + src1_cur, + dst_gate_cur, + shmem, + tgpig, + tiisg, + sgitg); + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_up_cur, + src1_cur, + dst_up_cur, + shmem, + tgpig, + tiisg, + sgitg); + + const short NSG = FC_mul_mv_nsg; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K; + device float *gate_f32 = (device float *)dst_gate_cur; + device float *up_f32 = (device float *)dst_up_cur; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride); + const float c = act.clamp_value; + const float route_weight = route_w[0]; + + if (tiisg == 0) { + for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) { + const uint out_row = first_row + row; + float g = gate_f32[out_row]; + float u = up_f32[out_row]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + +kernel void kernel_mul_mv_group8_q4_K_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const char * src0_gate0, + device const char * src0_gate1, + device const char * src0_gate2, + device const char * src0_gate3, + device const char * src0_gate4, + device const char * src0_gate5, + device const char * src0_gate6, + device const char * src0_gate7, + device const char * src0_up0, + device const char * src0_up1, + device const char * src0_up2, + device const char * src0_up3, + device const char * src0_up4, + device const char * src0_up5, + device const char * src0_up6, + device const char * src0_up7, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * ids, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr uint32_t expert_group_size = 48; + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t expert = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + if (expert < 0) { + return; + } + const uint32_t expert_u = (uint32_t)expert; + const uint32_t group_id = expert_u / expert_group_size; + if (group_id >= 8) { + return; + } + const uint32_t expert_local = expert_u - group_id * expert_group_size; + + device const char *src0_gate_cur = src0_gate0; + device const char *src0_up_cur = src0_up0; + switch (group_id) { + case 1: src0_gate_cur = src0_gate1; src0_up_cur = src0_up1; break; + case 2: src0_gate_cur = src0_gate2; src0_up_cur = src0_up2; break; + case 3: src0_gate_cur = src0_gate3; src0_up_cur = src0_up3; break; + case 4: src0_gate_cur = src0_gate4; src0_up_cur = src0_up4; break; + case 5: src0_gate_cur = src0_gate5; src0_up_cur = src0_up5; break; + case 6: src0_gate_cur = src0_gate6; src0_up_cur = src0_up6; break; + case 7: src0_gate_cur = src0_gate7; src0_up_cur = src0_up7; break; + default: break; + } + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + src0_gate_cur += (uint64_t)expert_local * args.nb02; + src0_up_cur += (uint64_t)expert_local * args.nb02; + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_gate_cur, + src1_cur, + dst_gate_cur, + shmem, + tgpig, + tiisg, + sgitg); + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_up_cur, + src1_cur, + dst_up_cur, + shmem, + tgpig, + tiisg, + sgitg); + + const short NSG = FC_mul_mv_nsg; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K; + device float *gate_f32 = (device float *)dst_gate_cur; + device float *up_f32 = (device float *)dst_up_cur; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride); + const float c = act.clamp_value; + const float route_weight = route_w[0]; + + if (tiisg == 0) { + for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) { + const uint out_row = first_row + row; + float g = gate_f32[out_row]; + float u = up_f32[out_row]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + +kernel void kernel_mul_mv_group24_q4_K_id_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src00, + device const char * src01, + device const char * src02, + device const char * src03, + device const char * src04, + device const char * src05, + device const char * src06, + device const char * src07, + device const char * src08, + device const char * src09, + device const char * src10, + device const char * src11, + device const char * src12, + device const char * src13, + device const char * src14, + device const char * src15, + device const char * src16, + device const char * src17, + device const char * src18, + device const char * src19, + device const char * src20, + device const char * src21, + device const char * src22, + device const char * src23, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr uint32_t expert_group_size = 16; + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t expert = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + if (expert < 0) { + return; + } + const uint32_t expert_u = (uint32_t)expert; + const uint32_t group_id = expert_u / expert_group_size; + if (group_id >= 24) { + return; + } + const uint32_t expert_local = expert_u - group_id * expert_group_size; + + device const char *src0_cur = ds4_q4_group24_select(group_id, + src00, src01, src02, src03, + src04, src05, src06, src07, + src08, src09, src10, src11, + src12, src13, src14, src15, + src16, src17, src18, src19, + src20, src21, src22, src23); + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + src0_cur += (uint64_t)expert_local * args.nb02; + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + device char *dst_cur = dst + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_cur, + src1_cur, + dst_cur, + shmem, + tgpig, + tiisg, + sgitg); + + (void)tiitg; +} + +kernel void kernel_mul_mv_group_q4_K_pair_swiglu_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + constant ds4_metal_moe_expert_group_args & group, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_gate, + device char * dst_up, + device char * dst_mid, + device const char * ids, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const int iid1 = tgpig.z / args.nei0; + const int idx = tgpig.z % args.nei0; + + tgpig.z = 0; + + const int32_t expert_global = ((device const int32_t *)(ids + iid1 * args.nbi1))[idx]; + if (expert_global < 0) { + return; + } + const uint32_t expert_u = (uint32_t)expert_global; + if (expert_u < group.expert_base || + expert_u >= group.expert_base + group.expert_count) { + return; + } + const uint32_t expert_local = expert_u - group.expert_base; + + const int64_t i11 = idx % args.ne11; + const int64_t i12 = iid1; + + device const char *src0_gate_cur = src0_gate + (uint64_t)expert_local * args.nb02; + device const char *src0_up_cur = src0_up + (uint64_t)expert_local * args.nb02; + device const char *src1_cur = src1 + i11 * args.nb11 + i12 * args.nb12; + + device char *dst_gate_cur = dst_gate + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + device char *dst_up_cur = dst_up + (idx * args.ne0 + i12 * args.ne1 * args.ne0) * sizeof(float); + + ds4_metal_args_mul_mv args0 = { + args.ne00, args.ne01, 1, + args.nb00, args.nb01, args.nb02, args.nb02, + args.ne10, 1, 1, + args.nb10, args.nb11, args.nb12, args.nb12, + args.ne0, 1, args.nr0, 1, 1, + }; + + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_gate_cur, + src1_cur, + dst_gate_cur, + shmem, + tgpig, + tiisg, + sgitg); + kernel_mul_mv_q4_K_f32_impl( + args0, + src0_up_cur, + src1_cur, + dst_up_cur, + shmem, + tgpig, + tiisg, + sgitg); + + const short NSG = FC_mul_mv_nsg; + const int first_row = (tgpig.x * NSG + sgitg) * N_R0_Q4_K; + device float *gate_f32 = (device float *)dst_gate_cur; + device float *up_f32 = (device float *)dst_up_cur; + const uint64_t pair_row = (uint64_t)i12 * (uint64_t)args.nei0 + (uint64_t)idx; + device float *mid_f32 = (device float *)(dst_mid + pair_row * act.mid_row_stride); + device const float *route_w = (device const float *)(weights + pair_row * act.weight_stride); + const float c = act.clamp_value; + const float route_weight = route_w[0]; + + if (tiisg == 0) { + for (int row = 0; row < N_R0_Q4_K && first_row + row < args.ne0; ++row) { + const uint out_row = first_row + row; + float g = gate_f32[out_row]; + float u = up_f32[out_row]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[out_row] = silu * u * route_weight; + } + } + + (void)tiitg; +} + + +/* IQ2_XXS down projection summed over the token's selected experts (the + * GLM 5.2 routed blobs are IQ2_XXS end to end, unlike DS4 Flash's Q2_K + * down). Same contract as the q2_K sum kernel: mid rows already carry + * silu(gate)*up*route_weight, expert ownership honors the TP split, and + * add_in folds the shared-expert partial when tp_addend is set. */ +kernel void kernel_mul_mv_id_iq2_xxs_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src0s, + device const char * src1, + device char * dst, + device const char * ids, + device const char * add_in, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_IQ2_XXS; + const int nb = args.ne00/QK_K; + const int nb32 = nb * (QK_K / 32); + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + float sumf[nr0] = {0.f}; + float yl[32]; + + threadgroup uint64_t *svalues = (threadgroup uint64_t *)(shmem); + threadgroup uint8_t *ssigns = (threadgroup uint8_t *)(svalues + 256); + { + int nval = 4; + int pos = (32 * sgitg + tiisg) * nval; + for (int i = 0; i < nval; ++i) svalues[pos + i] = ds4_metal_iq2xxs_grid[pos + i]; + nval = 2; + pos = (32 * sgitg + tiisg) * nval; + for (int i = 0; i < nval; ++i) ssigns[pos + i] = ds4_metal_ksigns_iq2xs[pos + i]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + const int ix = tiisg; + + for (int expert_slot = 0; expert_slot < args.nei0; expert_slot++) { + const int32_t expert = token_ids[expert_slot]; + if (!ds4_tp_owns_expert(expert, args.ne02, args.tp_rank, args.tp_world)) continue; + device const block_iq2_xxs *x = + (device const block_iq2_xxs *)(src0s + (int64_t)(expert - args.tp_expert_base)*args.nb02 + first_row*args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot*args.nb11); + device const float *y4 = y + 32 * ix; + + for (int ib32 = ix; ib32 < nb32; ib32 += 32) { + for (short i = 0; i < 32; ++i) { + yl[i] = y4[i]; + } + + const int ibl = ib32 / (QK_K / 32); + const int ib = ib32 % (QK_K / 32); + + device const block_iq2_xxs *xr = x + ibl; + device const uint16_t *q2 = xr->qs + 4 * ib; + device const half *dh = &xr->d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + device const uint8_t *aux8 = (device const uint8_t *)q2; + const uint32_t aux32 = q2[2] | (q2[3] << 16); + const float d = (float)dh[0] * (0.5f + (aux32 >> 28)); + + float s = 0; + for (short l = 0; l < 4; ++l) { + const threadgroup uint8_t *grid = (const threadgroup uint8_t *)(svalues + aux8[l]); + const uint8_t sign = ssigns[(aux32 >> 7 * l) & 127]; + for (short j = 0; j < 8; ++j) { + s += yl[8 * l + j] * grid[j] * + (sign & ds4_metal_kmask_iq2xs[j] ? -1.f : 1.f); + } + } + sumf[row] += d * s; + } + dh += args.nb01 / 2; + q2 += args.nb01 / 2; + } + + y4 += 32 * 32; + } + } + + device float * dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]) * 0.25f; + if (tiisg == 0) { + float outv = sum_all; + if (args.tp_addend) { + outv += ((device const float *) add_in)[first_row + row]; + } + dst_f32[first_row + row] = outv; + } + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_id_q2_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src0s, + device const char * src1, + device char * dst, + device const char * ids, + device const char * add_in, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q2_K; + const int nb = args.ne00/QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + float sumf[nr0] = {0.f}; + + const short ix = tiisg/8; + const short it = tiisg%8; + const short iq = it/4; + const short ir = it%4; + const short is = (8*ir)/16; + + for (int expert_slot = 0; expert_slot < args.nei0; expert_slot++) { + const int32_t expert = token_ids[expert_slot]; + if (!ds4_tp_owns_expert(expert, args.ne02, args.tp_rank, args.tp_world)) continue; + device const block_q2_K * x = (device const block_q2_K *)(src0s + (int64_t)(expert - args.tp_expert_base)*args.nb02 + first_row*args.nb01); + device const float * y = (device const float *)(token_src1 + expert_slot*args.nb11); + device const float * y4 = y + ix * QK_K + 128 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[32]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + for (short i = 0; i < 8; ++i) { + yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0]; + yl[i+ 8] = y4[i+32]; sumy[1] += yl[i+ 8]; + yl[i+16] = y4[i+64]; sumy[2] += yl[i+16]; + yl[i+24] = y4[i+96]; sumy[3] += yl[i+24]; + } + + device const uint8_t * sc = (device const uint8_t *)x[ib].scales + 8*iq + is; + device const uint16_t * qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half * dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + for (int i = 0; i < 8; i += 2) { + acc1[0] += yl[i+ 0] * (qs[i/2] & 0x0003); + acc2[0] += yl[i+ 1] * (qs[i/2] & 0x0300); + acc1[1] += yl[i+ 8] * (qs[i/2] & 0x000c); + acc2[1] += yl[i+ 9] * (qs[i/2] & 0x0c00); + acc1[2] += yl[i+16] * (qs[i/2] & 0x0030); + acc2[2] += yl[i+17] * (qs[i/2] & 0x3000); + acc1[3] += yl[i+24] * (qs[i/2] & 0x00c0); + acc2[3] += yl[i+25] * (qs[i/2] & 0xc000); + } + float dall = dh[0]; + float dmin = dh[1] * 1.f/16.f; + sumf[row] += dall * ((acc1[0] + 1.f/256.f * acc2[0]) * (sc[0] & 0xF) * 1.f/ 1.f + + (acc1[1] + 1.f/256.f * acc2[1]) * (sc[2] & 0xF) * 1.f/ 4.f + + (acc1[2] + 1.f/256.f * acc2[2]) * (sc[4] & 0xF) * 1.f/16.f + + (acc1[3] + 1.f/256.f * acc2[3]) * (sc[6] & 0xF) * 1.f/64.f) - + dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + + sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0)); + } + + qs += args.nb01/2; + sc += args.nb01; + dh += args.nb01/2; + } + + y4 += 4 * QK_K; + } + } + + device float * dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) { + float outv = sum_all; + if (args.tp_addend) { + outv += ((device const float *) add_in)[first_row + row]; + } + dst_f32[first_row + row] = outv; + } + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_slots6_q2_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src00, + device const char * src01, + device const char * src02, + device const char * src03, + device const char * src04, + device const char * src05, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q2_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + float sumf[nr0] = {0.f}; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const short is = (8 * ir) / 16; + + for (int expert_slot = 0; expert_slot < 6; expert_slot++) { + device const char *src0_cur = src00; + switch (expert_slot) { + case 1: src0_cur = src01; break; + case 2: src0_cur = src02; break; + case 3: src0_cur = src03; break; + case 4: src0_cur = src04; break; + case 5: src0_cur = src05; break; + default: break; + } + device const block_q2_K *x = + (device const block_q2_K *)(src0_cur + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[32]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16]; + yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24]; + } + + device const uint8_t *sc = (device const uint8_t *)x[ib].scales + 8 * iq + is; + device const uint16_t *qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + for (int i = 0; i < 8; i += 2) { + acc1[0] += yl[i + 0] * (qs[i / 2] & 0x0003); + acc2[0] += yl[i + 1] * (qs[i / 2] & 0x0300); + acc1[1] += yl[i + 8] * (qs[i / 2] & 0x000c); + acc2[1] += yl[i + 9] * (qs[i / 2] & 0x0c00); + acc1[2] += yl[i + 16] * (qs[i / 2] & 0x0030); + acc2[2] += yl[i + 17] * (qs[i / 2] & 0x3000); + acc1[3] += yl[i + 24] * (qs[i / 2] & 0x00c0); + acc2[3] += yl[i + 25] * (qs[i / 2] & 0xc000); + } + float dall = dh[0]; + float dmin = dh[1] * 1.f / 16.f; + sumf[row] += dall * ((acc1[0] + 1.f / 256.f * acc2[0]) * (sc[0] & 0xF) * 1.f / 1.f + + (acc1[1] + 1.f / 256.f * acc2[1]) * (sc[2] & 0xF) * 1.f / 4.f + + (acc1[2] + 1.f / 256.f * acc2[2]) * (sc[4] & 0xF) * 1.f / 16.f + + (acc1[3] + 1.f / 256.f * acc2[3]) * (sc[6] & 0xF) * 1.f / 64.f) - + dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + + sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0)); + } + + qs += args.nb01 / 2; + sc += args.nb01; + dh += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + } + + device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) dst_f32[first_row + row] = sum_all; + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_addr_q2_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const uint64_t * addrs, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q2_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + device const int32_t *token_ids = + (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + + float sumf[nr0] = {0.f}; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const short is = (8 * ir) / 16; + + for (int expert_slot = 0; expert_slot < 6; expert_slot++) { + const int32_t expert = token_ids[expert_slot]; + if (expert < 0 || expert >= args.ne02 || expert >= 384) { + continue; + } + const uint64_t addr = addrs[(uint)expert]; + if (addr == 0) { + continue; + } + device const char *src0_cur = + reinterpret_cast(addr); + device const block_q2_K *x = + (device const block_q2_K *)(src0_cur + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[32]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16]; + yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24]; + } + + device const uint8_t *sc = (device const uint8_t *)x[ib].scales + 8 * iq + is; + device const uint16_t *qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + for (int i = 0; i < 8; i += 2) { + acc1[0] += yl[i + 0] * (qs[i / 2] & 0x0003); + acc2[0] += yl[i + 1] * (qs[i / 2] & 0x0300); + acc1[1] += yl[i + 8] * (qs[i / 2] & 0x000c); + acc2[1] += yl[i + 9] * (qs[i / 2] & 0x0c00); + acc1[2] += yl[i + 16] * (qs[i / 2] & 0x0030); + acc2[2] += yl[i + 17] * (qs[i / 2] & 0x3000); + acc1[3] += yl[i + 24] * (qs[i / 2] & 0x00c0); + acc2[3] += yl[i + 25] * (qs[i / 2] & 0xc000); + } + float dall = dh[0]; + float dmin = dh[1] * 1.f / 16.f; + sumf[row] += dall * ((acc1[0] + 1.f / 256.f * acc2[0]) * (sc[0] & 0xF) * 1.f / 1.f + + (acc1[1] + 1.f / 256.f * acc2[1]) * (sc[2] & 0xF) * 1.f / 4.f + + (acc1[2] + 1.f / 256.f * acc2[2]) * (sc[4] & 0xF) * 1.f / 16.f + + (acc1[3] + 1.f / 256.f * acc2[3]) * (sc[6] & 0xF) * 1.f / 64.f) - + dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + + sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0)); + } + qs += args.nb01 / 2; + sc += args.nb01; + dh += args.nb01 / 2; + } + y4 += 4 * QK_K; + } + } + + device float * dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) dst_f32[first_row + row] = sum_all; + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_addr_q2_K_sum6_masked_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_stream_expert_split_args & split, + device const uint64_t * addrs, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q2_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + device const int32_t *token_ids = + (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + + float sumf[nr0] = {0.f}; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + const short is = (8 * ir) / 16; + + for (int expert_slot = 0; expert_slot < 6; expert_slot++) { + if ((split.active_mask & (1u << (uint)expert_slot)) == 0) { + continue; + } + const int32_t expert = token_ids[expert_slot]; + if (expert < 0 || expert >= args.ne02 || expert >= 384) { + continue; + } + const uint64_t addr = addrs[(uint)expert]; + if (addr == 0) { + continue; + } + device const char *src0_cur = + reinterpret_cast(addr); + device const block_q2_K *x = + (device const block_q2_K *)(src0_cur + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 128 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[32]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yl[i + 16] = y4[i + 64]; sumy[2] += yl[i + 16]; + yl[i + 24] = y4[i + 96]; sumy[3] += yl[i + 24]; + } + + device const uint8_t *sc = (device const uint8_t *)x[ib].scales + 8 * iq + is; + device const uint16_t *qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + for (int i = 0; i < 8; i += 2) { + acc1[0] += yl[i + 0] * (qs[i / 2] & 0x0003); + acc2[0] += yl[i + 1] * (qs[i / 2] & 0x0300); + acc1[1] += yl[i + 8] * (qs[i / 2] & 0x000c); + acc2[1] += yl[i + 9] * (qs[i / 2] & 0x0c00); + acc1[2] += yl[i + 16] * (qs[i / 2] & 0x0030); + acc2[2] += yl[i + 17] * (qs[i / 2] & 0x3000); + acc1[3] += yl[i + 24] * (qs[i / 2] & 0x00c0); + acc2[3] += yl[i + 25] * (qs[i / 2] & 0xc000); + } + float dall = dh[0]; + float dmin = dh[1] * 1.f / 16.f; + sumf[row] += dall * ((acc1[0] + 1.f / 256.f * acc2[0]) * (sc[0] & 0xF) * 1.f / 1.f + + (acc1[1] + 1.f / 256.f * acc2[1]) * (sc[2] & 0xF) * 1.f / 4.f + + (acc1[2] + 1.f / 256.f * acc2[2]) * (sc[4] & 0xF) * 1.f / 16.f + + (acc1[3] + 1.f / 256.f * acc2[3]) * (sc[6] & 0xF) * 1.f / 64.f) - + dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + + sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0)); + } + qs += args.nb01 / 2; + sc += args.nb01; + dh += args.nb01 / 2; + } + y4 += 4 * QK_K; + } + } + + device float * dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) { + if (split.accumulate) { + dst_f32[first_row + row] += sum_all; + } else { + dst_f32[first_row + row] = sum_all; + } + } + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_id_q4_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src0s, + device const char * src1, + device char * dst, + device const char * ids, + device const char * add_in, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q4_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + float sumf[nr0] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int expert_slot = 0; expert_slot < args.nei0; expert_slot++) { + const int32_t expert = token_ids[expert_slot]; + if (!ds4_tp_owns_expert(expert, args.ne02, args.tp_rank, args.tp_world)) continue; + device const block_q4_K *x = + (device const block_q4_K *)(src0s + (int64_t)(expert - args.tp_expert_base) * args.nb02 + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + } + + q1 += args.nb01 / 2; + sc += args.nb01 / 2; + dh += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + } + + device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) { + float outv = sum_all; + if (args.tp_addend) { + outv += ((device const float *) add_in)[first_row + row]; + } + dst_f32[first_row + row] = outv; + } + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_group_q4_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + constant ds4_metal_moe_expert_group_args & group, + device const char * src0s, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q4_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + float sumf[nr0] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int expert_slot = 0; expert_slot < 6; expert_slot++) { + const int32_t expert = token_ids[expert_slot]; + if (expert < 0) { + continue; + } + const uint32_t expert_u = (uint32_t)expert; + if (expert_u < group.expert_base || + expert_u >= group.expert_base + group.expert_count) { + continue; + } + const uint32_t expert_local = expert_u - group.expert_base; + + device const block_q4_K *x = + (device const block_q4_K *)(src0s + (uint64_t)expert_local * args.nb02 + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + } + + q1 += args.nb01 / 2; + sc += args.nb01 / 2; + dh += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + } + + device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) { + if (group.accumulate) { + dst_f32[first_row + row] += sum_all; + } else { + dst_f32[first_row + row] = sum_all; + } + } + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_table_q4_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const ds4_metal_q4_expert_table & table, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q4_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + float sumf[nr0] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int expert_slot = 0; expert_slot < 6; expert_slot++) { + const int32_t expert = token_ids[expert_slot]; + if (expert < 0 || expert >= args.ne02 || expert >= 384) { + return; + } + device const block_q4_K *x = + (device const block_q4_K *)(table.experts[(uint)expert] + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + } + + q1 += args.nb01 / 2; + sc += args.nb01 / 2; + dh += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + } + + device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) dst_f32[first_row + row] = sum_all; + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_addr_q4_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const ulong * addrs, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q4_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + float sumf[nr0] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int expert_slot = 0; expert_slot < 6; expert_slot++) { + const int32_t expert = token_ids[expert_slot]; + if (expert < 0 || expert >= args.ne02 || expert >= 384) { + return; + } + device const char *expert_base = + reinterpret_cast(addrs[(uint)expert]); + device const block_q4_K *x = + (device const block_q4_K *)(expert_base + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + } + + q1 += args.nb01 / 2; + sc += args.nb01 / 2; + dh += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + } + + device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) dst_f32[first_row + row] = sum_all; + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_slots6_q4_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src00, + device const char * src01, + device const char * src02, + device const char * src03, + device const char * src04, + device const char * src05, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q4_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + float sumf[nr0] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int expert_slot = 0; expert_slot < 6; expert_slot++) { + device const char *src0_cur = src00; + switch (expert_slot) { + case 1: src0_cur = src01; break; + case 2: src0_cur = src02; break; + case 3: src0_cur = src03; break; + case 4: src0_cur = src04; break; + case 5: src0_cur = src05; break; + default: break; + } + device const block_q4_K *x = + (device const block_q4_K *)(src0_cur + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + } + + q1 += args.nb01 / 2; + sc += args.nb01 / 2; + dh += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + } + + device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) dst_f32[first_row + row] = sum_all; + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_group6_q4_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src00, + device const char * src01, + device const char * src02, + device const char * src03, + device const char * src04, + device const char * src05, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr uint32_t expert_group_size = 64; + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q4_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + float sumf[nr0] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int expert_slot = 0; expert_slot < 6; expert_slot++) { + const int32_t expert = token_ids[expert_slot]; + if (expert < 0) { + continue; + } + const uint32_t expert_u = (uint32_t)expert; + const uint32_t group_id = expert_u / expert_group_size; + if (group_id >= 6) { + continue; + } + const uint32_t expert_local = expert_u - group_id * expert_group_size; + + device const char *src0_cur = src00; + switch (group_id) { + case 1: src0_cur = src01; break; + case 2: src0_cur = src02; break; + case 3: src0_cur = src03; break; + case 4: src0_cur = src04; break; + case 5: src0_cur = src05; break; + default: break; + } + + device const block_q4_K *x = + (device const block_q4_K *)(src0_cur + (uint64_t)expert_local * args.nb02 + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + } + + q1 += args.nb01 / 2; + sc += args.nb01 / 2; + dh += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + } + + device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) dst_f32[first_row + row] = sum_all; + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_group8_q4_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src00, + device const char * src01, + device const char * src02, + device const char * src03, + device const char * src04, + device const char * src05, + device const char * src06, + device const char * src07, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr uint32_t expert_group_size = 48; + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q4_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + float sumf[nr0] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int expert_slot = 0; expert_slot < 6; expert_slot++) { + const int32_t expert = token_ids[expert_slot]; + if (expert < 0) { + continue; + } + const uint32_t expert_u = (uint32_t)expert; + const uint32_t group_id = expert_u / expert_group_size; + if (group_id >= 8) { + continue; + } + const uint32_t expert_local = expert_u - group_id * expert_group_size; + + device const char *src0_cur = src00; + switch (group_id) { + case 1: src0_cur = src01; break; + case 2: src0_cur = src02; break; + case 3: src0_cur = src03; break; + case 4: src0_cur = src04; break; + case 5: src0_cur = src05; break; + case 6: src0_cur = src06; break; + case 7: src0_cur = src07; break; + default: break; + } + + device const block_q4_K *x = + (device const block_q4_K *)(src0_cur + (uint64_t)expert_local * args.nb02 + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + } + + q1 += args.nb01 / 2; + sc += args.nb01 / 2; + dh += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + } + + device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) dst_f32[first_row + row] = sum_all; + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +kernel void kernel_mul_mv_group24_q4_K_sum6_f32( + constant ds4_metal_args_mul_mv_id & args, + device const char * src00, + device const char * src01, + device const char * src02, + device const char * src03, + device const char * src04, + device const char * src05, + device const char * src06, + device const char * src07, + device const char * src08, + device const char * src09, + device const char * src10, + device const char * src11, + device const char * src12, + device const char * src13, + device const char * src14, + device const char * src15, + device const char * src16, + device const char * src17, + device const char * src18, + device const char * src19, + device const char * src20, + device const char * src21, + device const char * src22, + device const char * src23, + device const char * src1, + device char * dst, + device const char * ids, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr uint32_t expert_group_size = 16; + const short NSG = FC_mul_mv_nsg; + const short nr0 = N_R0_Q4_K; + const int nb = args.ne00 / QK_K; + const int first_row = (tgpig.x * NSG + sgitg) * nr0; + const uint token = tgpig.y; + device const int32_t *token_ids = (device const int32_t *)(ids + (uint64_t)token * args.nbi1); + device const char *token_src1 = src1 + (uint64_t)token * args.nb12; + + constexpr uint16_t kmask1 = 0x3f3f; + constexpr uint16_t kmask2 = 0x0f0f; + constexpr uint16_t kmask3 = 0xc0c0; + + const short ix = tiisg / 8; + const short it = tiisg % 8; + const short iq = it / 4; + const short ir = it % 4; + + float sumf[nr0] = {0.f}; + uint16_t sc16[4]; + thread const uint8_t *sc8 = (thread const uint8_t *)sc16; + + for (int expert_slot = 0; expert_slot < 6; expert_slot++) { + const int32_t expert = token_ids[expert_slot]; + if (expert < 0) { + continue; + } + const uint32_t expert_u = (uint32_t)expert; + const uint32_t group_id = expert_u / expert_group_size; + if (group_id >= 24) { + continue; + } + const uint32_t expert_local = expert_u - group_id * expert_group_size; + + device const char *src0_cur = ds4_q4_group24_select(group_id, + src00, src01, src02, src03, + src04, src05, src06, src07, + src08, src09, src10, src11, + src12, src13, src14, src15, + src16, src17, src18, src19, + src20, src21, src22, src23); + device const block_q4_K *x = + (device const block_q4_K *)(src0_cur + (uint64_t)expert_local * args.nb02 + first_row * args.nb01); + device const float *y = (device const float *)(token_src1 + expert_slot * args.nb11); + device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir; + + for (int ib = ix; ib < nb; ib += 4) { + float yl[16]; + float yh[16]; + float4 sumy = {0.f, 0.f, 0.f, 0.f}; + + for (short i = 0; i < 8; ++i) { + yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; + yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; + yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; + yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; + } + + device const uint16_t *sc = (device const uint16_t *)x[ib].scales + iq; + device const uint16_t *q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; + device const half *dh = &x[ib].d; + + for (short row = 0; row < nr0; row++) { + if (first_row + row < args.ne0) { + sc16[0] = sc[0] & kmask1; + sc16[1] = sc[2] & kmask1; + sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); + sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); + + device const uint16_t *q2 = q1 + 32; + + float4 acc1 = {0.f, 0.f, 0.f, 0.f}; + float4 acc2 = {0.f, 0.f, 0.f, 0.f}; + + FOR_UNROLL (short i = 0; i < 4; ++i) { + acc1[0] += yl[2 * i + 0] * (q1[i] & 0x000F); + acc1[1] += yl[2 * i + 1] * (q1[i] & 0x0F00); + acc1[2] += yl[2 * i + 8] * (q1[i] & 0x00F0); + acc1[3] += yl[2 * i + 9] * (q1[i] & 0xF000); + acc2[0] += yh[2 * i + 0] * (q2[i] & 0x000F); + acc2[1] += yh[2 * i + 1] * (q2[i] & 0x0F00); + acc2[2] += yh[2 * i + 8] * (q2[i] & 0x00F0); + acc2[3] += yh[2 * i + 9] * (q2[i] & 0xF000); + } + + sumf[row] += dh[0] * ((acc1[0] + 1.f / 256.f * acc1[1]) * sc8[0] + + (acc1[2] + 1.f / 256.f * acc1[3]) * sc8[1] * 1.f / 16.f + + (acc2[0] + 1.f / 256.f * acc2[1]) * sc8[4] + + (acc2[2] + 1.f / 256.f * acc2[3]) * sc8[5] * 1.f / 16.f) - + dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + + sumy[2] * sc8[6] + sumy[3] * sc8[7]); + } + + q1 += args.nb01 / 2; + sc += args.nb01 / 2; + dh += args.nb01 / 2; + } + + y4 += 4 * QK_K; + } + } + + device float *dst_f32 = (device float *)(dst + (uint64_t)token * args.nb1); + for (int row = 0; row < nr0 && first_row + row < args.ne0; row++) { + const float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) dst_f32[first_row + row] = sum_all; + } + + (void)shmem; + (void)tiitg; + (void)tgpig; +} + +#define QK_NL 16 + +// Builds the compact per-expert work map used by batched MoE matmul. DS4 routes +// each token to a small fixed top-k list, so this turns token-major ids into +// expert-major slices that the tiled matmul can consume. +template +kernel void kernel_mul_mm_id_map0( + constant ds4_metal_args_mul_mm_id_map0 & args, + device const char * src2, + device char * htpe, + device char * hids, + threadgroup char * shmem [[threadgroup(0)]], + ushort tpitg[[thread_position_in_threadgroup]], + ushort ntg[[threads_per_threadgroup]]) { + const short ide = tpitg; + + uint32_t n_all = 0; + + device int32_t * ids_i32 = (device int32_t *) hids + ide*args.ne21; + + for (int i21 = 0; i21 < args.ne21; i21 += ntg) { + if (i21 + tpitg < args.ne21) { + device const int32_t * src2_i32 = (device const int32_t *) (src2 + (i21 + tpitg)*args.nb21); + + threadgroup uint16_t * sids = (threadgroup uint16_t *) shmem + tpitg*ne20; + + #pragma unroll(ne20) + for (short i20 = 0; i20 < ne20; i20++) { + sids[i20] = src2_i32[i20]; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short t = 0; t < ntg; t++) { + if (i21 + t >= args.ne21) { + break; + } + + threadgroup const uint16_t * sids = (threadgroup const uint16_t *) shmem + t*ne20; + + short sel = 0; + #pragma unroll(ne20) + for (short i20 = 0; i20 < ne20; i20++) { + sel += (sids[i20] == ide)*(i20 + 1); + } + + ids_i32[n_all] = (i21 + t)*ne20 + sel - 1; + + n_all += sel > 0; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + device uint32_t * tpe_u32 = (device uint32_t *) (htpe); + tpe_u32[ide] = n_all; +} + +typedef decltype(kernel_mul_mm_id_map0<1>) kernel_mul_mm_id_map0_t; + +// Host-visible map builders for the routed-expert counts used by DS4 graph +// shapes. Some arities are generic leftovers retained for nearby batch sizes. +template [[host_name("kernel_mul_mm_id_map0_ne20_1" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<1>; +template [[host_name("kernel_mul_mm_id_map0_ne20_2" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<2>; +template [[host_name("kernel_mul_mm_id_map0_ne20_4" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<4>; +template [[host_name("kernel_mul_mm_id_map0_ne20_5" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<5>; +template [[host_name("kernel_mul_mm_id_map0_ne20_6" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<6>; +template [[host_name("kernel_mul_mm_id_map0_ne20_8" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<8>; +template [[host_name("kernel_mul_mm_id_map0_ne20_10")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<10>; +template [[host_name("kernel_mul_mm_id_map0_ne20_16")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<16>; +template [[host_name("kernel_mul_mm_id_map0_ne20_22")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<22>; + +// Batched routed-expert matmul. It reads the expert-major map produced above, +// loads selected expert weights, and writes results back to token-major slots +// so the DS4 FFN can apply SwiGLU, weighting, and the down projection. +template +kernel void kernel_mul_mm_id( + constant ds4_metal_args_mul_mm_id & args, + device const char * src0, + device const char * src1, + device const char * htpe, + device const char * hids, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr int NR0 = 64; + static_assert(NR1 == 32, "kernel_mul_mm_id accumulator layout supports only 32 routed rows"); + + constexpr int NK = 32; + constexpr int NL0 = NK/16; + constexpr int NL1 = NK/8; + constexpr int SA_BYTES = NR0 * NR1 * (int)sizeof(S0); + + threadgroup S0 * sa = (threadgroup S0 *)(shmem); + threadgroup S1 * sb = (threadgroup S1 *)(shmem + SA_BYTES); + + const int im = tgpig.z; + const int r0 = tgpig.y*NR0; + const int r1 = tgpig.x*NR1; + + device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe); + device const int32_t * ids_i32 = (device const int32_t *) (hids); + + const int32_t neh1 = tpe_u32[im]; + + if (r1 >= neh1) { + return; + } + + const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; + const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1; + + if (!ds4_tp_owns_expert(im, args.ne02, args.tp_rank, args.tp_world)) { + /* Unowned expert under the TP split: zero this tile's output rows so + * the downstream swiglu/sum stages stay unchanged. Each (token,slot) + * row belongs to exactly one expert, so nothing else writes them. */ + for (short j = sgitg; j < nr1; j += 4) { + const int idj = ids_i32[im*args.ne21 + r1 + j]; + + const short ide = idj % args.ne20; + const short idt = idj / args.ne20; + + device float * D = (device float *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0; + + for (int i = tiisg; i < nr0; i += 32) { + D[i] = 0.0f; + } + } + return; + } + + const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; + const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; + + const short il0 = (tiitg % NL0); + + short il = il0; + + const int id = ids_i32[im*args.ne21 + r1 + lr1]; + + const short i11 = (id % args.ne20) % args.ne11; + const short i12 = (id / args.ne20); + const short i13 = 0; + + const uint64_t offset0 = (uint64_t)(im - args.tp_expert_base)*args.nb02 + i13*args.nb03; + const short offset1 = il0/nl; + + device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1; + + const short iy = 8*(tiitg % NL1); + + device const T1 * y = (device const T1 *)(src1 + + args.nb13*i13 + + args.nb12*i12 + + args.nb11*i11 + + args.nb10*iy); + + S0_8x8 ma[4]; + S1_8x8 mb[2]; + + simdgroup_float8x8 mc[8]; + + for (short i = 0; i < 8; i++){ + mc[i] = make_filled_simdgroup_matrix(0.f); + } + + for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { + if (is_same::value && FC_mul_mm_bc_inp) { + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + *(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0; + } + } else { + S0_4x4 temp_a; + dequantize_func(x, il, temp_a); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + *(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4]; + } + } + + if (FC_mul_mm_bc_inp) { + for (short i = 0; i < 8; ++i) { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short lx = i; + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + + *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0; + } + } else { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + + *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y)); + } + + il = (il + 2 < nl) ? il + 2 : il % 2; + x = (il < 2) ? x + (2 + nl - 1)/nl : x; + + y += NK; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2)); + threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2)); + + FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma[i], lsma + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++){ + simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]); + } + + lsma += 8*64; + lsmb += 4*64; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short j = sgitg; j < nr1; j += 4) { + const int idj = ids_i32[im*args.ne21 + r1 + j]; + + const short ide = idj % args.ne20; + const short idt = idj / args.ne20; + + device float * D = (device float *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0; + device float4 * D4 = (device float4 *) D; + + threadgroup float * C = (threadgroup float *) shmem + j*NR0; + threadgroup float4 * C4 = (threadgroup float4 *) C; + + int i = tiisg; + for (; i < nr0/4; i += 32) { + *(D4 + i) = *(C4 + i); + } + + i = (4*(nr0/4)) + tiisg; + for (; i < nr0; i += 32) { + *(D + i) = *(C + i); + } + } +} + +// Address-table variant used by SSD streaming. The routing ids remain the +// model's original expert ids, but each expert's resident buffer is found via a +// GPU-address table instead of a contiguous full-layer tensor. +template +kernel void kernel_mul_mm_id_addr( + constant ds4_metal_args_mul_mm_id & args, + device const uint64_t * src0_addrs, + device const char * src1, + device const char * htpe, + device const char * hids, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr int NR0 = 64; + static_assert(NR1 == 32, "kernel_mul_mm_id_addr accumulator layout supports only 32 routed rows"); + + constexpr int NK = 32; + constexpr int NL0 = NK/16; + constexpr int NL1 = NK/8; + constexpr int SA_BYTES = NR0 * NR1 * (int)sizeof(S0); + + threadgroup S0 * sa = (threadgroup S0 *)(shmem); + threadgroup S1 * sb = (threadgroup S1 *)(shmem + SA_BYTES); + + const int im = tgpig.z; + const int r0 = tgpig.y*NR0; + const int r1 = tgpig.x*NR1; + + device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe); + device const int32_t * ids_i32 = (device const int32_t *) (hids); + + const int32_t neh1 = tpe_u32[im]; + + if (r1 >= neh1) { + return; + } + + const uint64_t base_addr = src0_addrs[im]; + if (base_addr == 0) { + return; + } + device const char * src0 = reinterpret_cast(base_addr); + + const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; + const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1; + + const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; + const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; + + const short il0 = (tiitg % NL0); + + short il = il0; + + const int id = ids_i32[im*args.ne21 + r1 + lr1]; + + const short i11 = (id % args.ne20) % args.ne11; + const short i12 = (id / args.ne20); + const short i13 = 0; + + const short offset1 = il0/nl; + + device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0)) + offset1; + + const short iy = 8*(tiitg % NL1); + + device const T1 * y = (device const T1 *)(src1 + + args.nb13*i13 + + args.nb12*i12 + + args.nb11*i11 + + args.nb10*iy); + + S0_8x8 ma[4]; + S1_8x8 mb[2]; + + simdgroup_float8x8 mc[8]; + + for (short i = 0; i < 8; i++){ + mc[i] = make_filled_simdgroup_matrix(0.f); + } + + for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { + if (is_same::value && FC_mul_mm_bc_inp) { + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + *(sa + 64*ib + 8*ly + lx) = loop_k + 16*il + i < args.ne00 ? *((device T0 *) x + i) : 0; + } + } else { + S0_4x4 temp_a; + dequantize_func(x, il, temp_a); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + + const short ib = 8*sx + sy; + + *(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4]; + } + } + + if (FC_mul_mm_bc_inp) { + for (short i = 0; i < 8; ++i) { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short lx = i; + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + + *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0; + } + } else { + const short sx = (tiitg%NL1); + const short sy = (tiitg/NL1)/8; + + const short ly = (tiitg/NL1)%8; + + const short ib = 4*sx + sy; + + *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y)); + } + + il = (il + 2 < nl) ? il + 2 : il % 2; + x = (il < 2) ? x + (2 + nl - 1)/nl : x; + + y += NK; + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2)); + threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2)); + + FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma[i], lsma + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++){ + simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]); + } + + lsma += 8*64; + lsmb += 4*64; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (short j = sgitg; j < nr1; j += 4) { + const int idj = ids_i32[im*args.ne21 + r1 + j]; + + const short ide = idj % args.ne20; + const short idt = idj / args.ne20; + + device float * D = (device float *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0; + device float4 * D4 = (device float4 *) D; + + threadgroup float * C = (threadgroup float *) shmem + j*NR0; + threadgroup float4 * C4 = (threadgroup float4 *) C; + + int i = tiisg; + for (; i < nr0/4; i += 32) { + *(D4 + i) = *(C4 + i); + } + + i = (4*(nr0/4)) + tiisg; + for (; i < nr0; i += 32) { + *(D + i) = *(C + i); + } + } +} + +// Fused routed gate+up grouped matmul with the SwiGLU epilogue, generic over +// the expert quant block. Both A streams share one staged B tile per k-step; +// each output keeps the exact MMA accumulation order of the separate GEMMs, +// and the epilogue matches kernel_dsv4_moe_swiglu_weight_f16, so the fused +// result is bit-identical to the unfused path. +template +kernel void kernel_mul_mm_id_pair_swiglu_f16_impl( + constant ds4_metal_args_mul_mm_id & args, + constant ds4_metal_dsv4_moe_swiglu_weight_args & act, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device const char * htpe, + device const char * hids, + device char * dst_mid, + device const char * weights, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + threadgroup half *sa_gate = (threadgroup half *)(shmem); + threadgroup half *sa_up = (threadgroup half *)(shmem + 4096); + threadgroup half *sb = (threadgroup half *)(shmem + 8192); + + constexpr int NR0 = 64; + constexpr int NR1 = 32; + constexpr int NK = 32; + constexpr int NL0 = NK/16; + constexpr int NL1 = NK/8; + + const int im = tgpig.z; + const int r0 = tgpig.y*NR0; + const int r1 = tgpig.x*NR1; + + device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe); + device const int32_t * ids_i32 = (device const int32_t *) (hids); + + const int32_t neh1 = tpe_u32[im]; + + if (r1 >= neh1) { + return; + } + + const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0; + const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1; + + const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; + const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; + + const short il0 = (tiitg % NL0); + short il = il0; + + const int id = ids_i32[im*args.ne21 + r1 + lr1]; + + const short i11 = (id % args.ne20) % args.ne11; + const short i12 = (id / args.ne20); + const short i13 = 0; + + const uint64_t offset0 = im*args.nb02 + i13*args.nb03; + const short offset1 = il0/nl; + + device const block_q * xg = + (device const block_q *)(src0_gate + args.nb01*(r0 + lr0) + offset0) + offset1; + device const block_q * xu = + (device const block_q *)(src0_up + args.nb01*(r0 + lr0) + offset0) + offset1; + + const short iy = 8*(tiitg % NL1); + + device const float * y = (device const float *)(src1 + + args.nb13*i13 + + args.nb12*i12 + + args.nb11*i11 + + args.nb10*iy); + + simdgroup_half8x8 ma_g[4]; + simdgroup_half8x8 ma_u[4]; + simdgroup_half8x8 mb[2]; + + simdgroup_float8x8 mc_gate[8]; + simdgroup_float8x8 mc_up[8]; + + for (short i = 0; i < 8; i++) { + mc_gate[i] = make_filled_simdgroup_matrix(0.f); + mc_up[i] = make_filled_simdgroup_matrix(0.f); + } + + for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) { + half4x4 temp_gate; + dequantize_func(xg, il, temp_gate); + half4x4 temp_up; + dequantize_func(xu, il, temp_up); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Stage the shared B tile and both expert A tiles in one pass: two + // threadgroup barriers per k-step total instead of two per matrix. + { + const short sx_b = (tiitg%NL1); + const short sy_b = (tiitg/NL1)/8; + const short ly_b = (tiitg/NL1)%8; + const short ib_b = 4*sx_b + sy_b; + *(threadgroup half2x4 *)(sb + 64*ib_b + 8*ly_b) = + (half2x4)(*((device float2x4 *) y)); + } + + FOR_UNROLL (short i = 0; i < 16; i++) { + const short sx = 2*il0 + i/8; + const short sy = (tiitg/NL0)/8; + const short lx = (tiitg/NL0)%8; + const short ly = i%8; + const short ib = 8*sx + sy; + *(sa_gate + 64*ib + 8*ly + lx) = temp_gate[i/4][i%4]; + *(sa_up + 64*ib + 8*ly + lx) = temp_up[i/4][i%4]; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup const half * lsma_gate = (sa_gate + 4*64*(sgitg%2)); + threadgroup const half * lsma_up = (sa_up + 4*64*(sgitg%2)); + threadgroup const half * lsmb = (sb + 2*64*(sgitg/2)); + + FOR_UNROLL (short ik = 0; ik < NK/8; ik++) { + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 4; i++) { + simdgroup_load(ma_g[i], lsma_gate + 64*i, 8, 0, false); + simdgroup_load(ma_u[i], lsma_up + 64*i, 8, 0, false); + } + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + simdgroup_barrier(mem_flags::mem_none); + + FOR_UNROLL (short i = 0; i < 8; i++) { + simdgroup_multiply_accumulate(mc_gate[i], mb[i/4], ma_g[i%4], mc_gate[i]); + simdgroup_multiply_accumulate(mc_up[i], mb[i/4], ma_u[i%4], mc_up[i]); + } + + lsma_gate += 8*64; + lsma_up += 8*64; + lsmb += 4*64; + } + + il = (il + 2 < nl) ? il + 2 : il % 2; + xg = (il < 2) ? xg + (2 + nl - 1)/nl : xg; + xu = (il < 2) ? xu + (2 + nl - 1)/nl : xu; + y += NK; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup float * temp_gate = (threadgroup float *) shmem; + threadgroup float * temp_up = temp_gate + NR0*NR1; + threadgroup float * temp_gate_str = + temp_gate + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0; + threadgroup float * temp_up_str = + temp_up + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0; + + for (short i = 0; i < 8; i++) { + simdgroup_store(mc_gate[i], temp_gate_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + simdgroup_store(mc_up[i], temp_up_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + const float c = act.clamp_value; + for (short j = sgitg; j < nr1; j += 4) { + const int idj = ids_i32[im*args.ne21 + r1 + j]; + + const short ide = idj % args.ne20; + const short idt = idj / args.ne20; + + device half *D = (device half *)(dst_mid + + ((uint64_t)idt*args.ne1 + (uint64_t)ide)*act.mid_row_stride) + r0; + device const float *w = (device const float *)(weights + (uint64_t)idj*act.weight_stride); + const float route_weight = w[0]; + + threadgroup float *Cg = temp_gate + j*NR0; + threadgroup float *Cu = temp_up + j*NR0; + + int i = tiisg; + for (; i < nr0; i += 32) { + float g = Cg[i]; + float u = Cu[i]; + if (c > 1.0e-6f) { + g = min(g, c); + u = clamp(u, -c, c); + } + const float silu = g / (1.0f + exp(-g)); + D[i] = (half)(silu * u * route_weight); + } + } +} + +typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl) mul_mm_id_pair_swiglu_f16_iq2; +typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl) mul_mm_id_pair_swiglu_f16_q4; +typedef decltype(kernel_mul_mm_id_pair_swiglu_f16_impl) mul_mm_id_pair_swiglu_f16_q5; + +// Host-visible fused routed pair matmuls for the DS4 expert quant formats. +template [[host_name("kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_iq2 kernel_mul_mm_id_pair_swiglu_f16_impl; +template [[host_name("kernel_mul_mm_id_q4_K_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_q4 kernel_mul_mm_id_pair_swiglu_f16_impl; +template [[host_name("kernel_mul_mm_id_q5_K_pair_swiglu_f16")]] kernel mul_mm_id_pair_swiglu_f16_q5 kernel_mul_mm_id_pair_swiglu_f16_impl; + +typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id; +typedef decltype(kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_f16_rhs; +typedef decltype(kernel_mul_mm_id<32, float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id_ff32; +typedef decltype(kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>) mul_mm_id_addr; +typedef decltype(kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>) mul_mm_id_addr_f16_rhs; + +// Host-visible batched MoE matmul variants for the DS4 quant formats. +template [[host_name("kernel_mul_mm_id_q8_0_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_q2_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_q4_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_q5_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_q6_K_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_iq2_xxs_f32")]] kernel mul_mm_id kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_q8_0_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q8_0, 2, dequantize_q8_0, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_q2_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_q4_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_q5_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q5_K, QK_NL, dequantize_q5_K, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_q6_K_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q6_K, QK_NL, dequantize_q6_K, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_iq2_xxs_f16")]] kernel mul_mm_id_f16_rhs kernel_mul_mm_id<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_iq2_xxs, QK_NL, dequantize_iq2_xxs, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_q4_K_ff32")]] kernel mul_mm_id_ff32 kernel_mul_mm_id<32, float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_q5_K_ff32")]] kernel mul_mm_id_ff32 kernel_mul_mm_id<32, float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, block_q5_K, QK_NL, dequantize_q5_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_q6_K_ff32")]] kernel mul_mm_id_ff32 kernel_mul_mm_id<32, float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, block_q6_K, QK_NL, dequantize_q6_K, float, float4x4, float, float2x4>; + +template [[host_name("kernel_mul_mm_id_addr_q2_K_f32")]] kernel mul_mm_id_addr kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_addr_q4_K_f32")]] kernel mul_mm_id_addr kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, float, float4x4, float, float2x4>; +template [[host_name("kernel_mul_mm_id_addr_q2_K_f16")]] kernel mul_mm_id_addr_f16_rhs kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q2_K, QK_NL, dequantize_q2_K, half, half4x4, half, half2x4>; +template [[host_name("kernel_mul_mm_id_addr_q4_K_f16")]] kernel mul_mm_id_addr_f16_rhs kernel_mul_mm_id_addr<32, half, half4x4, simdgroup_half8x8, half, half2x4, simdgroup_half8x8, block_q4_K, QK_NL, dequantize_q4_K, half, half4x4, half, half2x4>; + +#ifdef DS4_METAL_HAS_TENSOR +// Attention-output low-rank projection retained for Metal4 prefill. It uses +// the same direct-RHS idea as dense matmul: dequantize the Q8_0 low projection +// weights to a half tile, then let TensorOps read the dense head activations +// directly. Only the 64-token direct-RHS instantiation is exported because the +// staged-RHS and 32-token variants were benchmark-only experiments. +// +// Full tiles (the host dispatch guarantee for aligned batches) skip all bounds +// work. The weight tile is double-buffered: the next k-step's dequant +// overlaps the current cooperative matmul, so the k-loop needs one +// threadgroup barrier per step instead of two. +template +kernel void kernel_attn_out_low_q8_0_mpp_direct_rhs( + constant ds4_metal_args_mul_mm_id & args, + device const char * srcA, + device const char * srcB, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiitg [[thread_index_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + (void) sgitg; + + constexpr int NR0 = 64; + constexpr int NK = 32; + constexpr int NL = NK/16; + constexpr int NUM_THREADS = 128; + + const int K = args.ne00; + const int M = args.ne0; + const int N = args.ne21; + const int G = args.ne1; + const int group = tgpig.z; + const int r0 = tgpig.y*NR0; + const int r1 = tgpig.x*NR1; + const bool full_tile = r0 + NR0 <= M && r1 + NR1 <= N && (K % NK) == 0; + + threadgroup half *sa = (threadgroup half *)shmem; + auto tA0 = tensor(sa, dextents(NK, NR0)); + auto tA1 = tensor(sa + NR0*NK, dextents(NK, NR0)); + + device float *ptrB = (device float *)(srcB + args.nb11*group); + const int strideB = args.nb12/sizeof(float); + auto tB = tensor(ptrB, dextents(K, N), array({1, strideB})); + + matmul2d< + matmul2d_descriptor(NR1, NR0, NK, false, true, true, + matmul2d_descriptor::mode::multiply_accumulate), + execution_simdgroups<4>> mm; + + auto cT = mm.template get_destination_cooperative_tensor(); + + #pragma unroll + for (uint16_t i = 0; i < cT.get_capacity(); ++i) { + if (cT.is_valid_element(i)) { + cT[i] = 0.0f; + } + } + + auto stage_tile = [&](const int loop_k, threadgroup half *buf) { + for (int work = tiitg; work < NR0*NL; work += NUM_THREADS) { + const int row = work/NL; + const int k_chunk = work%NL; + const int k_pos = loop_k + k_chunk*16; + const short k_base = k_chunk*16; + + if (full_tile) { + device const block_q8_0 *row_ptr = + (device const block_q8_0 *)(srcA + args.nb01*(r0 + row) + group*args.nb02); + + half4x4 temp_a; + dequantize_q8_0_pairs(row_ptr + k_pos/32, (k_pos/16)%2, temp_a); + threadgroup half4 *dst4 = (threadgroup half4 *)(buf + row*NK + k_base); + dst4[0] = temp_a[0]; + dst4[1] = temp_a[1]; + dst4[2] = temp_a[2]; + dst4[3] = temp_a[3]; + } else if (r0 + row < M) { + const int block_idx = k_pos/32; + const short il = (k_pos/16)%2; + device const block_q8_0 *row_ptr = + (device const block_q8_0 *)(srcA + args.nb01*(r0 + row) + group*args.nb02); + + half4x4 temp_a; + dequantize_q8_0_pairs(row_ptr + block_idx, il, temp_a); + FOR_UNROLL (short i = 0; i < 16; i++) { + buf[row*NK + k_base + i] = (k_pos + i < K) ? temp_a[i/4][i%4] : (half)0; + } + } else { + FOR_UNROLL (short i = 0; i < 16; i++) { + buf[row*NK + k_base + i] = (half)0; + } + } + } + }; + + stage_tile(0, sa); + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint buf_sel = 0; + for (int loop_k = 0; loop_k < K; loop_k += NK) { + auto mA = (buf_sel ? tA1 : tA0).slice(0, 0); + auto mB = tB.slice(loop_k, r1); + mm.run(mB, mA, cT); + + const int next_k = loop_k + NK; + if (next_k < K) { + buf_sel ^= 1u; + stage_tile(next_k, buf_sel ? sa + NR0*NK : sa); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + device float *dst_group = (device float *)dst + group*M; + if (full_tile) { + device float *dst_tile = dst_group + r0 + (uint64_t)r1*G*M; + auto tD = tensor(dst_tile, dextents(NR0, NR1), array({1, G*M})); + cT.store(tD); + } else { + auto tD = tensor(dst_group, dextents(M, N), array({1, G*M})); + auto mD = tD.slice(r0, r1); + cT.store(mD); + } +} + +typedef decltype(kernel_attn_out_low_q8_0_mpp_direct_rhs<64>) attn_out_low_q8_0_mpp_direct_rhs_n64_t; + +template [[host_name("kernel_attn_out_low_q8_0_mpp_direct_rhs_n64")]] kernel attn_out_low_q8_0_mpp_direct_rhs_n64_t kernel_attn_out_low_q8_0_mpp_direct_rhs<64>; + +#endif + +#undef QK_NL +#undef kmask_iq2xs +#undef ksigns_iq2xs +#undef iq2xxs_grid +#undef QK_K +#undef N_R0_Q2_K +#undef N_R0_Q4_K +#undef N_R0_Q8_K +#undef N_R0_GLM_Q4_PAIR2_K +#undef N_R0_GLM_Q4_PAIR_K +#undef N_R0_Q5_PAIR_K +#undef N_R0_Q5_K +#undef N_R0_Q6_K +#undef N_R0_IQ2_XXS diff --git a/metal/norm.metal b/metal/norm.metal new file mode 100644 index 0000000..4bc72f2 --- /dev/null +++ b/metal/norm.metal @@ -0,0 +1,243 @@ +struct ds4_metal_args_norm { + int32_t ne00; + int32_t ne00_t; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + float eps; + int32_t nef1[3]; + int32_t nef2[3]; + int32_t nef3[3]; + uint64_t nbf1[3]; + uint64_t nbf2[3]; + uint64_t nbf3[3]; +}; + +// RMSNorm over one activation row, optionally fusing the learned weight +// multiply. DS4 calls this before attention, before the FFN, and for plain +// diagnostics that need normalized but unweighted rows. +template +kernel void kernel_rms_norm_fuse_impl( + constant ds4_metal_args_norm & args, + device const char * src0, + device const char * src1_0, + device const char * src1_1, + device char * dst, + threadgroup float * shmem_f32 [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + if (sgitg == 0) { + shmem_f32[tiisg] = 0.0f; + } + + const int i01 = tgpig.x; + const int i02 = tgpig.y; + const int i03 = tgpig.z; + + device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]); + + device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]); + device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]); + + float sumf = 0.0f; + + // parallel sum + for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) { + sumf += dot(x[i00], x[i00]); + } + sumf = simd_sum(sumf); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiisg == 0) { + shmem_f32[sgitg] = sumf; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + sumf = shmem_f32[tiisg]; + sumf = simd_sum(sumf); + + const float mean = sumf/args.ne00; + const float scale = 1.0f/sqrt(mean + args.eps); + + device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1); + for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) { + if (F == 1) { + y[i00] = (x[i00]*scale); + } + if (F == 2) { + y[i00] = (x[i00]*scale)*f0[i00]; + } + if (F == 3) { + y[i00] = (x[i00]*scale)*f0[i00] + f1[i00]; + } + } +} + +typedef decltype(kernel_rms_norm_fuse_impl) kernel_rms_norm_fuse_t; + +// Host-visible RMSNorm variants: plain norm and norm multiplied by weight. +template [[host_name("kernel_rms_norm_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl; +template [[host_name("kernel_rms_norm_mul_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl; + +kernel void kernel_add_rms_norm_mul_f32_4( + constant ds4_metal_args_norm & args, + device const char * src0, + device const char * src1, + device const char * weight, + device char * sum_dst, + device char * norm_dst, + threadgroup float * shmem_f32 [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + if (sgitg == 0) shmem_f32[tiisg] = 0.0f; + + const int i01 = tgpig.x; + const int i02 = tgpig.y; + const int i03 = tgpig.z; + device const float4 *a = (device const float4 *) + (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]); + device const float4 *b = (device const float4 *) + (src1 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]); + device const float4 *w = (device const float4 *) + (weight + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]); + device float4 *sum = (device float4 *) + (sum_dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1); + device float4 *norm = (device float4 *) + (norm_dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1); + + float sumf = 0.0f; + for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) { + const float4 v = a[i00] + b[i00]; + sum[i00] = v; + sumf += dot(v, v); + } + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) shmem_f32[sgitg] = sumf; + threadgroup_barrier(mem_flags::mem_threadgroup); + sumf = simd_sum(shmem_f32[tiisg]); + + const float mean = sumf / args.ne00; + const float scale = 1.0f / sqrt(mean + args.eps); + for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) { + norm[i00] = (sum[i00] * scale) * w[i00]; + } +} + +// RMSNorm reduction used when the following F16 matmul applies the row scale +// while staging its RHS tile. +kernel void kernel_rms_norm_scale_f32_4( + constant ds4_metal_args_norm & args, + device const char * src0, + device float * dst_scale, + threadgroup float * shmem_f32 [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + if (sgitg == 0) shmem_f32[tiisg] = 0.0f; + + const int i01 = tgpig.x; + const int i02 = tgpig.y; + const int i03 = tgpig.z; + device const float4 *x = (device const float4 *) + (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]); + + float sumf = 0.0f; + for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) { + sumf += dot(x[i00], x[i00]); + } + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) shmem_f32[sgitg] = sumf; + threadgroup_barrier(mem_flags::mem_threadgroup); + sumf = simd_sum(shmem_f32[tiisg]); + + const float mean = sumf / args.ne00; + const float scale = 1.0f / sqrt(mean + args.eps); + if (tpitg.x == 0) dst_scale[i01] = scale; +} + +struct ds4_metal_args_qkv_rms_norm { + int32_t q_n; + int32_t q_n4; + int32_t kv_n; + int32_t kv_n4; + uint64_t q_row_stride; + uint64_t kv_row_stride; + float eps; +}; + +// Normalizes DS4's q-lora row and KV row in one dispatch. The two reductions +// deliberately mirror kernel_rms_norm_mul_f32_4: Q uses the full 256-thread +// row shape for 1024 floats, while KV only has work in the first 128 lanes for +// its 512 floats. This keeps the q/kv normalization math aligned with the +// standalone kernels while removing one tiny launch from the attention setup. +kernel void kernel_dsv4_qkv_rms_norm_f32_4( + constant ds4_metal_args_qkv_rms_norm & args, + device const float4 * q_src, + device const float4 * q_weight, + device float4 * q_dst, + device const float4 * kv_src, + device const float4 * kv_weight, + device float4 * kv_dst, + threadgroup float * shmem_f32 [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + if (sgitg == 0) { + shmem_f32[tiisg] = 0.0f; + } + + const uint row = tgpig.x; + const bool kv_task = tgpig.y != 0; + const int n = kv_task ? args.kv_n : args.q_n; + const int n4 = kv_task ? args.kv_n4 : args.q_n4; + const uint64_t row_stride4 = (kv_task ? args.kv_row_stride : args.q_row_stride) / sizeof(float4); + + device const float4 * x = kv_task ? kv_src + row * row_stride4 : q_src + row * row_stride4; + device const float4 * w = kv_task ? kv_weight : q_weight; + device float4 * y = kv_task ? kv_dst + row * row_stride4 : q_dst + row * row_stride4; + + float sumf = 0.0f; + for (int i = tpitg.x; i < n4; i += ntg.x) { + const float4 v = x[i]; + sumf += dot(v, v); + } + sumf = simd_sum(sumf); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiisg == 0) { + shmem_f32[sgitg] = sumf; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + sumf = shmem_f32[tiisg]; + sumf = simd_sum(sumf); + +#ifdef DS4_METAL_NORM_RSQRT_DISABLE + // Match the formula used by kernel_rms_norm_fuse_impl above so both RMSNorm + // entry points produce bit-identical scales. Hardware rsqrt() and 1.0f/sqrt() + // can differ by ~1 ULP and that difference compounds across 43 layers. + const float scale = 1.0f / sqrt(sumf / float(n) + args.eps); +#else + const float scale = rsqrt(sumf / float(n) + args.eps); +#endif + + for (int i = tpitg.x; i < n4; i += ntg.x) { + y[i] = (x[i] * scale) * w[i]; + } +} diff --git a/metal/repeat.metal b/metal/repeat.metal new file mode 100644 index 0000000..98f68e9 --- /dev/null +++ b/metal/repeat.metal @@ -0,0 +1,52 @@ +// DS4 Metal repeat kernel used for HC embedding expansion. + +struct ds4_metal_args_repeat { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +}; + +// Repeats a source row into the HC channel dimension. DS4 uses this when the +// token embedding has to become an HC activation block before layer 0. +template +kernel void kernel_repeat( + constant ds4_metal_args_repeat & args, + device const char * src0, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + const int i3 = tgpig.z; + const int i2 = tgpig.y; + const int i1 = tgpig.x; + + const int i03 = i3%args.ne03; + const int i02 = i2%args.ne02; + const int i01 = i1%args.ne01; + + device const char * src0_ptr = src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01; + device char * dst_ptr = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1; + + for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) { + const int i00 = i0%args.ne00; + *((device T *)(dst_ptr + i0*args.nb0)) = *((device T *)(src0_ptr + i00*args.nb00)); + } +} + +typedef decltype(kernel_repeat) kernel_repeat_t; + +// Host-visible F32 repeat used for HC expansion of embeddings. +template [[host_name("kernel_repeat_f32")]] kernel kernel_repeat_t kernel_repeat; diff --git a/metal/set_rows.metal b/metal/set_rows.metal new file mode 100644 index 0000000..d3df6e0 --- /dev/null +++ b/metal/set_rows.metal @@ -0,0 +1,55 @@ +// DS4 Metal set-rows kernel used for KV writes. + +struct ds4_metal_args_set_rows { + int32_t nk0; + int32_t ne01; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +}; + +// Scatters rows into the KV cache by token position. DS4 uses this after Q/K/V +// preparation so decode and later prefill chunks can attend to previous tokens. +template +kernel void kernel_set_rows_f( + constant ds4_metal_args_set_rows & args, + device const char * src0, + device const char * src1, + device float * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint tiitg[[thread_index_in_threadgroup]], + uint3 tptg [[threads_per_threadgroup]]) { + const int32_t i03 = tgpig.z; + const int32_t i02 = tgpig.y; + + const int32_t i12 = i03%args.ne12; + const int32_t i11 = i02%args.ne11; + + const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x; + if (i01 >= args.ne01) { + return; + } + + const int32_t i10 = i01; + const TI i1 = ((const device TI *) (src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0]; + + device T * dst_row = ( device T *) ((device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3); + const device float * src_row = (const device float *) ( src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); + + for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) { + dst_row[ind] = (T) src_row[ind]; + } +} + +typedef decltype(kernel_set_rows_f) set_rows_f_t; + +// Host-visible F32/I32 scatter variant used by KV-cache writes. +template [[host_name("kernel_set_rows_f32_i32")]] kernel set_rows_f_t kernel_set_rows_f; diff --git a/metal/softmax.metal b/metal/softmax.metal new file mode 100644 index 0000000..7634705 --- /dev/null +++ b/metal/softmax.metal @@ -0,0 +1,241 @@ +// DS4 Metal softmax kernel used by the compressor pooling compatibility path. +// The single-compressed-row path is intentionally left as soft_max -> mul -> +// sum_rows instead of using the fused dsv4_softmax_pool kernel. + +struct ds4_metal_args_soft_max { + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + float scale; + float max_bias; + float m0; + float m1; + int32_t n_head_log2; +}; + +// Row softmax for score matrices. DS4 uses it in the literal one-compressor-row +// path where preserving the original graph operation boundary avoids drift. +template +kernel void kernel_soft_max( + constant ds4_metal_args_soft_max & args, + device const char * src0, + device const char * src1, + device const char * src2, + device char * dst, + threadgroup float * buf [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint sgitg[[simdgroup_index_in_threadgroup]], + uint tiisg[[thread_index_in_simdgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t i03 = tgpig.z; + const int32_t i02 = tgpig.y; + const int32_t i01 = tgpig.x; + + const int32_t i13 = i03%args.ne13; + const int32_t i12 = i02%args.ne12; + const int32_t i11 = i01; + + device const float * psrc0 = (device const float *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); + device const T * pmask = src1 != src0 ? (device const T * ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr; + device const float * psrc2 = src2 != src0 ? (device const float *) (src2) : nullptr; + device float * pdst = (device float *) (dst + i01*args.nb1 + i02*args.nb2 + i03*args.nb3); + + float slope = 1.0f; + + if (args.max_bias > 0.0f) { + const int32_t h = i02; + + const float base = h < args.n_head_log2 ? args.m0 : args.m1; + const int exp = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1; + + slope = pow(base, exp); + } + + float lmax = psrc2 ? psrc2[i02] : -INFINITY; + + for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) { + lmax = MAX(lmax, psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f)); + } + + float max_val = simd_max(lmax); + if (tptg.x > N_SIMDWIDTH) { + if (sgitg == 0) { + buf[tiisg] = -INFINITY; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiisg == 0) { + buf[sgitg] = max_val; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + max_val = buf[tiisg]; + max_val = simd_max(max_val); + } + + float lsum = 0.0f; + for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) { + const float exp_psrc0 = exp((psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f)) - max_val); + lsum += exp_psrc0; + pdst[i00] = exp_psrc0; + } + + threadgroup_barrier(mem_flags::mem_none); + + float sum = simd_sum(lsum); + + if (tptg.x > N_SIMDWIDTH) { + if (sgitg == 0) { + buf[tiisg] = 0.0f; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiisg == 0) { + buf[sgitg] = sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + sum = buf[tiisg]; + sum = simd_sum(sum); + } + + if (psrc2) { + sum += exp(psrc2[i02] - max_val); + } + + const float inv_sum = 1.0f/sum; + + for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) { + pdst[i00] *= inv_sum; + } +} + +// Vectorized float4 row softmax for contiguous score rows whose length is a +// multiple of four; used by the same DS4 compressor/indexer graph path. +template +kernel void kernel_soft_max_4( + constant ds4_metal_args_soft_max & args, + device const char * src0, + device const char * src1, + device const char * src2, + device char * dst, + threadgroup float * buf [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint sgitg[[simdgroup_index_in_threadgroup]], + uint tiisg[[thread_index_in_simdgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t i03 = tgpig.z; + const int32_t i02 = tgpig.y; + const int32_t i01 = tgpig.x; + + const int32_t i13 = i03%args.ne13; + const int32_t i12 = i02%args.ne12; + const int32_t i11 = i01; + + device const float4 * psrc4 = (device const float4 *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); + device const T * pmask = src1 != src0 ? (device const T * ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr; + device const float * psrc2 = src2 != src0 ? (device const float * ) (src2) : nullptr; + device float4 * pdst4 = (device float4 *) (dst + i01*args.nb1 + i02*args.nb2 + i03*args.nb3); + + float slope = 1.0f; + + if (args.max_bias > 0.0f) { + const int32_t h = i02; + + const float base = h < args.n_head_log2 ? args.m0 : args.m1; + const int exp = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1; + + slope = pow(base, exp); + } + + float4 lmax4 = psrc2 ? psrc2[i02] : -INFINITY; + + for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) { + lmax4 = fmax(lmax4, psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f))); + } + + const float lmax = MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3])); + + float max_val = simd_max(lmax); + if (tptg.x > N_SIMDWIDTH) { + if (sgitg == 0) { + buf[tiisg] = -INFINITY; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiisg == 0) { + buf[sgitg] = max_val; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + max_val = buf[tiisg]; + max_val = simd_max(max_val); + } + + float4 lsum4 = 0.0f; + for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) { + const float4 exp_psrc4 = exp((psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f))) - max_val); + lsum4 += exp_psrc4; + pdst4[i00] = exp_psrc4; + } + + const float lsum = lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3]; + + threadgroup_barrier(mem_flags::mem_none); + + float sum = simd_sum(lsum); + + if (tptg.x > N_SIMDWIDTH) { + if (sgitg == 0) { + buf[tiisg] = 0.0f; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiisg == 0) { + buf[sgitg] = sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + sum = buf[tiisg]; + sum = simd_sum(sum); + } + + if (psrc2) { + sum += exp(psrc2[i02] - max_val); + } + + const float inv_sum = 1.0f/sum; + + for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) { + pdst4[i00] *= inv_sum; + } +} + +typedef decltype(kernel_soft_max) kernel_soft_max_t; +typedef decltype(kernel_soft_max_4) kernel_soft_max_4_t; + +// Host-visible F32 softmax variants used by compressor pooling. +template [[host_name("kernel_soft_max_f32")]] kernel kernel_soft_max_t kernel_soft_max; +template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4; diff --git a/metal/sum_rows.metal b/metal/sum_rows.metal new file mode 100644 index 0000000..8efea9b --- /dev/null +++ b/metal/sum_rows.metal @@ -0,0 +1,102 @@ +// DS4 Metal row-sum kernel. + +#define FC_SUM_ROWS 1400 + +#define OP_SUM_ROWS_NUM_SUM_ROWS 10 +#define OP_SUM_ROWS_NUM_MEAN 11 + +struct ds4_metal_args_sum_rows { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int64_t ne0; + int64_t ne1; + int64_t ne2; + int64_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +}; + +static inline float sum(float x) { + return x; +} + +static inline float sum(float4 x) { + return x[0] + x[1] + x[2] + x[3]; +} + +constant short FC_sum_rows_op [[function_constant(FC_SUM_ROWS + 0)]]; + +// Reduces each row to a sum or mean. DS4 mainly uses the sum form to preserve +// the compressor-pooling graph boundary in the single-compressor-row case. +template +kernel void kernel_sum_rows_impl( + constant ds4_metal_args_sum_rows & args, + device const char * src0, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { +#define FC_OP FC_sum_rows_op + + const int i3 = tgpig.z; + const int i2 = tgpig.y; + const int i1 = tgpig.x; + + threadgroup T0 * shmem_t = (threadgroup T0 *) shmem; + + if (sgitg == 0) { + shmem_t[tiisg] = 0.0f; + } + + device const T0 * src_row = (device const T0 *) (src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03); + device T * dst_row = (device T *) (dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3); + + T0 sumf = T0(0.0f); + + for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) { + sumf += src_row[i0]; + } + + sumf = simd_sum(sumf); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiisg == 0) { + shmem_t[sgitg] = sumf; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + sumf = shmem_t[tiisg]; + sumf = simd_sum(sumf); + + if (tpitg.x == 0) { + if (FC_OP == OP_SUM_ROWS_NUM_MEAN) { + if (is_same::value) { + dst_row[0] = sum(sumf) / (4*args.ne00); + } else { + dst_row[0] = sum(sumf) / args.ne00; + } + } else { + dst_row[0] = sum(sumf); + } + } + +#undef FC_OP +} + +typedef decltype(kernel_sum_rows_impl) kernel_sum_rows_t; + +// Host-visible F32 row reduction used by compressor pooling. +template [[host_name("kernel_sum_rows_f32_f32")]] kernel kernel_sum_rows_t kernel_sum_rows_impl; diff --git a/metal/unary.metal b/metal/unary.metal new file mode 100644 index 0000000..717fb35 --- /dev/null +++ b/metal/unary.metal @@ -0,0 +1,312 @@ +#define FC_UNARY 1200 + +#define OP_UNARY_NUM_SCALE 10 +#define OP_UNARY_NUM_FILL 11 +#define OP_UNARY_NUM_CLAMP 12 +#define OP_UNARY_NUM_SQR 13 +#define OP_UNARY_NUM_SQRT 14 +#define OP_UNARY_NUM_SIN 15 +#define OP_UNARY_NUM_COS 16 +#define OP_UNARY_NUM_LOG 17 +#define OP_UNARY_NUM_LEAKY_RELU 18 + +#define OP_UNARY_NUM_TANH 100 +#define OP_UNARY_NUM_RELU 101 +#define OP_UNARY_NUM_SIGMOID 102 +#define OP_UNARY_NUM_GELU 103 +#define OP_UNARY_NUM_GELU_ERF 104 +#define OP_UNARY_NUM_GELU_QUICK 105 +#define OP_UNARY_NUM_SILU 106 +#define OP_UNARY_NUM_ELU 107 +#define OP_UNARY_NUM_NEG 108 +#define OP_UNARY_NUM_ABS 109 +#define OP_UNARY_NUM_SGN 110 +#define OP_UNARY_NUM_STEP 111 +#define OP_UNARY_NUM_HARDSWISH 112 +#define OP_UNARY_NUM_HARDSIGMOID 113 +#define OP_UNARY_NUM_EXP 114 +#define OP_UNARY_NUM_SOFTPLUS 115 +#define OP_UNARY_NUM_EXPM1 116 +#define OP_UNARY_NUM_FLOOR 117 +#define OP_UNARY_NUM_CEIL 118 +#define OP_UNARY_NUM_ROUND 119 +#define OP_UNARY_NUM_TRUNC 120 +#define OP_UNARY_NUM_XIELU 121 + +struct ds4_metal_args_unary { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + float slope; + float scale; + float bias; + float val; + float min; + float max; +}; + +constant float GELU_COEF_A = 0.044715f; +constant float GELU_QUICK_COEF = -1.702f; +constant float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; +constant float SQRT_2_INV = 0.70710678118654752440084436210484f; + +// based on Abramowitz and Stegun formula 7.1.26 or similar Hastings' approximation +// ref: https://www.johndcook.com/blog/python_erf/ +constant float p_erf = 0.3275911f; +constant float a1_erf = 0.254829592f; +constant float a2_erf = -0.284496736f; +constant float a3_erf = 1.421413741f; +constant float a4_erf = -1.453152027f; +constant float a5_erf = 1.061405429f; + +template +inline T erf_approx(T x) { + T sign_x = sign(x); + x = fabs(x); + T t = 1.0f / (1.0f + p_erf * x); + T y = 1.0f - (((((a5_erf * t + a4_erf) * t) + a3_erf) * t + a2_erf) * t + a1_erf) * t * exp(-x * x); + return sign_x * y; +} + +template T elu_approx(T x); + +template<> inline float elu_approx(float x) { + return (x > 0.f) ? x : (exp(x) - 1); +} + +template<> inline float4 elu_approx(float4 x) { + float4 res; + + res[0] = (x[0] > 0.0f) ? x[0] : (exp(x[0]) - 1.0f); + res[1] = (x[1] > 0.0f) ? x[1] : (exp(x[1]) - 1.0f); + res[2] = (x[2] > 0.0f) ? x[2] : (exp(x[2]) - 1.0f); + res[3] = (x[3] > 0.0f) ? x[3] : (exp(x[3]) - 1.0f); + + return res; +} + +constant short FC_unary_op [[function_constant(FC_UNARY + 0)]]; +constant bool FC_unary_cnt[[function_constant(FC_UNARY + 1)]]; + +// Generic unary elementwise op selected by function constant. DS4 only uses a +// small subset in inference, mainly sigmoid, SiLU, softplus, sqrt, clamp, +// scale, and fill. +template +kernel void kernel_unary_impl( + constant ds4_metal_args_unary & args, + device const char * src0, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { +#define FC_OP FC_unary_op +#define FC_CNT FC_unary_cnt + + device const T0 * src0_ptr; + device T * dst_ptr; + + int i0; + + if (FC_CNT) { + i0 = tgpig.x; + + src0_ptr = (device const T0 *) (src0); + dst_ptr = (device T *) (dst); + } else { + const int i03 = tgpig.z; + const int i02 = tgpig.y; + const int k0 = tgpig.x/args.ne01; + const int i01 = tgpig.x - k0*args.ne01; + + i0 = k0*ntg.x + tpitg.x; + + src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01); + dst_ptr = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1 ); + } + + { + if (!FC_CNT) { + if (i0 >= args.ne0) { + return; + } + } + + const TC x = (TC) src0_ptr[i0]; + + if (FC_OP == OP_UNARY_NUM_SCALE) { + dst_ptr[i0] = (T) (args.scale * x + args.bias); + } + + if (FC_OP == OP_UNARY_NUM_FILL) { + dst_ptr[i0] = (T) args.val; + } + + if (FC_OP == OP_UNARY_NUM_CLAMP) { + dst_ptr[i0] = (T) clamp(x, args.min, args.max); + } + + if (FC_OP == OP_UNARY_NUM_SQR) { + dst_ptr[i0] = (T) (x * x); + } + + if (FC_OP == OP_UNARY_NUM_SQRT) { + dst_ptr[i0] = (T) sqrt(x); + } + + if (FC_OP == OP_UNARY_NUM_SIN) { + dst_ptr[i0] = (T) sin(x); + } + + if (FC_OP == OP_UNARY_NUM_COS) { + dst_ptr[i0] = (T) cos(x); + } + + if (FC_OP == OP_UNARY_NUM_LOG) { + dst_ptr[i0] = (T) log(x); + } + + if (FC_OP == OP_UNARY_NUM_LEAKY_RELU) { + dst_ptr[i0] = (T) (TC(x > 0)*x + TC(x <= 0)*(x * args.slope)); + } + + if (FC_OP == OP_UNARY_NUM_TANH) { + dst_ptr[i0] = (T) precise::tanh(x); + } + + if (FC_OP == OP_UNARY_NUM_RELU) { + dst_ptr[i0] = (T) fmax(0, x); + } + + if (FC_OP == OP_UNARY_NUM_SIGMOID) { + dst_ptr[i0] = (T) (1 / (1 + exp(-x))); + } + + if (FC_OP == OP_UNARY_NUM_GELU) { + dst_ptr[i0] = (T) (0.5*x*(1 + precise::tanh(SQRT_2_OVER_PI*x*(1 + GELU_COEF_A*x*x)))); + } + + if (FC_OP == OP_UNARY_NUM_GELU_ERF) { + dst_ptr[i0] = (T) (0.5*x*(1 + erf_approx(SQRT_2_INV*x))); + } + + if (FC_OP == OP_UNARY_NUM_GELU_QUICK) { + dst_ptr[i0] = (T) (x * (1/(1 + exp(GELU_QUICK_COEF*x)))); + } + + if (FC_OP == OP_UNARY_NUM_SILU) { + dst_ptr[i0] = (T) (x / (1 + exp(-x))); + } + + if (FC_OP == OP_UNARY_NUM_ELU) { + dst_ptr[i0] = (T) elu_approx(x); + } + + if (FC_OP == OP_UNARY_NUM_NEG) { + dst_ptr[i0] = (T) -x; + } + + if (FC_OP == OP_UNARY_NUM_ABS) { + dst_ptr[i0] = (T) fabs(x); + } + + if (FC_OP == OP_UNARY_NUM_SGN) { + dst_ptr[i0] = T(x > 0) - T(x < 0); + } + + if (FC_OP == OP_UNARY_NUM_STEP) { + dst_ptr[i0] = T(x > 0); + } + + if (FC_OP == OP_UNARY_NUM_HARDSWISH) { + dst_ptr[i0] = (T) (x * fmax(0, fmin(1, x/6 + 0.5))); + } + + if (FC_OP == OP_UNARY_NUM_HARDSIGMOID) { + dst_ptr[i0] = (T) fmax(0, fmin(1, x/6 + 0.5)); + } + + if (FC_OP == OP_UNARY_NUM_EXP) { + dst_ptr[i0] = (T) exp(x); + } + + if (FC_OP == OP_UNARY_NUM_SOFTPLUS) { + dst_ptr[i0] = (T) select(log(1 + exp(x)), x, x > 20); + } + + if (FC_OP == OP_UNARY_NUM_EXPM1) { + // Metal target profiles used here do not all expose expm1(); this + // generic unary branch is not used by the DS4 inference graph. + dst_ptr[i0] = (T) (exp(x) - 1); + } + + if (FC_OP == OP_UNARY_NUM_FLOOR) { + dst_ptr[i0] = (T) floor(x); + } + + if (FC_OP == OP_UNARY_NUM_CEIL) { + dst_ptr[i0] = (T) ceil(x); + } + + if (FC_OP == OP_UNARY_NUM_ROUND) { + dst_ptr[i0] = (T) round(x); + } + + if (FC_OP == OP_UNARY_NUM_TRUNC) { + dst_ptr[i0] = (T) trunc(x); + } + + if (FC_OP == OP_UNARY_NUM_XIELU) { + const TC xi = x; + const TC gate = TC(xi > TC(0.0f)); + const TC clamped = fmin(xi, TC(args.val)); + const TC y_pos = TC(args.scale) * xi * xi + TC(args.bias) * xi; + const TC y_neg = (exp(clamped) - TC(1.0f) - xi) * TC(args.slope) + TC(args.bias) * xi; + dst_ptr[i0] = (T) (gate * y_pos + (TC(1.0f) - gate) * y_neg); + } + } + +#undef FC_OP +#undef FC_CNT +} + +typedef decltype(kernel_unary_impl) kernel_unary_t; + +// Decode router probability transform. The generic path applies softplus and +// sqrt as two elementwise kernels; DS4 decode always transforms one 256-wide +// expert-logit row, so this vectorized kernel does both in one pass. +kernel void kernel_dsv4_softplus_sqrt_f32_4( + constant ds4_metal_args_unary & args, + device const char *src, + device char *dst, + uint3 tgpig [[threadgroup_position_in_grid]], + ushort3 tpitg [[thread_position_in_threadgroup]], + ushort3 ntg [[threads_per_threadgroup]]) { + const int k0 = tgpig.x/args.ne01; + const int i01 = tgpig.x - k0*args.ne01; + const int i0 = k0*ntg.x + tpitg.x; + if (i0 >= args.ne0) return; + + device const float4 *s = (device const float4 *)(src + i01*args.nb01); + device float4 *d = (device float4 *)(dst + i01*args.nb1); + const float4 x = s[i0]; + const float4 sp = select(log(1.0f + exp(x)), x, x > 20.0f); + d[i0] = sqrt(sp); +} + +// Host-visible unary variants. Function constants select the actual DS4 op. +template [[host_name("kernel_unary_f32_f32")]] kernel kernel_unary_t kernel_unary_impl; +template [[host_name("kernel_unary_f32_f32_4")]] kernel kernel_unary_t kernel_unary_impl; +template [[host_name("kernel_unary_f16_f16")]] kernel kernel_unary_t kernel_unary_impl; diff --git a/native/metal/LICENSE b/native/metal/LICENSE new file mode 100644 index 0000000..5973a4c --- /dev/null +++ b/native/metal/LICENSE @@ -0,0 +1,22 @@ +MIT License + +Copyright (c) 2026 The ds4.c authors +Copyright (c) 2023-2026 The ggml authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/native/metal/README.md b/native/metal/README.md new file mode 100644 index 0000000..0f2a5f5 --- /dev/null +++ b/native/metal/README.md @@ -0,0 +1,15 @@ +# Vendored DS4 Metal boundary + +These files are a one-time snapshot of the DS4 Metal boundary from commit +`efdadd41e20134af4f3381e1ed90e96fe4faef6f`: + +- `ds4_metal.m` +- `ds4.h` +- `ds4_gpu.h` +- `ds4_ssd.h` +- `LICENSE` + +The matching Metal kernels live in the repository-level `metal/` directory. +DS4Server builds and bundles this local snapshot; it does not read a sibling +DS4 checkout. Rust owns the model, graph, session, sampling, and lifecycle. +Objective-C remains only at the platform Metal boundary. diff --git a/native/metal/ds4.h b/native/metal/ds4.h new file mode 100644 index 0000000..42bc116 --- /dev/null +++ b/native/metal/ds4.h @@ -0,0 +1,477 @@ +#ifndef DS4_H +#define DS4_H + +#include +#include +#include +#include + +#include "ds4_ssd.h" + +/* Public engine boundary. + * + * The CLI and server should treat ds4_engine as the loaded model and + * ds4_session as one mutable inference timeline. A session owns the live KV + * cache and logits; callers provide full token prefixes and let + * ds4_session_sync() reuse, extend, or rebuild the graph state. Keep this + * header narrow so HTTP/CLI code does not depend on tensor internals. */ + +typedef enum { + DS4_BACKEND_METAL, + DS4_BACKEND_CUDA, + DS4_BACKEND_CPU, +} ds4_backend; + +typedef enum { + DS4_THINK_NONE, + DS4_THINK_HIGH, + DS4_THINK_MAX, +} ds4_think_mode; + +typedef enum { + DS4_LOG_DEFAULT, + DS4_LOG_PREFILL, + DS4_LOG_GENERATION, + DS4_LOG_KVCACHE, + DS4_LOG_TOOL, + DS4_LOG_WARNING, + DS4_LOG_TIMING, + DS4_LOG_OK, + DS4_LOG_ERROR, +} ds4_log_type; + +typedef struct { + int *v; + int len; + int cap; +} ds4_tokens; + +typedef struct { + int id; + float logit; + float logprob; +} ds4_token_score; + +#define DS4_DEFAULT_TEMPERATURE 1.0f +#define DS4_DEFAULT_TOP_P 1.0f +#define DS4_DEFAULT_MIN_P 0.05f + +typedef struct ds4_engine ds4_engine; +typedef struct ds4_session ds4_session; + +typedef void (*ds4_session_progress_fn)(void *ud, const char *event, int current, int total); +typedef bool (*ds4_session_cancel_fn)(void *ud); + +#define DS4_SESSION_SYNC_INTERRUPTED 2 + +typedef enum { + DS4_DISTRIBUTED_NONE = 0, + DS4_DISTRIBUTED_COORDINATOR, + DS4_DISTRIBUTED_WORKER, +} ds4_distributed_role; + +typedef struct { + uint32_t start; + uint32_t end; + bool has_output; + bool set; +} ds4_distributed_layers; + +typedef struct { + ds4_distributed_role role; + ds4_distributed_layers layers; + const char *listen_host; + int listen_port; + const char *coordinator_host; + int coordinator_port; + uint32_t prefill_chunk; + uint32_t prefill_window; + uint32_t activation_bits; + bool replay_check; + bool debug; +} ds4_distributed_options; + +/* Tensor parallelism: two identical machines run the model in lockstep and + * split the heavy per-layer matvecs, exchanging partial sums at gates inside + * the graph (see misc/METAL_TENSOR_PARALLELISM.md). Each rank keeps one + * contiguous half of the routed experts resident; dense and shared weights + * remain replicated. The leader owns prompt/sampling and listens; the worker + * dials in and mirrors every session sync/eval. */ +typedef enum { + DS4_TP_NONE = 0, + DS4_TP_LEADER, + DS4_TP_WORKER, +} ds4_tp_role; + +typedef enum { + DS4_TP_TRANSPORT_AUTO = 0, + DS4_TP_TRANSPORT_RDMA, + DS4_TP_TRANSPORT_TCP, +} ds4_tp_transport; + +typedef struct { + ds4_tp_role role; + bool requested; /* --tensor-parallel with shared role options */ + const char *listen_host; /* leader listens here for the worker */ + int listen_port; + const char *leader_host; /* worker dials the leader */ + int leader_port; + ds4_tp_transport transport; + const char *rdma_device; + int rdma_gid_index; + bool rdma_gid_index_set; + bool glm_token_prefill; + int debug_hash; /* cross-check hidden state every N tokens */ +} ds4_tp_options; + +typedef struct { + const char *model_path; + const char *mtp_path; + ds4_backend backend; + int n_threads; + int context_size; + uint32_t prefill_chunk; + int mtp_draft_tokens; + float mtp_margin; + float dspark_confidence_threshold; + const char *directional_steering_file; + const char *expert_profile_path; + float directional_steering_attn; + float directional_steering_ffn; + int power_percent; + uint32_t ssd_streaming_cache_experts; + uint64_t ssd_streaming_cache_bytes; + uint32_t ssd_streaming_full_layers; + uint32_t ssd_streaming_preload_experts; + uint64_t simulate_used_memory_bytes; + bool warm_weights; + bool quality; + bool glm_mtp; + bool glm_mtp_timing; + bool dspark; + bool dspark_strict; + bool dspark_confidence_threshold_set; + bool cuda_tensor_parallel; + bool ssd_streaming; + bool ssd_streaming_cold; + bool ssd_streaming_full_layers_set; + bool inspect_only; + /* Multi-GPU placement uses this to price per-layer KV storage. */ + int placement_ctx_hint; + /* Server batch mode serializes execution and can share prefill scratch. */ + bool share_session_prefill_workspace; + bool first_token_test; + bool metal_graph_test; + bool load_slice; + uint32_t load_layer_start; + uint32_t load_layer_end; + bool load_output; + ds4_distributed_options distributed; + ds4_tp_options tp; +} ds4_engine_options; + +typedef void (*ds4_token_emit_fn)(void *ud, int token); +typedef void (*ds4_generation_done_fn)(void *ud); + +typedef struct { + uint64_t total_bytes; + uint64_t raw_bytes; + uint64_t compressed_bytes; + uint64_t scratch_bytes; + uint32_t prefill_cap; + uint32_t raw_cap; + uint32_t comp_cap; +} ds4_context_memory; + +typedef struct { + uint8_t *ptr; + uint64_t len; + uint64_t cap; +} ds4_session_snapshot; + +typedef struct { + char *path; + uint64_t bytes; +} ds4_session_payload_file; + +int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt); + +/* Multi-GPU pipeline-parallel entry point (wave 2). + * + * Accepts an optional ds4_gpu_config (defined in ds4_gpu_mgpu.h) that + * lets callers describe a multi-GPU placement target. Passing NULL is + * back-compatible with ds4_engine_open and produces identical engine + * state — bit-equivalent execution at runtime. + * + * When a non-NULL config is supplied AND the computed placement spans + * more than one tier (either multiple GPUs or any CPU-spill), this + * wave-2 implementation prints the layout and refuses to open: full + * multi-tier execution wiring lands in a follow-up task + * (mgpu-graph-session-execution). Callers receive a non-zero return + * and a documented stderr notice. */ +/* ds4_gpu_config is declared in ds4_gpu_mgpu.h, which callers should + * include separately. We forward-declare it here so this header can be + * used as-is (callers passing NULL don't need the struct definition). */ +struct ds4_gpu_config; +int ds4_engine_create_with_gpu_config(ds4_engine **out, + const ds4_engine_options *opt, + const struct ds4_gpu_config *gpu_cfg); +void ds4_engine_close(ds4_engine *e); +void ds4_engine_summary(ds4_engine *e); +int ds4_engine_vocab_size(ds4_engine *e); +uint32_t ds4_engine_prefill_chunk(ds4_engine *e); +int ds4_engine_power(ds4_engine *e); +int ds4_engine_set_power(ds4_engine *e, int power_percent); +const char *ds4_engine_model_name(ds4_engine *e); +int ds4_engine_layer_count(ds4_engine *e); +/* Decode gate schedule for the TP transport; see ds4_tp_identity. */ +void ds4_engine_tp_gate_schedule(ds4_engine *e, + uint32_t *start, + uint32_t *step, + uint32_t *per_token); +uint32_t ds4_engine_layer_compress_ratio(ds4_engine *e, uint32_t layer); +uint64_t ds4_engine_hidden_f32_values(ds4_engine *e); +int ds4_engine_embd_dim(ds4_engine *e); +uint64_t ds4_engine_model_bytes(ds4_engine *e); +int ds4_engine_tp_vocab_split(ds4_engine *e); +bool ds4_engine_glm_layer_payload_bytes(ds4_engine *e, + uint32_t layer, + uint32_t full_live, + uint32_t key_dim, + uint32_t value_dim, + uint32_t compact_live, + uint32_t index_live, + uint64_t *out); +/* Stable id for cache compatibility. 0 is the original Flash shape, so old + * KV files with the previously-zero reserved byte remain Flash-compatible; + * Pro and later shapes must use nonzero ids. */ +int ds4_engine_model_id(ds4_engine *e); +bool ds4_engine_is_glm_dsa(ds4_engine *e); +const char *ds4_backend_name(ds4_backend backend); +bool ds4_think_mode_enabled(ds4_think_mode mode); +const char *ds4_think_mode_name(ds4_think_mode mode); +const char *ds4_think_max_prefix(void); +const char *ds4_glm_reasoning_effort_text(ds4_think_mode mode); +uint32_t ds4_think_max_min_context(void); +ds4_think_mode ds4_think_mode_for_context(ds4_think_mode mode, int ctx_size); +/* Uses the active model shape selected by ds4_engine_open(); call after opening + * the GGUF so Flash/Pro dimensions are known. */ +ds4_context_memory ds4_context_memory_estimate(ds4_backend backend, int ctx_size); +ds4_context_memory ds4_context_memory_estimate_with_prefill( + ds4_backend backend, + int ctx_size, + uint32_t prefill_chunk); +ds4_context_memory ds4_context_memory_estimate_with_prefill_mode( + ds4_backend backend, + int ctx_size, + uint32_t prefill_chunk, + bool ssd_streaming); +bool ds4_log_is_tty(FILE *fp); +void ds4_log(FILE *fp, ds4_log_type type, const char *fmt, ...); +int ds4_engine_generate_argmax(ds4_engine *e, const ds4_tokens *prompt, + int n_predict, int ctx_size, + ds4_token_emit_fn emit, + ds4_generation_done_fn done, + void *emit_ud, + ds4_session_progress_fn progress, + void *progress_ud); +int ds4_engine_collect_imatrix(ds4_engine *e, + const char *dataset_path, + const char *output_path, + int ctx_size, + int max_prompts, + int max_tokens); +void ds4_engine_dump_tokens(ds4_engine *e, const ds4_tokens *tokens); +int ds4_dump_text_tokenization(const char *model_path, const char *text, FILE *fp); +int ds4_engine_head_test(ds4_engine *e, const ds4_tokens *prompt); +bool ds4_engine_is_glm_dsa(ds4_engine *e); +int ds4_engine_first_token_test(ds4_engine *e, const ds4_tokens *prompt); +int ds4_engine_metal_graph_test(ds4_engine *e, const ds4_tokens *prompt); +int ds4_engine_metal_graph_full_test(ds4_engine *e, const ds4_tokens *prompt); +int ds4_engine_metal_graph_prompt_test(ds4_engine *e, const ds4_tokens *prompt, int ctx_size); + +void ds4_tokens_push(ds4_tokens *tv, int token); +void ds4_tokens_free(ds4_tokens *tv); +void ds4_tokens_copy(ds4_tokens *dst, const ds4_tokens *src); +bool ds4_tokens_starts_with(const ds4_tokens *tokens, const ds4_tokens *prefix); + +void ds4_tokenize_text(ds4_engine *e, const char *text, ds4_tokens *out); +void ds4_tokenize_rendered_chat(ds4_engine *e, const char *text, ds4_tokens *out); +void ds4_chat_begin(ds4_engine *e, ds4_tokens *tokens); +void ds4_encode_chat_prompt( + ds4_engine *e, + const char *system, + const char *prompt, + ds4_think_mode think_mode, + ds4_tokens *out); +void ds4_chat_append_max_effort_prefix(ds4_engine *e, ds4_tokens *tokens); +void ds4_chat_append_message(ds4_engine *e, ds4_tokens *tokens, const char *role, const char *content); +void ds4_chat_append_assistant_prefix(ds4_engine *e, ds4_tokens *tokens, ds4_think_mode think_mode); + +char *ds4_token_text(ds4_engine *e, int token, size_t *len); +int ds4_token_eos(ds4_engine *e); +bool ds4_token_is_stop(ds4_engine *e, int token); +bool ds4_token_is_thinking_control(ds4_engine *e, int token); +bool ds4_token_is_stop_for_think_mode(ds4_engine *e, + int token, + ds4_think_mode mode); +int ds4_token_user(ds4_engine *e); +int ds4_token_assistant(ds4_engine *e); + +/* Tensor-parallel binding: allocates the GPU gate slab, registers it with + * the transport and arms the per-layer gate machinery. Call once, after + * ds4_tp_create() and before any session work. Transport lifecycle stays + * with the caller. */ +struct ds4_tp; +int ds4_engine_tp_bind(ds4_engine *e, struct ds4_tp *tp, char *err, size_t errlen); + +int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size); +void ds4_session_free(ds4_session *s); +int ds4_session_power(ds4_session *s); +int ds4_session_set_power(ds4_session *s, int power_percent); +bool ds4_session_is_distributed(ds4_session *s); +void ds4_session_set_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud); +/* UI-only progress. It may report fine-grained progress inside a prefill chunk; + * callers must not treat it as a durable KV checkpoint boundary. */ +void ds4_session_set_display_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud); +/* Optional cooperative cancellation. ds4_session_sync() checks it only at + * safe boundaries where the live checkpoint is either unchanged or represents a + * valid token prefix, and returns DS4_SESSION_SYNC_INTERRUPTED when it stops. */ +void ds4_session_set_cancel(ds4_session *s, ds4_session_cancel_fn fn, void *ud); +void ds4_session_report_progress(ds4_session *s, const char *event, int current, int total); +/* Distributed coordinator sessions return 1 when the full layer route is + * available, 0 when it is still incomplete, and -1 for a local API error. */ +int ds4_session_distributed_route_ready(ds4_session *s, char *err, size_t errlen); + +typedef enum { + DS4_SESSION_REWRITE_ERROR = -1, + DS4_SESSION_REWRITE_OK = 0, + /* The live backend state cannot be rewritten safely in place. The caller should + * restore an older checkpoint if it has one, then sync to the prompt. */ + DS4_SESSION_REWRITE_REBUILD_NEEDED = 1, +} ds4_session_rewrite_result; + +/* Synchronize the live session to a full prompt token prefix. If the current + * checkpoint is a prefix, only the suffix is evaluated; otherwise the backend + * state is refilled from scratch. */ +#define DS4_SESSION_SYNC_INTERRUPTED 2 +int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen); +bool ds4_session_rewrite_requires_rebuild(int live_len, int canonical_len, int common); +ds4_session_rewrite_result ds4_session_rewrite_from_common( + ds4_session *s, const ds4_tokens *prompt, int common, + char *err, size_t errlen); +int ds4_session_common_prefix(ds4_session *s, const ds4_tokens *prompt); +int ds4_session_argmax(ds4_session *s); +int ds4_session_argmax_excluding(ds4_session *s, int excluded_id); +int ds4_sample_logits(const float *logits, int n_vocab, float temperature, + int top_k, float top_p, float min_p, uint64_t *rng); +int ds4_session_sample(ds4_session *s, float temperature, int top_k, float top_p, float min_p, uint64_t *rng); +#ifdef DS4_TEST_HOOKS +int ds4_test_sample_logits(const float *logits, uint32_t n_vocab, + float temperature, int top_k, + float top_p, float min_p, uint64_t *rng, + float *prob_scratch); +uint64_t ds4_test_mixed_native_count(void); +#endif +int ds4_session_top_logprobs(ds4_session *s, ds4_token_score *out, int k); +int ds4_session_token_logprob(ds4_session *s, int token, ds4_token_score *out); +int ds4_session_copy_logits(ds4_session *s, float *out, int cap); +int ds4_session_set_logits(ds4_session *s, const float *logits, int n); +/* Pay the one-time first-submission GPU cost outside any measured window; + * used by the TP worker right after session create (no-op on CPU/GLM). */ +void ds4_session_gpu_warmup(ds4_session *s); +int ds4_session_eval(ds4_session *s, int token, char *err, size_t errlen); + +typedef struct { + ds4_session *session; + int token; +} ds4_decode_item; + +/* Advance independent sessions by one token each. Batch size one is exactly + * ds4_session_eval(). Backends without native batching use a correctness-first + * sequential fallback. */ +int ds4_sessions_eval_batch(ds4_decode_item *items, int count, + char *err, size_t errlen); +/* Advance one resumed prefill suffix and an independent decode batch as one + * scheduling step. Unsupported combinations use the ordinary serialized + * session operations. */ +int ds4_sessions_eval_batch_with_prefill( + ds4_decode_item *items, int count, + ds4_session *prefill_session, const ds4_tokens *prefill_prompt, + char *err, size_t errlen); +int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, + int max_tokens, int eos_token, + int *accepted, int accepted_cap, + char *err, size_t errlen); +/* TP worker side of a mirrored speculative-verify block: run its half of the + * batch verify for KV side effects, then obey the leader's commit frame + * (keep, or roll back and replay). Only called from ds4_tp_worker_run. */ +int ds4_session_tp_spec_cycle(ds4_session *s, const int *drafts, int draft_n, + char *err, size_t errlen); +void ds4_session_invalidate(ds4_session *s); +void ds4_session_rewind(ds4_session *s, int pos); +int ds4_session_pos(ds4_session *s); +int ds4_session_ctx(ds4_session *s); +int ds4_session_prefill_cap(ds4_session *s); +int ds4_engine_routed_quant_bits(ds4_engine *e); +bool ds4_engine_has_output_head(ds4_engine *e); +bool ds4_engine_has_mtp(ds4_engine *e); +int ds4_engine_mtp_draft_tokens(ds4_engine *e); +const ds4_tokens *ds4_session_tokens(ds4_session *s); + +/* Low-level graph slice entry points used by distributed inference. The + * transport/session routing logic lives in ds4_distributed.c. */ +int ds4_session_layer_slice_reset(ds4_session *s, char *err, size_t errlen); +int ds4_session_eval_layer_slice(ds4_session *s, + const int *tokens, + uint32_t n_tokens, + uint32_t pos0, + uint32_t layer_start, + uint32_t layer_end, + const float *input_hc, + float *output_hc, + bool output_logits, + float *logits, + char *err, + size_t errlen); +int ds4_session_eval_output_head_from_hc(ds4_session *s, + const float *hidden_hc, + uint32_t n_tokens, + float *logits, + char *err, + size_t errlen); + +/* Disk KV payload helpers. HTTP/agent code owns the outer file header and + * persistence policy; the engine owns the DS4-specific serialized graph state. */ +#define DS4_SESSION_PAYLOAD_MAGIC UINT32_C(0x34565344) /* "DSV4" */ +#define DS4_SESSION_PAYLOAD_VERSION UINT32_C(2) +#define DS4_SESSION_PAYLOAD_U32_FIELDS 13u +#define DS4_SESSION_LAYER_PAYLOAD_MAGIC UINT32_C(0x4c565344) /* "DSVL" */ +#define DS4_SESSION_LAYER_PAYLOAD_VERSION UINT32_C(1) +#define DS4_SESSION_LAYER_PAYLOAD_U32_FIELDS 14u + +uint64_t ds4_session_payload_bytes(ds4_session *s); +int ds4_session_stage_payload(ds4_session *s, ds4_session_payload_file *out, + char *err, size_t errlen); +int ds4_session_write_staged_payload(const ds4_session_payload_file *payload, + FILE *fp, char *err, size_t errlen); +void ds4_session_payload_file_free(ds4_session_payload_file *payload); +int ds4_session_save_payload(ds4_session *s, FILE *fp, char *err, size_t errlen); +int ds4_session_load_payload(ds4_session *s, FILE *fp, uint64_t payload_bytes, char *err, size_t errlen); +int ds4_session_save_snapshot(ds4_session *s, ds4_session_snapshot *snap, char *err, size_t errlen); +int ds4_session_load_snapshot(ds4_session *s, const ds4_session_snapshot *snap, char *err, size_t errlen); +void ds4_session_snapshot_free(ds4_session_snapshot *snap); + +uint64_t ds4_session_layer_payload_bytes(ds4_session *s, + uint32_t layer_start, + uint32_t layer_end); +int ds4_session_save_layer_payload(ds4_session *s, FILE *fp, + uint32_t layer_start, uint32_t layer_end, + char *err, size_t errlen); +int ds4_session_load_layer_payload(ds4_session *s, FILE *fp, + uint64_t payload_bytes, + const int *tokens, uint32_t n_tokens, + uint32_t layer_start, uint32_t layer_end, + char *err, size_t errlen); + +#endif diff --git a/native/metal/ds4_gpu.h b/native/metal/ds4_gpu.h new file mode 100644 index 0000000..2000bba --- /dev/null +++ b/native/metal/ds4_gpu.h @@ -0,0 +1,2603 @@ +#ifndef DS4_GPU_H +#define DS4_GPU_H + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* ========================================================================= + * GPU Tensor and Command Lifetime. + * ========================================================================= + * + * Opaque device tensor used by the DS4-specific GPU executor. + * + * The public GPU API is tensor-resident: activations, KV state, and scratch + * buffers stay device-owned across the whole prefill/decode command sequence. + */ +#ifndef DS4_GPU_TENSOR_DEFINED +#define DS4_GPU_TENSOR_DEFINED +typedef struct ds4_gpu_tensor ds4_gpu_tensor; +#endif + +#ifndef DS4_GPU_ATTENTION_DECODE_ROW_DEFINED +#define DS4_GPU_ATTENTION_DECODE_ROW_DEFINED +#define DS4_GPU_ATTENTION_DECODE_BATCH_MAX 32u +typedef struct { + uint64_t raw_kv; + uint64_t comp_kv; + uint64_t topk; + uint32_t pos; + uint32_t n_raw; + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_comp; + uint32_t top_k; + uint32_t window; + uint32_t ratio; + uint32_t indexed; +} ds4_gpu_attention_decode_row; +#endif + +int ds4_gpu_init(void); +void ds4_gpu_cleanup(void); + +ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes); +ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes); +ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes); +void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor); +uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor); +void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor); +int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count); +int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes); +int ds4_gpu_tensor_read(const ds4_gpu_tensor *tensor, uint64_t offset, void *data, uint64_t bytes); +int ds4_gpu_tensor_copy(ds4_gpu_tensor *dst, uint64_t dst_offset, + const ds4_gpu_tensor *src, uint64_t src_offset, + uint64_t bytes); +int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset, + const ds4_gpu_tensor *src, uint64_t src_offset, + uint64_t count); +int ds4_gpu_moe_handoff_pack_tensor( + ds4_gpu_tensor *packed, + const ds4_gpu_tensor *ffn_norm, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_embd, + uint32_t n_expert); +int ds4_gpu_pack_slot_rows_f32_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *slots, + uint32_t n_rows, + uint32_t width, + uint32_t n_slots, + uint32_t slot_cap); + +int ds4_gpu_begin_commands(void); +int ds4_gpu_flush_encoder(void); +int ds4_gpu_flush_commands(void); +int ds4_gpu_commands_active(void); +int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value); +int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label); +int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label); +#ifdef DS4_ROCM_BUILD +int ds4_gpu_tensor_read_after_selected_event(const ds4_gpu_tensor *tensor, + uint64_t offset, + void *data, + uint64_t bytes, + uint64_t event_value, + const char *label); +#endif +int ds4_gpu_end_commands(void); +int ds4_gpu_synchronize(void); + +int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size); +int ds4_gpu_set_model_fd(int fd); +int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map); +int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes); +int ds4_gpu_set_model_map_spans(const void *model_map, uint64_t model_size, const uint64_t *offsets, const uint64_t *sizes, uint32_t count, uint64_t max_tensor_bytes); +int ds4_gpu_cache_model_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, const char *label); +int ds4_gpu_cache_q8_f16_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, uint64_t in_dim, uint64_t out_dim, const char *label); +int ds4_gpu_q8_cache_suppressed(void); +void ds4_gpu_set_q8_cache_suppressed(int suppressed); +#ifdef DS4_ROCM_BUILD +void ds4_gpu_release_q8_f16_cache(void); +#endif + +/* Model-file ranges assigned to CUDA devices by the multi-GPU placement + * planner. Metal keeps these declarations for the shared engine interface. */ +#ifndef DS4_MAX_GPUS +#define DS4_MAX_GPUS 16 +#endif +typedef struct { + uint64_t source_offset; + uint64_t bytes; + int target_device; +} ds4_tensor_range; + +int ds4_gpu_device_cache_tensors(int device_id, + const ds4_tensor_range *ranges, + int n_ranges); +int ds4_gpu_register_support_map(const void *map, uint64_t size, uint64_t bias); +int ds4_gpu_device_cache_support_tensors(int device_id, + int entry_device_id, + const ds4_tensor_range *ranges, + int n_ranges, + int from_main_map); +uint64_t ds4_gpu_tier_free_vram(int logical_tier); +int ds4_gpu_lookup_cache(uint64_t source_offset, uint64_t bytes, + int *out_device_id, void **out_device_ptr); +int ds4_gpu_lookup_cache_device(uint64_t source_offset, uint64_t bytes); + +int ds4_gpu_pro_q4_expert_table_auto_available(void); +int ds4_gpu_preload_q4_expert_tables(const void *model_map, uint64_t model_size, + uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, + uint64_t gate_expert_bytes, uint64_t down_expert_bytes, + uint32_t n_total_expert); +int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes); +void ds4_gpu_set_quality(bool quality); +void ds4_gpu_set_glm_model(bool enabled); +void ds4_gpu_set_ssd_streaming(bool enabled); +void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled); +#ifdef __APPLE__ +void ds4_gpu_release_zero_prefix_prefill_mask_cache(void); +#endif +void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts); +void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes); +uint64_t ds4_gpu_recommended_working_set_size(void); +uint32_t ds4_gpu_stream_expert_cache_configured_count(void); +uint32_t ds4_gpu_stream_expert_cache_current_count(void); +typedef struct ds4_gpu_stream_expert_table { + const void *model_map; + uint64_t model_size; + uint32_t layer; + uint32_t n_total_expert; + uint64_t gate_offset; + uint64_t up_offset; + uint64_t down_offset; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; +} ds4_gpu_stream_expert_table; +/* Reset only the prompt-local eviction heuristic. The resident SSD expert + * cache itself is intentionally kept warm across sessions. */ +void ds4_gpu_stream_expert_cache_reset_route_hotness(void); +void ds4_gpu_stream_expert_cache_release_resident(void); +uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes); +int ds4_gpu_stream_expert_cache_seed_selected( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_selected); +int ds4_gpu_stream_expert_cache_begin_selected_load( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_selected); +int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( + const ds4_gpu_stream_expert_table *table, + const ds4_gpu_tensor *selected, + uint32_t n_selected); +#ifdef __APPLE__ +/* The async selected-load worker registers itself so Metal cache paths never + * wait on command buffers from that thread (they fail the load instead and + * the caller retries synchronously). */ +void ds4_gpu_stream_expert_cache_note_service_thread(void); +#endif +#if defined(DS4_ROCM_BUILD) || (!defined(DS4_NO_GPU) && !defined(__APPLE__)) +int ds4_gpu_stream_expert_cache_prepare_selected_batch( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_tokens, + uint32_t n_selected); +#endif +#ifdef DS4_ROCM_BUILD +int ds4_gpu_stream_expert_cache_load_layer( + const ds4_gpu_stream_expert_table *table); +int ds4_gpu_stream_expert_cache_seed_from_layer_selected( + const ds4_gpu_stream_expert_table *table, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_seed_tokens, + uint32_t n_selected); +int ds4_gpu_stream_expert_cache_release_layer_cache(void); +#endif +int ds4_gpu_stream_expert_cache_seed_experts( + const ds4_gpu_stream_expert_table *table, + const int32_t *expert_ids, + const uint32_t *expert_priorities, + uint32_t n_experts); +void ds4_gpu_print_memory_report(const char *label); + +/* Tensor-parallel per-layer gates (Metal only). The encoder calls + * ds4_gpu_tp_gate_encode() right after the kernels that produce a partial + * block output in the TP slab: it closes the current encoder, makes the GPU + * signal a shared event, queues the exchange on a service thread, and makes + * the GPU wait for the CPU-signaled release before the combine kernel runs. + * Sequence values are assigned internally and increase monotonically; both + * ranks encode the identical gate sequence so values pair up by + * construction. The exchange callback runs on the service thread and must + * return nonzero on success. */ +typedef int (*ds4_gpu_tp_exchange_fn)(void *ud, uint32_t layer, uint32_t gate, uint64_t seq); +/* Bind one rank of the two-way split. slab is the transport slab tensor and + * gpu_flags_off is the offset of its GPU-written gate-ready flag words. */ +int ds4_gpu_tp_init(uint32_t rank, + ds4_gpu_tensor *slab, uint64_t gpu_flags_off, + ds4_gpu_tp_exchange_fn fn, void *ud); +void ds4_gpu_tp_shutdown(void); +/* Multi-session TP reuses slab slots across several encoded graph tapes. + * Shared-event arrival is required in that mode to make each partial vector + * CPU-visible before the transport thread reads it. */ +void ds4_gpu_tp_set_session_batch_mode(int enabled); +/* The coordinator-only DSpark support model does not participate in TP. + * Suspend ownership only while encoding it; base-model verification remains + * split across both ranks. */ +void ds4_gpu_tp_suspend_expert_sharding(int suspend); +int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate); +/* Verify-block batch gates: one exchange per layer moving `rows` partial + * rows at once (speculative verify). The callback runs on the gate service + * thread with the same ud as the row-gate exchange fn. */ +typedef int (*ds4_gpu_tp_batch_exchange_fn)(void *ud, uint32_t layer, + uint32_t rows, uint64_t seq); +void ds4_gpu_tp_set_batch_exchange(ds4_gpu_tp_batch_exchange_fn fn); +int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows); +/* Prefill batch gates: the service thread exchanges `bytes` between two + * CPU-visible bounce tensors directly (payloads far beyond slab slots). */ +typedef int (*ds4_gpu_tp_big_exchange_fn)(void *ud, uint32_t layer, + uint64_t seq, const void *out, + void *in, uint64_t bytes); +void ds4_gpu_tp_set_big_exchange(ds4_gpu_tp_big_exchange_fn fn); +int ds4_gpu_tp_big_gate_encode(uint32_t layer, uint32_t rows, + const ds4_gpu_tensor *out_t, + ds4_gpu_tensor *in_t, + uint64_t bytes); +/* Split big gate: kick publishes the GPU arrival marker (batch shared + * event, whose completion semantics make the bounce payload visible to + * the exchange thread) and queues the exchange, returning the gate seq + * (0 on failure); wait encodes the release. Multiple kicks may be in + * flight; waiting on the last seq covers all earlier kicks (monotonic + * release event, in-order service thread). */ +uint64_t ds4_gpu_tp_big_gate_kick(uint32_t layer, uint32_t rows, + const ds4_gpu_tensor *out_t, + ds4_gpu_tensor *in_t, + uint64_t bytes); +int ds4_gpu_tp_big_gate_wait(uint64_t seq); +/* Pause/resume the DVFS keep-alive around work that keeps the GPU busy. + * No-op when TP is not bound. */ +void ds4_gpu_tp_keepalive_pause(int paused); +/* Split attention heads across the two TP ranks in the GLM batch-prefill + * attention kernels (qk-low, attention-lora, value-project). The caller + * zeroes the unowned head range of the heads buffer and combines the + * attn-output partials over the TP big-gate exchange. */ +void ds4_gpu_tp_set_attn_head_split(int enabled); +/* Skip the whole-file model residency set (TP sharding: only the + * owned ranges are warmed; the rest must never be paged in). Call before + * the model is mapped. */ +void ds4_gpu_model_residency_skip(int skip); +/* Nonzero after any gate exchange failed; the eval must abort. */ +int ds4_gpu_tp_failed(void); + +/* Tensor-parallel sliced projections (Metal decode path only). + * + * ds4_gpu_matmul_q8_0_kslice_tensor computes a k-range partial matvec: + * out[out_dim] = W[:, k_off : k_off + k_cnt] @ x[x_elem_off : +k_cnt] where + * W rows span full_in_dim quantized Q8_0 elements. k offsets/counts must be + * multiples of 32 (Q8_0 block). Partial results from both ranks sum to the + * full projection. + * + * ds4_gpu_attention_output_q8_tp_tensor is the group-sliced attention output + * pair: low projection for groups [group0, group0+group_cnt) plus the + * matching k-slice of the expand projection, producing this rank's partial + * attention block output (n_tokens == 1 only). */ +int ds4_gpu_matmul_q8_0_kslice_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t full_in_dim, + uint64_t k_off, + uint64_t k_cnt, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t x_elem_off); +/* CUDA multi-row variant. Each input row contains only the owned contiguous + * K slice, while each output row spans the full projection width. */ +int ds4_gpu_matmul_q8_0_kslice_rows_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t full_in_dim, + uint64_t out_dim, + uint64_t k_off, + uint64_t k_cnt, + const ds4_gpu_tensor *x, + uint64_t n_rows); +int ds4_gpu_matmul_quant_kslice_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint64_t full_in_dim, + uint64_t k_off, + uint64_t k_cnt, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t x_elem_off); +int ds4_gpu_attention_output_q8_tp_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups_total, + uint32_t group0, + uint32_t group_cnt, + uint64_t out_dim, + const ds4_gpu_tensor *heads); + +/* ========================================================================= + * Embeddings and Indexer Helpers. + * ========================================================================= + * + * These kernels seed HC state from token embeddings and implement the ratio-4 + * compressed-attention indexer that chooses visible compressed rows. + */ + +int ds4_gpu_embed_token_hc_tensor( + ds4_gpu_tensor *out_hc, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t token, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_embed_tokens_hc_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *tokens, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_embed_token_q8_0_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t token, + uint32_t n_embd); + +int ds4_gpu_embed_tokens_q8_0_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *tokens, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd); + +int ds4_gpu_embed_token_quant_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_vocab, + uint32_t token, + uint32_t n_embd); + +int ds4_gpu_embed_tokens_quant_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *tokens, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd); + +int ds4_gpu_indexer_score_one_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_head, + uint32_t head_dim, + float scale); + +int ds4_gpu_indexer_scores_prefill_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale); + +int ds4_gpu_indexer_scores_decode_batch_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale); + +int ds4_gpu_dspark_markov_argmax_tensor(ds4_gpu_tensor *out_idx, + const ds4_gpu_tensor *logits_row, + const void *model_map, + uint64_t model_size, + uint64_t w1_offset, + uint64_t w2_offset, + uint32_t prev_token, + uint32_t vocab, + uint32_t rank); +int ds4_gpu_indexer_topk_tensor( + ds4_gpu_tensor *selected, + const ds4_gpu_tensor *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k); + +int ds4_gpu_indexer_top1_value_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *values, + const ds4_gpu_tensor *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t index_offset); + +int ds4_gpu_matmul_q8_0_top1_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *values, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t index_offset); + +int ds4_gpu_set_decode_fast_attention(int enabled); +int ds4_gpu_set_decode_score_vec4(int enabled); + +/* GPU argmax over n_vocab F32 logits. Writes the winning index as int32 at + * out_idx[0]. Tie-break: lower index wins (matches host sample_argmax). */ +int ds4_gpu_argmax_tensor( + ds4_gpu_tensor *out_idx, + const ds4_gpu_tensor *logits, + uint32_t n_vocab); + +int ds4_gpu_dsv4_topk_mask_tensor( + ds4_gpu_tensor *mask, + const ds4_gpu_tensor *topk, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k); + +/* ========================================================================= + * Dense Projections, Norms, RoPE, and KV Rounding. + * ========================================================================= + * + * The graph uses these primitives for Q/KV projections, HC/output projections, + * attention output projections, and DS4's tail-only RoPE. + */ + +int ds4_gpu_matmul_q8_0_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +int ds4_gpu_matmul_q8_0_decode_mpp_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +int ds4_gpu_matmul_q8_0_decode_mpp_model_view_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +int ds4_gpu_matmul_q8_0_rows_scalar_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +int ds4_gpu_matmul_quant_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +int ds4_gpu_matmul_quant_decode_mpp_model_view_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +int ds4_gpu_matmul_quant_rows_scalar_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +/* Optional fused GPU operations. + * + * These are acceleration hooks, not required backend primitives. A backend + * that does not provide the fused kernel must still define the symbol and + * return 0. Callers then use the portable sequence of required primitives. + * Backends that return nonzero from a fused half-output operation must also + * implement the matching half-input HC expansion helpers below. + */ +int ds4_gpu_matmul_q8_0_pair_tensor( + ds4_gpu_tensor *out0, + ds4_gpu_tensor *out1, + const void *model_map, + uint64_t model_size, + uint64_t weight0_offset, + uint64_t weight1_offset, + uint64_t in_dim, + uint64_t out0_dim, + uint64_t out1_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +/* Multi-row decode projections that preserve the one-row reduction order. */ +int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_rows); +int ds4_gpu_matmul_q8_0_pair_decode_rows_exact_tensor( + ds4_gpu_tensor *out0, + ds4_gpu_tensor *out1, + const void *model_map, + uint64_t model_size, + uint64_t weight0_offset, + uint64_t weight1_offset, + uint64_t in_dim, + uint64_t out0_dim, + uint64_t out1_dim, + const ds4_gpu_tensor *x, + uint32_t n_rows); + +int ds4_gpu_matmul_q8_0_f16_out_tensor( + ds4_gpu_tensor *out_h, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp); +int ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor( + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *prequant, + uint32_t expert_split, + bool home_rank); + +int ds4_gpu_shared_mid_swiglu_q8_0_tensor( + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp); + +int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp); + +int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + float clamp); + +int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + float clamp); + +int ds4_gpu_matmul_f16_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +/* Exact multi-row form of the DeepSeek 4096x256 F16 router projection. */ +int ds4_gpu_matmul_f16_router_rows_exact_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + const ds4_gpu_tensor *x, + uint32_t n_rows); + +int ds4_gpu_matmul_f16_pair_tensor( + ds4_gpu_tensor *out_a, + ds4_gpu_tensor *out_b, + const void *model_map, + uint64_t model_size, + uint64_t weight_a_offset, + uint64_t weight_b_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +/* Optional Metal decode fusion. Returns 1 when the paired projection and + * recurrent compressor-state store were encoded, 0 when the optimized path + * is unavailable, and -1 on an attempted-path error. */ +int ds4_gpu_matmul_f16_pair_compressor_store_tensor( + ds4_gpu_tensor *out_kv, + ds4_gpu_tensor *out_score, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const void *model_map, + uint64_t model_size, + uint64_t weight_kv_offset, + uint64_t weight_score_offset, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t in_dim, + uint32_t width, + const ds4_gpu_tensor *x, + uint32_t ratio, + uint32_t pos); + +int ds4_gpu_matmul_f32_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok); + +int ds4_gpu_repeat_hc_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *row, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_repeat_hc_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *rows, + uint32_t n_tokens, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_rms_norm_plain_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + uint32_t n, + float eps); + +int ds4_gpu_rms_norm_plain_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + uint32_t n, + uint32_t rows, + float eps); + +int ds4_gpu_rms_norm_weight_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + float eps); + +int ds4_gpu_rms_norm_weight_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + uint32_t rows, + float eps); + +int ds4_gpu_add_rms_norm_weight_tensor( + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *sum_out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + float eps); + +int ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( + ds4_gpu_tensor *q_out, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t q_weight_offset, + uint32_t q_n, + ds4_gpu_tensor *kv_out, + const ds4_gpu_tensor *kv, + uint64_t kv_weight_offset, + uint32_t kv_n, + uint32_t rows, + float eps); + +int ds4_gpu_dsv4_qkv_rms_norm_rows_kv_rope_tensor( + ds4_gpu_tensor *q_out, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t q_weight_offset, + uint32_t q_n, + ds4_gpu_tensor *kv_out, + const ds4_gpu_tensor *kv, + uint64_t kv_weight_offset, + uint32_t kv_n, + uint32_t rows, + uint32_t kv_n_head, + uint32_t kv_head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps); + +int ds4_gpu_head_rms_norm_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + float eps); + +int ds4_gpu_head_rms_norm_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps); + +int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *q_half, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps); + +int ds4_gpu_dsv4_fp8_kv_quantize_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t head_dim, + uint32_t n_rot); + +int ds4_gpu_dsv4_indexer_qat_tensor( + ds4_gpu_tensor *x, + uint32_t n_rows, + uint32_t head_dim); + +int ds4_gpu_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_kv_lora_rms_norm_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *kv_raw, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + float eps); + +int ds4_gpu_glm_k_b_project_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *kv_norm, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t n_head); + +int ds4_gpu_glm_k_b_project_typed_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *kv_norm, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_tokens, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t n_head); + +int ds4_gpu_glm_store_compact_kv_tensor( + ds4_gpu_tensor *kv_lora_cache, + ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *kv_norm, + const ds4_gpu_tensor *kv_raw, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + bool cache_f16); + +int ds4_gpu_glm_qkv_norm_store_compact_kv_tensor( + ds4_gpu_tensor *q_out, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t q_weight_offset, + uint32_t q_n, + ds4_gpu_tensor *kv_lora_cache, + ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *kv_raw, + uint64_t kv_weight_offset, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + bool cache_f16, + float eps); + +int ds4_gpu_glm_store_indexer_k_tensor( + ds4_gpu_tensor *indexer_key_cache, + const ds4_gpu_tensor *raw_k, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t bias_offset, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t n_ctx_orig, + float eps, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool cache_f16); + +int ds4_gpu_glm_build_kv_cache_tensor( + ds4_gpu_tensor *key_cache, + ds4_gpu_tensor *value_cache, + const ds4_gpu_tensor *kv_raw, + const ds4_gpu_tensor *k_nope, + const ds4_gpu_tensor *value, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool cache_f16); + +int ds4_gpu_glm_build_kv_cache_flash_tensor( + ds4_gpu_tensor *key_cache, + ds4_gpu_tensor *value_cache, + const ds4_gpu_tensor *kv_raw, + const ds4_gpu_tensor *k_nope, + const ds4_gpu_tensor *value, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool cache_f16); + +int ds4_gpu_glm_attention_full_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16); + +int ds4_gpu_glm_fill_selected_range_tensor( + ds4_gpu_tensor *selected, + uint32_t n_selected); + +int ds4_gpu_glm_fill_selected_range_batch_tensor( + ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + uint32_t pad_row); + +int ds4_gpu_glm_indexer_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_indexer_score_one_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16); + +int ds4_gpu_glm_indexer_scores_batch_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16); + +int ds4_gpu_glm_qk_lowrank_q8_0_tensor( + ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim); + +int ds4_gpu_glm_qk_lowrank_q8_0_batch_tensor( + ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim); + +int ds4_gpu_glm_qk_lowrank_typed_tensor( + ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim); + +int ds4_gpu_glm_qk_lowrank_typed_batch_tensor( + ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim); + +int ds4_gpu_glm_value_project_q8_0_batch_heads_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *lora, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t value_dim); + +int ds4_gpu_glm_value_project_typed_batch_heads_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *lora, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t value_dim); + +int ds4_gpu_glm_attention_indexed_decode_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_rope_tail_decode_rows_tensor( + ds4_gpu_tensor *x, + const ds4_gpu_attention_decode_row *rows, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_attention_indexed_decode_typed_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + uint32_t value_weight_type, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *partial_lora, + ds4_gpu_tensor *partial_ms, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + bool selected_rows_valid, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + uint32_t block_rows, + uint32_t n_blocks, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *partial_lora, + ds4_gpu_tensor *partial_ms, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + uint32_t value_weight_type, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + bool selected_rows_valid, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + uint32_t block_rows, + uint32_t n_blocks, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_attention_indexed_batch_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_attention_indexed_batch_typed_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + uint32_t value_weight_type, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_sort_i32_rows_asc_tensor( + ds4_gpu_tensor *dst, + const ds4_gpu_tensor *src, + uint32_t row_width, + uint32_t n_rows); + +int ds4_gpu_glm_attention_indexed_batch_lora_tensor( + ds4_gpu_tensor *lora_out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( + ds4_gpu_tensor *lora_out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( + ds4_gpu_tensor *lora_out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + +int ds4_gpu_glm_attention_flash_staged_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16); + +int ds4_gpu_glm_attention_flash_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16); + +/* Release decode fused KV finalizer: after the standalone RoPE kernel, this + * performs DS4's FP8 non-RoPE KV round trip and writes the F16-rounded raw + * attention cache row in one dispatch. */ +int ds4_gpu_kv_fp8_store_raw_tensor( + ds4_gpu_tensor *kv, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t row, + uint32_t head_dim, + uint32_t n_rot); + +/* Exact multi-session form of the decode KV finalizer. KV rows are + * contiguous, while each output row is written to its session-private cache. */ +int ds4_gpu_kv_fp8_store_raw_decode_rows_tensor( + ds4_gpu_tensor *kv, + ds4_gpu_tensor *const *raw_caches, + const uint32_t *raw_caps, + const uint32_t *raw_rows, + uint32_t n_rows, + uint32_t head_dim, + uint32_t n_rot); + +/* Reference/raw-cache primitive kept for prefill and diagnostics. Decode uses + * ds4_gpu_kv_fp8_store_raw_tensor unless a diagnostic reference path is + * explicitly selected by the graph driver. */ +int ds4_gpu_store_raw_kv_tensor( + ds4_gpu_tensor *raw_cache, + const ds4_gpu_tensor *kv, + uint32_t raw_cap, + uint32_t row, + uint32_t head_dim); + +int ds4_gpu_store_raw_kv_batch_tensor( + ds4_gpu_tensor *raw_cache, + const ds4_gpu_tensor *kv, + uint32_t raw_cap, + uint32_t pos0, + uint32_t n_tokens, + uint32_t head_dim); + +/* ========================================================================= + * KV Compression and Attention. + * ========================================================================= + * + * Compressed layers maintain rolling score/KV state and append pooled rows at + * ratio boundaries. Attention kernels consume raw SWA rows, compressed rows, + * and optional indexer masks. + */ + +int ds4_gpu_compressor_update_tensor( + const ds4_gpu_tensor *kv_cur, + const ds4_gpu_tensor *sc_cur, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + ds4_gpu_tensor *comp_cache, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos, + uint32_t comp_row, + uint32_t n_rot, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps, + bool state_already_stored); + +int ds4_gpu_compressor_store_batch_tensor( + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens); + +int ds4_gpu_compressor_prefill_tensor( + ds4_gpu_tensor *comp_cache, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens, + uint32_t n_rot, + uint32_t n_ctx_orig, + bool quantize_fp8, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps); + +int ds4_gpu_compressor_prefill_ratio4_replay_tensor( + ds4_gpu_tensor *comp_cache, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t pos0, + uint32_t n_tokens, + uint32_t n_rot, + uint32_t n_ctx_orig, + bool quantize_fp8, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps); + +int ds4_gpu_compressor_prefill_state_ratio4_tensor( + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv_tail, + const ds4_gpu_tensor *sc_tail, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t head_dim, + uint32_t pos0); + +int ds4_gpu_attention_decode_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_comp, + const ds4_gpu_tensor *comp_mask, + uint32_t use_mask, + uint32_t n_head, + uint32_t head_dim); + +int ds4_gpu_attention_decode_heads_rope_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_comp, + const ds4_gpu_tensor *comp_mask, + uint32_t use_mask, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + int *fused_inv_rope); + +/* Multi-session decode over contiguous Q/head rows and private KV caches. + * The row table is copied into CUDA launch parameters, so no device-side + * descriptor upload or synchronization is required. */ +int ds4_gpu_attention_decode_rows_rope_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_attention_decode_row *rows, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); +/* Diagnostic/public form of the dk=512 gathered decode-attention KV staging + * step. The compressed source must be F16; dst writes chronological raw-ring + * rows followed by compressed rows and must not overlap either source. */ +int ds4_gpu_flash_kv_stage_f16_tensor( + ds4_gpu_tensor *dst, + const ds4_gpu_tensor *raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_raw, + const ds4_gpu_tensor *comp, + uint32_t comp_is_f16, + uint32_t n_comp, + uint32_t head_dim); + +int ds4_gpu_attention_prefill_raw_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t window, + uint32_t n_head, + uint32_t head_dim); + +/* Rectangular raw prefill attention: q is a view of the n_q query rows at + * token positions [q_row0, q_row0 + n_q) of the chunk, raw_kv keeps all + * n_kv rows, heads receives n_q output rows. Used by the TP prefill row + * split; the square entry above is the q_row0 = 0, n_q = n_kv case. */ +int ds4_gpu_attention_prefill_raw_heads_range_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_kv, + uint32_t window, + uint32_t n_head, + uint32_t head_dim); + +int ds4_gpu_attention_decode_raw_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t window, + uint32_t n_head, + uint32_t head_dim); + +int ds4_gpu_attention_noncausal_raw_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_head, + uint32_t head_dim); + +int ds4_gpu_attention_decode_mixed_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim); + +int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *topk, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t top_k, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim); + +int ds4_gpu_attention_prefill_static_mixed_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim); + +/* Rectangular static-mixed prefill attention: q is a view of the n_q query + * rows at token positions [q_row0, q_row0 + n_q) of the chunk, while raw_kv + * keeps all n_tokens rows and comp_kv all n_comp compressed keys. Used by + * the TP prefill row split; the square entry above is q_row0 = 0, + * n_q = n_tokens. */ +int ds4_gpu_attention_prefill_static_mixed_heads_range_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim); + +int ds4_gpu_attention_prefill_masked_mixed_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim); + +int ds4_gpu_attention_output_q8_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, + ds4_gpu_tensor *low_tmp, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens); +int ds4_gpu_attention_output_q4_K_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, + ds4_gpu_tensor *low_tmp, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint32_t out_b_type, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens); + +int ds4_gpu_attention_output_q8_batch_f16_tensor( + ds4_gpu_tensor *out_h, + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens); + +int ds4_gpu_attention_output_low_q8_tensor( + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + const ds4_gpu_tensor *heads); +int ds4_gpu_attention_output_low_q4_K_slice_tensor( + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t group0, + uint32_t group_cnt, + const ds4_gpu_tensor *heads); + +int ds4_gpu_attention_output_low_q8_rows_exact_tensor( + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups_total, + uint32_t group0, + uint32_t group_cnt, + const ds4_gpu_tensor *heads, + uint32_t n_rows); + +int ds4_gpu_attention_output_q8_tp_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups_total, + uint32_t group0, + uint32_t group_cnt, + uint64_t out_dim, + const ds4_gpu_tensor *heads); + +/* ========================================================================= + * Router, Shared Expert, and Routed MoE. + * ========================================================================= + * + * These kernels implement the FFN body: router probabilities/top-k or hash + * routing, shared SwiGLU, and the IQ2_XXS/Q2_K/Q4_K routed experts. + */ + +int ds4_gpu_swiglu_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *gate, + const ds4_gpu_tensor *up, + uint32_t n, + float clamp, + float weight); + +int ds4_gpu_add_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + uint32_t n); + +int ds4_gpu_add3_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + const ds4_gpu_tensor *c, + uint32_t n); + +int ds4_gpu_directional_steering_project_tensor( + ds4_gpu_tensor *x, + const ds4_gpu_tensor *directions, + uint32_t layer, + uint32_t width, + uint32_t rows, + float scale); + +int ds4_gpu_router_select_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + uint64_t hash_offset, + uint32_t hash_rows, + uint32_t token, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_expert_groups, + uint32_t n_group_used, + bool has_bias, + bool hash_mode, + const ds4_gpu_tensor *logits); + +int ds4_gpu_router_select_batch_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + uint64_t hash_offset, + uint32_t hash_rows, + uint32_t n_expert_groups, + uint32_t n_group_used, + bool has_bias, + bool hash_mode, + const ds4_gpu_tensor *logits, + const ds4_gpu_tensor *tokens, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_tokens); + +int ds4_gpu_glm_router_select_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + const ds4_gpu_tensor *logits, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale); + +int ds4_gpu_glm_router_select_batch_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + const ds4_gpu_tensor *logits, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_tokens); + +int ds4_gpu_glm_routed_moe_one_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + bool force_resident); + +int ds4_gpu_glm_routed_moe_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride, + bool force_resident); + +int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride); + +int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected); +void ds4_gpu_set_glm_mtp_verify_mode(bool enabled); + +int ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *block_out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + uint64_t in_start, + uint64_t in_count, + const ds4_gpu_tensor *x, + const ds4_gpu_tensor *block_add, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_routed_moe_one_owned_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *experts, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t resident_expert_base, + uint32_t resident_expert_count, + float clamp, + const ds4_gpu_tensor *x, + ds4_gpu_tensor *down_output, + bool pack_fixed3, + ds4_gpu_tensor *shared_prequant); + +int ds4_gpu_routed_moe_batch_owned_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *experts, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t resident_expert_base, + uint32_t resident_expert_count, + float clamp, + const ds4_gpu_tensor *x, + uint32_t layer_index, + uint32_t n_tokens, + bool *mid_is_f16); + +int ds4_gpu_routed_moe_owned_slots_combine_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *home_slots, + const ds4_gpu_tensor *peer_slots, + const ds4_gpu_tensor *selected, + uint32_t out_dim, + uint32_t expert_split); + +int ds4_gpu_routed_moe_owned_slots_combine_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *home_slots, + const ds4_gpu_tensor *peer_slots, + const ds4_gpu_tensor *selected, + uint32_t out_dim, + uint32_t expert_split, + uint32_t rows); + +int ds4_gpu_routed_moe_owned_packed_combine_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *home_slots, + const ds4_gpu_tensor *peer_packed, + const ds4_gpu_tensor *selected, + uint32_t out_dim, + uint32_t expert_split); + +int ds4_gpu_routed_moe_one_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *experts, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + float clamp, + const ds4_gpu_tensor *x, + const ds4_gpu_tensor *add_in, + uint32_t layer_index, + bool force_resident); + +int ds4_gpu_routed_moe_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *experts, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + float clamp, + const ds4_gpu_tensor *x, + uint32_t layer_index, + uint32_t n_tokens, + bool *mid_is_f16, + bool force_resident); + +/* ========================================================================= + * Hyper-Connection Kernels. + * ========================================================================= + * + * HC kernels reduce four residual streams before a sublayer and expand the + * sublayer output back into four streams afterward. + */ + +int ds4_gpu_hc_split_sinkhorn_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *mix, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps); + +int ds4_gpu_hc_weighted_sum_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *weights, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_hc_weighted_sum_norm_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *weights, + const void *model_map, + uint64_t model_size, + uint64_t norm_weight_offset, + uint32_t n_embd, + uint32_t n_hc, + float norm_eps); + +int ds4_gpu_hc_weighted_sum_split_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); + +/* Release decode fused HC pre-sublayer operation: split the HC mixer and + * immediately reduce four HC streams into the active 4096-wide sublayer row. */ +int ds4_gpu_hc_split_weighted_sum_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *mix, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps); + +int ds4_gpu_hc_split_weighted_sum_norm_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *mix, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps, + float norm_eps); + +/* Batched HC RMSNorm followed by its narrow F16 mixer projection. On the + * tuned Metal path, scale_scratch stores one float per row instead of the + * full normalized HC tensor; other shapes retain the established fallback. */ +int ds4_gpu_hc_rms_scale_project_f16_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *scale_scratch, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t in_dim, + uint32_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_rows, + float eps); + +int ds4_gpu_output_hc_weights_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *pre, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_hc, + float eps); + +int ds4_gpu_hc_expand_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *post, + const ds4_gpu_tensor *comb, + uint32_t n_embd, + uint32_t n_hc); +int ds4_gpu_hc_expand_add_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *block_add, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *post, + const ds4_gpu_tensor *comb, + uint32_t n_embd, + uint32_t n_hc); + + +int ds4_gpu_hc_expand_add_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *block_add, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *post, + const ds4_gpu_tensor *comb, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_hc_expand_split_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_hc_expand_split_half_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out_h, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_hc_expand_add_split_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *block_add, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_hc_expand_add_split_half_add_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *block_add_h, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_shared_down_hc_expand_q8_0_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *shared_out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *shared_mid, + const ds4_gpu_tensor *routed_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_shared_down_hc_expand_add_q8_0_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *shared_out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *shared_mid, + const ds4_gpu_tensor *routed_out, + const ds4_gpu_tensor *routed_add, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_shared_down_hc_expand_owned_q8_0_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *shared_out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *shared_mid, + const ds4_gpu_tensor *home_slots, + const ds4_gpu_tensor *peer_packed, + const ds4_gpu_tensor *selected, + uint32_t expert_split, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); + +int ds4_gpu_matmul_q8_0_hc_expand_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *block_out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/native/metal/ds4_metal.m b/native/metal/ds4_metal.m new file mode 100644 index 0000000..40d2b74 --- /dev/null +++ b/native/metal/ds4_metal.m @@ -0,0 +1,39615 @@ +#import +#import + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ds4.h" +#include "ds4_gpu.h" + +/* + * Objective-C Metal glue for the Rust engine. + * + * Rust owns model semantics and graph scheduling. This file owns only + * Metal objects: device/queue/library setup, mmap-backed weight views, command + * batching, persistent tensors, scratch buffers, and thin wrappers around the + * kernel files in the metal directory. Keeping this boundary narrow makes the + * inference path readable from Rust while still using Objective-C where Metal + * requires it. + */ + +enum { + DS4_METAL_TENSOR_Q4_0 = 2, + DS4_METAL_TENSOR_Q8_0 = 8, + DS4_METAL_TENSOR_Q2_K = 10, + DS4_METAL_TENSOR_Q4_K = 12, + DS4_METAL_TENSOR_Q5_K = 13, + DS4_METAL_TENSOR_Q6_K = 14, + DS4_METAL_TENSOR_Q8_K = 15, + DS4_METAL_TENSOR_IQ2_XXS = 16, +}; + +@class DS4MetalQ4ExpertTable; + +static id g_device; +static id g_queue; +static id g_library; +static id g_batch_cb; +static id g_batch_enc; +static BOOL g_batch_has_work; +static NSMutableArray> *g_pending_cbs; +static id g_selected_readback_event; +static uint64_t g_selected_readback_event_value; +static id g_set_rows_f32_i32_pipeline; +static id g_get_rows_f32_pipeline; +static id g_get_rows_f16_pipeline; +static id g_get_rows_i32_pipeline; +static id g_get_rows_q8_0_pipeline; +static id g_get_rows_q4_0_pipeline; +static id g_get_rows_q4_K_pipeline; +static id g_repeat_f32_pipeline; +static id g_concat_pipeline; +static id g_cpy_f32_f32_pipeline; +static id g_cpy_f32_f16_pipeline; +static id g_cpy_contig_f32_f16_pipeline; +static id g_cpy_f16_f32_pipeline; +static id g_cpy_f16_f16_pipeline; +static id g_cpy_contig_f16_f32_pipeline; +static id g_cpy_contig_f16_f16_pipeline; +static id g_flash_kv_stage_f16_pipeline; +static id g_swiglu_pipeline; +static id g_swiglu_flat_pipeline; +static id g_add_pipeline; +static id g_add2_pipeline; +static id g_add3_pipeline; +static id g_moe_sum6_pipeline; +static id g_moe_sum8_pipeline; +static id g_mul_pipeline; +static id g_rms_norm_pipeline; +static id g_rms_norm_plain_pipeline; +static id g_add_rms_norm_pipeline; +static id g_rms_norm_scale_pipeline; +static id g_dsv4_qkv_rms_norm_pipeline; +static id g_hc_split_sinkhorn_pipeline; +static id g_hc_split_weighted_sum_pipeline; +static id g_hc_split_weighted_sum_norm_pipeline; +static id g_hc_weighted_sum_pipeline; +static id g_hc_weighted_sum_norm_pipeline; +static id g_output_hc_weights4_pipeline; +static id g_hc_expand_pipeline; +static id g_unary_sigmoid_pipeline; +static id g_unary_silu_pipeline; +static id g_unary_softplus_pipeline; +static id g_unary_sqrt_pipeline; +static id g_unary_clamp_pipeline; +static id g_unary_scale_pipeline; +static id g_unary_fill_pipeline; +static id g_unary_fill_f16_pipeline; +static id g_bin_mul_scalar_pipeline; +static id g_bin_div_row_pipeline; +static id g_moe_mul_mv_id_iq2_xxs_pipeline; +static id g_moe_mul_mv_id_iq2_xxs_pair_pipeline; +static id g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; +static id g_moe_mul_mv_id_q2_k_pipeline; +static id g_moe_mul_mv_id_q2_k_sum6_pipeline; +static id g_moe_mul_mv_id_iq2_xxs_sum6_pipeline; +static id g_moe_mul_mv_id_q4_k_pipeline; +static id g_moe_mul_mv_id_q4_k_pair_pipeline; +static id g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_id_q4_k_sum6_pipeline; +static id g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_group_q4_k_sum6_pipeline; +static id g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_group6_q4_k_sum6_pipeline; +static id g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_group8_q4_k_sum6_pipeline; +static id g_moe_mul_mv_group24_q4_k_id_pipeline; +static id g_moe_mul_mv_group24_q4_k_sum6_pipeline; +static id g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline; +static id g_moe_mul_mv_slots6_q2_k_sum6_pipeline; +static id g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_slots6_q4_k_sum6_pipeline; +static id g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline; +static id g_moe_mul_mv_addr_iq2_xxs_pipeline; +static id g_moe_mul_mv_addr_q2_k_sum6_pipeline; +static id g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline; +static id g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline; +static id g_moe_stream_expert_cache_validate_pipeline; +static id g_moe_q4_gather_slots6_pipeline; +static id g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_table_q4_k_sum6_pipeline; +static id g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline; +static id g_moe_mul_mv_addr_q4_k_sum6_pipeline; +static id g_moe_table_q4_pair_gate_encoder; +static id g_moe_table_q4_pair_up_encoder; +static id g_moe_table_q4_sum_down_encoder; +static id g_rope_tail_batch_pipeline; +static id g_rope_tail_inplace_pair_pipeline; +static id g_rope_tail_inplace_pair_shared4_pipeline; +static id g_rope_tail_inplace_pair_affine_pipeline; +static id g_dsv4_fp8_kv_quantize_pipeline; +static id g_dsv4_indexer_qat_pipeline; +static id g_dsv4_kv_fp8_store_pipeline; +static id g_dsv4_ratio4_shift_pipeline; +static id g_dsv4_compressor_pack_ratio4_pipeline; +static id g_dsv4_softmax_pool_ratio4_direct_pipeline; +static id g_dsv4_softmax_pool_pipeline; +static id g_soft_max_f32_pipeline; +static id g_soft_max_f32_4_pipeline; +static id g_argsort_f32_i32_desc_pipeline; +static id g_argsort_merge_f32_i32_desc_pipeline; +static id g_sum_rows_f32_f32_pipeline; +static id g_dsv4_topk_mask_pipeline; +static id g_dsv4_topk_mask_scatter_pipeline; +static id g_dsv4_indexer_weighted_sum_pipeline; +static id g_dsv4_indexer_score_one_direct_pipeline; +static id g_dsv4_compressor_store_one_pipeline; +static id g_dsv4_sort_i32_rows_asc_pipeline; +static id g_dsv4_indexed_attention_heads8_pipeline; +static id g_dsv4_indexed_attention_heads8_rb16_pipeline; +static id g_dsv4_softplus_sqrt_pipeline; +static id g_dsv4_router_finalize_one_pipeline; +static id g_dsv4_router_finalize_one_simd_pipeline; +static id g_dsv4_router_finalize_weights_one_simd_pipeline; +static id g_dsv4_router_weights_one_pipeline; +static id g_glm_router_select_one_pipeline; +static id g_glm_kv_lora_rms_norm_pipeline; +static id g_glm_k_b_project_pipeline; +static id g_glm_store_compact_kv_pipeline; +static id g_glm_qkv_norm_store_compact_kv_pipeline; +static id g_glm_store_indexer_k_pipeline; +static id g_glm_build_kv_cache_pipeline; +static id g_glm_build_kv_cache_decode_group4_pipeline; +static id g_glm_build_kv_cache_flash_pipeline; +static id g_glm_attention_full_pipeline; +static id g_glm_fill_selected_range_pipeline; +static id g_glm_fill_selected_range_batch_pipeline; +static id g_glm_indexer_rope_tail_pipeline; +static id g_glm_indexer_score_one_pipeline; +static id g_glm_indexer_score_one_direct_pipeline; +static id g_glm_indexer_scores_batch_pipeline; +static id g_glm_indexer_scores_tiled_pipeline; +static id g_glm_indexer_scores_tiled_f32_pipeline; +static id g_glm_qk_lowrank_pipeline; +static id g_glm_qk_lowrank_glm52_pipeline; +static id g_glm_qk_lowrank_glm52_sg_pipeline; +static id g_glm_qk_lowrank_batch_pipeline; +static id g_glm_qk_lowrank_batch_glm52_t4_pipeline; +static id g_glm_value_project_q8_0_pipeline; +static id g_glm_value_project_q8_0_batch_heads_pipeline; +static id g_glm_value_project_q8_0_batch_heads_mma_pipeline; +static id g_glm_attention_indexed_decode_pipeline; +static id g_glm_attention_indexed_decode_split_group8_partial_pipeline; +static id g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline; +static id g_glm_attention_indexed_decode_split_group8_reduce_pipeline; +static id g_glm_attention_indexed_decode_split_group8_reduce16_pipeline; +static id g_glm_attention_indexed_batch_pipeline; +static id g_glm_attention_indexed_batch_group2_pipeline; +static id g_glm_attention_indexed_batch_q2_group4_pipeline; +static id g_glm_attention_indexed_batch_group8_pipeline; +static id g_glm_attention_indexed_batch_lora_group8_vec_pipeline; +static id g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline; +static id g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline; +static id g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline; +static id g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline; +static id g_glm_q4_k_pair_swiglu_f32_pipeline; +static id g_glm_q4_k_pair_swiglu2_f32_pipeline; +static id g_glm_q4_k_pair_swiglu4_f32_pipeline; +static id g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline; +static id g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline; +static id g_glm_q2_k_pair_swiglu_f32_pipeline; +static id g_glm_q2_k_addr_pair_swiglu2_f32_pipeline; +static id g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline; +static id g_glm_q4_k_addr_pair_swiglu_f32_pipeline; +static id g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline; +static id g_glm_q2_k_down_f32_pipeline; +static id g_glm_q4_k_down_f32_pipeline; +static id g_glm_q2_k_addr_down_f32_pipeline; +static id g_glm_q4_k_addr_down_f32_pipeline; +static id g_glm_q5_k_pair_swiglu_f32_pipeline; +static id g_glm_q5_k_pair_swiglu_mapped_f32_pipeline; +static id g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline; +static id g_glm_q5_k_down_f32_pipeline; +static id g_glm_q6_k_down_f32_pipeline; +static id g_dsv4_router_weights_batch_pipeline; +static id g_dsv4_hc_expand4_pipeline; +static NSMutableDictionary> *g_pipeline_cache; +static NSMutableDictionary> *g_model_buffer_cache; +static NSMutableDictionary *g_q4_expert_table_cache; +static NSMutableDictionary *g_q4_expert_layer_residency_cache; +static NSMutableArray> *g_transient_buffers; +static id g_model_residency_set; + +typedef struct { + id __strong mask; + id __strong blk; + NSUInteger mask_bytes; + NSUInteger blk_bytes; + uint32_t kind; + uint32_t n_tokens; + uint32_t n_comp; + uint32_t n_keys; + uint32_t window; + uint32_t ratio; + uint32_t nqptg; + uint32_t ncpsg; + bool has_kvpad; + bool bc_mask; + bool valid; + bool blk_ready; +} ds4_gpu_zero_prefix_prefill_mask_cache_entry; + +enum { + DS4_GPU_PREFILL_MASK_CACHE_RAW = 1, + DS4_GPU_PREFILL_MASK_CACHE_RATIO4 = 2, + DS4_GPU_PREFILL_MASK_CACHE_RATIO128 = 3, + DS4_GPU_PREFILL_MASK_CACHE_SLOTS = 3, +}; + +static ds4_gpu_zero_prefix_prefill_mask_cache_entry + g_zero_prefix_prefill_mask_cache[DS4_GPU_PREFILL_MASK_CACHE_SLOTS]; +static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void); +static id g_flash_attn_mask_buffer; +static id g_flash_attn_zero_mask_buffer; +static id g_flash_attn_pad_buffer; +static id g_flash_attn_tmp_buffer; +static id g_flash_attn_blk_buffer; +static id g_flash_attn_ring_buffer; +static id g_flash_attn_kv_buffer; +static id g_glm_flash_attn_mask_buffer; +static id g_compressor_pool_kv_buffer; +static id g_compressor_pool_score_buffer; +static id g_compressor_pool_score_cont_buffer; +static id g_compressor_pool_softmax_buffer; +static id g_compressor_pool_product_buffer; +static id g_compressor_store_ape_buffer; +static id g_compressor_store_score_buffer; +static id g_embed_rows_buffer; +static id g_router_selection_buffer; +static id g_router_weight_sum_buffer; +static id g_indexer_head_scores_buffer; +static id g_indexer_topk_buffer; +static id g_indexed_topk_buffer; +static id g_f16_round_scratch_buffer; +static id g_raw_store_round_buffer; +static id g_moe_gate_scratch_buffer; +static id g_moe_down_scratch_buffer; +static id g_moe_id_map_buffer; +static id g_moe_q4_gate_slots_buffer; +static id g_moe_q4_up_slots_buffer; +static id g_moe_q4_down_slots_buffer; +static id g_attn_out_group_ids_buffer; +static int g_model_fd = -1; +static const void *g_model_map_ptr; +static uint64_t g_model_map_size; +static uint64_t g_model_mapped_offset; +static uint64_t g_model_mapped_size; +static uint64_t g_model_mapped_max_tensor_bytes; +static uint64_t g_tensor_alloc_live_bytes; +static uint64_t g_tensor_alloc_peak_bytes; +static pthread_mutex_t g_tensor_mu = PTHREAD_MUTEX_INITIALIZER; +static uintptr_t *g_tensor_live_slots; +static size_t g_tensor_live_cap; +static size_t g_tensor_live_count; +static size_t g_tensor_live_tombs; +static uint64_t g_model_wrap_count; +static uint64_t g_model_wrap_bytes; +static uint64_t g_model_wrap_max_bytes; +static uint64_t g_model_buffer_cache_bytes; +static uint64_t g_model_buffer_cache_evictions; +static int g_model_buffer_cache_over_limit; +static uint64_t g_stream_expert_cache_bytes; +static uint64_t g_stream_expert_cache_expert_bytes; +static uint32_t g_stream_expert_cache_entry_count; +static uint32_t g_stream_expert_cache_budget_override; +static uint64_t g_stream_expert_cache_hits; +static uint64_t g_stream_expert_cache_misses; +static uint64_t g_stream_expert_cache_evictions; +static uint64_t g_stream_expert_cache_wraps; +static uint64_t g_stream_expert_cache_clock; +static uint64_t g_stream_expert_cache_evict_advise_bytes; +static uint64_t g_stream_expert_cache_willneed_advise_bytes; +static uint64_t g_stream_expert_cache_pread_bytes; +static double g_stream_expert_cache_pread_ms; +static uint64_t g_stream_expert_cache_buffer_allocs; +static uint64_t g_stream_expert_cache_buffer_reuses; +static uint64_t g_stream_expert_cache_decode_tokens; +static uint64_t g_stream_expert_cache_hotness_decay_token; +static uint64_t g_stream_expert_timing_selected_calls; +static double g_stream_expert_timing_selected_read_ms; +static double g_stream_expert_timing_selected_sync_ms; +static double g_stream_expert_timing_selected_copy_ms; +static double g_stream_expert_timing_selected_bind_ms; +static uint64_t g_stream_expert_timing_split_layers; +static uint64_t g_stream_expert_timing_split_resident_experts; +static uint64_t g_stream_expert_timing_split_missing_experts; +static double g_stream_expert_timing_split_resident_ms; +static double g_stream_expert_timing_split_missing_ms; +static double g_stream_expert_timing_split_missing_load_ms; +static double g_stream_expert_timing_split_missing_slot_ms; +static double g_stream_expert_timing_split_missing_prune_ms; +static double g_stream_expert_timing_split_missing_addr_ms; +static double g_stream_expert_timing_split_missing_wait_ms; +static uint64_t g_stream_expert_timing_load_calls; +static double g_stream_expert_timing_load_prepare_ms; +static double g_stream_expert_timing_load_pread_ms; +static double g_stream_expert_timing_load_modify_ms; +static double g_stream_expert_timing_load_install_ms; +static uint64_t g_stream_expert_timing_prepare_batch_reuse_calls; +static double g_stream_expert_timing_prepare_batch_reuse_ms; +static uint64_t g_stream_expert_timing_prepare_buffer_calls; +static double g_stream_expert_timing_prepare_buffer_ms; +static uint64_t g_stream_expert_timing_prepare_task_experts; +static double g_stream_expert_timing_prepare_task_ms; +static uint64_t g_stream_expert_timing_reuse_scan_calls; +static uint64_t g_stream_expert_timing_reuse_scan_entries; +static double g_stream_expert_timing_reuse_scan_ms; +static double g_stream_expert_timing_reuse_clear_ms; +static uint64_t g_stream_expert_timing_readahead_calls; +static uint64_t g_stream_expert_timing_readahead_bytes; +static double g_stream_expert_timing_readahead_ms; +static uint64_t g_stream_expert_timing_cache_all_resident_layers; +static uint64_t g_stream_expert_timing_cache_all_missing_layers; +static uint64_t g_stream_expert_timing_cache_mixed_layers; +static uint64_t g_stream_expert_timing_cache_resident_experts; +static uint64_t g_stream_expert_timing_cache_missing_experts; +typedef struct { + uint64_t selected_calls; + double selected_read_ms; + double selected_sync_ms; + double selected_copy_ms; + double selected_bind_ms; + uint64_t split_layers; + uint64_t split_resident_experts; + uint64_t split_missing_experts; + double split_resident_ms; + double split_missing_ms; + double split_missing_load_ms; + double split_missing_slot_ms; + double split_missing_prune_ms; + double split_missing_addr_ms; + double split_missing_wait_ms; + uint64_t load_calls; + double load_prepare_ms; + double load_pread_ms; + double load_modify_ms; + double load_install_ms; + uint64_t prepare_batch_reuse_calls; + double prepare_batch_reuse_ms; + uint64_t prepare_buffer_calls; + double prepare_buffer_ms; + uint64_t prepare_task_experts; + double prepare_task_ms; + uint64_t reuse_scan_calls; + uint64_t reuse_scan_entries; + double reuse_scan_ms; + double reuse_clear_ms; + uint64_t readahead_calls; + uint64_t readahead_bytes; + double readahead_ms; + uint64_t cache_all_resident_layers; + uint64_t cache_all_missing_layers; + uint64_t cache_mixed_layers; + uint64_t cache_resident_experts; + uint64_t cache_missing_experts; +} ds4_gpu_stream_expert_timing_snapshot; +static ds4_gpu_stream_expert_timing_snapshot g_stream_expert_timing_last_report; +static int g_stream_prefill_batch_selected_addr_building; +static int g_glm_stream_expert_addr_table_building; +static uint64_t g_model_residency_count; +static int g_model_residency_added_to_queue; +static int g_glm_model_mode; +static int g_ssd_streaming_mode; +static int g_glm_streaming_prefill_full_layer_runtime; +static int g_metal4_runtime_available; +static int g_metal4_family_supported; +static int g_metal4_queue_supported; +static int g_metal4_m5_neural_accelerators_hint; +static int g_metal4_tensor_api_enabled; +static int g_metal4_tensor_api_compile_supported; +static char g_metal_device_name[128]; +static int ds4_gpu_model_map_log_enabled(void); +static int ds4_gpu_stream_expert_cache_note_expert_size( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes); +static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void); +static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats); +static void ds4_gpu_stream_expert_pending_load_clear(void); +static void ds4_gpu_stream_expert_pread_pool_shutdown(void); +static int ds4_gpu_stream_expert_timing_summary_enabled(void); +static int ds4_gpu_stream_expert_cache_entry_protected( + uint32_t layer, + uint32_t expert, + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect); + +/* The async selected-load worker registers itself so cache paths that would + * flush/wait on command buffers (a race against the encoding thread) fail + * the load instead; the caller then retries on the main thread. */ +static pthread_t g_stream_expert_service_thread; +static int g_stream_expert_service_thread_set; + +void ds4_gpu_stream_expert_cache_note_service_thread(void) { + g_stream_expert_service_thread = pthread_self(); + g_stream_expert_service_thread_set = 1; +} + +static int ds4_gpu_stream_expert_cache_on_service_thread(void) { + return g_stream_expert_service_thread_set && + pthread_equal(pthread_self(), g_stream_expert_service_thread); +} +static NSUInteger g_flash_attn_mask_bytes; +static NSUInteger g_flash_attn_zero_mask_bytes; +static NSUInteger g_flash_attn_pad_bytes; +static NSUInteger g_flash_attn_tmp_bytes; +static NSUInteger g_flash_attn_blk_bytes; +static NSUInteger g_flash_attn_ring_bytes; +static NSUInteger g_flash_attn_kv_bytes; +static NSUInteger g_glm_flash_attn_mask_bytes; +static uint32_t g_glm_flash_attn_mask_pos0; +static uint32_t g_glm_flash_attn_mask_tokens; +static uint32_t g_glm_flash_attn_mask_cache_len; +static int g_glm_flash_attn_mask_valid; +static NSUInteger g_compressor_pool_kv_bytes; +static NSUInteger g_compressor_pool_score_bytes; +static NSUInteger g_compressor_pool_score_cont_bytes; +static NSUInteger g_compressor_pool_softmax_bytes; +static NSUInteger g_compressor_pool_product_bytes; +static NSUInteger g_compressor_store_ape_bytes; +static NSUInteger g_compressor_store_score_bytes; +static NSUInteger g_embed_rows_bytes; +static NSUInteger g_router_selection_bytes; +static NSUInteger g_router_weight_sum_bytes; +static NSUInteger g_indexer_head_scores_bytes; +static NSUInteger g_indexer_topk_bytes; +static NSUInteger g_indexed_topk_bytes; +static NSUInteger g_f16_round_scratch_bytes; +static NSUInteger g_raw_store_round_bytes; +static NSUInteger g_moe_gate_scratch_bytes; +static NSUInteger g_moe_down_scratch_bytes; +static NSUInteger g_moe_id_map_bytes; +static NSUInteger g_moe_q4_gate_slots_bytes; +static NSUInteger g_moe_q4_up_slots_bytes; +static NSUInteger g_moe_q4_down_slots_bytes; +static NSUInteger g_attn_out_group_ids_bytes; +static int g_initialized; +static int g_quality_mode; +static int g_mpp_invalid_env_reported; +#define DS4_METAL_MAX_ROUTED_EXPERT_USED 8 +static int32_t g_routed_moe_selected_override[DS4_METAL_MAX_ROUTED_EXPERT_USED]; +static uint32_t g_routed_moe_selected_override_n; +static int g_moe_selected_trace_record_initialized; +static FILE *g_moe_selected_trace_record_fp; +static uint64_t g_moe_selected_trace_record_count; +static int g_moe_selected_trace_replay_initialized; +static int32_t *g_moe_selected_trace_replay_ids; +static uint64_t g_moe_selected_trace_replay_count; +static uint64_t g_moe_selected_trace_replay_pos; + +static double ds4_gpu_gib(uint64_t bytes); + +static uint64_t ds4_gpu_system_memory_bytes(void) { + uint64_t bytes = 0; + size_t len = sizeof(bytes); + if (sysctlbyname("hw.memsize", &bytes, &len, NULL, 0) != 0) return 0; + return len == sizeof(bytes) ? bytes : 0; +} + +static void ds4_gpu_print_device_summary(void) { + const char *name = g_device.name ? [g_device.name UTF8String] : "unknown Metal device"; + uint64_t mem = ds4_gpu_system_memory_bytes(); + if (mem) { + double gib = (double)mem / 1024.0 / 1024.0 / 1024.0; + fprintf(stderr, "ds4: Metal device %s, %.2f GiB RAM\n", name, gib); + } else { + fprintf(stderr, "ds4: Metal device %s\n", name); + } +} + +#define DS4_METAL_MAX_MODEL_VIEWS 4096 +/* Compatibility fallback for callers that cannot provide a parsed GGUF tensor + * span. The normal DS4 engine passes the exact maximum tensor byte size. */ +#define DS4_METAL_FALLBACK_MAX_TENSOR_BYTES (4ull * 1024ull * 1024ull * 1024ull) + +typedef struct { + __strong id buffer; + const void *model_map; + uint64_t model_size; + uint64_t model_offset; + uint64_t bytes; +} ds4_gpu_model_view; + +static ds4_gpu_model_view g_model_views[DS4_METAL_MAX_MODEL_VIEWS]; +static uint32_t g_model_view_count; + +enum { + DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER = 80, + DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT = 384, + DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED = DS4_METAL_MAX_ROUTED_EXPERT_USED, + DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES = + DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER * + DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT, + DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS = 256, + DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS = 16, + DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS = 16, +}; + +typedef struct { + uint32_t layer; + uint32_t expert; + uint64_t hits; +} ds4_gpu_moe_selected_hotlist_entry; + +static int g_moe_selected_hotlist_initialized; +static uint64_t + g_moe_selected_hotlist_counts[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; +static uint64_t g_moe_selected_hotlist_records; +static uint64_t g_moe_selected_hotlist_selections; + +typedef struct { + __strong id gate_buffer; + __strong id up_buffer; + __strong id down_buffer; + const void *model_map; + uint64_t model_size; + uint64_t gate_abs_offset; + uint64_t up_abs_offset; + uint64_t down_abs_offset; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + uint64_t logical_bytes; + uint64_t last_used; + uint64_t use_count; + NSUInteger gate_inner; + NSUInteger up_inner; + NSUInteger down_inner; + uint64_t inflight_seq; + uint32_t slab_slot; + uint8_t valid; + uint8_t slab_backed; +} ds4_gpu_stream_expert_cache_entry; + +typedef struct { + __strong id gate_buffer; + __strong id up_buffer; + __strong id down_buffer; + NSUInteger gate_inner; + NSUInteger up_inner; + NSUInteger down_inner; +} ds4_gpu_stream_expert_reusable_buffers; + +static ds4_gpu_stream_expert_cache_entry + g_stream_expert_cache[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; +static ds4_gpu_stream_expert_cache_entry + g_stream_full_expert_addr_entry[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint32_t g_stream_expert_cache_layer_count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_hits[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_misses[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_evictions[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_pread_bytes[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static double g_stream_expert_cache_layer_pread_ms[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_last_hits[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_last_misses[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_last_evictions[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint64_t g_stream_expert_cache_layer_last_pread_bytes[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static double g_stream_expert_cache_layer_last_pread_ms[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static uint32_t + g_stream_expert_cache_route_hotness[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; +static id g_stream_expert_cache_gate_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_expert_cache_up_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_expert_cache_down_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_expert_cache_slabs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; +static uint32_t g_stream_expert_cache_slab_start_slot[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; +static uint32_t g_stream_expert_cache_slab_slot_count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; +static uint32_t g_stream_expert_cache_slab_slots_used[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS]; +static uint32_t g_stream_expert_cache_slab_count; +static uint32_t g_stream_expert_cache_slab_total_slots; +static uint32_t g_stream_expert_cache_free_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES]; +static uint32_t g_stream_expert_cache_free_slot_count; +static uint64_t g_stream_expert_cache_slab_slot_bytes; +static uint64_t g_stream_expert_cache_cb_seq; +static uint64_t g_stream_expert_cache_done_seq; +static uint64_t g_stream_expert_cache_batch_seq; +static uint64_t g_stream_expert_cache_owned_seq; +static uint64_t g_stream_expert_cache_pending_max_seq; +static id g_stream_compact_gate_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_compact_up_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_compact_down_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_compact_selected_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_selected_id_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER]; +static id g_stream_expert_validate_status_buffer; + +@interface DS4MetalTensor : NSObject +@property(nonatomic, strong) id buffer; +@property(nonatomic, assign) uint64_t offset; +@property(nonatomic, assign) uint64_t bytes; +@property(nonatomic, assign) uint8_t owner; +@end + +@implementation DS4MetalTensor +@end + +@interface DS4MetalQ4ExpertTable : NSObject +@property(nonatomic, strong) id argumentBuffer; +@property(nonatomic, strong) id addressBuffer; +@property(nonatomic, strong) NSMutableArray> *expertBuffers; +@property(nonatomic, strong) id residencySet; +@property(nonatomic, assign) BOOL residencySetAddedToQueue; +@property(nonatomic, assign) uint32_t nExpert; +@property(nonatomic, assign) uint64_t expertBytes; +@end + +@implementation DS4MetalQ4ExpertTable +- (void)dealloc { +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + if (_residencySet) { + if (_residencySetAddedToQueue && + g_queue && + [g_queue respondsToSelector:@selector(removeResidencySet:)]) { + [g_queue removeResidencySet:_residencySet]; + } + [_residencySet endResidency]; + } + } +#endif +} +@end + +@interface DS4MetalQ4LayerResidency : NSObject +@property(nonatomic, strong) id residencySet; +@property(nonatomic, assign) BOOL addedToQueue; +@end + +@implementation DS4MetalQ4LayerResidency +- (void)dealloc { +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + if (_residencySet) { + if (_addedToQueue && + g_queue && + [g_queue respondsToSelector:@selector(removeResidencySet:)]) { + [g_queue removeResidencySet:_residencySet]; + } + [_residencySet endResidency]; + } + } +#endif +} +@end + +static DS4MetalTensor *ds4_gpu_tensor_obj(ds4_gpu_tensor *tensor) { + return (__bridge DS4MetalTensor *)tensor; +} + +static const DS4MetalTensor *ds4_gpu_tensor_const_obj(const ds4_gpu_tensor *tensor) { + return (__bridge const DS4MetalTensor *)tensor; +} + +/* C code owns ds4_gpu_tensor handles as retained Objective-C objects. Freeing + * the same opaque handle twice would make the second __bridge_transfer release + * an already-deallocated object, which macOS reports as malloc corruption. The + * live table lets free validate a handle before touching Objective-C state; the + * same mutex also serializes the diagnostic allocation counters. */ +static uint64_t ds4_gpu_tensor_ptr_hash(uintptr_t ptr) { + uint64_t x = (uint64_t)(ptr >> 4); + x ^= x >> 33; + x *= UINT64_C(0xff51afd7ed558ccd); + x ^= x >> 33; + x *= UINT64_C(0xc4ceb9fe1a85ec53); + x ^= x >> 33; + return x; +} + +static int ds4_gpu_tensor_live_resize_locked(size_t min_cap) { + size_t new_cap = 1024; + while (new_cap < min_cap) new_cap <<= 1; + + uintptr_t *new_slots = calloc(new_cap, sizeof(new_slots[0])); + if (!new_slots) return 0; + + for (size_t i = 0; i < g_tensor_live_cap; i++) { + const uintptr_t key = g_tensor_live_slots[i]; + if (key == 0 || key == UINTPTR_MAX) continue; + + size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (new_cap - 1); + while (new_slots[idx] != 0) idx = (idx + 1) & (new_cap - 1); + new_slots[idx] = key; + } + + free(g_tensor_live_slots); + g_tensor_live_slots = new_slots; + g_tensor_live_cap = new_cap; + g_tensor_live_tombs = 0; + return 1; +} + +static int ds4_gpu_tensor_live_insert_locked(const void *ptr) { + if (!ptr || (uintptr_t)ptr == UINTPTR_MAX) return 0; + if ((g_tensor_live_count + g_tensor_live_tombs + 1) * 10 >= + g_tensor_live_cap * 7) + { + const size_t min_cap = g_tensor_live_cap ? g_tensor_live_cap * 2 : 1024; + if (!ds4_gpu_tensor_live_resize_locked(min_cap)) return 0; + } + + const uintptr_t key = (uintptr_t)ptr; + size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (g_tensor_live_cap - 1); + size_t tomb = (size_t)-1; + for (;;) { + const uintptr_t cur = g_tensor_live_slots[idx]; + if (cur == key) return 0; + if (cur == UINTPTR_MAX) { + if (tomb == (size_t)-1) tomb = idx; + } else if (cur == 0) { + if (tomb != (size_t)-1) { + idx = tomb; + g_tensor_live_tombs--; + } + g_tensor_live_slots[idx] = key; + g_tensor_live_count++; + return 1; + } + idx = (idx + 1) & (g_tensor_live_cap - 1); + } +} + +static int ds4_gpu_tensor_live_remove_locked(const void *ptr) { + if (!ptr || g_tensor_live_cap == 0) return 0; + + const uintptr_t key = (uintptr_t)ptr; + size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (g_tensor_live_cap - 1); + for (;;) { + const uintptr_t cur = g_tensor_live_slots[idx]; + if (cur == 0) return 0; + if (cur == key) { + g_tensor_live_slots[idx] = UINTPTR_MAX; + g_tensor_live_count--; + g_tensor_live_tombs++; + return 1; + } + idx = (idx + 1) & (g_tensor_live_cap - 1); + } +} + +static int ds4_gpu_tensor_track_alloc_locked( + const void *ptr, + uint64_t bytes, + uint64_t *live_snap, + uint64_t *peak_snap) +{ + if (!ds4_gpu_tensor_live_insert_locked(ptr)) return 0; + + g_tensor_alloc_live_bytes += bytes; + if (g_tensor_alloc_live_bytes > g_tensor_alloc_peak_bytes) { + g_tensor_alloc_peak_bytes = g_tensor_alloc_live_bytes; + } + if (live_snap) *live_snap = g_tensor_alloc_live_bytes; + if (peak_snap) *peak_snap = g_tensor_alloc_peak_bytes; + return 1; +} + +static int ds4_gpu_tensor_track_view_locked(const void *ptr) { + return ds4_gpu_tensor_live_insert_locked(ptr); +} + +static int ds4_gpu_tensor_prepare_free( + ds4_gpu_tensor *tensor, + uint8_t *owner, + uint64_t *bytes, + uint64_t *live_snap, + uint64_t *peak_snap) +{ + pthread_mutex_lock(&g_tensor_mu); + if (!ds4_gpu_tensor_live_remove_locked(tensor)) { + pthread_mutex_unlock(&g_tensor_mu); + fprintf(stderr, + "ds4: Metal tensor free ignored for unknown handle %p\n", + (void *)tensor); + return 0; + } + + DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor); + const uint8_t obj_owner = obj.owner; + const uint64_t obj_bytes = obj.bytes; + if (obj_owner) { + if (obj_bytes <= g_tensor_alloc_live_bytes) { + g_tensor_alloc_live_bytes -= obj_bytes; + } else { + g_tensor_alloc_live_bytes = 0; + } + } + if (owner) *owner = obj_owner; + if (bytes) *bytes = obj_bytes; + if (live_snap) *live_snap = g_tensor_alloc_live_bytes; + if (peak_snap) *peak_snap = g_tensor_alloc_peak_bytes; + pthread_mutex_unlock(&g_tensor_mu); + return 1; +} + +static void ds4_gpu_tensor_tracking_reset(void) { + pthread_mutex_lock(&g_tensor_mu); + if (g_tensor_live_count != 0) { + fprintf(stderr, + "ds4: Metal cleanup discarded %zu live tensor handles\n", + g_tensor_live_count); + } + free(g_tensor_live_slots); + g_tensor_live_slots = NULL; + g_tensor_live_cap = 0; + g_tensor_live_count = 0; + g_tensor_live_tombs = 0; + g_tensor_alloc_live_bytes = 0; + g_tensor_alloc_peak_bytes = 0; + pthread_mutex_unlock(&g_tensor_mu); +} + +static id ds4_gpu_tensor_buffer(const ds4_gpu_tensor *tensor) { + if (!tensor) return nil; + const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); + return obj.buffer; +} + +static NSUInteger ds4_gpu_tensor_offset(const ds4_gpu_tensor *tensor) { + if (!tensor) return 0; + const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); + return (NSUInteger)obj.offset; +} + +static id ds4_gpu_new_command_buffer(void); +static void ds4_gpu_stream_expert_cache_note_owned_created(void); + +static id ds4_gpu_command_buffer(int *owned) { + if (g_batch_cb) { + *owned = 0; + return g_batch_cb; + } + *owned = 1; + id cb = ds4_gpu_new_command_buffer(); + if (cb) ds4_gpu_stream_expert_cache_note_owned_created(); + return cb; +} + +static id ds4_gpu_compute_encoder(id cb) { + if (g_batch_cb && cb == g_batch_cb) { + g_batch_has_work = YES; + if (!g_batch_enc) g_batch_enc = [cb computeCommandEncoder]; + return g_batch_enc; + } + return [cb computeCommandEncoder]; +} + +static void ds4_gpu_end_compute_encoder(id cb, id enc) { + if (!enc) return; + if (g_batch_cb && cb == g_batch_cb && enc == g_batch_enc) return; + [enc endEncoding]; +} + +static void ds4_gpu_close_batch_encoder(void) { + if (!g_batch_enc) return; + [g_batch_enc endEncoding]; + g_batch_enc = nil; +} + +static double g_gpu_busy_accum; +static uint64_t g_gpu_busy_cbs; + +static int ds4_gpu_wait_command_buffer(id cb, const char *label) { + [cb waitUntilCompleted]; + if (getenv("DS4_METAL_GPU_BUSY_PROFILE")) { + const double busy = cb.GPUEndTime - cb.GPUStartTime; + if (busy > 0) g_gpu_busy_accum += busy; + if ((++g_gpu_busy_cbs % 64u) == 0u) { + fprintf(stderr, "ds4: gpu busy accum %.1f ms over %llu cbs\n", + g_gpu_busy_accum * 1000.0, + (unsigned long long)g_gpu_busy_cbs); + } + } + if (cb.status == MTLCommandBufferStatusError) { + fprintf(stderr, "ds4: Metal %s failed: %s\n", + label, [[cb.error localizedDescription] UTF8String]); + return 0; + } + return 1; +} + +static id ds4_gpu_new_command_buffer(void) { + static int initialized; + static int use_unretained; + if (!initialized) { + use_unretained = getenv("DS4_METAL_UNRETAINED_COMMAND_BUFFERS") != NULL; + initialized = 1; + } + if (use_unretained) { + return [g_queue commandBufferWithUnretainedReferences]; + } + return [g_queue commandBuffer]; +} + +static uint64_t ds4_gpu_exact_view_cache_limit_bytes(void) { + static int initialized; + static uint64_t limit_bytes; + if (initialized) return limit_bytes; + + const uint64_t mib = 1024ull * 1024ull; + const uint64_t gib = 1024ull * mib; + limit_bytes = 64ull * gib; + + const char *gib_env = getenv("DS4_METAL_EXACT_VIEW_CACHE_GIB"); + if (gib_env && gib_env[0]) { + char *end = NULL; + unsigned long long v = strtoull(gib_env, &end, 10); + if (end != gib_env && *end == '\0') { + limit_bytes = v > UINT64_MAX / gib ? UINT64_MAX : (uint64_t)v * gib; + } + } + + const char *mib_env = getenv("DS4_METAL_EXACT_VIEW_CACHE_MIB"); + if (mib_env && mib_env[0]) { + char *end = NULL; + unsigned long long v = strtoull(mib_env, &end, 10); + if (end != mib_env && *end == '\0') { + limit_bytes = v > UINT64_MAX / mib ? UINT64_MAX : (uint64_t)v * mib; + } + } + + initialized = 1; + return limit_bytes; +} + +static void ds4_gpu_model_buffer_cache_note_insert(uint64_t bytes) { + if (g_model_buffer_cache_bytes > UINT64_MAX - bytes) { + g_model_buffer_cache_bytes = UINT64_MAX; + } else { + g_model_buffer_cache_bytes += bytes; + } + + const uint64_t limit = ds4_gpu_exact_view_cache_limit_bytes(); + if (limit != 0 && g_model_buffer_cache_bytes > limit) { + g_model_buffer_cache_over_limit = 1; + } +} + +static void ds4_gpu_model_buffer_cache_clear(const char *reason) { + if (!g_model_buffer_cache) { + g_model_buffer_cache_bytes = 0; + g_model_buffer_cache_over_limit = 0; + return; + } + + const NSUInteger entries = [g_model_buffer_cache count]; + if (entries != 0) { + if (getenv("DS4_METAL_EXACT_VIEW_CACHE_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal exact model view cache evict reason=%s entries=%lu bytes=%.2f GiB limit=%.2f GiB\n", + reason ? reason : "unknown", + (unsigned long)entries, + ds4_gpu_gib(g_model_buffer_cache_bytes), + ds4_gpu_gib(ds4_gpu_exact_view_cache_limit_bytes())); + } + [g_model_buffer_cache removeAllObjects]; + g_model_buffer_cache_evictions++; + } + g_model_buffer_cache_bytes = 0; + g_model_buffer_cache_over_limit = 0; +} + +static void ds4_gpu_model_buffer_cache_maybe_evict(const char *reason) { + if (g_model_buffer_cache_over_limit) { + ds4_gpu_model_buffer_cache_clear(reason); + } +} + +static uint64_t ds4_gpu_stream_expert_cache_next_cb_seq(void) { + if (g_stream_expert_cache_cb_seq == UINT64_MAX) { + /* + * A real wrap would require an astronomical number of command buffers. + * Resetting the epoch space is still safer than letting zero become a + * valid in-flight marker. + */ + g_stream_expert_cache_cb_seq = 0; + g_stream_expert_cache_done_seq = 0; + g_stream_expert_cache_batch_seq = 0; + g_stream_expert_cache_owned_seq = 0; + g_stream_expert_cache_pending_max_seq = 0; + for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) { + for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) { + g_stream_expert_cache[layer][expert].inflight_seq = 0; + } + } + } + return ++g_stream_expert_cache_cb_seq; +} + +static void ds4_gpu_stream_expert_cache_note_batch_created(void) { + g_stream_expert_cache_batch_seq = + ds4_gpu_stream_expert_cache_next_cb_seq(); +} + +static void ds4_gpu_stream_expert_cache_note_batch_committed(void) { + if (g_stream_expert_cache_batch_seq > g_stream_expert_cache_pending_max_seq) { + g_stream_expert_cache_pending_max_seq = g_stream_expert_cache_batch_seq; + } + g_stream_expert_cache_batch_seq = 0; +} + +static void ds4_gpu_stream_expert_cache_note_owned_created(void) { + g_stream_expert_cache_owned_seq = + ds4_gpu_stream_expert_cache_next_cb_seq(); +} + +static void ds4_gpu_stream_expert_cache_note_pending_completed(void) { + if (g_stream_expert_cache_pending_max_seq > g_stream_expert_cache_done_seq) { + g_stream_expert_cache_done_seq = g_stream_expert_cache_pending_max_seq; + } + g_stream_expert_cache_pending_max_seq = 0; +} + +static void ds4_gpu_stream_expert_cache_note_owned_completed(void) { + if (g_stream_expert_cache_owned_seq > g_stream_expert_cache_done_seq) { + g_stream_expert_cache_done_seq = g_stream_expert_cache_owned_seq; + } + g_stream_expert_cache_owned_seq = 0; +} + +static int ds4_gpu_stream_expert_cache_entry_inflight( + const ds4_gpu_stream_expert_cache_entry *e) { + return e && e->valid && e->inflight_seq > g_stream_expert_cache_done_seq; +} + +static int ds4_gpu_stream_expert_cache_mark_inflight( + ds4_gpu_stream_expert_cache_entry *e) { + if (!e || !e->valid) return 0; + const uint64_t seq = g_stream_expert_cache_batch_seq ? + g_stream_expert_cache_batch_seq : + g_stream_expert_cache_owned_seq; + if (seq == 0) return 0; + e->inflight_seq = seq; + return 1; +} + +static int ds4_gpu_stream_expert_cache_mark_entries_inflight( + ds4_gpu_stream_expert_cache_entry * const *entries, + uint32_t n_entries, + uint32_t active_mask) { + if (!entries || n_entries == 0) return 0; + for (uint32_t i = 0; i < n_entries; i++) { + if (active_mask != 0 && (active_mask & (1u << i)) == 0) continue; + if (!ds4_gpu_stream_expert_cache_mark_inflight(entries[i])) return 0; + } + return 1; +} + +static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label); + +static int ds4_gpu_wait_pending_command_buffers(const char *label) { + int ok = 1; + for (id pending in g_pending_cbs) { + if (!ds4_gpu_wait_command_buffer(pending, label)) ok = 0; + } + [g_pending_cbs removeAllObjects]; + ds4_gpu_stream_expert_cache_note_pending_completed(); + if (!ok) ds4_gpu_invalidate_zero_prefix_prefill_block_maps(); + return ok; +} + +static int ds4_gpu_finish_command_buffer(id cb, int owned, const char *label) { + if (!owned) return 1; + + [cb commit]; + int ok = ds4_gpu_wait_pending_command_buffers(label); + if (!ds4_gpu_wait_command_buffer(cb, label)) { + ok = 0; + ds4_gpu_invalidate_zero_prefix_prefill_block_maps(); + } + ds4_gpu_stream_expert_cache_note_owned_completed(); + [g_transient_buffers removeAllObjects]; + ds4_gpu_model_buffer_cache_maybe_evict(label); + return ok; +} + +static int ds4_gpu_device_name_contains(const char *needle); + +static int ds4_gpu_use_m5_private_scratch(void) { + static int initialized; + static int enabled; + if (!initialized) { + enabled = ds4_gpu_device_name_contains("M5"); + initialized = 1; + } + return enabled; +} + +static int ds4_gpu_scratch_needs_cpu_access(const char *label) { + if (!label) return 0; + return strstr(label, "mask") != NULL || + strcmp(label, "ds4_attention_output_group_ids") == 0; +} + +static MTLResourceOptions ds4_gpu_model_resource_options(void) { + MTLResourceOptions options = MTLResourceStorageModeShared; + if (getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) { + options |= MTLResourceHazardTrackingModeUntracked; + } + return options; +} + +static int ds4_gpu_ensure_scratch_buffer( + id __strong *buffer, + NSUInteger *capacity, + NSUInteger bytes, + const char *label) { + if (*buffer && *capacity >= bytes) return 1; + if (bytes == 0) bytes = 1; + if (bytes > NSUIntegerMax) return 0; + + MTLResourceOptions options = MTLResourceStorageModeShared; + if (ds4_gpu_use_m5_private_scratch() && + !ds4_gpu_scratch_needs_cpu_access(label)) { + /* + * M5 scratch buffers that only flow between Metal kernels do not need + * CPU-visible shared storage. This reduces shared-memory traffic and + * residency pressure for the long prefill scratch pools without + * changing the public buffer lifetime model. Keep default hazard + * tracking because the graph reuses these buffers across dependent + * compute encoders. + */ + options = MTLResourceStorageModePrivate; + } + + *buffer = [g_device newBufferWithLength:bytes options:options]; + if (!*buffer && options != MTLResourceStorageModeShared) { + *buffer = [g_device newBufferWithLength:bytes options:MTLResourceStorageModeShared]; + } + if (!*buffer) { + fprintf(stderr, "ds4: failed to allocate Metal scratch buffer %s (%llu bytes)\n", + label, (unsigned long long)bytes); + *capacity = 0; + return 0; + } + (*buffer).label = [NSString stringWithUTF8String:label]; + *capacity = bytes; + return 1; +} + +static int ds4_gpu_ensure_zero_attention_mask(NSUInteger bytes) { + const NSUInteger capacity = 8192u * sizeof(uint16_t); + if (bytes > capacity) return 0; + if (g_flash_attn_zero_mask_buffer && + g_flash_attn_zero_mask_bytes >= capacity) { + return 1; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_zero_mask_buffer, + &g_flash_attn_zero_mask_bytes, + capacity, + "ds4_flash_attn_zero_mask")) { + return 0; + } + void *contents = [g_flash_attn_zero_mask_buffer contents]; + if (!contents) return 0; + memset(contents, 0, g_flash_attn_zero_mask_bytes); + return 1; +} + +static uint64_t round_up_u64(uint64_t v, uint64_t align) { + return (v + align - 1) & ~(align - 1); +} + +static uint64_t ds4_gpu_effective_model_max_tensor_bytes(uint64_t map_size, uint64_t max_tensor_bytes) { + if (max_tensor_bytes != 0) return max_tensor_bytes; + return map_size < DS4_METAL_FALLBACK_MAX_TENSOR_BYTES ? + map_size : DS4_METAL_FALLBACK_MAX_TENSOR_BYTES; +} + +static id ds4_gpu_get_pipeline(const char *function_name); +static int ds4_gpu_warm_model_views(void); +static double ds4_gpu_gib(uint64_t bytes); + +static double ds4_gpu_now_ms(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return ts.tv_sec * 1000.0 + ts.tv_nsec / 1000000.0; +} + +static int ds4_gpu_moe_selected_hotlist_cmp(const void *a, const void *b) { + const ds4_gpu_moe_selected_hotlist_entry *ea = a; + const ds4_gpu_moe_selected_hotlist_entry *eb = b; + if (ea->hits < eb->hits) return 1; + if (ea->hits > eb->hits) return -1; + if (ea->layer != eb->layer) return ea->layer < eb->layer ? -1 : 1; + if (ea->expert != eb->expert) return ea->expert < eb->expert ? -1 : 1; + return 0; +} + +static int ds4_gpu_moe_selected_hotlist_merge_requested(void) { + return getenv("DS4_MOE_RECORD_SELECTED_HOTLIST_MERGE") != NULL && + getenv("DS4_MOE_RECORD_SELECTED_HOTLIST_FRESH") == NULL; +} + +static int ds4_gpu_moe_selected_hotlist_load_existing(const char *path) { + if (!path || !path[0] || !ds4_gpu_moe_selected_hotlist_merge_requested()) { + return 1; + } + + FILE *fp = fopen(path, "rb"); + if (!fp) { + if (errno == ENOENT) return 1; + fprintf(stderr, "ds4: failed to open selected hotlist merge file %s\n", path); + return 0; + } + + char line[256]; + uint64_t lineno = 0; + uint64_t loaded_entries = 0; + uint64_t loaded_hits = 0; + uint64_t header_records = UINT64_MAX; + uint64_t header_selections = UINT64_MAX; + while (fgets(line, sizeof(line), fp)) { + lineno++; + char *p = line; + while (*p && isspace((unsigned char)*p)) p++; + if (*p == '\0') continue; + if (*p == '#') { + unsigned long long value = 0; + if (sscanf(p, "# layer_records %llu", &value) == 1) { + header_records = (uint64_t)value; + } else if (sscanf(p, "# selections %llu", &value) == 1) { + header_selections = (uint64_t)value; + } + continue; + } + + errno = 0; + char *end = NULL; + unsigned long layer = strtoul(p, &end, 10); + if (end == p || errno != 0) goto bad_line; + p = end; + while (*p && isspace((unsigned char)*p)) p++; + + errno = 0; + unsigned long expert = strtoul(p, &end, 10); + if (end == p || errno != 0) goto bad_line; + p = end; + while (*p && isspace((unsigned char)*p)) p++; + + errno = 0; + unsigned long long hits = strtoull(p, &end, 10); + if (end == p || errno != 0) goto bad_line; + if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && + hits != 0) { + uint64_t *dst = &g_moe_selected_hotlist_counts[layer][expert]; + if (*dst > UINT64_MAX - (uint64_t)hits) { + *dst = UINT64_MAX; + } else { + *dst += (uint64_t)hits; + } + loaded_entries++; + if (loaded_hits > UINT64_MAX - (uint64_t)hits) { + loaded_hits = UINT64_MAX; + } else { + loaded_hits += (uint64_t)hits; + } + } + continue; + +bad_line: + fprintf(stderr, + "ds4: invalid selected hotlist merge line %" PRIu64 " in %s\n", + lineno, + path); + fclose(fp); + return 0; + } + if (ferror(fp)) { + fprintf(stderr, "ds4: failed to read selected hotlist merge file %s\n", path); + fclose(fp); + return 0; + } + fclose(fp); + + g_moe_selected_hotlist_records = + header_records != UINT64_MAX ? header_records : loaded_hits / 6u; + g_moe_selected_hotlist_selections = + header_selections != UINT64_MAX ? header_selections : loaded_hits; + fprintf(stderr, + "ds4: merged selected-id hotlist %s " + "(%" PRIu64 " entries, %" PRIu64 " hits)\n", + path, + loaded_entries, + loaded_hits); + return 1; +} + +static void ds4_gpu_moe_selected_hotlist_close(void) { + const char *path = getenv("DS4_MOE_RECORD_SELECTED_HOTLIST"); + if (!g_moe_selected_hotlist_initialized || !path || !path[0]) return; + + const size_t cap = + (size_t)DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER * + DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + ds4_gpu_moe_selected_hotlist_entry *entries = + malloc(cap * sizeof(entries[0])); + if (!entries) { + fprintf(stderr, "ds4: failed to allocate selected hotlist entries\n"); + return; + } + + size_t n = 0; + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + const uint64_t hits = + g_moe_selected_hotlist_counts[layer][expert]; + if (hits == 0) continue; + entries[n++] = (ds4_gpu_moe_selected_hotlist_entry) { + .layer = layer, + .expert = expert, + .hits = hits, + }; + } + } + qsort(entries, n, sizeof(entries[0]), ds4_gpu_moe_selected_hotlist_cmp); + + FILE *fp = fopen(path, "wb"); + if (!fp) { + fprintf(stderr, "ds4: failed to open selected hotlist file %s\n", path); + free(entries); + return; + } + fprintf(fp, + "# ds4 selected-id hotlist v1\n" + "# layer_records %" PRIu64 "\n" + "# selections %" PRIu64 "\n" + "# columns: layer expert hits weight\n", + g_moe_selected_hotlist_records, + g_moe_selected_hotlist_selections); + for (size_t i = 0; i < n; i++) { + fprintf(fp, + "%u %u %" PRIu64 " 0\n", + entries[i].layer, + entries[i].expert, + entries[i].hits); + } + free(entries); + + if (fclose(fp) != 0) { + fprintf(stderr, "ds4: failed to close selected hotlist file %s\n", path); + } else { + fprintf(stderr, + "ds4: wrote selected-id hotlist to %s " + "(%" PRIu64 " layer records, %" PRIu64 " selections)\n", + path, + g_moe_selected_hotlist_records, + g_moe_selected_hotlist_selections); + } +} + +static int ds4_gpu_moe_selected_hotlist_record( + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_selected, + uint32_t n_total_expert) { + const char *path = getenv("DS4_MOE_RECORD_SELECTED_HOTLIST"); + if (!path || !path[0]) return 1; + if (!selected_ids || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { + return 0; + } + if (!g_moe_selected_hotlist_initialized) { + g_moe_selected_hotlist_initialized = 1; + if (!ds4_gpu_moe_selected_hotlist_load_existing(path)) return 0; + atexit(ds4_gpu_moe_selected_hotlist_close); + } + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return 1; + + g_moe_selected_hotlist_records++; + for (uint32_t i = 0; i < n_selected; i++) { + if (selected_ids[i] < 0) continue; + const uint32_t expert = (uint32_t)selected_ids[i]; + if (expert >= n_total_expert || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + continue; + } + g_moe_selected_hotlist_counts[layer][expert]++; + g_moe_selected_hotlist_selections++; + } + return 1; +} + +static void ds4_gpu_moe_selected_trace_record_close(void) { + if (g_moe_selected_trace_record_fp) { + const char *path = getenv("DS4_MOE_RECORD_SELECTED_IDS"); + fclose(g_moe_selected_trace_record_fp); + g_moe_selected_trace_record_fp = NULL; + fprintf(stderr, + "ds4: recorded %" PRIu64 " routed-MoE selected-id entries to %s\n", + g_moe_selected_trace_record_count, + path && path[0] ? path : "(unknown)"); + } +} + +static int ds4_gpu_moe_selected_trace_record( + const int32_t selected_ids[6], + uint32_t n_selected) { + const char *path = getenv("DS4_MOE_RECORD_SELECTED_IDS"); + if (!path || !path[0]) return 1; + if (n_selected != 6) { + fprintf(stderr, "ds4: selected-id recording expects exactly 6 selected experts\n"); + return 0; + } + + if (!g_moe_selected_trace_record_initialized) { + g_moe_selected_trace_record_initialized = 1; + g_moe_selected_trace_record_fp = fopen(path, "wb"); + if (!g_moe_selected_trace_record_fp) { + fprintf(stderr, "ds4: failed to open selected-id record file %s\n", path); + return 0; + } + setvbuf(g_moe_selected_trace_record_fp, NULL, _IOFBF, 1u << 20); + atexit(ds4_gpu_moe_selected_trace_record_close); + } + + if (fwrite(selected_ids, sizeof(selected_ids[0]), n_selected, g_moe_selected_trace_record_fp) != n_selected) { + fprintf(stderr, "ds4: failed to write selected-id record file %s\n", path); + return 0; + } + if (fflush(g_moe_selected_trace_record_fp) != 0) { + fprintf(stderr, "ds4: failed to flush selected-id record file %s\n", path); + return 0; + } + g_moe_selected_trace_record_count++; + return 1; +} + +static int ds4_gpu_moe_selected_trace_replay( + int32_t selected_ids[6], + uint32_t n_selected) { + const char *path = getenv("DS4_MOE_REPLAY_SELECTED_IDS"); + if (!path || !path[0]) return 0; + if (n_selected != 6) { + fprintf(stderr, "ds4: selected-id replay expects exactly 6 selected experts\n"); + return -1; + } + + if (!g_moe_selected_trace_replay_initialized) { + g_moe_selected_trace_replay_initialized = 1; + FILE *fp = fopen(path, "rb"); + if (!fp) { + fprintf(stderr, "ds4: failed to open selected-id replay file %s\n", path); + return -1; + } + if (fseeko(fp, 0, SEEK_END) != 0) { + fprintf(stderr, "ds4: failed to seek selected-id replay file %s\n", path); + fclose(fp); + return -1; + } + const off_t end = ftello(fp); + if (end < 0) { + fprintf(stderr, "ds4: failed to size selected-id replay file %s\n", path); + fclose(fp); + return -1; + } + if (fseeko(fp, 0, SEEK_SET) != 0) { + fprintf(stderr, "ds4: failed to rewind selected-id replay file %s\n", path); + fclose(fp); + return -1; + } + + const uint64_t bytes = (uint64_t)end; + const uint64_t entry_bytes = (uint64_t)n_selected * sizeof(selected_ids[0]); + if (bytes == 0 || (bytes % entry_bytes) != 0) { + fprintf(stderr, + "ds4: selected-id replay file %s has invalid size %" PRIu64 "\n", + path, + bytes); + fclose(fp); + return -1; + } + if (bytes > SIZE_MAX) { + fprintf(stderr, "ds4: selected-id replay file %s is too large\n", path); + fclose(fp); + return -1; + } + g_moe_selected_trace_replay_count = bytes / entry_bytes; + g_moe_selected_trace_replay_ids = malloc((size_t)bytes); + if (!g_moe_selected_trace_replay_ids) { + fprintf(stderr, "ds4: failed to allocate selected-id replay buffer\n"); + fclose(fp); + return -1; + } + if (fread(g_moe_selected_trace_replay_ids, 1, (size_t)bytes, fp) != (size_t)bytes) { + fprintf(stderr, "ds4: failed to read selected-id replay file %s\n", path); + fclose(fp); + free(g_moe_selected_trace_replay_ids); + g_moe_selected_trace_replay_ids = NULL; + return -1; + } + fclose(fp); + fprintf(stderr, + "ds4: loaded %" PRIu64 " routed-MoE selected-id entries from %s\n", + g_moe_selected_trace_replay_count, + path); + } + + if (g_moe_selected_trace_replay_pos >= g_moe_selected_trace_replay_count) { + fprintf(stderr, + "ds4: selected-id replay exhausted after %" PRIu64 " entries\n", + g_moe_selected_trace_replay_pos); + return -1; + } + memcpy(selected_ids, + g_moe_selected_trace_replay_ids + g_moe_selected_trace_replay_pos * n_selected, + (size_t)n_selected * sizeof(selected_ids[0])); + g_moe_selected_trace_replay_pos++; + return 1; +} + +static int ds4_gpu_progress_enabled(void) { + return ds4_log_is_tty(stderr); +} + +static void ds4_gpu_progress_begin(const char *what) { + if (!ds4_gpu_progress_enabled()) return; + fprintf(stderr, "ds4: %s...", what); + fflush(stderr); +} + +static void ds4_gpu_progress_done(void) { + if (!ds4_gpu_progress_enabled()) return; + fputs(" done\n", stderr); + fflush(stderr); +} + +static void ds4_gpu_progress_failed(void) { + if (!ds4_gpu_progress_enabled()) return; + fputs(" failed\n", stderr); + fflush(stderr); +} + +static void ds4_gpu_model_views_clear(void) { + for (uint32_t i = 0; i < g_model_view_count; i++) { + g_model_views[i].buffer = nil; + g_model_views[i].model_map = NULL; + g_model_views[i].model_size = 0; + g_model_views[i].model_offset = 0; + g_model_views[i].bytes = 0; + } + g_model_view_count = 0; +} + +static void ds4_gpu_model_residency_clear(void) { +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + if (g_model_residency_set) { + if (g_model_residency_added_to_queue && + g_queue && + [g_queue respondsToSelector:@selector(removeResidencySet:)]) { + [g_queue removeResidencySet:g_model_residency_set]; + } + [g_model_residency_set endResidency]; + [g_model_residency_set removeAllAllocations]; + g_model_residency_set = nil; + } + } +#endif + g_model_residency_count = 0; + g_model_residency_added_to_queue = 0; +} + +/* TP sharding keeps only this rank's expert ranges warm, + * so whole-view residency requests (which would page in the full file) + * must be skipped; pages fault in lazily through the same view buffers, + * exactly like ssd-streaming mode. */ +static int g_model_residency_skipped; + +void ds4_gpu_model_residency_skip(int skip) { + g_model_residency_skipped = skip; +} + +static int ds4_gpu_model_residency_request_views(void) { + if (g_model_view_count == 0 || + g_ssd_streaming_mode || + g_model_residency_skipped || + getenv("DS4_METAL_NO_RESIDENCY") != NULL) { + return 1; + } + +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + /* + * Register all model views as one residency set before inference. This + * is a GPU residency/budgeting hint, not a request to fault the whole + * 80+ GB file into memory. Its purpose is to make the driver see the + * complete set of large shared allocations during setup instead of + * discovering them lazily from the first measured graph command, where + * VM validation and residency accounting would look like model compute. + */ + MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init]; + desc.label = @"ds4_model"; + desc.initialCapacity = g_model_view_count; + + NSError *error = nil; + g_model_residency_set = [g_device newResidencySetWithDescriptor:desc error:&error]; + if (!g_model_residency_set) { + fprintf(stderr, "ds4: Metal model residency set creation failed: %s\n", + [[error localizedDescription] UTF8String]); + return 0; + } + + for (uint32_t i = 0; i < g_model_view_count; i++) { + [g_model_residency_set addAllocation:g_model_views[i].buffer]; + } + [g_model_residency_set commit]; + [g_model_residency_set requestResidency]; + if (getenv("DS4_METAL_DISABLE_QUEUE_RESIDENCY_SET") == NULL && + g_queue && + [g_queue respondsToSelector:@selector(addResidencySet:)]) { + [g_queue addResidencySet:g_model_residency_set]; + g_model_residency_added_to_queue = 1; + } + g_model_residency_count = g_model_view_count; + } +#endif + + return 1; +} + +static int ds4_gpu_add_model_view_range( + const void *model_map, + uint64_t model_size, + uint64_t map_offset, + uint64_t map_size, + uint64_t max_tensor_bytes, + bool use_default_view_cap, + uint64_t *mapped_model_size_out) { + const uint64_t page = (uint64_t)getpagesize(); + const uintptr_t model_addr = (uintptr_t)model_map; + + if ((model_addr & (uintptr_t)(page - 1)) != 0) { + fprintf(stderr, "ds4: Metal model mmap base is not page aligned\n"); + return 0; + } + if (map_offset > model_size || map_size > model_size - map_offset) { + fprintf(stderr, "ds4: Metal model mapped range is outside the GGUF mapping\n"); + return 0; + } + const uint64_t page_model_offset = map_offset & ~(page - 1); + const uint64_t leading = map_offset - page_model_offset; + if (map_size > UINT64_MAX - leading || + leading + map_size > UINT64_MAX - (page - 1)) + { + fprintf(stderr, "ds4: Metal model mapped range overflows page alignment\n"); + return 0; + } + const uint64_t mapped_model_size = round_up_u64(leading + map_size, page); + uint64_t max_buffer = (uint64_t)[g_device maxBufferLength]; + max_buffer &= ~(page - 1); + + /* + * Wrap only the tensor-data part of the GGUF file. Metadata is parsed by the + * CPU and is never dereferenced by kernels, so exposing it to Metal only + * grows the residency set and the VM range the driver must validate. + * + * Metal buffers have a device-specific maximum length, and this model is + * larger than that maximum on the target machines. Creating one no-copy + * buffer per tensor would avoid the length limit, but it would also move a + * lot of VM-object creation and residency bookkeeping into graph setup. The + * stable shape here is a tiny number of page-aligned views created once. + * + * Adjacent views intentionally overlap by more than the largest tensor, plus + * one page for alignment. That invariant guarantees every tensor lies wholly + * inside at least one view, so hot paths pass one buffer and one inner byte + * offset. We never split a weight tensor across command encoders. + */ + if (max_tensor_bytes > map_size) { + fprintf(stderr, "ds4: Metal model max tensor span is larger than a mapped tensor span\n"); + return 0; + } + if (max_tensor_bytes > UINT64_MAX - (page - 1)) { + fprintf(stderr, "ds4: Metal model max tensor span overflows page alignment\n"); + return 0; + } + const uint64_t max_tensor_rounded = round_up_u64(max_tensor_bytes, page); + if (max_tensor_rounded > UINT64_MAX - page) { + fprintf(stderr, "ds4: Metal model view overlap overflows page slack\n"); + return 0; + } + const uint64_t overlap = max_tensor_rounded + page; + if (max_buffer == 0 || max_buffer <= overlap) { + fprintf(stderr, + "ds4: Metal maxBufferLength is too small for DS4 model views " + "(max tensor %.2f GiB, max buffer %.2f GiB)\n", + ds4_gpu_gib(max_tensor_bytes), + ds4_gpu_gib(max_buffer)); + return 0; + } + + uint64_t view_limit = max_buffer; + const char *view_limit_env = getenv("DS4_METAL_MODEL_VIEW_MAX_GIB"); + if (view_limit_env && view_limit_env[0]) { + char *end = NULL; + unsigned long long gib = strtoull(view_limit_env, &end, 10); + if (end != view_limit_env && gib > 0) { + uint64_t env_limit = gib * 1024ull * 1024ull * 1024ull; + env_limit &= ~(page - 1); + if (env_limit > 0) view_limit = env_limit; + } + } else if (use_default_view_cap && mapped_model_size > max_buffer) { + /* + * Very large no-copy buffers can make Metal's VM validation dominate + * startup or the first graph command on multi-hundred-GiB slices. Keep + * ordinary contiguous model mappings unchanged, but let distributed + * span maps use smaller overlapping views when a range already has to + * be split. + */ + const uint64_t default_limit = 128ull * 1024ull * 1024ull * 1024ull; + if (view_limit > default_limit) view_limit = default_limit; + } + if (view_limit > max_buffer) view_limit = max_buffer; + view_limit &= ~(page - 1); + if (view_limit == 0 || view_limit <= overlap) { + fprintf(stderr, + "ds4: Metal model view cap is too small for DS4 model views " + "(cap %.2f GiB, max tensor %.2f GiB)\n", + ds4_gpu_gib(view_limit), + ds4_gpu_gib(max_tensor_bytes)); + return 0; + } + + const uint64_t step = view_limit - overlap; + uint64_t off = 0; + while (off < mapped_model_size) { + if (g_model_view_count == DS4_METAL_MAX_MODEL_VIEWS) { + fprintf(stderr, "ds4: Metal model needs more mapped views than expected\n"); + return 0; + } + + uint64_t view_bytes = mapped_model_size - off; + if (view_bytes > view_limit) view_bytes = view_limit; + + id buffer = [g_device newBufferWithBytesNoCopy:(void *)(model_addr + page_model_offset + off) + length:(NSUInteger)view_bytes + options:ds4_gpu_model_resource_options() + deallocator:nil]; + if (!buffer) { + fprintf(stderr, + "ds4: Metal could not wrap mmaped model view at %.2f GiB, size %.2f GiB\n", + (double)(page_model_offset + off) / (1024.0 * 1024.0 * 1024.0), + (double)view_bytes / (1024.0 * 1024.0 * 1024.0)); + return 0; + } + buffer.label = [NSString stringWithFormat:@"ds4_model_view_%u", g_model_view_count]; + + g_model_views[g_model_view_count].buffer = buffer; + g_model_views[g_model_view_count].model_map = model_map; + g_model_views[g_model_view_count].model_size = model_size; + g_model_views[g_model_view_count].model_offset = page_model_offset + off; + g_model_views[g_model_view_count].bytes = view_bytes; + g_model_view_count++; + + g_model_wrap_count++; + g_model_wrap_bytes += view_bytes; + if (view_bytes > g_model_wrap_max_bytes) g_model_wrap_max_bytes = view_bytes; + + if (off + view_bytes >= mapped_model_size) break; + off += step; + } + + if (mapped_model_size_out) *mapped_model_size_out += mapped_model_size; + return 1; +} + +static int ds4_gpu_finish_model_views( + double t0, + uint64_t mapped_model_size, + uint64_t display_offset) { + const double t_mapped = ds4_gpu_now_ms(); + const int request_residency = + !g_ssd_streaming_mode && + getenv("DS4_METAL_NO_RESIDENCY") == NULL; + if (request_residency) ds4_gpu_progress_begin("requesting Metal residency (may take tens of seconds)"); + if (!ds4_gpu_model_residency_request_views()) { + if (request_residency) ds4_gpu_progress_failed(); + return 0; + } + if (request_residency) ds4_gpu_progress_done(); + const double t_resident = ds4_gpu_now_ms(); + int warmed = 1; + const double t_warm0 = ds4_gpu_now_ms(); + const int warm_model_views = !g_ssd_streaming_mode && + getenv("DS4_METAL_NO_RESIDENCY") == NULL && + getenv("DS4_METAL_NO_MODEL_WARMUP") == NULL; + if (warm_model_views) { + /* + * The first GPU command touching no-copy mmap storage can pay command + * queue setup, page-table validation, and shared-allocation residency + * costs. Sample each model view here so timed graph execution starts + * after that one-time work. The stride is intentionally coarse: this is + * a validation touch over the VM ranges, not a full model prefetch. A + * dense prefetch would create exactly the kind of memory pressure and + * startup stalls this path is designed to avoid. + */ + if (g_model_residency_skipped) { + /* TP sharding: a single command buffer binding every + * view demands residency of them all and OOMs; the engine's + * CPU-side sharded warm pre-faults the owned bytes instead. */ + warmed = 1; + } else { + ds4_gpu_progress_begin("warming Metal model views"); + warmed = ds4_gpu_warm_model_views(); + if (warmed) ds4_gpu_progress_done(); + else ds4_gpu_progress_failed(); + } + } + const double t_warm = ds4_gpu_now_ms(); + if (ds4_gpu_model_map_log_enabled()) { + fprintf(stderr, + "ds4: Metal model views created in %.3f ms, residency requested in %.3f ms, warmup %.3f ms (mapped %.2f MiB from offset %.2f MiB)\n", + t_mapped - t0, + t_resident - t_mapped, + t_warm - t_warm0, + mapped_model_size / 1024.0 / 1024.0, + display_offset / 1024.0 / 1024.0); + } + if (!warmed) return 0; + return 1; +} + +static int ds4_gpu_map_model_views( + const void *model_map, + uint64_t model_size, + uint64_t map_offset, + uint64_t map_size, + uint64_t max_tensor_bytes) { + const double t0 = ds4_gpu_now_ms(); + uint64_t mapped_model_size = 0; + if (!ds4_gpu_add_model_view_range(model_map, + model_size, + map_offset, + map_size, + max_tensor_bytes, + false, + &mapped_model_size)) { + return 0; + } + return ds4_gpu_finish_model_views(t0, mapped_model_size, map_offset); +} + +static id ds4_gpu_new_transient_buffer(NSUInteger bytes, const char *label) { + if (bytes == 0) bytes = 1; + + id buffer = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!buffer) { + fprintf(stderr, "ds4: failed to allocate Metal transient buffer %s (%llu bytes)\n", + label ? label : "(unnamed)", (unsigned long long)bytes); + return nil; + } + if (label) buffer.label = [NSString stringWithUTF8String:label]; + + /* + * CPU-filled buffers must survive until their command buffer completes. + * A local ObjC strong variable is not enough when the encoder function + * returns before the caller commits the command buffer. + */ + [g_transient_buffers addObject:buffer]; + return buffer; +} + +static int ds4_gpu_zero_prefix_prefill_mask_cache_enabled(void) { + if (getenv("DS4_METAL_DISABLE_M3_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL || + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL) { + return 0; + } + return ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL; +} + +static ds4_gpu_zero_prefix_prefill_mask_cache_entry * +ds4_gpu_get_zero_prefix_prefill_mask_cache( + uint32_t kind, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t n_keys, + uint32_t window, + uint32_t ratio, + uint32_t nqptg, + uint32_t ncpsg, + bool has_kvpad, + bool bc_mask, + NSUInteger mask_bytes, + NSUInteger blk_bytes, + bool *created) { + if (created) *created = false; + if (!created || !ds4_gpu_zero_prefix_prefill_mask_cache_enabled() || + mask_bytes == 0 || blk_bytes == 0) { + return NULL; + } + + uint32_t slot = UINT32_MAX; + switch (kind) { + case DS4_GPU_PREFILL_MASK_CACHE_RAW: slot = 0; break; + case DS4_GPU_PREFILL_MASK_CACHE_RATIO4: slot = 1; break; + case DS4_GPU_PREFILL_MASK_CACHE_RATIO128: slot = 2; break; + default: return NULL; + } + + ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = + &g_zero_prefix_prefill_mask_cache[slot]; + if (entry->valid && entry->mask && entry->blk && + entry->mask_bytes == mask_bytes && entry->blk_bytes == blk_bytes && + entry->kind == kind && entry->n_tokens == n_tokens && + entry->n_comp == n_comp && entry->n_keys == n_keys && + entry->window == window && entry->ratio == ratio && + entry->nqptg == nqptg && entry->ncpsg == ncpsg && + entry->has_kvpad == has_kvpad && entry->bc_mask == bc_mask) { + return entry; + } + + id mask = [g_device newBufferWithLength:mask_bytes + options:MTLResourceStorageModeShared]; + id blk = [g_device newBufferWithLength:blk_bytes + options:MTLResourceStorageModePrivate]; + if (!blk) { + blk = [g_device newBufferWithLength:blk_bytes + options:MTLResourceStorageModeShared]; + } + if (!mask || !blk) return NULL; + + mask.label = [NSString stringWithFormat:@"ds4_prefill_mask_cache_%u", kind]; + blk.label = [NSString stringWithFormat:@"ds4_prefill_blk_cache_%u", kind]; + + /* A replaced entry may still be referenced by an uncommitted or in-flight + * command buffer. Keep its resources alive with the other batch-scoped + * buffers until command completion instead of mutating or releasing them. */ + if (entry->mask) [g_transient_buffers addObject:entry->mask]; + if (entry->blk) [g_transient_buffers addObject:entry->blk]; + + entry->mask = mask; + entry->blk = blk; + entry->mask_bytes = mask_bytes; + entry->blk_bytes = blk_bytes; + entry->kind = kind; + entry->n_tokens = n_tokens; + entry->n_comp = n_comp; + entry->n_keys = n_keys; + entry->window = window; + entry->ratio = ratio; + entry->nqptg = nqptg; + entry->ncpsg = ncpsg; + entry->has_kvpad = has_kvpad; + entry->bc_mask = bc_mask; + /* The caller publishes the entry only after synchronously filling mask. */ + entry->valid = false; + entry->blk_ready = false; + *created = true; + return entry; +} + +static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void) { + for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { + g_zero_prefix_prefill_mask_cache[i].blk_ready = false; + } +} + +static void ds4_gpu_clear_zero_prefix_prefill_mask_cache(void) { + for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { + ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = + &g_zero_prefix_prefill_mask_cache[i]; + entry->mask = nil; + entry->blk = nil; + entry->mask_bytes = 0; + entry->blk_bytes = 0; + entry->kind = 0; + entry->n_tokens = 0; + entry->n_comp = 0; + entry->n_keys = 0; + entry->window = 0; + entry->ratio = 0; + entry->nqptg = 0; + entry->ncpsg = 0; + entry->has_kvpad = false; + entry->bc_mask = false; + entry->valid = false; + entry->blk_ready = false; + } +} + +void ds4_gpu_release_zero_prefix_prefill_mask_cache(void) { + /* Layer-major prefill waits each layer before advancing, so its final + * release point has no outstanding cache users. Keep the guard here for + * diagnostic callers that may have an open or asynchronously flushed CB. */ + if (!g_initialized || g_batch_cb || + (g_pending_cbs && [g_pending_cbs count] != 0)) { + return; + } + ds4_gpu_clear_zero_prefix_prefill_mask_cache(); +} + +static id ds4_gpu_get_mul_mm_pipeline( + const char *function_name, + bool bc_inp, + bool bc_out) { + NSString *key = [NSString stringWithFormat:@"%s_bci=%d_bco=%d", + function_name, bc_inp ? 1 : 0, bc_out ? 1 : 0]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&bc_inp type:MTLDataTypeBool atIndex:700]; + [constants setConstantValue:&bc_out type:MTLDataTypeBool atIndex:701]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_mul_mm_id_pipeline( + const char *function_name, + bool bc_inp) { + NSString *key = [NSString stringWithFormat:@"%s_bci=%d", + function_name, bc_inp ? 1 : 0]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&bc_inp type:MTLDataTypeBool atIndex:700]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_pipeline( + const char *function_name) { + NSString *key = [NSString stringWithFormat:@"%s", function_name]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found\n", function_name); + return nil; + } + + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static int ds4_gpu_disable_hot_pipeline_statics(void) { + static int initialized; + static int disabled; + if (!initialized) { + disabled = getenv("DS4_METAL_DISABLE_HOT_PIPELINE_STATICS") != NULL; + initialized = 1; + } + return disabled; +} + +static id ds4_gpu_hot_pipeline( + id pipeline, + const char *fallback_name) { + if (!ds4_gpu_disable_hot_pipeline_statics()) return pipeline; + return ds4_gpu_get_pipeline(fallback_name); +} + +static int ds4_gpu_use_compressor_pair_nr4(void) { + static int initialized; + static int enabled; + if (!initialized) { + enabled = getenv("DS4_METAL_COMPRESSOR_PAIR_NR4") != NULL; + initialized = 1; + } + return enabled; +} + +static int ds4_gpu_device_name_contains(const char *needle); + +static int ds4_gpu_env_value_eq(const char *v, size_t n, const char *literal) { + size_t m = strlen(literal); + if (n != m) return 0; + for (size_t i = 0; i < n; i++) { + if (tolower((unsigned char)v[i]) != tolower((unsigned char)literal[i])) return 0; + } + return 1; +} + +static int ds4_gpu_env_bool(const char *name) { + const char *v = getenv(name); + if (!v) return -1; + + while (isspace((unsigned char)*v)) v++; + size_t n = strlen(v); + while (n > 0 && isspace((unsigned char)v[n - 1])) n--; + if (n == 0) return 1; + + if (ds4_gpu_env_value_eq(v, n, "1") || + ds4_gpu_env_value_eq(v, n, "true") || + ds4_gpu_env_value_eq(v, n, "yes") || + ds4_gpu_env_value_eq(v, n, "on")) { + return 1; + } + if (ds4_gpu_env_value_eq(v, n, "0") || + ds4_gpu_env_value_eq(v, n, "false") || + ds4_gpu_env_value_eq(v, n, "no") || + ds4_gpu_env_value_eq(v, n, "off")) { + return 0; + } + + if (!g_mpp_invalid_env_reported) { + fprintf(stderr, + "ds4: invalid Metal boolean environment value %s=%.*s; treating presence as enabled\n", + name, (int)n, v); + g_mpp_invalid_env_reported = 1; + } + return 1; +} + +static uint64_t ds4_gpu_env_u64(const char *name, + uint64_t fallback, + uint64_t min_value, + uint64_t max_value) { + const char *v = getenv(name); + if (!v) return fallback; + while (isspace((unsigned char)*v)) v++; + if (!*v) return fallback; + + errno = 0; + char *end = NULL; + unsigned long long parsed = strtoull(v, &end, 10); + if (end == v || errno == ERANGE) return fallback; + while (isspace((unsigned char)*end)) end++; + if (*end) return fallback; + + if (parsed < min_value) return fallback; + uint64_t value = (uint64_t)parsed; + if (value > max_value) value = max_value; + return value; +} + +static uint32_t ds4_gpu_glm_full_attention_max_cache_len(void) { + /* + * kernel_glm_attention_full stores one score per visible token in + * threadgroup memory, plus 256 reduction slots. Keep the default under + * the 32 KiB envelope used by current Apple GPUs while allowing short + * dense decode to step past the 4096-token prefill boundary. + */ + return 7680u; +} + +static uint32_t ds4_gpu_glm_flash_attention_max_cache_len(void) { + /* + * Staged FlashAttention uses fixed-size threadgroup scratch and separate + * KV staging buffers, so it is not bound by the legacy full-attention + * kernel's one-score-per-token threadgroup-memory envelope. + */ + return 8192u; +} + +static int ds4_gpu_mpp_available(void) { + return g_metal4_tensor_api_enabled && !g_quality_mode; +} + +/* + * Retained Metal4 defaults live here instead of behind user-visible options. + * The public runtime has one automatic accelerated path plus the global + * DS4_METAL_DISABLE_METAL4 comparison switch. Benchmark-only alternatives that + * lost during M5 work are removed or kept out of the dispatch path so future + * changes do not accidentally turn old experiments into new modes. + */ +static int ds4_gpu_use_mpp_attn_out_low_matmul(void) { + return ds4_gpu_mpp_available(); +} + +enum { + DS4_METAL_ATTN_OUT_MPP_TILE_N = 64, +}; + +static void ds4_gpu_warn_mpp_fallback(void) { + static int warned; + if (!warned) { + fprintf(stderr, "ds4: accelerated Metal prefill matmul unavailable; falling back to legacy kernel\n"); + warned = 1; + } +} + +static int ds4_gpu_device_name_contains(const char *needle) { + return g_metal_device_name[0] != '\0' && strstr(g_metal_device_name, needle) != NULL; +} + +static int ds4_gpu_compile_tensor_probe(void) { +#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000 + if (!g_device) return 0; + if (@available(macOS 26.0, *)) { + const char *src = + "#include \n" + "#include \n" + "#include \n" + "using namespace metal;\n" + "using namespace mpp::tensor_ops;\n" + "kernel void ds4_tensor_probe(\n" + " tensor> A [[buffer(0)]],\n" + " tensor> B [[buffer(1)]],\n" + " device float *C [[buffer(2)]],\n" + " uint2 tgid [[threadgroup_position_in_grid]]) {\n" + " auto tA = A.slice(0, (int)tgid.y);\n" + " auto tB = B.slice((int)tgid.x, 0);\n" + " matmul2d> mm;\n" + " auto cT = mm.get_destination_cooperative_tensor();\n" + " auto sA = tA.slice(0, 0);\n" + " auto sB = tB.slice(0, 0);\n" + " mm.run(sB, sA, cT);\n" + " auto tC = tensor, tensor_inline>(C, dextents(16, 16));\n" + " cT.store(tC);\n" + "}\n"; + + NSError *error = nil; + NSString *source = [NSString stringWithUTF8String:src]; + id probe_library = [g_device newLibraryWithSource:source options:[MTLCompileOptions new] error:&error]; + if (!probe_library) { + fprintf(stderr, "ds4: Metal 4 tensor API probe compile failed: %s\n", + error ? [[error localizedDescription] UTF8String] : "(unknown)"); + return 0; + } + id fn = [probe_library newFunctionWithName:@"ds4_tensor_probe"]; + if (!fn) { + fprintf(stderr, "ds4: Metal 4 tensor API probe function missing\n"); + return 0; + } + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal 4 tensor API probe pipeline failed: %s\n", + error ? [[error localizedDescription] UTF8String] : "(unknown)"); + return 0; + } + return 1; + } +#endif + return 0; +} + +static void ds4_gpu_detect_metal4_features(void) { + g_metal4_runtime_available = 0; + g_metal4_family_supported = 0; + g_metal4_queue_supported = 0; + g_metal4_m5_neural_accelerators_hint = 0; + g_metal4_tensor_api_enabled = 0; + g_metal4_tensor_api_compile_supported = 0; + g_metal_device_name[0] = '\0'; + + if (!g_device) return; + + const char *name = [[g_device name] UTF8String]; + if (name) { + snprintf(g_metal_device_name, sizeof(g_metal_device_name), "%s", name); + } + + const int metal4_disabled = ds4_gpu_env_bool("DS4_METAL_DISABLE_METAL4") > 0; + +#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000 + if (@available(macOS 26.0, *)) { + g_metal4_runtime_available = 1; + g_metal4_family_supported = + !metal4_disabled && [g_device supportsFamily:MTLGPUFamilyMetal4] ? 1 : 0; + g_metal4_queue_supported = [g_device respondsToSelector:@selector(newMTL4CommandQueue)] ? 1 : 0; + + /* + * Apple does not currently expose a separate "Neural Accelerator" bit + * through Metal. On public M5 systems the hardware signal is the device + * generation plus Metal 4 support, so keep this as a conservative hint. + */ + if (g_metal4_family_supported && ds4_gpu_device_name_contains("M5")) { + g_metal4_m5_neural_accelerators_hint = 1; + } + + if (g_metal4_family_supported) { + const int default_enable = + ds4_gpu_device_name_contains("M5") || + ds4_gpu_device_name_contains("M6") || + ds4_gpu_device_name_contains("A19") || + ds4_gpu_device_name_contains("A20"); + + /* + * Metal 4 TensorOps are portable in source, but on pre-M5 hardware + * they can map to ordinary shader fallbacks. Keep the automatic + * fast path restricted to hardware generations where the Neural + * Accelerator/TensorOps path is expected to pay off; older Metal + * machines continue to use the established kernels unless a future + * device is explicitly added here. + */ + if (default_enable) { + g_metal4_tensor_api_compile_supported = ds4_gpu_compile_tensor_probe(); + g_metal4_tensor_api_enabled = g_metal4_tensor_api_compile_supported; + if (!g_metal4_tensor_api_enabled) { + fprintf(stderr, "ds4: Metal 4 tensor API probe failed; using legacy Metal kernels\n"); + } + } else { + fprintf(stderr, "ds4: Metal 4 tensor API disabled for pre-M5/pre-A19 devices\n"); + } + } + } +#endif +} + +static int ds4_gpu_warm_model_views(void) { + if (g_model_view_count == 0) return 1; + + id pipeline = ds4_gpu_get_pipeline("kernel_touch_u8_stride"); + if (!pipeline) return 0; + + uint64_t stride = 1024ull * 1024ull; + const char *stride_env = getenv("DS4_METAL_MODEL_WARMUP_STRIDE_MB"); + if (stride_env && stride_env[0]) { + char *end = NULL; + unsigned long long mb = strtoull(stride_env, &end, 10); + if (end != stride_env && mb > 0 && mb <= 1024) { + stride = mb * 1024ull * 1024ull; + } + } + const char *stride_kb_env = getenv("DS4_METAL_MODEL_WARMUP_STRIDE_KB"); + if (stride_kb_env && stride_kb_env[0]) { + char *end = NULL; + unsigned long long kb = strtoull(stride_kb_env, &end, 10); + if (end != stride_kb_env && kb > 0 && kb <= 1024ull * 1024ull) { + stride = kb * 1024ull; + const uint64_t page = (uint64_t)getpagesize(); + if (stride < page) stride = page; + } + } + + uint64_t total_touches = 0; + for (uint32_t i = 0; i < g_model_view_count; i++) { + total_touches += (g_model_views[i].bytes + stride - 1) / stride; + } + if (total_touches == 0 || total_touches > (uint64_t)NSUIntegerMax) return 0; + + const NSUInteger out_bytes = (NSUInteger)total_touches; + id out = [g_device newBufferWithLength:out_bytes + options:MTLResourceStorageModeShared]; + if (!out) { + fprintf(stderr, "ds4: Metal model warmup scratch allocation failed\n"); + return 0; + } + out.label = @"ds4_model_warmup"; + + id cb = ds4_gpu_new_command_buffer(); + if (!cb) { + fprintf(stderr, "ds4: Metal model warmup command buffer allocation failed\n"); + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + uint64_t dst_offset = 0; + for (uint32_t i = 0; i < g_model_view_count; i++) { + const uint64_t bytes = g_model_views[i].bytes; + const uint64_t n = (bytes + stride - 1) / stride; + [enc setBuffer:g_model_views[i].buffer offset:0 atIndex:0]; + [enc setBuffer:out offset:0 atIndex:1]; + [enc setBytes:&stride length:sizeof(stride) atIndex:2]; + [enc setBytes:&bytes length:sizeof(bytes) atIndex:3]; + [enc setBytes:&dst_offset length:sizeof(dst_offset) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)((n + 255) / 256), 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + dst_offset += n; + } + ds4_gpu_end_compute_encoder(cb, enc); + + [cb commit]; + [cb waitUntilCompleted]; + + if (cb.status == MTLCommandBufferStatusError) { + fprintf(stderr, "ds4: Metal model warmup failed: %s\n", + [[cb.error localizedDescription] UTF8String]); + return 0; + } + + return 1; +} + +static const char *ds4_gpu_mul_mm_id_map0_name(uint32_t ne20) { + switch (ne20) { + case 1: return "kernel_mul_mm_id_map0_ne20_1"; + case 2: return "kernel_mul_mm_id_map0_ne20_2"; + case 4: return "kernel_mul_mm_id_map0_ne20_4"; + case 5: return "kernel_mul_mm_id_map0_ne20_5"; + case 6: return "kernel_mul_mm_id_map0_ne20_6"; + case 8: return "kernel_mul_mm_id_map0_ne20_8"; + case 10: return "kernel_mul_mm_id_map0_ne20_10"; + case 16: return "kernel_mul_mm_id_map0_ne20_16"; + case 22: return "kernel_mul_mm_id_map0_ne20_22"; + default: return NULL; + } +} + +static id ds4_gpu_get_mul_mv_pipeline( + const char *function_name, + int16_t nsg) { + NSString *key = [NSString stringWithFormat:@"%s_nsg=%d", function_name, (int)nsg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_mul_mv_ext_pipeline( + const char *function_name, + int16_t nsg, + int16_t nxpsg) { + NSString *key = [NSString stringWithFormat:@"%s_nsg=%d_nxpsg=%d", + function_name, (int)nsg, (int)nxpsg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600]; + [constants setConstantValue:&nxpsg type:MTLDataTypeShort atIndex:601]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_flash_attn_pad_pipeline( + bool has_mask, + int32_t ncpsg) { + NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_pad_mask=%d_ncpsg=%d", + has_mask ? 1 : 0, (int)ncpsg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:100]; + [constants setConstantValue:&ncpsg type:MTLDataTypeInt atIndex:125]; + + NSError *error = nil; + id fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_pad" + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_pad function not found: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_pad pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_flash_attn_blk_pipeline( + int32_t nqptg, + int32_t ncpsg) { + NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_blk_nqptg=%d_ncpsg=%d", + (int)nqptg, (int)ncpsg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&nqptg type:MTLDataTypeInt atIndex:224]; + [constants setConstantValue:&ncpsg type:MTLDataTypeInt atIndex:225]; + + NSError *error = nil; + id fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_blk" + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_blk function not found: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_blk pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_flash_attn_pipeline( + const char *function_name, + bool has_mask, + bool has_sinks, + bool has_bias, + bool has_scap, + bool has_kvpad, + bool bc_mask, + int32_t ns10, + int32_t ns20, + int32_t nsg) { + NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_ns10=%d_ns20=%d_nsg=%d", + function_name, + has_mask ? 1 : 0, + has_sinks ? 1 : 0, + has_bias ? 1 : 0, + has_scap ? 1 : 0, + has_kvpad ? 1 : 0, + bc_mask ? 1 : 0, + (int)ns10, + (int)ns20, + (int)nsg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) return cached; + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:300]; + [constants setConstantValue:&has_sinks type:MTLDataTypeBool atIndex:301]; + [constants setConstantValue:&has_bias type:MTLDataTypeBool atIndex:302]; + [constants setConstantValue:&has_scap type:MTLDataTypeBool atIndex:303]; + [constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:304]; + [constants setConstantValue:&bc_mask type:MTLDataTypeBool atIndex:310]; + [constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:320]; + [constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:321]; + [constants setConstantValue:&nsg type:MTLDataTypeInt atIndex:322]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + return pipeline; +} + +static id ds4_gpu_get_flash_attn_vec_pipeline( + const char *function_name, + bool has_mask, + bool has_sinks, + bool has_bias, + bool has_scap, + bool has_kvpad, + bool shared_kvpad, + int32_t ns10, + int32_t ns20, + int32_t nsg, + int32_t nwg) { + /* + * Decode calls this once per layer with identical arguments, so memoize + * the last hit and skip the NSString key + dictionary lookup on the hot + * path. The generic cache below remains the fallback for new variants. + */ + static struct { + const char *fn; + bool m, s, b, c, k, sp; + int32_t n10, n20, sg, wg; + id pipeline; + } memo; + if (memo.pipeline && memo.fn != NULL && strcmp(memo.fn, function_name) == 0 && + memo.m == has_mask && memo.s == has_sinks && memo.b == has_bias && + memo.c == has_scap && memo.k == has_kvpad && memo.sp == shared_kvpad && + memo.n10 == ns10 && memo.n20 == ns20 && memo.sg == nsg && memo.wg == nwg) { + return memo.pipeline; + } + + NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_sharedpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d", + function_name, + has_mask ? 1 : 0, + has_sinks ? 1 : 0, + has_bias ? 1 : 0, + has_scap ? 1 : 0, + has_kvpad ? 1 : 0, + shared_kvpad ? 1 : 0, + (int)ns10, + (int)ns20, + (int)nsg, + (int)nwg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) { + memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias, + has_scap, has_kvpad, shared_kvpad, ns10, ns20, + nsg, nwg, cached }; + return cached; + } + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:400]; + [constants setConstantValue:&has_sinks type:MTLDataTypeBool atIndex:401]; + [constants setConstantValue:&has_bias type:MTLDataTypeBool atIndex:402]; + [constants setConstantValue:&has_scap type:MTLDataTypeBool atIndex:403]; + [constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:404]; + [constants setConstantValue:&shared_kvpad type:MTLDataTypeBool atIndex:405]; + [constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:420]; + [constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:421]; + [constants setConstantValue:&nsg type:MTLDataTypeInt atIndex:422]; + [constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:423]; + + NSError *error = nil; + NSString *name = [NSString stringWithUTF8String:function_name]; + id fn = [g_library newFunctionWithName:name + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal %s function not found: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n", + function_name, [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias, + has_scap, has_kvpad, shared_kvpad, ns10, ns20, + nsg, nwg, pipeline }; + return pipeline; +} + +static id ds4_gpu_get_flash_attn_reduce_pipeline( + int32_t dv, + int32_t nwg) { + /* Same per-layer memo pattern as the vec getter above. */ + static int32_t memo_dv, memo_nwg; + static id memo_pipeline; + if (memo_pipeline && memo_dv == dv && memo_nwg == nwg) { + return memo_pipeline; + } + + NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_vec_reduce_dv=%d_nwg=%d", + (int)dv, (int)nwg]; + id cached = [g_pipeline_cache objectForKey:key]; + if (cached) { + memo_dv = dv; memo_nwg = nwg; memo_pipeline = cached; + return cached; + } + + MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init]; + [constants setConstantValue:&dv type:MTLDataTypeInt atIndex:500]; + [constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:501]; + + NSError *error = nil; + id fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_vec_reduce" + constantValues:constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce function not found: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + error = nil; + id pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!pipeline) { + fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + [g_pipeline_cache setObject:pipeline forKey:key]; + memo_dv = dv; memo_nwg = nwg; memo_pipeline = pipeline; + return pipeline; +} + +static uint32_t ds4_gpu_flash_attn_vec_nsg(uint32_t n_keys, uint32_t nwg, uint32_t ncpsg) { + uint32_t nsg = 1; + while (2u * nwg * nsg * ncpsg < n_keys && nsg < 4u) { + nsg *= 2u; + } + return nsg; +} + +static int ds4_gpu_trace_allocs(void) { + static int initialized; + static int enabled; + if (!initialized) { + enabled = getenv("DS4_METAL_TRACE_ALLOCS") != NULL; + initialized = 1; + } + return enabled; +} + +static double ds4_gpu_mib(uint64_t bytes) { + return (double)bytes / (1024.0 * 1024.0); +} + +static double ds4_gpu_gib(uint64_t bytes) { + return (double)bytes / (1024.0 * 1024.0 * 1024.0); +} + +static ds4_gpu_stream_expert_timing_snapshot +ds4_gpu_stream_expert_timing_current(void) { + return (ds4_gpu_stream_expert_timing_snapshot) { + .selected_calls = g_stream_expert_timing_selected_calls, + .selected_read_ms = g_stream_expert_timing_selected_read_ms, + .selected_sync_ms = g_stream_expert_timing_selected_sync_ms, + .selected_copy_ms = g_stream_expert_timing_selected_copy_ms, + .selected_bind_ms = g_stream_expert_timing_selected_bind_ms, + .split_layers = g_stream_expert_timing_split_layers, + .split_resident_experts = g_stream_expert_timing_split_resident_experts, + .split_missing_experts = g_stream_expert_timing_split_missing_experts, + .split_resident_ms = g_stream_expert_timing_split_resident_ms, + .split_missing_ms = g_stream_expert_timing_split_missing_ms, + .split_missing_load_ms = + g_stream_expert_timing_split_missing_load_ms, + .split_missing_slot_ms = + g_stream_expert_timing_split_missing_slot_ms, + .split_missing_prune_ms = + g_stream_expert_timing_split_missing_prune_ms, + .split_missing_addr_ms = + g_stream_expert_timing_split_missing_addr_ms, + .split_missing_wait_ms = + g_stream_expert_timing_split_missing_wait_ms, + .load_calls = g_stream_expert_timing_load_calls, + .load_prepare_ms = g_stream_expert_timing_load_prepare_ms, + .load_pread_ms = g_stream_expert_timing_load_pread_ms, + .load_modify_ms = g_stream_expert_timing_load_modify_ms, + .load_install_ms = g_stream_expert_timing_load_install_ms, + .prepare_batch_reuse_calls = + g_stream_expert_timing_prepare_batch_reuse_calls, + .prepare_batch_reuse_ms = + g_stream_expert_timing_prepare_batch_reuse_ms, + .prepare_buffer_calls = + g_stream_expert_timing_prepare_buffer_calls, + .prepare_buffer_ms = + g_stream_expert_timing_prepare_buffer_ms, + .prepare_task_experts = + g_stream_expert_timing_prepare_task_experts, + .prepare_task_ms = + g_stream_expert_timing_prepare_task_ms, + .reuse_scan_calls = + g_stream_expert_timing_reuse_scan_calls, + .reuse_scan_entries = + g_stream_expert_timing_reuse_scan_entries, + .reuse_scan_ms = + g_stream_expert_timing_reuse_scan_ms, + .reuse_clear_ms = + g_stream_expert_timing_reuse_clear_ms, + .readahead_calls = + g_stream_expert_timing_readahead_calls, + .readahead_bytes = + g_stream_expert_timing_readahead_bytes, + .readahead_ms = + g_stream_expert_timing_readahead_ms, + .cache_all_resident_layers = + g_stream_expert_timing_cache_all_resident_layers, + .cache_all_missing_layers = + g_stream_expert_timing_cache_all_missing_layers, + .cache_mixed_layers = g_stream_expert_timing_cache_mixed_layers, + .cache_resident_experts = + g_stream_expert_timing_cache_resident_experts, + .cache_missing_experts = + g_stream_expert_timing_cache_missing_experts, + }; +} + +static uint64_t ds4_gpu_stream_expert_timing_delta_u64( + uint64_t current, + uint64_t previous) { + return current >= previous ? current - previous : current; +} + +static double ds4_gpu_stream_expert_timing_delta_f64( + double current, + double previous) { + return current >= previous ? current - previous : current; +} + +static ds4_gpu_stream_expert_timing_snapshot +ds4_gpu_stream_expert_timing_delta( + ds4_gpu_stream_expert_timing_snapshot current, + ds4_gpu_stream_expert_timing_snapshot previous) { + return (ds4_gpu_stream_expert_timing_snapshot) { + .selected_calls = + ds4_gpu_stream_expert_timing_delta_u64(current.selected_calls, + previous.selected_calls), + .selected_read_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.selected_read_ms, + previous.selected_read_ms), + .selected_sync_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.selected_sync_ms, + previous.selected_sync_ms), + .selected_copy_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.selected_copy_ms, + previous.selected_copy_ms), + .selected_bind_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.selected_bind_ms, + previous.selected_bind_ms), + .split_layers = + ds4_gpu_stream_expert_timing_delta_u64(current.split_layers, + previous.split_layers), + .split_resident_experts = + ds4_gpu_stream_expert_timing_delta_u64(current.split_resident_experts, + previous.split_resident_experts), + .split_missing_experts = + ds4_gpu_stream_expert_timing_delta_u64(current.split_missing_experts, + previous.split_missing_experts), + .split_resident_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.split_resident_ms, + previous.split_resident_ms), + .split_missing_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.split_missing_ms, + previous.split_missing_ms), + .split_missing_load_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.split_missing_load_ms, + previous.split_missing_load_ms), + .split_missing_slot_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.split_missing_slot_ms, + previous.split_missing_slot_ms), + .split_missing_prune_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.split_missing_prune_ms, + previous.split_missing_prune_ms), + .split_missing_addr_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.split_missing_addr_ms, + previous.split_missing_addr_ms), + .split_missing_wait_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.split_missing_wait_ms, + previous.split_missing_wait_ms), + .load_calls = + ds4_gpu_stream_expert_timing_delta_u64(current.load_calls, + previous.load_calls), + .load_prepare_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.load_prepare_ms, + previous.load_prepare_ms), + .load_pread_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.load_pread_ms, + previous.load_pread_ms), + .load_modify_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.load_modify_ms, + previous.load_modify_ms), + .load_install_ms = + ds4_gpu_stream_expert_timing_delta_f64(current.load_install_ms, + previous.load_install_ms), + .prepare_batch_reuse_calls = + ds4_gpu_stream_expert_timing_delta_u64( + current.prepare_batch_reuse_calls, + previous.prepare_batch_reuse_calls), + .prepare_batch_reuse_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.prepare_batch_reuse_ms, + previous.prepare_batch_reuse_ms), + .prepare_buffer_calls = + ds4_gpu_stream_expert_timing_delta_u64( + current.prepare_buffer_calls, + previous.prepare_buffer_calls), + .prepare_buffer_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.prepare_buffer_ms, + previous.prepare_buffer_ms), + .prepare_task_experts = + ds4_gpu_stream_expert_timing_delta_u64( + current.prepare_task_experts, + previous.prepare_task_experts), + .prepare_task_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.prepare_task_ms, + previous.prepare_task_ms), + .reuse_scan_calls = + ds4_gpu_stream_expert_timing_delta_u64( + current.reuse_scan_calls, + previous.reuse_scan_calls), + .reuse_scan_entries = + ds4_gpu_stream_expert_timing_delta_u64( + current.reuse_scan_entries, + previous.reuse_scan_entries), + .reuse_scan_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.reuse_scan_ms, + previous.reuse_scan_ms), + .reuse_clear_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.reuse_clear_ms, + previous.reuse_clear_ms), + .readahead_calls = + ds4_gpu_stream_expert_timing_delta_u64( + current.readahead_calls, + previous.readahead_calls), + .readahead_bytes = + ds4_gpu_stream_expert_timing_delta_u64( + current.readahead_bytes, + previous.readahead_bytes), + .readahead_ms = + ds4_gpu_stream_expert_timing_delta_f64( + current.readahead_ms, + previous.readahead_ms), + .cache_all_resident_layers = + ds4_gpu_stream_expert_timing_delta_u64( + current.cache_all_resident_layers, + previous.cache_all_resident_layers), + .cache_all_missing_layers = + ds4_gpu_stream_expert_timing_delta_u64( + current.cache_all_missing_layers, + previous.cache_all_missing_layers), + .cache_mixed_layers = + ds4_gpu_stream_expert_timing_delta_u64(current.cache_mixed_layers, + previous.cache_mixed_layers), + .cache_resident_experts = + ds4_gpu_stream_expert_timing_delta_u64( + current.cache_resident_experts, + previous.cache_resident_experts), + .cache_missing_experts = + ds4_gpu_stream_expert_timing_delta_u64( + current.cache_missing_experts, + previous.cache_missing_experts), + }; +} + +static int ds4_gpu_stream_expert_timing_has_data( + ds4_gpu_stream_expert_timing_snapshot s) { + return s.selected_calls != 0 || + s.split_layers != 0 || + s.load_calls != 0 || + s.cache_all_resident_layers != 0 || + s.cache_all_missing_layers != 0 || + s.cache_mixed_layers != 0; +} + +static void ds4_gpu_stream_expert_timing_print( + const char *scope, + ds4_gpu_stream_expert_timing_snapshot s) { + if (!ds4_gpu_stream_expert_timing_has_data(s)) return; + const double selected_calls = (double)s.selected_calls; + const double split_layers = (double)s.split_layers; + const double selected_read_avg = + selected_calls != 0.0 ? s.selected_read_ms / selected_calls : 0.0; + const double selected_sync_avg = + selected_calls != 0.0 ? s.selected_sync_ms / selected_calls : 0.0; + const double selected_copy_avg = + selected_calls != 0.0 ? s.selected_copy_ms / selected_calls : 0.0; + const double selected_bind_avg = + selected_calls != 0.0 ? s.selected_bind_ms / selected_calls : 0.0; + const double split_resident_avg = + split_layers != 0.0 ? s.split_resident_ms / split_layers : 0.0; + const double split_missing_avg = + split_layers != 0.0 ? s.split_missing_ms / split_layers : 0.0; + const double split_missing_load_avg = + split_layers != 0.0 ? + s.split_missing_load_ms / split_layers : 0.0; + const double split_missing_slot_avg = + split_layers != 0.0 ? + s.split_missing_slot_ms / split_layers : 0.0; + const double split_missing_prune_avg = + split_layers != 0.0 ? + s.split_missing_prune_ms / split_layers : 0.0; + const double split_missing_addr_avg = + split_layers != 0.0 ? + s.split_missing_addr_ms / split_layers : 0.0; + const double split_missing_wait_avg = + split_layers != 0.0 ? + s.split_missing_wait_ms / split_layers : 0.0; + const double load_calls = (double)s.load_calls; + const double load_prepare_avg = + load_calls != 0.0 ? s.load_prepare_ms / load_calls : 0.0; + const double load_pread_avg = + load_calls != 0.0 ? s.load_pread_ms / load_calls : 0.0; + const double load_modify_avg = + load_calls != 0.0 ? s.load_modify_ms / load_calls : 0.0; + const double load_install_avg = + load_calls != 0.0 ? s.load_install_ms / load_calls : 0.0; + const double prepare_batch_reuse_avg = + s.prepare_batch_reuse_calls != 0 ? + s.prepare_batch_reuse_ms / + (double)s.prepare_batch_reuse_calls : 0.0; + const double prepare_buffer_avg = + s.prepare_buffer_calls != 0 ? + s.prepare_buffer_ms / (double)s.prepare_buffer_calls : 0.0; + const double prepare_task_avg = + s.prepare_task_experts != 0 ? + s.prepare_task_ms / (double)s.prepare_task_experts : 0.0; + const double reuse_scan_avg = + s.reuse_scan_calls != 0 ? + s.reuse_scan_ms / (double)s.reuse_scan_calls : 0.0; + const double reuse_scan_entries_avg = + s.reuse_scan_calls != 0 ? + (double)s.reuse_scan_entries / (double)s.reuse_scan_calls : 0.0; + const double readahead_avg = + s.readahead_calls != 0 ? + s.readahead_ms / (double)s.readahead_calls : 0.0; + const double split_resident_experts_avg = + split_layers != 0.0 ? + (double)s.split_resident_experts / split_layers : 0.0; + const double split_missing_experts_avg = + split_layers != 0.0 ? + (double)s.split_missing_experts / split_layers : 0.0; + const uint64_t cache_layers = + s.cache_all_resident_layers + + s.cache_all_missing_layers + + s.cache_mixed_layers; + const double cache_layer_count = (double)cache_layers; + const double cache_resident_experts_avg = + cache_layer_count != 0.0 ? + (double)s.cache_resident_experts / cache_layer_count : 0.0; + const double cache_missing_experts_avg = + cache_layer_count != 0.0 ? + (double)s.cache_missing_experts / cache_layer_count : 0.0; + fprintf(stderr, + "ds4: streaming expert timing %s selected_calls=%llu read_avg=%.3f ms sync_avg=%.3f ms copy_avg=%.3f ms bind_avg=%.3f ms read_total=%.3f ms sync_total=%.3f ms copy_total=%.3f ms bind_total=%.3f ms split_layers=%llu resident_experts_avg=%.2f missing_experts_avg=%.2f resident_submit_avg=%.3f ms missing_bind_avg=%.3f ms resident_submit_total=%.3f ms missing_bind_total=%.3f ms missing_load_avg=%.3f ms missing_slot_avg=%.3f ms missing_prune_avg=%.3f ms missing_addr_avg=%.3f ms missing_wait_avg=%.3f ms missing_wait_total=%.3f ms load_calls=%llu load_prepare_avg=%.3f ms load_pread_avg=%.3f ms load_modify_avg=%.3f ms load_install_avg=%.3f ms prepare_batch_reuse_calls=%llu prepare_batch_reuse_avg=%.3f ms prepare_batch_reuse_total=%.3f ms prepare_buffer_calls=%llu prepare_buffer_avg=%.3f ms prepare_buffer_total=%.3f ms prepare_task_experts=%llu prepare_task_avg=%.3f ms prepare_task_total=%.3f ms reuse_scan_calls=%llu reuse_scan_entries_avg=%.1f reuse_scan_avg=%.3f ms reuse_scan_total=%.3f ms reuse_clear_total=%.3f ms readahead_calls=%llu readahead_avg=%.3f ms readahead_total=%.3f ms readahead_gib=%.2f cache_all_resident=%llu cache_all_missing=%llu cache_mixed=%llu cache_resident_avg=%.2f cache_missing_avg=%.2f\n", + scope ? scope : "total", + (unsigned long long)s.selected_calls, + selected_read_avg, + selected_sync_avg, + selected_copy_avg, + selected_bind_avg, + s.selected_read_ms, + s.selected_sync_ms, + s.selected_copy_ms, + s.selected_bind_ms, + (unsigned long long)s.split_layers, + split_resident_experts_avg, + split_missing_experts_avg, + split_resident_avg, + split_missing_avg, + s.split_resident_ms, + s.split_missing_ms, + split_missing_load_avg, + split_missing_slot_avg, + split_missing_prune_avg, + split_missing_addr_avg, + split_missing_wait_avg, + s.split_missing_wait_ms, + (unsigned long long)s.load_calls, + load_prepare_avg, + load_pread_avg, + load_modify_avg, + load_install_avg, + (unsigned long long)s.prepare_batch_reuse_calls, + prepare_batch_reuse_avg, + s.prepare_batch_reuse_ms, + (unsigned long long)s.prepare_buffer_calls, + prepare_buffer_avg, + s.prepare_buffer_ms, + (unsigned long long)s.prepare_task_experts, + prepare_task_avg, + s.prepare_task_ms, + (unsigned long long)s.reuse_scan_calls, + reuse_scan_entries_avg, + reuse_scan_avg, + s.reuse_scan_ms, + s.reuse_clear_ms, + (unsigned long long)s.readahead_calls, + readahead_avg, + s.readahead_ms, + ds4_gpu_gib(s.readahead_bytes), + (unsigned long long)s.cache_all_resident_layers, + (unsigned long long)s.cache_all_missing_layers, + (unsigned long long)s.cache_mixed_layers, + cache_resident_experts_avg, + cache_missing_experts_avg); +} + +static void ds4_gpu_print_task_memory_report(void) { + task_vm_info_data_t info; + mach_msg_type_number_t count = TASK_VM_INFO_COUNT; + const kern_return_t kr = task_info(mach_task_self(), + TASK_VM_INFO, + (task_info_t)&info, + &count); + if (kr != KERN_SUCCESS) return; + + fprintf(stderr, + "ds4: macOS task memory footprint %.2f GiB, resident %.2f GiB, virtual %.2f GiB\n", + ds4_gpu_gib((uint64_t)info.phys_footprint), + ds4_gpu_gib((uint64_t)info.resident_size), + ds4_gpu_gib((uint64_t)info.virtual_size)); +} + +void ds4_gpu_print_memory_report(const char *label) { + uint64_t cached_prefill_mask_bytes = 0; + uint64_t cached_prefill_blk_bytes = 0; + for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) { + const ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry = + &g_zero_prefix_prefill_mask_cache[i]; + if (entry->mask) cached_prefill_mask_bytes += entry->mask_bytes; + if (entry->blk) cached_prefill_blk_bytes += entry->blk_bytes; + } + const uint64_t scratch = + (uint64_t)g_flash_attn_mask_bytes + + (uint64_t)g_flash_attn_zero_mask_bytes + + cached_prefill_mask_bytes + + (uint64_t)g_flash_attn_pad_bytes + + (uint64_t)g_flash_attn_tmp_bytes + + (uint64_t)g_flash_attn_blk_bytes + + cached_prefill_blk_bytes + + (uint64_t)g_flash_attn_ring_bytes + + (uint64_t)g_flash_attn_kv_bytes + + (uint64_t)g_glm_flash_attn_mask_bytes + + (uint64_t)g_compressor_pool_kv_bytes + + (uint64_t)g_compressor_pool_score_bytes + + (uint64_t)g_compressor_pool_score_cont_bytes + + (uint64_t)g_compressor_pool_softmax_bytes + + (uint64_t)g_compressor_pool_product_bytes + + (uint64_t)g_compressor_store_ape_bytes + + (uint64_t)g_compressor_store_score_bytes + + (uint64_t)g_embed_rows_bytes + + (uint64_t)g_router_selection_bytes + + (uint64_t)g_router_weight_sum_bytes + + (uint64_t)g_indexer_head_scores_bytes + + (uint64_t)g_indexer_topk_bytes + + (uint64_t)g_indexed_topk_bytes + + (uint64_t)g_f16_round_scratch_bytes + + (uint64_t)g_raw_store_round_bytes + + (uint64_t)g_moe_gate_scratch_bytes + + (uint64_t)g_moe_down_scratch_bytes + + (uint64_t)g_moe_id_map_bytes + + (uint64_t)g_moe_q4_gate_slots_bytes + + (uint64_t)g_moe_q4_up_slots_bytes + + (uint64_t)g_moe_q4_down_slots_bytes; + + pthread_mutex_lock(&g_tensor_mu); + const uint64_t tensor_live_snap = g_tensor_alloc_live_bytes; + const uint64_t tensor_peak_snap = g_tensor_alloc_peak_bytes; + pthread_mutex_unlock(&g_tensor_mu); + + uint64_t tracked_live = tensor_live_snap; + if (tracked_live > UINT64_MAX - g_stream_expert_cache_bytes) { + tracked_live = UINT64_MAX; + } else { + tracked_live += g_stream_expert_cache_bytes; + } + + const bool color = ds4_log_is_tty(stderr); + const char *green = color ? "\x1b[32m" : ""; + const char *bright_green = color ? "\x1b[1;32m" : ""; + const char *reset = color ? "\x1b[0m" : ""; + fprintf(stderr, + "%sds4: Metal memory%s%s: runtime %.2f GiB + streaming experts %.2f GiB = %s%.2f GiB tracked live%s\n", + green, + label && label[0] ? " " : "", + label && label[0] ? label : "", + ds4_gpu_gib(tensor_live_snap), + ds4_gpu_gib(g_stream_expert_cache_bytes), + bright_green, + ds4_gpu_gib(tracked_live), + reset); + if (color) fputs(green, stderr); + fprintf(stderr, + "ds4: runtime tensors live %.2f MiB peak %.2f MiB\n", + ds4_gpu_mib(tensor_live_snap), + ds4_gpu_mib(tensor_peak_snap)); + ds4_gpu_print_task_memory_report(); + fprintf(stderr, + "ds4: mmap model wrapper spans %llu buffers %.2f GiB total, %.2f GiB max (not copied)\n", + (unsigned long long)g_model_wrap_count, + ds4_gpu_gib(g_model_wrap_bytes), + ds4_gpu_gib(g_model_wrap_max_bytes)); + if (g_model_buffer_cache && [g_model_buffer_cache count] != 0) { + const uint64_t limit = ds4_gpu_exact_view_cache_limit_bytes(); + if (limit == 0) { + fprintf(stderr, + "ds4: exact model view cache %lu buffers %.2f GiB unlimited, %llu evictions (not copied)\n", + (unsigned long)[g_model_buffer_cache count], + ds4_gpu_gib(g_model_buffer_cache_bytes), + (unsigned long long)g_model_buffer_cache_evictions); + } else { + fprintf(stderr, + "ds4: exact model view cache %lu buffers %.2f GiB / %.2f GiB, %llu evictions (not copied)\n", + (unsigned long)[g_model_buffer_cache count], + ds4_gpu_gib(g_model_buffer_cache_bytes), + ds4_gpu_gib(limit), + (unsigned long long)g_model_buffer_cache_evictions); + } + } + if (g_stream_expert_cache_hits != 0 || + g_stream_expert_cache_misses != 0 || + g_stream_expert_cache_bytes != 0) { + const uint64_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + uint64_t target_bytes = 0; + if (budget != 0 && g_stream_expert_cache_expert_bytes != 0) { + target_bytes = + budget > UINT64_MAX / g_stream_expert_cache_expert_bytes ? + UINT64_MAX : + budget * g_stream_expert_cache_expert_bytes; + } + const uint64_t lookups = g_stream_expert_cache_hits + g_stream_expert_cache_misses; + const double hit_rate = lookups ? + (double)g_stream_expert_cache_hits / (double)lookups : 0.0; + if (g_stream_expert_cache_evict_advise_bytes != 0 || + g_stream_expert_cache_willneed_advise_bytes != 0 || + g_stream_expert_cache_pread_bytes != 0) { + fprintf(stderr, + "ds4: streaming expert cache budget=%llu experts entries=%u expert=%.2f MiB target=%.2f GiB live=%.2f GiB, hits=%llu misses=%llu hit_rate=%.3f wraps=%llu evictions=%llu buffer_allocs=%llu buffer_reuses=%llu evict_dontneed=%.2f GiB miss_willneed=%.2f GiB miss_pread=%.2f GiB pread_ms=%.3f\n", + (unsigned long long)budget, + g_stream_expert_cache_entry_count, + ds4_gpu_mib(g_stream_expert_cache_expert_bytes), + ds4_gpu_gib(target_bytes), + ds4_gpu_gib(g_stream_expert_cache_bytes), + (unsigned long long)g_stream_expert_cache_hits, + (unsigned long long)g_stream_expert_cache_misses, + hit_rate, + (unsigned long long)g_stream_expert_cache_wraps, + (unsigned long long)g_stream_expert_cache_evictions, + (unsigned long long)g_stream_expert_cache_buffer_allocs, + (unsigned long long)g_stream_expert_cache_buffer_reuses, + ds4_gpu_gib(g_stream_expert_cache_evict_advise_bytes), + ds4_gpu_gib(g_stream_expert_cache_willneed_advise_bytes), + ds4_gpu_gib(g_stream_expert_cache_pread_bytes), + g_stream_expert_cache_pread_ms); + } else { + fprintf(stderr, + "ds4: streaming expert cache budget=%llu experts entries=%u expert=%.2f MiB target=%.2f GiB live=%.2f GiB, hits=%llu misses=%llu hit_rate=%.3f wraps=%llu evictions=%llu buffer_allocs=%llu buffer_reuses=%llu\n", + (unsigned long long)budget, + g_stream_expert_cache_entry_count, + ds4_gpu_mib(g_stream_expert_cache_expert_bytes), + ds4_gpu_gib(target_bytes), + ds4_gpu_gib(g_stream_expert_cache_bytes), + (unsigned long long)g_stream_expert_cache_hits, + (unsigned long long)g_stream_expert_cache_misses, + hit_rate, + (unsigned long long)g_stream_expert_cache_wraps, + (unsigned long long)g_stream_expert_cache_evictions, + (unsigned long long)g_stream_expert_cache_buffer_allocs, + (unsigned long long)g_stream_expert_cache_buffer_reuses); + } + if (ds4_gpu_stream_expert_timing_summary_enabled()) { + const ds4_gpu_stream_expert_timing_snapshot total = + ds4_gpu_stream_expert_timing_current(); + if (ds4_gpu_stream_expert_timing_has_data(total)) { + const ds4_gpu_stream_expert_timing_snapshot delta = + ds4_gpu_stream_expert_timing_delta( + total, + g_stream_expert_timing_last_report); + ds4_gpu_stream_expert_timing_print("total", total); + ds4_gpu_stream_expert_timing_print("delta", delta); + g_stream_expert_timing_last_report = total; + } + } + if (getenv("DS4_METAL_STREAMING_EXPERT_LAYER_STATS") != NULL) { + fprintf(stderr, "ds4: streaming expert cache per-layer stats:\n"); + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + const uint64_t hits = g_stream_expert_cache_layer_hits[layer]; + const uint64_t misses = g_stream_expert_cache_layer_misses[layer]; + const uint64_t lookups = hits + misses; + const uint64_t evictions = + g_stream_expert_cache_layer_evictions[layer]; + const uint64_t pread_bytes = + g_stream_expert_cache_layer_pread_bytes[layer]; + const double pread_ms = + g_stream_expert_cache_layer_pread_ms[layer]; + const uint32_t cached = + g_stream_expert_cache_layer_count[layer]; + const uint32_t layer_slots = + ds4_gpu_stream_expert_cache_configured_count() != 0 ? + DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT : 0; + if (lookups == 0 && evictions == 0 && pread_bytes == 0 && + cached == 0) { + continue; + } + const double layer_hit_rate = lookups ? + (double)hits / (double)lookups : 0.0; + fprintf(stderr, + "ds4: layer=%u layer_slots=%u cached=%u hits=%llu misses=%llu hit_rate=%.3f evictions=%llu miss_pread=%.2f GiB pread_ms=%.3f\n", + layer, + layer_slots, + cached, + (unsigned long long)hits, + (unsigned long long)misses, + layer_hit_rate, + (unsigned long long)evictions, + ds4_gpu_gib(pread_bytes), + pread_ms); + } + if (getenv("DS4_METAL_STREAMING_EXPERT_LAYER_STATS_DELTA") != NULL) { + fprintf(stderr, "ds4: streaming expert cache per-layer delta:\n"); + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + const uint64_t hits = + g_stream_expert_cache_layer_hits[layer]; + const uint64_t misses = + g_stream_expert_cache_layer_misses[layer]; + const uint64_t evictions = + g_stream_expert_cache_layer_evictions[layer]; + const uint64_t pread_bytes = + g_stream_expert_cache_layer_pread_bytes[layer]; + const double pread_ms = + g_stream_expert_cache_layer_pread_ms[layer]; + const uint64_t delta_hits = + hits - g_stream_expert_cache_layer_last_hits[layer]; + const uint64_t delta_misses = + misses - g_stream_expert_cache_layer_last_misses[layer]; + const uint64_t delta_evictions = + evictions - + g_stream_expert_cache_layer_last_evictions[layer]; + const uint64_t delta_pread_bytes = + pread_bytes - + g_stream_expert_cache_layer_last_pread_bytes[layer]; + const double delta_pread_ms = + pread_ms - + g_stream_expert_cache_layer_last_pread_ms[layer]; + const uint64_t lookups = delta_hits + delta_misses; + if (lookups == 0 && delta_evictions == 0 && + delta_pread_bytes == 0) { + continue; + } + const double hit_rate = lookups ? + (double)delta_hits / (double)lookups : 0.0; + fprintf(stderr, + "ds4: layer=%u cached=%u hits=%llu misses=%llu hit_rate=%.3f evictions=%llu miss_pread=%.2f GiB pread_ms=%.3f\n", + layer, + g_stream_expert_cache_layer_count[layer], + (unsigned long long)delta_hits, + (unsigned long long)delta_misses, + hit_rate, + (unsigned long long)delta_evictions, + ds4_gpu_gib(delta_pread_bytes), + delta_pread_ms); + g_stream_expert_cache_layer_last_hits[layer] = hits; + g_stream_expert_cache_layer_last_misses[layer] = misses; + g_stream_expert_cache_layer_last_evictions[layer] = + evictions; + g_stream_expert_cache_layer_last_pread_bytes[layer] = + pread_bytes; + g_stream_expert_cache_layer_last_pread_ms[layer] = + pread_ms; + } + } + } + } + fprintf(stderr, + "ds4: model residency requests %llu%s\n", + (unsigned long long)g_model_residency_count, + g_ssd_streaming_mode ? " (ssd-streaming)" : + (getenv("DS4_METAL_NO_RESIDENCY") != NULL ? " (disabled)" : "")); + fprintf(stderr, + "ds4: device %s, Metal 4 runtime %s, family %s, MTL4 queue %s, tensor API %s, M5 neural accelerators %s\n", + g_metal_device_name[0] ? g_metal_device_name : "(unknown)", + g_metal4_runtime_available ? "yes" : "no", + g_metal4_family_supported ? "yes" : "no", + g_metal4_queue_supported ? "yes" : "no", + g_metal4_tensor_api_enabled ? "enabled" : + (g_metal4_tensor_api_compile_supported ? "available" : "disabled"), + g_metal4_m5_neural_accelerators_hint ? "likely" : "not detected"); + fprintf(stderr, + "ds4: accelerated Metal path %s%s\n", + ds4_gpu_mpp_available() ? "enabled" : "disabled", + g_quality_mode ? " by --quality" : + (!g_metal4_tensor_api_enabled ? " (tensor API unavailable)" : "")); + fprintf(stderr, + "ds4: device %s, Metal 4 runtime %s, family %s, MTL4 queue %s, tensor API %s, M5 neural accelerators %s\n", + g_metal_device_name[0] ? g_metal_device_name : "(unknown)", + g_metal4_runtime_available ? "yes" : "no", + g_metal4_family_supported ? "yes" : "no", + g_metal4_queue_supported ? "yes" : "no", + g_metal4_tensor_api_enabled ? "enabled" : + (g_metal4_tensor_api_compile_supported ? "available" : "disabled"), + g_metal4_m5_neural_accelerators_hint ? "likely" : "not detected"); + fprintf(stderr, + "ds4: scratch %.2f MiB (flash mask %.2f, pad %.2f, tmp %.2f, blk %.2f, ring %.2f, kv %.2f, compressor %.2f, router %.2f, indexer %.2f, moe %.2f, f16 %.2f, raw-store %.2f)\n", + ds4_gpu_mib(scratch), + ds4_gpu_mib((uint64_t)g_flash_attn_mask_bytes + + (uint64_t)g_glm_flash_attn_mask_bytes + + (uint64_t)g_flash_attn_zero_mask_bytes + + cached_prefill_mask_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_pad_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_tmp_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_blk_bytes + + cached_prefill_blk_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_ring_bytes), + ds4_gpu_mib((uint64_t)g_flash_attn_kv_bytes), + ds4_gpu_mib((uint64_t)g_compressor_pool_kv_bytes + + (uint64_t)g_compressor_pool_score_bytes + + (uint64_t)g_compressor_pool_score_cont_bytes + + (uint64_t)g_compressor_pool_softmax_bytes + + (uint64_t)g_compressor_pool_product_bytes + + (uint64_t)g_compressor_store_ape_bytes + + (uint64_t)g_compressor_store_score_bytes + + (uint64_t)g_embed_rows_bytes), + ds4_gpu_mib((uint64_t)g_router_selection_bytes + + (uint64_t)g_router_weight_sum_bytes), + ds4_gpu_mib((uint64_t)g_indexer_head_scores_bytes + + (uint64_t)g_indexer_topk_bytes + + (uint64_t)g_indexed_topk_bytes), + ds4_gpu_mib((uint64_t)g_moe_gate_scratch_bytes + + (uint64_t)g_moe_down_scratch_bytes + + (uint64_t)g_moe_id_map_bytes + + (uint64_t)g_moe_q4_gate_slots_bytes + + (uint64_t)g_moe_q4_up_slots_bytes + + (uint64_t)g_moe_q4_down_slots_bytes), + ds4_gpu_mib((uint64_t)g_f16_round_scratch_bytes), + ds4_gpu_mib((uint64_t)g_raw_store_round_bytes)); + if (color) fputs(reset, stderr); +} + +void ds4_gpu_set_quality(bool quality) { + g_quality_mode = quality ? 1 : 0; +} + +void ds4_gpu_set_glm_model(bool enabled) { + g_glm_model_mode = enabled ? 1 : 0; +} + +void ds4_gpu_set_ssd_streaming(bool enabled) { + g_ssd_streaming_mode = enabled ? 1 : 0; + ds4_gpu_stream_expert_cache_clear_all(1); + if (g_ssd_streaming_mode) { + fprintf(stderr, + "ds4: Metal SSD streaming mode enabled; full model residency and warmup are skipped\n"); + } +} + +void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled) { + g_glm_streaming_prefill_full_layer_runtime = enabled ? 1 : 0; +} + +void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) { + if (experts > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { + experts = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES; + } + g_stream_expert_cache_budget_override = experts; + ds4_gpu_stream_expert_cache_clear_all(1); +} + +void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) { + /* + * Pre-seed the cache's single slab size class with the model's uniform + * per-expert bytes (first routed layer). With a mixed-precision GGUF this + * pins the class to the majority layers so the boosted ones are rejected + * deterministically from startup, instead of depending on which layer + * happens to touch the cache first. + */ + g_stream_expert_cache_expert_bytes = bytes; +} + +uint64_t ds4_gpu_recommended_working_set_size(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_device) return 0; + return (uint64_t)[g_device recommendedMaxWorkingSetSize]; +} + +static int ds4_gpu_model_map_log_enabled(void) { + if (!g_ssd_streaming_mode) return 1; + const char *trace = getenv("DS4_METAL_STREAMING_MAP_TRACE"); + return trace && trace[0] && strcmp(trace, "0") != 0; +} + +static id ds4_gpu_wrap_model_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset); + +static id ds4_gpu_wrap_model_exact_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset); + +static const char *ds4_gpu_source = +"#include \n" +"#ifdef DS4_METAL_HAS_TENSOR\n" +"#include \n" +"#include \n" +"#endif\n" +"using namespace metal;\n" +"#ifdef DS4_METAL_HAS_TENSOR\n" +"using namespace mpp::tensor_ops;\n" +"#endif\n" +"\n" +"#define MAX(x, y) ((x) > (y) ? (x) : (y))\n" +"#define MIN(x, y) ((x) < (y) ? (x) : (y))\n" +"#define SWAP(x, y) { auto tmp = (x); (x) = (y); (y) = tmp; }\n" +"#define QK8_0 32\n" +"#ifndef QK_K\n" +"#define QK_K 256\n" +"#endif\n" +"#define N_SIMDWIDTH 32\n" +"#define N_R0_Q8_0 2\n" +"#define N_SG_Q8_0 4\n" +"#define FC_MUL_MV 600\n" +"#define FC_MUL_MM 700\n" +"#define FC_BIN 1300\n" +"#define FOR_UNROLL(x) _Pragma(\"clang loop unroll(full)\") for (x)\n" +"#define M_PI_F 3.14159265358979323846f\n" +"\n" +"// Reads one byte per stride to warm model-backed pages without copying the\n" +"// model. This is outside inference and exists only to reduce first-use stalls.\n" +"kernel void kernel_touch_u8_stride(\n" +" device const uchar *src [[buffer(0)]],\n" +" device uchar *dst [[buffer(1)]],\n" +" constant ulong &stride [[buffer(2)]],\n" +" constant ulong &bytes [[buffer(3)]],\n" +" constant ulong &dst_offset [[buffer(4)]],\n" +" uint gid [[thread_position_in_grid]]) {\n" +" ulong off = (ulong)gid * stride;\n" +" if (off >= bytes) return;\n" +" dst[dst_offset + (ulong)gid] = src[off];\n" +"}\n" +"\n" +"enum ds4_sort_order {\n" +" DS4_SORT_ORDER_ASC,\n" +" DS4_SORT_ORDER_DESC,\n" +"};\n" +"\n" +"struct block_q8_0 {\n" +" half d;\n" +" int8_t qs[QK8_0];\n" +"};\n" +"\n" +"struct block_q8_K {\n" +" float d;\n" +" int8_t qs[QK_K];\n" +" int16_t bsums[QK_K / 16];\n" +"};\n" +"\n" +"\n"; + +static NSString *ds4_gpu_full_source(void) { + NSString *base = [NSString stringWithUTF8String:ds4_gpu_source]; + NSFileManager *fm = [NSFileManager defaultManager]; + /* + * Kernels are kept as separate files for review, then concatenated into one + * Metal library. Environment overrides are still honored so a diagnostic + * run can swap one source file without changing the executable. + */ + NSArray *> *required_sources = @[ + @[@"DS4_METAL_FLASH_ATTN_SOURCE", @"metal/flash_attn.metal"], + @[@"DS4_METAL_DENSE_SOURCE", @"metal/dense.metal"], + @[@"DS4_METAL_MOE_SOURCE", @"metal/moe.metal"], + @[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"], + @[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"], + @[@"DS4_METAL_DSV4_KV_SOURCE", @"metal/dsv4_kv.metal"], + @[@"DS4_METAL_DSV4_ROPE_SOURCE", @"metal/dsv4_rope.metal"], + @[@"DS4_METAL_DSV4_MISC_SOURCE", @"metal/dsv4_misc.metal"], + @[@"DS4_METAL_ARGSORT_SOURCE", @"metal/argsort.metal"], + @[@"DS4_METAL_CPY_SOURCE", @"metal/cpy.metal"], + @[@"DS4_METAL_CONCAT_SOURCE", @"metal/concat.metal"], + @[@"DS4_METAL_GET_ROWS_SOURCE", @"metal/get_rows.metal"], + @[@"DS4_METAL_SUM_ROWS_SOURCE", @"metal/sum_rows.metal"], + @[@"DS4_METAL_SOFTMAX_SOURCE", @"metal/softmax.metal"], + @[@"DS4_METAL_REPEAT_SOURCE", @"metal/repeat.metal"], + @[@"DS4_METAL_GLU_SOURCE", @"metal/glu.metal"], + @[@"DS4_METAL_NORM_SOURCE", @"metal/norm.metal"], + @[@"DS4_METAL_BIN_SOURCE", @"metal/bin.metal"], + @[@"DS4_METAL_SET_ROWS_SOURCE", @"metal/set_rows.metal"], + ]; + + NSMutableString *source = [NSMutableString stringWithString:base]; + for (NSArray *spec in required_sources) { + const char *override_path = getenv([spec[0] UTF8String]); + NSMutableArray *paths = [NSMutableArray array]; + if (override_path && override_path[0]) { + [paths addObject:[NSString stringWithUTF8String:override_path]]; + } + [paths addObject:spec[1]]; + [paths addObject:[@"./" stringByAppendingString:spec[1]]]; + + NSString *loaded = nil; + NSString *loaded_path = nil; + for (NSString *path in paths) { + if (![fm fileExistsAtPath:path]) continue; + + NSError *error = nil; + loaded = [NSString stringWithContentsOfFile:path + encoding:NSUTF8StringEncoding + error:&error]; + if (!loaded) { + fprintf(stderr, "ds4: failed to read Metal source %s: %s\n", + [path UTF8String], [[error localizedDescription] UTF8String]); + return nil; + } + loaded_path = path; + break; + } + + if (!loaded) { + fprintf(stderr, + "ds4: Metal source %s not found (set %s to override)\n", + [spec[1] UTF8String], [spec[0] UTF8String]); + return nil; + } + [source appendFormat:@"\n// appended %@\n%@\n", loaded_path, loaded]; + } + return source; +} + +typedef struct { + int32_t ne00t; + int32_t ne00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +} ds4_gpu_get_rows_args; + +typedef struct { + int32_t n_embd; + int32_t n_vocab; + int32_t n_tokens; + uint64_t src_row_bytes; + uint64_t dst_row_bytes; + uint64_t token_stride; +} ds4_gpu_get_rows_q8_0_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +} ds4_gpu_repeat_args; + +typedef struct { + int32_t nk0; + int32_t ne01; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +} ds4_gpu_set_rows_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + int32_t dim; +} ds4_gpu_concat_args; + +typedef struct { + int64_t nk0; + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int64_t ne0; + int64_t ne1; + int64_t ne2; + int64_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +} ds4_gpu_cpy_args; + +static ds4_gpu_cpy_args ds4_gpu_make_cpy_1d_args( + uint32_t n, + uint64_t src_elem, + uint64_t dst_elem) { + return (ds4_gpu_cpy_args) { + .nk0 = (int64_t)n, + .ne00 = (int64_t)n, + .ne01 = 1, + .ne02 = 1, + .ne03 = 1, + .nb00 = src_elem, + .nb01 = (uint64_t)n * src_elem, + .nb02 = (uint64_t)n * src_elem, + .nb03 = (uint64_t)n * src_elem, + .ne0 = (int64_t)n, + .ne1 = 1, + .ne2 = 1, + .ne3 = 1, + .nb0 = dst_elem, + .nb1 = (uint64_t)n * dst_elem, + .nb2 = (uint64_t)n * dst_elem, + .nb3 = (uint64_t)n * dst_elem, + }; +} + +static NSUInteger ds4_gpu_cpy_threads(uint32_t n, id pipeline) { + NSUInteger nth = 32u; + const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; + while (nth < (NSUInteger)n && nth < max_threads) nth *= 2u; + if (nth > max_threads) nth = max_threads; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + return nth ? nth : 1u; +} + +static float ds4_gpu_negative_infinity(void) { + union { uint32_t u; float f; } v = { 0xff800000u }; + return v.f; +} + +static float ds4_gpu_positive_infinity(void) { + union { uint32_t u; float f; } v = { 0x7f800000u }; + return v.f; +} + +static int ds4_gpu_encode_cpy_f32_f32_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n); + +static int ds4_gpu_encode_cpy_f32_f32_3d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride); + +static int ds4_gpu_encode_cpy_f32_f32_3d_src_strided( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_col_stride, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride); + +static int ds4_gpu_encode_cpy_f32_f16_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n); + +static int ds4_gpu_encode_cpy_f32_f16_2d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint64_t src_row_stride, + uint64_t dst_row_stride); + +static int ds4_gpu_encode_cpy_f16_f32_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n); + +static int ds4_gpu_encode_fill_f32_rows( + id cb, + id buf, + NSUInteger offset, + uint32_t width, + uint32_t rows, + float value); + +static int ds4_gpu_encode_add_f32_1d( + id cb, + id a, + NSUInteger a_off, + id b, + NSUInteger b_off, + id out, + NSUInteger out_off, + uint32_t n); + +typedef struct { + int32_t ne00; + uint64_t nb01; + int32_t ne10; + uint64_t nb11; + int32_t ne0; + uint64_t nb1; + int32_t i00; + int32_t i10; + float alpha; + float limit; +} ds4_gpu_glu_args; + +typedef struct { + uint32_t n; +} ds4_gpu_add_flat_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + uint64_t offs; + uint64_t o1[8]; +} ds4_gpu_bin_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + float slope; + float scale; + float bias; + float val; + float min; + float max; +} ds4_gpu_unary_args; + +static ds4_gpu_bin_args ds4_gpu_make_bin_rows_args(uint32_t n, uint32_t rows, uint32_t rhs_n) { + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + const uint64_t rhs_row_bytes = (uint64_t)rhs_n * sizeof(float); + return (ds4_gpu_bin_args) { + .ne00 = (int32_t)n, + .ne01 = (int32_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = row_bytes, + .nb03 = row_bytes, + .ne10 = (int32_t)rhs_n, + .ne11 = 1, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = rhs_row_bytes, + .nb12 = rhs_row_bytes, + .nb13 = rhs_row_bytes, + .ne0 = (int32_t)n, + .ne1 = (int32_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = row_bytes, + .nb2 = row_bytes, + .nb3 = row_bytes, + .offs = 0, + .o1 = { 0 }, + }; +} + +static ds4_gpu_unary_args ds4_gpu_make_unary_rows_args( + uint32_t n, + uint32_t rows, + int c4, + float scale, + float bias) { + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + const uint32_t n_kernel = c4 ? n / 4u : n; + return (ds4_gpu_unary_args) { + .ne00 = (int32_t)n_kernel, + .ne01 = (int32_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = row_bytes, + .nb03 = row_bytes, + .ne0 = (int32_t)n_kernel, + .ne1 = (int32_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = row_bytes, + .nb2 = row_bytes, + .nb3 = row_bytes, + .slope = 0.0f, + .scale = scale, + .bias = bias, + .val = 0.0f, + .min = 0.0f, + .max = 0.0f, + }; +} + +static ds4_gpu_bin_args ds4_gpu_make_bin_same_rows_args(uint32_t n, uint32_t rows) { + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + return (ds4_gpu_bin_args) { + .ne00 = (int32_t)n, + .ne01 = (int32_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = (uint64_t)rows * row_bytes, + .nb03 = (uint64_t)rows * row_bytes, + .ne10 = (int32_t)n, + .ne11 = (int32_t)rows, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = row_bytes, + .nb12 = (uint64_t)rows * row_bytes, + .nb13 = (uint64_t)rows * row_bytes, + .ne0 = (int32_t)n, + .ne1 = (int32_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = row_bytes, + .nb2 = (uint64_t)rows * row_bytes, + .nb3 = (uint64_t)rows * row_bytes, + .offs = 0, + .o1 = { 0 }, + }; +} + +static int ds4_gpu_encode_bin_f32_rows( + id cb, + id pipeline, + const ds4_gpu_bin_args *args, + id a, + NSUInteger a_off, + id b, + NSUInteger b_off, + id out, + NSUInteger out_off); + +static int ds4_gpu_encode_sum_rows_f32( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t width, + uint32_t rows); + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int32_t nr0; + int16_t r2; + int16_t r3; +} ds4_gpu_q8_0_matvec_args; + +typedef struct { + int32_t ne00; + int32_t ne02; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int16_t r2; + int16_t r3; +} ds4_gpu_mul_mm_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; + int32_t ne11; + int32_t ne12; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne0; + int32_t ne1; + int16_t r2; + int16_t r3; +} ds4_gpu_mul_mv_ext_args; + +typedef ds4_gpu_q8_0_matvec_args ds4_gpu_f16_matvec_args; + +static ds4_gpu_q8_0_matvec_args ds4_gpu_make_q8_0_mv_args(uint64_t in_dim, uint64_t out_dim) { + const uint64_t row_bytes = (in_dim / 32u) * 34u; + return (ds4_gpu_q8_0_matvec_args) { + .ne00 = (int32_t)in_dim, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = 34, + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)in_dim, + .ne11 = 1, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * sizeof(float), + .nb13 = in_dim * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = 1, + .nr0 = 2, + .r2 = 1, + .r3 = 1, + }; +} + +static ds4_gpu_f16_matvec_args ds4_gpu_make_f16_mv_args(uint64_t in_dim, uint64_t out_dim) { + const uint64_t row_bytes = in_dim * sizeof(uint16_t); + return (ds4_gpu_f16_matvec_args) { + .ne00 = (int32_t)in_dim, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = sizeof(uint16_t), + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)in_dim, + .ne11 = 1, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * sizeof(float), + .nb13 = in_dim * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = 1, + .nr0 = 2, + .r2 = 1, + .r3 = 1, + }; +} + +static ds4_gpu_q8_0_matvec_args ds4_gpu_make_f32_mv_args( + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_vec) { + const uint64_t row_bytes = in_dim * sizeof(float); + return (ds4_gpu_q8_0_matvec_args) { + .ne00 = (int32_t)in_dim, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)in_dim, + .ne11 = (int32_t)n_vec, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * n_vec * sizeof(float), + .nb13 = in_dim * n_vec * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = (int32_t)n_vec, + .nr0 = 2, + .r2 = 1, + .r3 = 1, + }; +} + +typedef struct { + const char *function_name; + int16_t nsg; + int32_t nr0; + NSUInteger smem; +} ds4_gpu_mv_dispatch; + +static int ds4_gpu_tp_world_is_two(void); + +static ds4_gpu_mv_dispatch ds4_gpu_make_q8_0_mv_dispatch(void) { + const uint64_t default_nsg = ds4_gpu_tp_world_is_two() ? 2u : 4u; + const int16_t nsg = + (int16_t)ds4_gpu_env_u64("DS4_METAL_Q8_MV_NSG", default_nsg, 1u, 8u); + const uint64_t rows = ds4_gpu_env_u64("DS4_METAL_Q8_MV_ROWS", 2u, 2u, 4u); + if (rows >= 4u) { + return (ds4_gpu_mv_dispatch) { + .function_name = "kernel_mul_mv_q8_0_f32_r4", + .nsg = nsg, + .nr0 = 4, + .smem = 32u * 4u * sizeof(float), + }; + } + return (ds4_gpu_mv_dispatch) { + .function_name = "kernel_mul_mv_q8_0_f32", + .nsg = nsg, + .nr0 = 2, + .smem = 32u * 2u * sizeof(float), + }; +} + +static ds4_gpu_mv_dispatch ds4_gpu_make_plain_mv_dispatch( + uint64_t in_dim, + int f32_weights) { + if (in_dim < 32) { + return (ds4_gpu_mv_dispatch) { + .function_name = f32_weights ? "kernel_mul_mv_f32_f32_short" : "kernel_mul_mv_f16_f32_short", + .nsg = 1, + .nr0 = 32, + .smem = 0, + }; + } + + const int16_t nsg = (int16_t)((in_dim + 127u) / 128u > 8u ? 8u : (in_dim + 127u) / 128u); + const int use_4 = (in_dim % 4u) == 0; + return (ds4_gpu_mv_dispatch) { + .function_name = f32_weights + ? (use_4 ? "kernel_mul_mv_f32_f32_4" : "kernel_mul_mv_f32_f32") + : (use_4 ? "kernel_mul_mv_f16_f32_4" : "kernel_mul_mv_f16_f32"), + .nsg = nsg, + .nr0 = 2, + .smem = 32u * 2u * sizeof(float), + }; +} + +static ds4_gpu_mul_mm_args ds4_gpu_make_mm_args( + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + uint64_t row_bytes) { + return (ds4_gpu_mul_mm_args) { + .ne00 = (int32_t)in_dim, + .ne02 = 1, + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * n_tok * sizeof(float), + .nb13 = in_dim * n_tok * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = (int32_t)n_tok, + .r2 = 1, + .r3 = 1, + }; +} + +static ds4_gpu_mul_mv_ext_args ds4_gpu_make_mv_ext_args( + uint64_t in_dim, + uint64_t out_dim, + uint64_t n_tok, + uint64_t elem_bytes, + uint64_t row_bytes) { + return (ds4_gpu_mul_mv_ext_args) { + .ne00 = (int32_t)in_dim, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = elem_bytes, + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)in_dim, + .ne11 = (int32_t)n_tok, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * n_tok * sizeof(float), + .nb13 = in_dim * n_tok * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = (int32_t)n_tok, + .r2 = 1, + .r3 = 1, + }; +} + +static int16_t ds4_gpu_mv_ext_nxpsg(uint64_t in_dim, uint64_t n_tok) { + if ((in_dim % 256u) == 0 && n_tok < 3) return 16; + if ((in_dim % 128u) == 0) return 8; + return 4; +} + +static int16_t ds4_gpu_mv_ext_r1ptg(uint64_t n_tok) { + switch (n_tok) { + case 2: return 2; + case 3: + case 6: return 3; + case 4: + case 7: + case 8: return 4; + case 5: return 5; + default: return n_tok > 8 ? 4 : 0; + } +} + +static const char *ds4_gpu_mv_ext_name(int q8, int16_t r1ptg) { + if (q8) { + switch (r1ptg) { + case 2: return "kernel_mul_mv_ext_q8_0_f32_r1_2"; + case 3: return "kernel_mul_mv_ext_q8_0_f32_r1_3"; + case 4: return "kernel_mul_mv_ext_q8_0_f32_r1_4"; + case 5: return "kernel_mul_mv_ext_q8_0_f32_r1_5"; + default: return NULL; + } + } + + switch (r1ptg) { + case 2: return "kernel_mul_mv_ext_f16_f32_r1_2"; + case 3: return "kernel_mul_mv_ext_f16_f32_r1_3"; + case 4: return "kernel_mul_mv_ext_f16_f32_r1_4"; + case 5: return "kernel_mul_mv_ext_f16_f32_r1_5"; + default: return NULL; + } +} + +static const char *ds4_gpu_mv_ext_f32_name(int16_t r1ptg) { + switch (r1ptg) { + case 2: return "kernel_mul_mv_ext_f32_f32_r1_2"; + case 3: return "kernel_mul_mv_ext_f32_f32_r1_3"; + case 4: return "kernel_mul_mv_ext_f32_f32_r1_4"; + case 5: return "kernel_mul_mv_ext_f32_f32_r1_5"; + default: return NULL; + } +} + +static const char *ds4_gpu_mv_ext_q8_pair_swiglu_name(int16_t r1ptg) { + switch (r1ptg) { + case 2: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_2"; + case 3: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_3"; + case 4: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_4"; + case 5: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_5"; + default: return NULL; + } +} + +typedef struct { + int32_t ne00; + int32_t ne00_t; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + float eps; + int32_t nef1[3]; + int32_t nef2[3]; + int32_t nef3[3]; + uint64_t nbf1[3]; + uint64_t nbf2[3]; + uint64_t nbf3[3]; +} ds4_gpu_rms_norm_args; + +typedef struct { + int32_t q_n; + int32_t q_n4; + int32_t kv_n; + int32_t kv_n4; + uint64_t q_row_stride; + uint64_t kv_row_stride; + float eps; +} ds4_gpu_qkv_rms_norm_args; + +static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_args(uint32_t n, uint32_t rows, float eps) { + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + return (ds4_gpu_rms_norm_args) { + .ne00 = (int32_t)n, + .ne00_t = (int32_t)(n / 4u), + .nb1 = row_bytes, + .nb2 = row_bytes * rows, + .nb3 = row_bytes * rows, + .eps = eps, + .nef1 = { (int32_t)rows, 1, 1 }, + .nef2 = { 1, 1, 1 }, + .nef3 = { 1, 1, 1 }, + .nbf1 = { row_bytes, row_bytes, row_bytes }, + .nbf2 = { row_bytes * rows, row_bytes, row_bytes }, + .nbf3 = { row_bytes * rows, row_bytes, row_bytes }, + }; +} + +static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_3d_args( + uint32_t n0, + uint32_t n1, + uint32_t n2, + float eps) { + const uint64_t row_bytes = (uint64_t)n0 * sizeof(float); + const uint64_t plane_bytes = row_bytes * n1; + return (ds4_gpu_rms_norm_args) { + .ne00 = (int32_t)n0, + .ne00_t = (int32_t)(n0 / 4u), + .nb1 = row_bytes, + .nb2 = plane_bytes, + .nb3 = plane_bytes * n2, + .eps = eps, + .nef1 = { (int32_t)n1, 1, 1 }, + .nef2 = { (int32_t)n2, 1, 1 }, + .nef3 = { 1, 1, 1 }, + .nbf1 = { row_bytes, row_bytes, row_bytes }, + .nbf2 = { plane_bytes, row_bytes, row_bytes }, + .nbf3 = { plane_bytes * n2, row_bytes, row_bytes }, + }; +} + +static NSUInteger ds4_gpu_rms_norm_threads(uint32_t n) { + NSUInteger ne00_t = n / 4u; + NSUInteger nth = 32u; + while (nth < ne00_t && nth < 1024u) nth *= 2u; + if (nth > ne00_t) nth = ne00_t; + return nth ? nth : 1u; +} + +static NSUInteger ds4_gpu_rms_norm_pipeline_threads( + uint32_t n, + id pipeline) { + NSUInteger ne00_t = n / 4u; + NSUInteger max_threads = pipeline ? [pipeline maxTotalThreadsPerThreadgroup] : 1024u; + NSUInteger nth = 32u; + while (nth < ne00_t && nth < max_threads) nth *= 2u; + if (nth > max_threads) nth = max_threads; + if (nth > ne00_t) nth = ne00_t; + return nth ? nth : 1u; +} + +typedef struct { + int32_t n_hc; + int32_t sinkhorn_iters; + int64_t n_rows; + int64_t mix_hc; + uint64_t nb01; + uint64_t nb1; + float eps; +} ds4_gpu_hc_split_args; + +typedef struct { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb_w0; + uint64_t nb_w1; + uint64_t nb0; + uint64_t nb1; +} ds4_gpu_hc_weighted_sum_args; + +typedef struct { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb_w0; + uint64_t nb_w1; + uint64_t nb0; + uint64_t nb1; + uint64_t nb_norm1; + float norm_eps; +} ds4_gpu_hc_weighted_sum_norm_args; + +typedef struct { + float post_scale; + float eps; +} ds4_gpu_output_hc_weights4_args; + +typedef struct { + int64_t n_embd; + int32_t n_hc; + int32_t sinkhorn_iters; + int64_t n_rows; + int64_t mix_hc; + uint64_t nb_mix1; + uint64_t nb_split1; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb0; + uint64_t nb1; + float eps; +} ds4_gpu_hc_split_weighted_sum_args; + +typedef struct { + int64_t n_embd; + int32_t n_hc; + int32_t sinkhorn_iters; + int64_t n_rows; + int64_t mix_hc; + uint64_t nb_mix1; + uint64_t nb_split1; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb0; + uint64_t nb1; + uint64_t nb_norm1; + float eps; + float norm_eps; +} ds4_gpu_hc_split_weighted_sum_norm_args; + +typedef struct { + int64_t n_embd; + int64_t n_hc; + int64_t n_tokens; + uint64_t nb_block0; + uint64_t nb_block1; + uint64_t nb_add0; + uint64_t nb_add1; + uint64_t nb_res0; + uint64_t nb_res1; + uint64_t nb_res2; + uint64_t nb_post0; + uint64_t nb_post1; + uint64_t nb_comb0; + uint64_t nb_comb1; + uint64_t nb_comb2; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + int32_t has_add; +} ds4_gpu_hc_expand_args; + +typedef struct { + int32_t nei0; + int32_t nei1; + uint64_t nbi1; + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + int32_t ne10; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + int32_t ne0; + int32_t ne1; + uint64_t nb1; + int32_t nr0; + /* Tensor-parallel expert ownership; see ds4_metal_args_mul_mv_id in + * metal/moe.metal. Zero (from struct literals) means no split. */ + int32_t tp_rank; + int32_t tp_world; + int32_t tp_addend; + int32_t tp_expert_base; +} ds4_gpu_mul_mv_id_args; + +typedef struct { + uint32_t n_total_expert; + uint32_t n_expert; +} ds4_gpu_stream_expert_validate_args; + +typedef struct { + uint32_t active_mask; + uint32_t accumulate; +} ds4_gpu_stream_expert_split_args; + +typedef struct { + int32_t ne02; + int32_t ne10; + int32_t ne11; + uint64_t nb11; + uint64_t nb12; + int32_t ne21; + int32_t ne20; + uint64_t nb21; +} ds4_gpu_mul_mm_id_map_args; + +typedef struct { + int32_t ne00; + int32_t ne02; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne20; + int32_t ne21; + int32_t ne0; + int32_t ne1; + int16_t r2; + int16_t r3; + int32_t tp_rank; + int32_t tp_world; + int32_t tp_expert_base; +} ds4_gpu_mul_mm_id_args; + +static int ds4_gpu_encode_mul_mv_id( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0); + +static int ds4_gpu_encode_attn_out_low_q8_direct( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0); + +static int ds4_gpu_encode_attn_out_low_q8_mpp( + id cb, + id pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off); + +static int ds4_gpu_encode_attn_out_low_q8_mpp( + id cb, + id pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off); + +static ds4_gpu_mul_mm_id_map_args ds4_gpu_make_mul_mm_id_map_args( + uint32_t src0_cols, + uint32_t src0_experts, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens); + +static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args( + uint32_t src0_cols, + uint32_t src0_rows, + uint32_t src0_experts, + uint64_t src0_row_bytes, + uint64_t src0_expert_bytes, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens); +static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args_src1_size( + uint32_t src0_cols, + uint32_t src0_rows, + uint32_t src0_experts, + uint64_t src0_row_bytes, + uint64_t src0_expert_bytes, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens, + uint32_t src1_elem_size); + +static int ds4_gpu_encode_mul_mm_id( + id cb, + id map_pipeline, + id mm_pipeline, + const ds4_gpu_mul_mm_id_map_args *map_args, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off); + +static int ds4_gpu_encode_mul_mm_id_map( + id cb, + id map_pipeline, + const ds4_gpu_mul_mm_id_map_args *map_args, + const ds4_gpu_mul_mm_id_args *mm_args, + id ids, + NSUInteger ids_off); + +static int ds4_gpu_encode_mul_mm_id_mapped( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off); +static int ds4_gpu_encode_mul_mm_id_mapped_tile( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes); +static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0_addrs, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes, + ds4_gpu_stream_expert_cache_entry * const *resources, + uint32_t resource_count, + uint32_t resource_kind, + id overflow_resource); + +typedef struct { + int32_t ne11; + int32_t ne_12_2; + int32_t ne_12_3; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + uint64_t nb21; + uint64_t nb22; + uint64_t nb23; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; +} ds4_gpu_flash_attn_pad_args; + +typedef struct { + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_raw; + uint32_t n_comp; + uint32_t pad_rows; + uint32_t shared_pad; +} ds4_gpu_flash_kv_stage_f16_args; + +typedef struct { + int32_t ne01; + int32_t ne30; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; +} ds4_gpu_flash_attn_blk_args; + +typedef struct { + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + int32_t ne_12_2; + int32_t ne_12_3; + int32_t ns10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ns20; + uint64_t nb21; + uint64_t nb22; + uint64_t nb23; + int32_t ne31; + int32_t ne32; + int32_t ne33; + uint64_t nb31; + uint64_t nb32; + uint64_t nb33; + int32_t ne1; + int32_t ne2; + int32_t ne3; + float scale; + float max_bias; + float m0; + float m1; + int32_t n_head_log2; + float logit_softcap; +} ds4_gpu_flash_attn_vec_args; + +typedef struct { + int32_t nrows; +} ds4_gpu_flash_attn_reduce_args; + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + int32_t n_dims; + int32_t mode; + int32_t n_ctx_orig; + int32_t inverse; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + bool src2; +} ds4_gpu_rope_tail_batch_args; + +typedef struct { + uint64_t row_bytes; + uint64_t token_bytes; + int32_t head_dim; + int32_t n_dims; + int32_t n_ctx_orig; + int32_t inverse; + uint32_t pos0; + uint32_t pos_step; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +} ds4_gpu_rope_affine_pair_args; + +_Static_assert(sizeof(ds4_gpu_rope_affine_pair_args) == 64, + "Metal affine RoPE argument ABI changed"); + +static ds4_gpu_rope_tail_batch_args ds4_gpu_make_rope_tail_args( + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); + const uint64_t tok_bytes = (uint64_t)n_head * row_bytes; + return (ds4_gpu_rope_tail_batch_args) { + .ne00 = head_dim, + .ne01 = n_head, + .ne02 = n_tok, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = tok_bytes, + .nb03 = (uint64_t)n_tok * tok_bytes, + .nb0 = sizeof(float), + .nb1 = row_bytes, + .nb2 = tok_bytes, + .nb3 = (uint64_t)n_tok * tok_bytes, + .n_dims = (int32_t)n_rot, + .mode = 0, + .n_ctx_orig = (int32_t)n_ctx_orig, + .inverse = inverse ? 1 : 0, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .src2 = false, + }; +} + +static int ds4_gpu_encode_rope_tail_inplace( + id cb, + id xbuf, + NSUInteger xoff, + const ds4_gpu_rope_tail_batch_args *args, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t pos0, + uint32_t pos_step) { + const uint32_t tail_threads = args->n_dims > 0 ? (uint32_t)args->n_dims : 0u; + const bool lane_compatible = + tail_threads <= head_dim && ((head_dim - tail_threads) & 31u) == 0u; + const bool force_affine_position = + getenv("DS4_METAL_ENABLE_AFFINE_ROPE_PAIR") != NULL; + const bool use_inplace_pair = + g_rope_tail_inplace_pair_pipeline != nil && + args->mode == 0 && !args->src2 && + lane_compatible && + (ds4_gpu_device_name_contains("M3") || + (ds4_gpu_device_name_contains("M5") && n_tok == 1u) || + getenv("DS4_METAL_ENABLE_INPLACE_ROPE_PAIR") != NULL || + force_affine_position) && + getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL; + const bool use_shared_coeff = + use_inplace_pair && !force_affine_position && + g_rope_tail_inplace_pair_shared4_pipeline != nil && + /* The 256-thread grouped schedule helps long prefill, but reduces the + * per-head parallelism that short batches and decode rely on. */ + tail_threads == 64u && n_head >= 4u && n_tok >= 32u && + getenv("DS4_METAL_DISABLE_M3_SHARED_ROPE_COEFF") == NULL; + if (force_affine_position && + g_rope_tail_inplace_pair_affine_pipeline == nil && + getenv("DS4_METAL_DISABLE_M3_INPLACE_ROPE_PAIR") == NULL && + getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL) { + fprintf(stderr, + "ds4: forced affine-position RoPE pipeline is unavailable\n"); + return 0; + } + /* Keep long prefill on the proven shared4 kernel. Reconstructing affine + * positions inside its coefficient cohort perturbs YaRN fast-math codegen; + * the compact affine specialization is exact for the decode pair schedule. */ + const bool use_affine_position = + use_inplace_pair && !use_shared_coeff && + g_rope_tail_inplace_pair_affine_pipeline != nil && + (n_tok == 1u || force_affine_position) && + (ds4_gpu_device_name_contains("M3") || + ds4_gpu_device_name_contains("M5") || + force_affine_position) && + getenv("DS4_METAL_DISABLE_M3_AFFINE_ROPE_PAIR") == NULL; + + int32_t pos_stack[256]; + int32_t *pos = NULL; + id posbuf = nil; + const NSUInteger pos_bytes = (NSUInteger)n_tok * sizeof(int32_t); + if (!use_affine_position) { + pos = pos_stack; + if (n_tok > (uint32_t)(sizeof(pos_stack) / sizeof(pos_stack[0]))) { + pos = malloc((size_t)n_tok * sizeof(*pos)); + if (!pos) { + fprintf(stderr, "ds4: failed to allocate Metal RoPE position buffer\n"); + return 0; + } + } + for (uint32_t t = 0; t < n_tok; t++) { + pos[t] = (int32_t)(pos0 + t * pos_step); + } + + if (pos_bytes > 4096u) { + /* + * Metal inline setBytes data is meant for small constants. Long + * prefill RoPE calls need thousands of positions; passing that much + * inline can make the Apple driver abort the process. + */ + posbuf = ds4_gpu_new_transient_buffer( + pos_bytes, "ds4_rope_positions"); + if (!posbuf) { + if (pos != pos_stack) free(pos); + return 0; + } + memcpy([posbuf contents], pos, pos_bytes); + } + } + + ds4_gpu_rope_affine_pair_args affine_args; + if (use_affine_position) { + const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); + affine_args = (ds4_gpu_rope_affine_pair_args) { + .row_bytes = row_bytes, + .token_bytes = (uint64_t)n_head * row_bytes, + .head_dim = (int32_t)head_dim, + .n_dims = args->n_dims, + .n_ctx_orig = args->n_ctx_orig, + .inverse = args->inverse, + .pos0 = pos0, + .pos_step = pos_step, + .freq_base = args->freq_base, + .freq_scale = args->freq_scale, + .ext_factor = args->ext_factor, + .attn_factor = args->attn_factor, + .beta_fast = args->beta_fast, + .beta_slow = args->beta_slow, + }; + } + const NSUInteger reference_nth = + (NSUInteger)(head_dim < 256u ? head_dim : 256u); + const NSUInteger pair_nth = + (NSUInteger)(tail_threads < 256u ? tail_threads : 256u); + const NSUInteger nth = use_shared_coeff ? 256u : + (use_inplace_pair ? pair_nth : reference_nth); + const NSUInteger head_groups = use_shared_coeff ? + (NSUInteger)((n_head + 3u) / 4u) : (NSUInteger)n_head; + id enc = ds4_gpu_compute_encoder(cb); + id pipeline = use_affine_position ? + g_rope_tail_inplace_pair_affine_pipeline : + (use_shared_coeff ? + g_rope_tail_inplace_pair_shared4_pipeline : + (use_inplace_pair ? + g_rope_tail_inplace_pair_pipeline : g_rope_tail_batch_pipeline)); + [enc setComputePipelineState:pipeline]; + if (use_affine_position) { + [enc setBytes:&affine_args length:sizeof(affine_args) atIndex:0]; + } else { + [enc setBytes:args length:sizeof(*args) atIndex:0]; + } + [enc setBuffer:xbuf offset:xoff atIndex:1]; + if (!use_affine_position) { + if (posbuf) { + [enc setBuffer:posbuf offset:0 atIndex:2]; + } else { + [enc setBytes:pos length:pos_bytes atIndex:2]; + } + [enc setBuffer:xbuf offset:xoff atIndex:3]; + } + [enc setBuffer:xbuf offset:xoff atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(head_groups, n_tok, 1) + threadsPerThreadgroup:MTLSizeMake(nth ? nth : 1u, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (pos && pos != pos_stack) free(pos); + return 1; +} + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + int32_t n_rot; +} ds4_gpu_dsv4_fp8_kv_quantize_args; + +typedef struct { + int32_t head_dim; + int32_t n_rot; + int32_t raw_row; +} ds4_gpu_dsv4_kv_fp8_store_args; + +typedef struct { + uint32_t n_rows; + uint32_t head_dim; + uint64_t row_stride; +} ds4_gpu_dsv4_indexer_qat_args; + +typedef struct { + uint32_t width; +} ds4_gpu_dsv4_ratio4_shift_args; + +typedef struct { + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t n_threads; +} ds4_gpu_dsv4_compressor_pack_ratio4_args; + +typedef struct { + int64_t n_rows; + uint32_t head_dim; + uint32_t n_comp; + uint32_t replay; + uint32_t pad; +} ds4_gpu_dsv4_softmax_pool_ratio4_direct_args; + +typedef struct { + uint32_t width; + uint32_t ratio; + uint32_t pos; + uint32_t ape_type; +} ds4_gpu_dsv4_compressor_store_one_args; + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; +} ds4_gpu_dsv4_softmax_pool_args; + +typedef struct { + uint32_t width; + uint32_t ratio; + uint32_t pos0; + uint32_t n_tokens; +} ds4_gpu_dsv4_compressor_score_ape_args; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + int32_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + int32_t top_k; +} ds4_gpu_kargs_argsort; + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne0; + int32_t ne1; + int32_t ne2; + int32_t ne3; + int32_t top_k; + int32_t len; +} ds4_gpu_kargs_argsort_merge; + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + int64_t ne03; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int64_t ne0; + int64_t ne1; + int64_t ne2; + int64_t ne3; + uint64_t nb0; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; +} ds4_gpu_kargs_sum_rows; + +typedef struct { + int32_t ne00; + int32_t ne01; + int32_t ne02; + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne11; + int32_t ne12; + int32_t ne13; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + float scale; + float max_bias; + float m0; + float m1; + int32_t n_head_log2; +} ds4_gpu_softmax_args; + +typedef struct { + int64_t ne00; + int64_t ne01; + uint64_t nb00; + uint64_t nb01; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; +} ds4_gpu_dsv4_topk_mask_args; + +typedef struct { + int64_t ne00; + int64_t ne01; + int64_t ne02; + uint64_t nb00; + uint64_t nb01; + uint64_t nb02; + int64_t ne10; + int64_t ne11; + uint64_t nb10; + uint64_t nb11; + int64_t ne0; + int64_t ne1; + uint64_t nb0; + uint64_t nb1; + float scale; +} ds4_gpu_dsv4_indexer_weighted_sum_args; + +typedef struct { + uint32_t has_bias; + uint32_t hash_mode; + uint32_t use_token_buffer; + uint32_t token; + uint32_t hash_rows; +} ds4_gpu_dsv4_router_select_one_args; + +typedef struct { + uint32_t n_expert; + uint32_t n_expert_used; + float expert_weight_scale; + uint32_t pad0; +} ds4_gpu_glm_router_select_one_args; + +typedef struct { + uint32_t n_tokens; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + float eps; +} ds4_gpu_glm_kv_lora_rms_norm_args; + +typedef struct { + uint32_t n_tokens; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t n_head; + uint32_t row_bytes; + uint32_t weight_type; + uint32_t pad1; + uint32_t pad2; +} ds4_gpu_glm_k_b_project_args; + +typedef struct { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t qk_rope; + uint32_t cache_f16; + uint32_t pad1; +} ds4_gpu_glm_store_compact_kv_args; + +typedef struct { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t q_n; + uint32_t q_n4; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t kv_lora_n4; + uint32_t qk_rope; + uint32_t cache_f16; + float eps; + uint32_t pad0; +} ds4_gpu_glm_qkv_norm_store_compact_kv_args; + +typedef struct { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t head_dim; + uint32_t rot_dim; + uint32_t n_ctx_orig; + uint32_t cache_f16; + uint32_t pad0; + float eps; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + float pad1; +} ds4_gpu_glm_store_indexer_k_args; + +typedef struct { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_cap; + uint32_t n_head; + uint32_t kv_raw_dim; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_rope; + uint32_t value_dim; + uint32_t n_ctx_orig; + uint32_t cache_f16; + uint32_t pad0; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +} ds4_gpu_glm_build_kv_cache_args; + +typedef struct { + uint32_t pos0; + uint32_t n_tokens; + uint32_t cache_len; + uint32_t cache_cap; + uint32_t n_head; + uint32_t qk_dim; + uint32_t value_dim; + uint32_t pad0; + uint32_t cache_f16; + uint32_t pad1; + uint32_t pad2; + float scale; +} ds4_gpu_glm_attention_full_args; + +typedef struct { + uint32_t n_selected; +} ds4_gpu_glm_fill_selected_range_args; + +typedef struct { + uint32_t n_tokens; + uint32_t pos0; + uint32_t n_selected; + uint32_t pad_row; +} ds4_gpu_glm_fill_selected_range_batch_args; + +typedef struct { + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t rot_dim; + uint32_t rot_offset; + uint32_t pos0; + uint32_t n_ctx_orig; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; +} ds4_gpu_glm_rope_tail_args; + +typedef struct { + uint32_t n_rows; + uint32_t n_head; + uint32_t head_dim; + uint32_t cache_f16; + float scale; +} ds4_gpu_glm_indexer_score_one_args; + +typedef struct { + uint32_t n_rows; + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t pos0; + uint32_t cache_f16; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t weights_token_stride; + uint64_t score_token_stride; + float scale; +} ds4_gpu_glm_indexer_scores_batch_args; + +typedef struct { + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_dim; + uint32_t row_bytes; + uint32_t weight_type; + uint32_t pad1; + uint32_t pad2; +} ds4_gpu_glm_qk_lowrank_args; + +typedef struct { + uint32_t n_tokens; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_dim; + uint32_t row_bytes; + uint32_t weight_type; + uint32_t head_base; +} ds4_gpu_glm_qk_lowrank_batch_args; + +typedef struct { + uint32_t n_selected; + uint32_t cache_cap; + uint32_t cache_f16; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_rope; + uint32_t value_dim; + uint32_t n_ctx_orig; + uint32_t value_row_bytes; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t value_type; +} ds4_gpu_glm_attention_indexed_decode_args; + +typedef struct { + uint32_t n_selected; + uint32_t cache_cap; + uint32_t cache_f16; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_rope; + uint32_t value_dim; + uint32_t n_ctx_orig; + uint32_t value_row_bytes; + uint32_t block_rows; + uint32_t n_blocks; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t value_type; +} ds4_gpu_glm_attention_indexed_decode_split_args; + +typedef struct { + uint32_t n_tokens; + uint32_t n_selected; + uint32_t cache_cap; + uint32_t cache_f16; + uint32_t n_head; + uint32_t kv_lora_dim; + uint32_t qk_nope; + uint32_t qk_rope; + uint32_t value_dim; + uint32_t n_ctx_orig; + uint32_t value_row_bytes; + uint32_t value_type; + uint32_t pos0; + float scale; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint32_t head_base; +} ds4_gpu_glm_attention_indexed_batch_args; + +typedef struct { + uint32_t in_dim; + uint32_t mid_dim; + uint32_t out_dim; + uint32_t n_total_expert; + uint32_t n_expert_used; + uint32_t n_tokens; + uint32_t mid_token_stride; + uint32_t down_type; + int32_t tp_rank; + int32_t tp_world; + int32_t tp_expert_base; + uint64_t gate_expert_bytes; + uint64_t gate_row_bytes; + uint64_t up_expert_bytes; + uint64_t up_row_bytes; + uint64_t down_expert_bytes; + uint64_t down_row_bytes; +} ds4_gpu_glm_routed_moe_args; + +typedef struct { + uint32_t n_tokens; + uint32_t n_head; + uint32_t n_raw; + uint32_t raw_cap; + uint32_t raw_start; + uint32_t n_comp; + uint32_t top_k; + uint32_t pos0; + uint32_t window; + uint32_t ratio; + uint32_t comp_kv_f16; + uint32_t pad0; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t raw_row_stride; + uint64_t comp_row_stride; + uint64_t topk_token_stride; + uint64_t dst_token_stride; + uint64_t dst_head_stride; + float scale; +} ds4_gpu_dsv4_indexed_attention_args; + +typedef struct { + uint32_t n_comp; + uint32_t n_tokens; + uint32_t n_head; + uint32_t head_dim; + uint32_t pos0; + uint32_t ratio; + uint64_t q_token_stride; + uint64_t q_head_stride; + uint64_t weights_token_stride; + uint64_t index_row_stride; + uint64_t score_token_stride; + float scale; +} ds4_gpu_dsv4_indexer_scores_fused_args; + +typedef struct { + uint32_t width; + uint32_t rows; + uint64_t gate_row_stride; + uint64_t up_row_stride; + uint64_t mid_row_stride; + uint64_t weight_stride; + uint32_t write_clamped; + float clamp_value; +} ds4_gpu_dsv4_moe_swiglu_weight_args; + +typedef struct { + uint32_t expert_base; + uint32_t expert_count; + uint32_t accumulate; + uint32_t pad0; +} ds4_gpu_moe_expert_group_args; + +typedef struct { + uint64_t expert_bytes; + uint32_t group_size; + uint32_t n_slots; +} ds4_gpu_q4_gather_slots6_args; + +typedef struct { + uint32_t width; + uint32_t tokens; + uint64_t src_token_stride; + uint64_t dst_token_stride; +} ds4_gpu_dsv4_moe_sum6_args; + +/* Compile the single in-repo Metal source and create the pipelines that every + * session uses. Shape-dependent kernels with function constants are built + * lazily by the small ds4_gpu_get_* caches, so startup stays predictable + * while long-context prefill and decode can still pick specialized variants. */ +int ds4_gpu_init(void) { + if (g_initialized) return 1; + + @autoreleasepool { + g_device = MTLCreateSystemDefaultDevice(); + if (!g_device) { + fprintf(stderr, "ds4: Metal device not available\n"); + return 0; + } + ds4_gpu_print_device_summary(); + ds4_gpu_detect_metal4_features(); + + g_queue = [g_device newCommandQueue]; + if (!g_queue) { + fprintf(stderr, "ds4: failed to create Metal command queue\n"); + g_device = nil; + return 0; + } + g_model_buffer_cache = [NSMutableDictionary dictionary]; + g_model_buffer_cache_bytes = 0; + g_model_buffer_cache_evictions = 0; + g_model_buffer_cache_over_limit = 0; + g_q4_expert_table_cache = [NSMutableDictionary dictionary]; + g_q4_expert_layer_residency_cache = [NSMutableDictionary dictionary]; + g_pipeline_cache = [NSMutableDictionary dictionary]; + g_transient_buffers = [NSMutableArray array]; + g_pending_cbs = [NSMutableArray array]; + if (!g_model_buffer_cache || !g_q4_expert_table_cache || + !g_q4_expert_layer_residency_cache || + !g_pipeline_cache || !g_transient_buffers || !g_pending_cbs) { + fprintf(stderr, "ds4: Metal bookkeeping allocation failed\n"); + g_pending_cbs = nil; + g_transient_buffers = nil; + g_pipeline_cache = nil; + g_q4_expert_layer_residency_cache = nil; + g_q4_expert_table_cache = nil; + g_model_buffer_cache = nil; + g_queue = nil; + g_device = nil; + return 0; + } + + NSError *error = nil; + NSString *source = ds4_gpu_full_source(); + if (!source) { + g_queue = nil; + g_device = nil; + return 0; + } + MTLCompileOptions *options = [MTLCompileOptions new]; + NSMutableDictionary *macros = [NSMutableDictionary new]; + if (g_metal4_tensor_api_enabled) { + macros[@"DS4_METAL_HAS_TENSOR"] = @"1"; + fprintf(stderr, "ds4: Metal 4 tensor API enabled for Tensor kernels\n"); + } + + const int drift_hc_stable = ds4_gpu_env_bool("DS4_METAL_HC_STABLE") != 0; // default ON + const int drift_norm_unify = ds4_gpu_env_bool("DS4_METAL_NORM_RSQRT_DISABLE") != 0; // default ON + const int drift_kv_raw_f32 = ds4_gpu_env_bool("DS4_METAL_KV_RAW_F32") > 0; // default OFF + const int drift_rope_exp2_log2 = ds4_gpu_env_bool("DS4_METAL_ROPE_EXP2_LOG2") > 0; // default OFF + const int drift_math_safe = ds4_gpu_env_bool("DS4_METAL_MATH_SAFE") > 0; // default OFF + + if (drift_math_safe) { + // MTLCompileOptions.fastMathEnabled defaults to YES and Apple's + // headers explicitly say this "may violate the IEEE 754 standard". + // Different fast-math optimizations get applied across the + // matmul2d cooperative-tensor path and the legacy + // simdgroup_multiply_accumulate path on M5, amplifying the + // mismatch. MTLMathModeSafe pins the entire library to strict + // IEEE-754 semantics. Diagnostic-only: useful to localize drift + // sources but not to ship as a default. + if (@available(macOS 15.0, *)) { + options.mathMode = MTLMathModeSafe; + fprintf(stderr, "ds4: Metal shader library math mode = safe (strict IEEE-754) by DS4_METAL_MATH_SAFE\n"); + } else { +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wdeprecated-declarations" + options.fastMathEnabled = NO; +#pragma clang diagnostic pop + fprintf(stderr, "ds4: Metal shader library fast-math disabled by DS4_METAL_MATH_SAFE (pre-macOS 15)\n"); + } + } + + if (drift_hc_stable) macros[@"DS4_METAL_HC_STABLE"] = @"1"; + if (drift_norm_unify) macros[@"DS4_METAL_NORM_RSQRT_DISABLE"] = @"1"; + if (drift_kv_raw_f32) macros[@"DS4_METAL_KV_RAW_F32"] = @"1"; + if (drift_rope_exp2_log2) macros[@"DS4_METAL_ROPE_EXP2_LOG2"] = @"1"; + fprintf(stderr, + "ds4: drift-patch flags hc_stable=%s norm_unify=%s kv_raw_f32=%s rope_exp2_log2=%s math_safe=%s tensor_matmul=%s\n", + drift_hc_stable ? "on" : "off", + drift_norm_unify ? "on" : "off", + drift_kv_raw_f32 ? "on" : "off", + drift_rope_exp2_log2 ? "on" : "off", + drift_math_safe ? "on" : "off", + g_metal4_tensor_api_enabled ? "on" : "off"); + options.preprocessorMacros = macros; + id library = [g_device newLibraryWithSource:source options:options error:&error]; + if (!library) { + fprintf(stderr, "ds4: Metal shader compilation failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_library = library; + + id fn = [library newFunctionWithName:@"kernel_get_rows_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_get_rows_f16"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_f16 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_f16 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_get_rows_i32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_i32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_i32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_i32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_i32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_get_rows_q8_0_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q8_0_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_q8_0_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_q8_0_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q8_0_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_get_rows_q4_0_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q4_0_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_q4_0_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_q4_0_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q4_0_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_get_rows_q4_K_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q4_K_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_get_rows_q4_K_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_get_rows_q4_K_pipeline) { + fprintf(stderr, "ds4: Metal kernel_get_rows_q4_K_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_repeat_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_repeat_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_repeat_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_repeat_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_repeat_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_set_rows_f32_i32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_set_rows_f32_i32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_set_rows_f32_i32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_set_rows_f32_i32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_set_rows_f32_i32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_concat"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_concat function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_concat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_concat_pipeline) { + fprintf(stderr, "ds4: Metal kernel_concat pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_f32_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_f32_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_f32_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_f32_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_f32_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_f32_f16"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_f32_f16 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_f32_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_f32_f16 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_contig_f32_f16_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_contig_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_contig_f32_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_f16_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_f16_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_f16_f16"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_f16_f16 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_f16_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_f16_f16 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_contig_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_contig_f16_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f16_bits_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_cpy_contig_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_cpy_contig_f16_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_flash_kv_stage_f16"]; + if (fn) { + g_flash_kv_stage_f16_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_flash_kv_stage_f16_pipeline) { + fprintf(stderr, + "ds4: optional Metal gathered KV staging pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal gathered KV staging kernel unavailable\n"); + } + + fn = [library newFunctionWithName:@"kernel_dsv4_fp8_kv_quantize_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_fp8_kv_quantize_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_fp8_kv_quantize_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_indexer_hadamard_fp4_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_hadamard_fp4_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_indexer_qat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_indexer_qat_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_hadamard_fp4_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_kv_fp8_store_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_kv_fp8_store_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_kv_fp8_store_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_kv_fp8_store_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_kv_fp8_store_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_ratio4_shift_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_ratio4_shift_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_ratio4_shift_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_swiglu_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_swiglu_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_swiglu_flat_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_swiglu_flat_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + g_swiglu_flat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_swiglu_flat_pipeline) { + fprintf(stderr, "ds4: Metal kernel_swiglu_flat_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_moe_sum6_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum6_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_moe_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_moe_sum8_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum8_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + g_moe_sum8_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_sum8_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum8_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *bin_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t bin_op = 0; + int16_t bin_f = 1; + bool bin_rb = false; + bool bin_cb = false; + [bin_constants setConstantValue:&bin_op type:MTLDataTypeShort atIndex:1300]; + [bin_constants setConstantValue:&bin_f type:MTLDataTypeShort atIndex:1301]; + [bin_constants setConstantValue:&bin_rb type:MTLDataTypeBool atIndex:1302]; + [bin_constants setConstantValue:&bin_cb type:MTLDataTypeBool atIndex:1303]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" + constantValues:bin_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + g_add_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_add_pipeline) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_add2_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_add2_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + g_add2_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_add2_pipeline) { + fprintf(stderr, "ds4: Metal kernel_add2_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_add3_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_add3_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + g_add3_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_add3_pipeline) { + fprintf(stderr, "ds4: Metal kernel_add3_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *bin_mul_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t bin_mul_plain_op = 2; + int16_t bin_mul_plain_f = 1; + bool bin_mul_plain_rb = false; + bool bin_mul_plain_cb = false; + [bin_mul_constants setConstantValue:&bin_mul_plain_op type:MTLDataTypeShort atIndex:1300]; + [bin_mul_constants setConstantValue:&bin_mul_plain_f type:MTLDataTypeShort atIndex:1301]; + [bin_mul_constants setConstantValue:&bin_mul_plain_rb type:MTLDataTypeBool atIndex:1302]; + [bin_mul_constants setConstantValue:&bin_mul_plain_cb type:MTLDataTypeBool atIndex:1303]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" + constantValues:bin_mul_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + g_mul_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_mul_pipeline) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *bin_mul_scalar_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t bin_mul_op = 2; + int16_t bin_mul_f = 1; + bool bin_mul_rb = false; + bool bin_mul_cb = true; + [bin_mul_scalar_constants setConstantValue:&bin_mul_op type:MTLDataTypeShort atIndex:1300]; + [bin_mul_scalar_constants setConstantValue:&bin_mul_f type:MTLDataTypeShort atIndex:1301]; + [bin_mul_scalar_constants setConstantValue:&bin_mul_rb type:MTLDataTypeBool atIndex:1302]; + [bin_mul_scalar_constants setConstantValue:&bin_mul_cb type:MTLDataTypeBool atIndex:1303]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" + constantValues:bin_mul_scalar_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul-scalar function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + g_bin_mul_scalar_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_bin_mul_scalar_pipeline) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul-scalar pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *bin_div_row_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t bin_div_op = 3; + int16_t bin_div_f = 1; + bool bin_div_rb = false; + bool bin_div_cb = true; + [bin_div_row_constants setConstantValue:&bin_div_op type:MTLDataTypeShort atIndex:1300]; + [bin_div_row_constants setConstantValue:&bin_div_f type:MTLDataTypeShort atIndex:1301]; + [bin_div_row_constants setConstantValue:&bin_div_rb type:MTLDataTypeBool atIndex:1302]; + [bin_div_row_constants setConstantValue:&bin_div_cb type:MTLDataTypeBool atIndex:1303]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32" + constantValues:bin_div_row_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 div-row function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + g_bin_div_row_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_bin_div_row_pipeline) { + fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 div-row pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_rms_norm_mul_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_rms_norm_mul_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rms_norm_pipeline) { + fprintf(stderr, "ds4: Metal kernel_rms_norm_mul_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_rms_norm_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_rms_norm_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_rms_norm_plain_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rms_norm_plain_pipeline) { + fprintf(stderr, "ds4: Metal kernel_rms_norm_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_add_rms_norm_mul_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_add_rms_norm_mul_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_add_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_add_rms_norm_pipeline) { + fprintf(stderr, "ds4: Metal kernel_add_rms_norm_mul_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_qkv_rms_norm_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_qkv_rms_norm_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_qkv_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_qkv_rms_norm_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_qkv_rms_norm_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *moe_mv_id_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t moe_mv_id_nsg = 2; + [moe_mv_id_constants setConstantValue:&moe_mv_id_nsg type:MTLDataTypeShort atIndex:600]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_iq2_xxs_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_iq2_xxs_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_iq2_xxs_pair_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_iq2_xxs_pair_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q2_k_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q2_k_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q2_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_iq2_xxs_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q2_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q2_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_slots6_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_slots6_q2_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q2_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_addr_iq2_xxs_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_iq2_xxs_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q2_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_addr_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_q2_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q2_K_sum6_masked_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_masked_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_masked_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_stream_expert_cache_validate"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_stream_expert_cache_validate function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_stream_expert_cache_validate_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_stream_expert_cache_validate_pipeline) { + fprintf(stderr, "ds4: Metal kernel_stream_expert_cache_validate pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q4_k_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q4_k_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_pair_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q4_k_pair_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q4_k_pair_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_id_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_id_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group6_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group6_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group6_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group6_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group8_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group8_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group8_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group8_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group24_q4_K_id_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_id_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group24_q4_k_id_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group24_q4_k_id_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_id_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_group24_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_group24_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_group24_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_slots6_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_slots6_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_q4_gather_slots6"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_q4_gather_slots6 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_q4_gather_slots6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_q4_gather_slots6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_q4_gather_slots6 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_table_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_pair_swiglu_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_table_q4_pair_gate_encoder = [fn newArgumentEncoderWithBufferIndex:2]; + g_moe_table_q4_pair_up_encoder = [fn newArgumentEncoderWithBufferIndex:3]; + if (!g_moe_table_q4_pair_gate_encoder || !g_moe_table_q4_pair_up_encoder) { + fprintf(stderr, "ds4: Metal Q4 expert-table pair argument encoder creation failed\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_pair_swiglu_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_table_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_sum6_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_table_q4_sum_down_encoder = [fn newArgumentEncoderWithBufferIndex:1]; + if (!g_moe_table_q4_sum_down_encoder) { + fprintf(stderr, "ds4: Metal Q4 expert-table down argument encoder creation failed\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_moe_mul_mv_table_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_table_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_sum6_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q4_K_pair_swiglu_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (fn) { + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_pair_swiglu_f32 pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_pair_swiglu_f32 function unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q4_K_sum6_f32" + constantValues:moe_mv_id_constants + error:&error]; + if (fn) { + g_moe_mul_mv_addr_q4_k_sum6_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_moe_mul_mv_addr_q4_k_sum6_pipeline) { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_sum6_f32 pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_sum6_f32 function unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_rope_tail_batch_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_batch_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_rope_tail_inplace_pair_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_inplace_pair_pipeline) { + fprintf(stderr, + "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_shared4"]; + if (!fn) { + fprintf(stderr, + "ds4: optional Metal shared-head RoPE kernel unavailable; using per-head path\n"); + } else { + g_rope_tail_inplace_pair_shared4_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_inplace_pair_shared4_pipeline) { + fprintf(stderr, + "ds4: optional Metal shared-head RoPE pipeline unavailable; using per-head path: %s\n", + [[error localizedDescription] UTF8String]); + } + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_affine"]; + if (fn) { + g_rope_tail_inplace_pair_affine_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_rope_tail_inplace_pair_affine_pipeline) { + fprintf(stderr, + "ds4: optional Metal affine-position RoPE pair pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal affine-position RoPE pair kernel unavailable\n"); + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_softmax_pool"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_softmax_pool_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_softmax_pool_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_soft_max_f32"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_soft_max_f32 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_soft_max_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_soft_max_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_soft_max_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_soft_max_f32_4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_soft_max_f32_4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_soft_max_f32_4_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_soft_max_f32_4_pipeline) { + fprintf(stderr, "ds4: Metal kernel_soft_max_f32_4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_argsort_f32_i32_desc"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_argsort_f32_i32_desc function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_argsort_f32_i32_desc_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_argsort_f32_i32_desc_pipeline) { + fprintf(stderr, "ds4: Metal kernel_argsort_f32_i32_desc pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_argsort_merge_f32_i32_desc"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_argsort_merge_f32_i32_desc function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_argsort_merge_f32_i32_desc_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_argsort_merge_f32_i32_desc_pipeline) { + fprintf(stderr, "ds4: Metal kernel_argsort_merge_f32_i32_desc pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *sum_rows_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t sum_rows_op = 10; + [sum_rows_constants setConstantValue:&sum_rows_op type:MTLDataTypeShort atIndex:1400]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_sum_rows_f32_f32" + constantValues:sum_rows_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_sum_rows_f32_f32 function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_sum_rows_f32_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_sum_rows_f32_f32_pipeline) { + fprintf(stderr, "ds4: Metal kernel_sum_rows_f32_f32 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_topk_mask"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_topk_mask_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_topk_mask_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_topk_mask_scatter"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask_scatter function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_topk_mask_scatter_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_topk_mask_scatter_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask_scatter pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_indexer_weighted_sum"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_weighted_sum function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_dsv4_indexer_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_dsv4_indexer_weighted_sum_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_weighted_sum pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_sinkhorn"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_sinkhorn function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_hc_split_sinkhorn_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_split_sinkhorn_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_sinkhorn pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_weighted_sum"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_hc_split_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_split_weighted_sum_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_weighted_sum_norm4"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum_norm4 function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_hc_split_weighted_sum_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_split_weighted_sum_norm_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum_norm4 pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_weighted_sum function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_hc_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_weighted_sum_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_weighted_sum pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum_norm4"]; + if (fn) { + g_hc_weighted_sum_norm_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_weighted_sum_norm_pipeline) { + fprintf(stderr, + "ds4: optional Metal output HC sum/RMSNorm pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal output HC sum/RMSNorm kernel unavailable\n"); + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_dsv4_output_hc_weights4"]; + if (fn) { + g_output_hc_weights4_pipeline = + [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_output_hc_weights4_pipeline) { + fprintf(stderr, + "ds4: optional Metal output HC weights4 pipeline unavailable: %s\n", + [[error localizedDescription] UTF8String]); + } + } else { + fprintf(stderr, + "ds4: optional Metal output HC weights4 kernel unavailable\n"); + } + + MTLFunctionConstantValues *unary_sigmoid_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_sigmoid_op = 102; + bool unary_cnt = false; + [unary_sigmoid_constants setConstantValue:&unary_sigmoid_op type:MTLDataTypeShort atIndex:1200]; + [unary_sigmoid_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_sigmoid_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sigmoid function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_sigmoid_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_sigmoid_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sigmoid pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_silu_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_silu_op = 106; + [unary_silu_constants setConstantValue:&unary_silu_op type:MTLDataTypeShort atIndex:1200]; + [unary_silu_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_silu_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 silu function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_silu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_silu_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 silu pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_softplus_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_softplus_op = 115; + [unary_softplus_constants setConstantValue:&unary_softplus_op type:MTLDataTypeShort atIndex:1200]; + [unary_softplus_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_softplus_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 softplus function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_softplus_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_softplus_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 softplus pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_sqrt_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_sqrt_op = 14; + [unary_sqrt_constants setConstantValue:&unary_sqrt_op type:MTLDataTypeShort atIndex:1200]; + [unary_sqrt_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_sqrt_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sqrt function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_sqrt_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_sqrt_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sqrt pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_clamp_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_clamp_op = 12; + [unary_clamp_constants setConstantValue:&unary_clamp_op type:MTLDataTypeShort atIndex:1200]; + [unary_clamp_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32" + constantValues:unary_clamp_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32 clamp function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_clamp_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_clamp_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32 clamp pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_scale_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_scale_op = 10; + [unary_scale_constants setConstantValue:&unary_scale_op type:MTLDataTypeShort atIndex:1200]; + [unary_scale_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_scale_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 scale function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_scale_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_scale_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 scale pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + MTLFunctionConstantValues *unary_fill_constants = [[MTLFunctionConstantValues alloc] init]; + int16_t unary_fill_op = 11; + [unary_fill_constants setConstantValue:&unary_fill_op type:MTLDataTypeShort atIndex:1200]; + [unary_fill_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201]; + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4" + constantValues:unary_fill_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 fill function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_fill_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_fill_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 fill pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + error = nil; + fn = [library newFunctionWithName:@"kernel_unary_f16_f16" + constantValues:unary_fill_constants + error:&error]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_unary_f16_f16 fill function not found: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + g_unary_fill_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_unary_fill_f16_pipeline) { + fprintf(stderr, "ds4: Metal kernel_unary_f16_f16 fill pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + fn = [library newFunctionWithName:@"kernel_dsv4_hc_expand"]; + if (!fn) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_expand function not found\n"); + g_queue = nil; + g_device = nil; + return 0; + } + g_hc_expand_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error]; + if (!g_hc_expand_pipeline) { + fprintf(stderr, "ds4: Metal kernel_dsv4_hc_expand pipeline failed: %s\n", + [[error localizedDescription] UTF8String]); + g_queue = nil; + g_device = nil; + return 0; + } + + g_dsv4_indexer_score_one_direct_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexer_score_one_direct"); + g_dsv4_compressor_store_one_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_compressor_store_one"); + g_dsv4_compressor_pack_ratio4_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_compressor_pack_ratio4"); + g_dsv4_softmax_pool_ratio4_direct_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_softmax_pool_ratio4_direct"); + g_rms_norm_scale_pipeline = + ds4_gpu_get_pipeline("kernel_rms_norm_scale_f32_4"); + g_dsv4_sort_i32_rows_asc_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_sort_i32_rows_asc"); + g_dsv4_indexed_attention_heads8_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8"); + g_dsv4_indexed_attention_heads8_rb16_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_rb16"); + g_dsv4_softplus_sqrt_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_softplus_sqrt_f32_4"); + g_dsv4_router_finalize_one_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one"); + g_dsv4_router_finalize_one_simd_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one_simd"); + g_dsv4_router_finalize_weights_one_simd_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_weights_one_simd"); + g_dsv4_router_weights_one_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_weights_one"); + g_glm_router_select_one_pipeline = + ds4_gpu_get_pipeline("kernel_glm_router_select_one"); + g_glm_kv_lora_rms_norm_pipeline = + ds4_gpu_get_pipeline("kernel_glm_kv_lora_rms_norm"); + g_glm_k_b_project_pipeline = + ds4_gpu_get_pipeline("kernel_glm_k_b_project_q8_0"); + g_glm_store_compact_kv_pipeline = + ds4_gpu_get_pipeline("kernel_glm_store_compact_kv"); + g_glm_qkv_norm_store_compact_kv_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qkv_norm_store_compact_kv"); + g_glm_store_indexer_k_pipeline = + ds4_gpu_get_pipeline("kernel_glm_store_indexer_k"); + g_glm_build_kv_cache_pipeline = + ds4_gpu_get_pipeline("kernel_glm_build_kv_cache"); + g_glm_build_kv_cache_decode_group4_pipeline = + ds4_gpu_get_pipeline("kernel_glm_build_kv_cache_decode_group4"); + g_glm_build_kv_cache_flash_pipeline = + ds4_gpu_get_pipeline("kernel_glm_build_kv_cache_flash"); + g_glm_attention_full_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_full"); + g_glm_fill_selected_range_pipeline = + ds4_gpu_get_pipeline("kernel_glm_fill_selected_range"); + g_glm_fill_selected_range_batch_pipeline = + ds4_gpu_get_pipeline("kernel_glm_fill_selected_range_batch"); + g_glm_indexer_rope_tail_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_rope_tail_f32"); + g_glm_indexer_score_one_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_score_one"); + g_glm_indexer_score_one_direct_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_score_one_direct"); + g_glm_indexer_scores_batch_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_scores_batch"); + g_glm_indexer_scores_tiled_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_scores_tiled"); + g_glm_indexer_scores_tiled_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_indexer_scores_tiled_f32"); + g_glm_qk_lowrank_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0"); + g_glm_qk_lowrank_glm52_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_glm52"); + g_glm_qk_lowrank_glm52_sg_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_glm52_sg"); + g_glm_qk_lowrank_batch_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch"); + g_glm_qk_lowrank_batch_glm52_t4_pipeline = + ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch_glm52_t4"); + g_glm_value_project_q8_0_pipeline = + ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0"); + g_glm_value_project_q8_0_batch_heads_pipeline = + ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0_batch_heads"); + g_glm_value_project_q8_0_batch_heads_mma_pipeline = + ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0_batch_heads_mma"); + g_glm_attention_indexed_decode_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode"); + g_glm_attention_indexed_decode_split_group8_partial_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_partial"); + g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads"); + g_glm_attention_indexed_decode_split_group8_reduce_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce"); + g_glm_attention_indexed_decode_split_group8_reduce16_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce16"); + g_glm_attention_indexed_batch_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch"); + g_glm_attention_indexed_batch_group2_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_group2"); + g_glm_attention_indexed_batch_q2_group4_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_q2_group4"); + g_glm_attention_indexed_batch_group8_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_group8"); + g_glm_attention_indexed_batch_lora_group8_vec_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec"); + g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid"); + g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads"); + g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal"); + g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline = + ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads"); + g_glm_q4_k_pair_swiglu_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu_f32"); + g_glm_q4_k_pair_swiglu2_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_f32"); + g_glm_q4_k_pair_swiglu4_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu4_f32"); + g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_mapped_f32"); + g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_mapped_row_f32"); + g_glm_q2_k_pair_swiglu_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q2_K_pair_swiglu_f32"); + g_glm_q2_k_addr_pair_swiglu2_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_pair_swiglu2_f32"); + g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_pair_swiglu2_f32_masked"); + g_glm_q4_k_addr_pair_swiglu_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_pair_swiglu_f32"); + g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_pair_swiglu_f32_masked"); + g_glm_q2_k_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q2_K_down_f32"); + g_glm_q4_k_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_down_simd_f32"); + g_glm_q2_k_addr_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_down_f32"); + g_glm_q4_k_addr_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_down_simd_f32"); + g_glm_q5_k_pair_swiglu_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_f32"); + g_glm_q5_k_pair_swiglu_mapped_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_mapped_f32"); + g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_mapped_row_f32"); + g_glm_q5_k_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q5_K_down_f32"); + g_glm_q6_k_down_f32_pipeline = + ds4_gpu_get_pipeline("kernel_glm_q6_K_down_f32"); + g_dsv4_router_weights_batch_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_router_weights_batch"); + g_dsv4_hc_expand4_pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_hc_expand4"); + if (!g_dsv4_indexer_score_one_direct_pipeline || + !g_dsv4_compressor_store_one_pipeline || + !g_dsv4_sort_i32_rows_asc_pipeline || + !g_dsv4_indexed_attention_heads8_pipeline || + !g_dsv4_indexed_attention_heads8_rb16_pipeline || + !g_dsv4_softplus_sqrt_pipeline || + !g_dsv4_router_finalize_one_pipeline || + !g_dsv4_router_weights_one_pipeline || + !g_glm_router_select_one_pipeline || + !g_glm_kv_lora_rms_norm_pipeline || + !g_glm_k_b_project_pipeline || + !g_glm_store_compact_kv_pipeline || + !g_glm_qkv_norm_store_compact_kv_pipeline || + !g_glm_store_indexer_k_pipeline || + !g_glm_build_kv_cache_pipeline || + !g_glm_build_kv_cache_decode_group4_pipeline || + !g_glm_build_kv_cache_flash_pipeline || + !g_glm_attention_full_pipeline || + !g_glm_fill_selected_range_pipeline || + !g_glm_fill_selected_range_batch_pipeline || + !g_glm_indexer_rope_tail_pipeline || + !g_glm_indexer_score_one_pipeline || + !g_glm_indexer_score_one_direct_pipeline || + !g_glm_indexer_scores_batch_pipeline || + !g_glm_indexer_scores_tiled_pipeline || + !g_glm_indexer_scores_tiled_f32_pipeline || + !g_glm_qk_lowrank_pipeline || + !g_glm_qk_lowrank_glm52_pipeline || + !g_glm_qk_lowrank_glm52_sg_pipeline || + !g_glm_qk_lowrank_batch_pipeline || + !g_glm_qk_lowrank_batch_glm52_t4_pipeline || + !g_glm_value_project_q8_0_pipeline || + !g_glm_value_project_q8_0_batch_heads_pipeline || + !g_glm_value_project_q8_0_batch_heads_mma_pipeline || + !g_glm_attention_indexed_decode_pipeline || + !g_glm_attention_indexed_decode_split_group8_partial_pipeline || + !g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline || + !g_glm_attention_indexed_decode_split_group8_reduce_pipeline || + !g_glm_attention_indexed_decode_split_group8_reduce16_pipeline || + !g_glm_attention_indexed_batch_pipeline || + !g_glm_attention_indexed_batch_group2_pipeline || + !g_glm_attention_indexed_batch_q2_group4_pipeline || + !g_glm_attention_indexed_batch_group8_pipeline || + !g_glm_attention_indexed_batch_lora_group8_vec_pipeline || + !g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline || + !g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline || + !g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline || + !g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline || + !g_glm_q4_k_pair_swiglu_f32_pipeline || + !g_glm_q4_k_pair_swiglu2_f32_pipeline || + !g_glm_q4_k_pair_swiglu4_f32_pipeline || + !g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline || + !g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline || + !g_glm_q2_k_pair_swiglu_f32_pipeline || + !g_glm_q2_k_addr_pair_swiglu2_f32_pipeline || + !g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline || + !g_glm_q4_k_addr_pair_swiglu_f32_pipeline || + !g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline || + !g_glm_q2_k_down_f32_pipeline || + !g_glm_q4_k_down_f32_pipeline || + !g_glm_q2_k_addr_down_f32_pipeline || + !g_glm_q4_k_addr_down_f32_pipeline || + !g_glm_q5_k_pair_swiglu_f32_pipeline || + !g_glm_q5_k_pair_swiglu_mapped_f32_pipeline || + !g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline || + !g_glm_q5_k_down_f32_pipeline || + !g_glm_q6_k_down_f32_pipeline || + !g_dsv4_hc_expand4_pipeline) { + g_queue = nil; + g_device = nil; + return 0; + } + + g_initialized = 1; + } + + return 1; +} + +ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) { + if (!g_initialized && !ds4_gpu_init()) return NULL; + if (bytes == 0 || bytes > (uint64_t)NSUIntegerMax) return NULL; + + @autoreleasepool { + DS4MetalTensor *tensor = [DS4MetalTensor new]; + tensor.buffer = [g_device newBufferWithLength:(NSUInteger)bytes + options:MTLResourceStorageModeShared]; + if (!tensor.buffer) { + return NULL; + } + tensor.offset = 0; + tensor.bytes = bytes; + tensor.owner = 1; + uint64_t live_snap = 0; + uint64_t peak_snap = 0; + pthread_mutex_lock(&g_tensor_mu); + const int tracked = ds4_gpu_tensor_track_alloc_locked( + (__bridge const void *)tensor, + bytes, + &live_snap, + &peak_snap); + pthread_mutex_unlock(&g_tensor_mu); + if (!tracked) { + fprintf(stderr, "ds4: failed to track Metal tensor allocation\n"); + tensor.buffer = nil; + return NULL; + } + if (ds4_gpu_trace_allocs()) { + fprintf(stderr, + "ds4: Metal tensor alloc %.3f MiB live %.3f MiB peak %.3f MiB\n", + (double)bytes / (1024.0 * 1024.0), + (double)live_snap / (1024.0 * 1024.0), + (double)peak_snap / (1024.0 * 1024.0)); + } + return (__bridge_retained ds4_gpu_tensor *)tensor; + } +} + +ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) { + return ds4_gpu_tensor_alloc(bytes); +} + +int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes) { + (void)kv_cache_bytes; + (void)context_bytes; + return 0; +} + +ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes) { + if (!base) return NULL; + const DS4MetalTensor *base_obj = ds4_gpu_tensor_const_obj(base); + if (offset > base_obj.bytes || bytes > base_obj.bytes - offset) return NULL; + if (base_obj.offset > UINT64_MAX - offset) return NULL; + const uint64_t absolute_offset = base_obj.offset + offset; + if (absolute_offset > (uint64_t)NSUIntegerMax) return NULL; + + @autoreleasepool { + DS4MetalTensor *view = [DS4MetalTensor new]; + view.buffer = base_obj.buffer; + view.offset = absolute_offset; + view.bytes = bytes; + view.owner = 0; + pthread_mutex_lock(&g_tensor_mu); + const int tracked = ds4_gpu_tensor_track_view_locked((__bridge const void *)view); + pthread_mutex_unlock(&g_tensor_mu); + if (!tracked) { + fprintf(stderr, "ds4: failed to track Metal tensor view\n"); + view.buffer = nil; + return NULL; + } + return (__bridge_retained ds4_gpu_tensor *)view; + } +} + +void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor) { + if (!tensor) return; + @autoreleasepool { + uint8_t owner = 0; + uint64_t bytes = 0; + uint64_t live_snap = 0; + uint64_t peak_snap = 0; + if (!ds4_gpu_tensor_prepare_free(tensor, + &owner, + &bytes, + &live_snap, + &peak_snap)) { + return; + } + DS4MetalTensor *obj = (__bridge_transfer DS4MetalTensor *)tensor; + if (owner) { + if (ds4_gpu_trace_allocs()) { + fprintf(stderr, + "ds4: Metal tensor free %.3f MiB live %.3f MiB peak %.3f MiB\n", + (double)bytes / (1024.0 * 1024.0), + (double)live_snap / (1024.0 * 1024.0), + (double)peak_snap / (1024.0 * 1024.0)); + } + } + obj.buffer = nil; + obj.offset = 0; + obj.bytes = 0; + obj.owner = 0; + } +} + +uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor) { + if (!tensor) return 0; + const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); + return obj.bytes; +} + +void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor) { + if (!tensor) return NULL; + DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor); + return (uint8_t *)[obj.buffer contents] + obj.offset; +} + +int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count) { + if (!tensor || count > ds4_gpu_tensor_bytes(tensor) / sizeof(float)) return 0; + float *p = ds4_gpu_tensor_contents(tensor); + if (!p && count != 0) return 0; + for (uint64_t i = 0; i < count; i++) p[i] = value; + return 1; +} + +int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes) { + if (!tensor || (!data && bytes != 0)) return 0; + DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor); + if (offset > obj.bytes || bytes > obj.bytes - offset) return 0; + if (bytes != 0) { + memcpy((uint8_t *)[obj.buffer contents] + obj.offset + offset, data, (size_t)bytes); + } + return 1; +} + +int ds4_gpu_tensor_read(const ds4_gpu_tensor *tensor, uint64_t offset, void *data, uint64_t bytes) { + if (!tensor || (!data && bytes != 0)) return 0; + const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor); + if (offset > obj.bytes || bytes > obj.bytes - offset) return 0; + if (bytes != 0) { + memcpy(data, (const uint8_t *)[obj.buffer contents] + obj.offset + offset, (size_t)bytes); + } + return 1; +} + +int ds4_gpu_tensor_copy(ds4_gpu_tensor *dst, uint64_t dst_offset, + const ds4_gpu_tensor *src, uint64_t src_offset, + uint64_t bytes) { + if (!dst || !src) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + DS4MetalTensor *d = ds4_gpu_tensor_obj(dst); + const DS4MetalTensor *s = ds4_gpu_tensor_const_obj(src); + if (dst_offset > d.bytes || bytes > d.bytes - dst_offset) return 0; + if (src_offset > s.bytes || bytes > s.bytes - src_offset) return 0; + if (bytes == 0) return 1; + if (!g_batch_cb) return 0; + + ds4_gpu_close_batch_encoder(); + g_batch_has_work = YES; + id blit = [g_batch_cb blitCommandEncoder]; + if (!blit) return 0; + [blit copyFromBuffer:s.buffer + sourceOffset:(NSUInteger)(s.offset + src_offset) + toBuffer:d.buffer + destinationOffset:(NSUInteger)(d.offset + dst_offset) + size:(NSUInteger)bytes]; + [blit endEncoding]; + return 1; +} + +int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset, + const ds4_gpu_tensor *src, uint64_t src_offset, + uint64_t count) { + if (!dst || !src) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + DS4MetalTensor *d = ds4_gpu_tensor_obj(dst); + const DS4MetalTensor *s = ds4_gpu_tensor_const_obj(src); + if (count == 0) return 1; + if (count > UINT64_MAX / sizeof(float) || + count > UINT64_MAX / sizeof(uint16_t)) { + return 0; + } + const uint64_t src_bytes = count * sizeof(float); + const uint64_t dst_bytes = count * sizeof(uint16_t); + if (src_offset > s.bytes || src_bytes > s.bytes - src_offset || + dst_offset > d.bytes || dst_bytes > d.bytes - dst_offset) { + return 0; + } + + @autoreleasepool { + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + uint64_t done = 0; + int ok = 1; + while (done < count && ok) { + uint64_t chunk64 = count - done; + if (chunk64 > UINT32_MAX) chunk64 = UINT32_MAX; + const uint32_t chunk = (uint32_t)chunk64; + ok = ds4_gpu_encode_cpy_f32_f16_1d( + cb, + s.buffer, + (NSUInteger)(s.offset + src_offset + done * sizeof(float)), + d.buffer, + (NSUInteger)(d.offset + dst_offset + done * sizeof(uint16_t)), + chunk); + done += chunk; + } + if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "tensor f32 to f16 copy"); + return ok; + } +} + +int ds4_gpu_pack_slot_rows_f32_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *slots, + uint32_t n_rows, + uint32_t width, + uint32_t n_slots, + uint32_t slot_cap) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !slots || n_rows == 0 || width == 0 || n_slots == 0 || + slot_cap == 0 || n_rows > slot_cap) { + return 0; + } + + @autoreleasepool { + id slotsbuf = ds4_gpu_tensor_buffer(slots); + id outbuf = ds4_gpu_tensor_buffer(out); + uint64_t row_bytes = 0; + uint64_t slot_plane_bytes = 0; + uint64_t slots_bytes = 0; + uint64_t out_rows = 0; + uint64_t out_bytes = 0; + if ((uint64_t)width > UINT64_MAX / sizeof(float)) return 0; + row_bytes = (uint64_t)width * sizeof(float); + if ((uint64_t)slot_cap > UINT64_MAX / row_bytes) return 0; + slot_plane_bytes = (uint64_t)slot_cap * row_bytes; + if ((uint64_t)n_slots > UINT64_MAX / slot_plane_bytes) return 0; + slots_bytes = (uint64_t)n_slots * slot_plane_bytes; + if ((uint64_t)n_rows > UINT64_MAX / n_slots) return 0; + out_rows = (uint64_t)n_rows * n_slots; + if (out_rows > UINT64_MAX / row_bytes) return 0; + out_bytes = out_rows * row_bytes; + if (!slotsbuf || !outbuf || + ds4_gpu_tensor_bytes(slots) < slots_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal slot-row pack received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb, + slotsbuf, + ds4_gpu_tensor_offset(slots), + outbuf, + ds4_gpu_tensor_offset(out), + width, + n_slots, + n_rows, + sizeof(float), + slot_plane_bytes, + row_bytes, + row_bytes, + (uint64_t)n_slots * row_bytes)) { + return 0; + } + if (!ds4_gpu_finish_command_buffer(cb, owned, "slot-row pack")) return 0; + } + + return 1; +} + +int ds4_gpu_begin_commands(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (g_batch_cb) return 0; + g_batch_cb = ds4_gpu_new_command_buffer(); + g_batch_has_work = NO; + if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); + return g_batch_cb != nil; +} + +int ds4_gpu_flush_encoder(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb) return 0; + ds4_gpu_close_batch_encoder(); + return 1; +} + +int ds4_gpu_flush_commands(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb) return 0; + + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + g_batch_cb = nil; + g_batch_has_work = NO; + [cb commit]; + [g_pending_cbs addObject:cb]; + ds4_gpu_stream_expert_cache_note_batch_committed(); + + g_batch_cb = ds4_gpu_new_command_buffer(); + g_batch_has_work = NO; + if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); + if (!g_batch_cb) { + (void)ds4_gpu_wait_pending_command_buffers("command batch"); + [g_transient_buffers removeAllObjects]; + return 0; + } + return 1; +} + +int ds4_gpu_commands_active(void) { + return g_batch_cb != nil; +} + +static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label) { + const char *what = label ? label : "streaming expert cache in-flight"; + if (g_batch_cb && ds4_gpu_flush_commands() == 0) return 0; + if ([g_pending_cbs count] != 0 && + ds4_gpu_wait_pending_command_buffers(what) == 0) { + return 0; + } + return 1; +} + +int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) { + if (!event_value) return 0; + *event_value = 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb) return 0; + + if (@available(macOS 12.0, *)) { + if (!g_selected_readback_event) { + g_selected_readback_event = [g_device newSharedEvent]; + if (!g_selected_readback_event) { + fprintf(stderr, "ds4: failed to create Metal shared event for selected-id overlap\n"); + return 0; + } + } + + ds4_gpu_close_batch_encoder(); + const uint64_t value = ++g_selected_readback_event_value; + [g_batch_cb encodeSignalEvent:g_selected_readback_event value:value]; + g_batch_has_work = YES; + *event_value = value; + return 1; + } + + fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n"); + return 0; +} + +int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb || event_value == 0) return 0; + + if (@available(macOS 12.0, *)) { + if (!g_selected_readback_event) return 0; + + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + g_batch_cb = nil; + g_batch_has_work = NO; + [cb commit]; + ds4_gpu_stream_expert_cache_note_batch_committed(); + + const char *what = label ? label : "selected-id overlap"; + const BOOL signaled = + [g_selected_readback_event waitUntilSignaledValue:event_value timeoutMS:60000]; + [g_pending_cbs addObject:cb]; + if (!signaled) { + fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", what); + (void)ds4_gpu_wait_pending_command_buffers(what); + [g_transient_buffers removeAllObjects]; + return 0; + } + if (cb.status == MTLCommandBufferStatusError) { + fprintf(stderr, "ds4: Metal %s failed: %s\n", + what, + [[cb.error localizedDescription] UTF8String]); + (void)ds4_gpu_wait_pending_command_buffers(what); + [g_transient_buffers removeAllObjects]; + return 0; + } + + g_batch_cb = ds4_gpu_new_command_buffer(); + g_batch_has_work = NO; + if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); + if (!g_batch_cb) { + (void)ds4_gpu_wait_pending_command_buffers(what); + [g_transient_buffers removeAllObjects]; + return 0; + } + return 1; + } + + fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n"); + return 0; +} + +/* + * Tensor-parallel gates. + * + * A TP gate is a mid-command-stream rendezvous with the peer machine: the + * kernels ahead of the gate leave a partial block output in a slab slot, + * the GPU signals g_tp_gpu_event, and the pre-encoded combine kernel waits + * on g_tp_cpu_event. A dedicated service thread bridges the two: it spins + * until the GPU reaches the gate, runs the transport exchange (RDMA WRITE + * plus flag poll, or a TCP write/read pair — behind the callback), and + * CPU-signals the release. On exchange failure the release is signaled + * anyway so the GPU never deadlocks; the failure latches in g_tp_failed + * and the eval aborts at the next command-buffer boundary. + * + * Gate sequence values increase monotonically per encoded gate. Both ranks + * encode the identical graph, so the values agree by construction and slots + * never need resetting between tokens. + */ +typedef struct { + uint32_t layer; + uint32_t gate; + uint32_t rows; /* 0 = row gate; >0 = verify-block batch gate */ + uint32_t event_arrival; + uint64_t seq; + /* Big batch gates (prefill): exchange big_bytes from big_out into + * big_in directly (CPU-visible bounce buffers), bypassing the slab. */ + const void *big_out; + void *big_in; + uint64_t big_bytes; +} ds4_gpu_tp_request; + +enum { DS4_GPU_TP_QUEUE = 1024 }; + +static id g_tp_gpu_event; /* GPU -> service thread */ +static id g_tp_cpu_event; /* service thread -> GPU */ +/* Batch (verify-block) gates run on their own sequence space and release + * event: the row-gate seq feeds the RDMA pre-posted recv accounting, which + * requires consecutive values, and a shared release event would make a + * small batch value satisfy waits armed against the larger row seq. */ +static id g_tp_batch_gpu_event; +static id g_tp_batch_cpu_event; +static uint64_t g_tp_batch_seq; +/* Batch flag values are tagged so a stale row-gate seq in the reused FFN + * flag word can never satisfy a batch arrival spin (and vice versa). */ +#define DS4_TP_BATCH_FLAG_TAG 0x80000000u +/* Expert-ownership split parameters for routed kernels. World 1 means TP is + * not bound; world 2 assigns each rank one contiguous expert range. */ +static int32_t g_tp_split_rank; +static int32_t g_tp_split_world = 1; +static int32_t g_tp_session_batch_mode; + +static int ds4_gpu_tp_world_is_two(void) { + return g_tp_split_world == 2; +} + +/* Return the contiguous routed-expert range backed by this process. Rank 1 + * owns the high range and receives any odd-count remainder. */ +static void ds4_gpu_tp_expert_range(uint32_t n_total_expert, + uint32_t *first_expert, + uint32_t *n_expert) { + *first_expert = 0; + *n_expert = n_total_expert; + if (g_tp_split_world != 2) return; + + const uint32_t low_experts = n_total_expert / 2u; + if (g_tp_split_rank == 1) { + *first_expert = low_experts; + *n_expert = n_total_expert - low_experts; + } else { + *n_expert = low_experts; + } +} + +/* Attention head split for GLM batch prefill: each rank computes a + * contiguous half of the heads in the qk-low / attention-lora / + * value-project batch kernels; the caller zeroes the unowned head range + * of the heads buffer and combines the attn-output partials over the + * TP big-gate exchange. */ +static int32_t g_tp_attn_head_split; + +void ds4_gpu_tp_set_attn_head_split(int enabled) { + g_tp_attn_head_split = enabled ? 1 : 0; +} + +static void ds4_gpu_tp_attn_head_range(uint32_t n_head, + uint32_t group, + uint32_t *head_base, + uint32_t *head_count) { + *head_base = 0; + *head_count = n_head; + if (!g_tp_attn_head_split || g_tp_split_world != 2) return; + const uint32_t half = n_head / 2u; + if (half == 0u || (half % group) != 0u || (n_head % 2u) != 0u) return; + *head_count = half; + *head_base = g_tp_split_rank == 1 ? half : 0u; +} +/* Flag gates (DS4_TP_FLAG_GATES): the GPU publishes gate arrival by storing + * the sequence number into a slab word instead of signaling the shared + * event; the service thread spin-reads it from shared memory, which wakes + * hundreds of microseconds earlier than signaledValue polling. The + * CPU->GPU release direction stays on the shared event. */ +static bool g_tp_flag_gates; +static id g_tp_slab_buffer; +static NSUInteger g_tp_slab_buffer_off; +static volatile uint32_t *g_tp_gpu_flags; /* CPU view of the flag words */ +static uint64_t g_tp_gpu_flags_off; +static uint64_t g_tp_seq; +static ds4_gpu_tp_exchange_fn g_tp_exchange_fn; +static ds4_gpu_tp_batch_exchange_fn g_tp_batch_exchange_fn; +static ds4_gpu_tp_big_exchange_fn g_tp_big_exchange_fn; + +void ds4_gpu_tp_set_big_exchange(ds4_gpu_tp_big_exchange_fn fn) { + g_tp_big_exchange_fn = fn; +} + +static void *g_tp_exchange_ud; +static pthread_t g_tp_thread; +static int g_tp_thread_running; +static int g_tp_shutdown; +static int g_tp_failed_flag; +static pthread_mutex_t g_tp_mutex = PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t g_tp_cond = PTHREAD_COND_INITIALIZER; +static ds4_gpu_tp_request g_tp_queue[DS4_GPU_TP_QUEUE]; +static uint32_t g_tp_queue_head; +static uint32_t g_tp_queue_count; + +static uint64_t g_tp_stat_gates; +static double g_tp_stat_gpu_wait_ms; +static double g_tp_stat_exchange_ms; + +/* GPU keep-alive (see kernel_dsv4_tp_keepalive): its own queue and thread, + * alive exactly as long as the TP gate machinery. */ +static id g_tp_keepalive_queue; +static id g_tp_keepalive_buffer; +static pthread_t g_tp_keepalive_thread; +static int g_tp_keepalive_running; + +/* Nonzero while a verify block runs: the GPU is genuinely busy + * there, so the keep-alive is a pure parasite (~2.3ms per 5-row block + * measured against the single-machine verify). */ +static volatile int g_tp_keepalive_paused; + +void ds4_gpu_tp_keepalive_pause(int paused) { + g_tp_keepalive_paused = paused; +} + +void ds4_gpu_tp_set_session_batch_mode(int enabled) { + g_tp_session_batch_mode = enabled ? 1 : 0; +} + +static uint32_t ds4_gpu_tp_keepalive_tgs_from_env(void) { + uint32_t ka_tgs = 1; + const char *tgs_env = getenv("DS4_TP_KEEPALIVE_TGS"); + if (tgs_env) { + int v = atoi(tgs_env); + if (v > 0 && v <= 2048) ka_tgs = (uint32_t)v; + } + return ka_tgs; +} + +static void *ds4_gpu_tp_keepalive_thread(void *arg) { + (void)arg; + /* Swept on the M5 Max pair: too few iterations lets clocks sag. Current + * TP split-resident Flash runs show 1.2M is a small Q4/Q2 decode win over + * 800k, while two threadgroups waste work. */ + uint32_t iters = 1200000; + const char *env = getenv("DS4_TP_KEEPALIVE_ITERS"); + if (env) iters = (uint32_t)atoi(env); + /* One ALU-only threadgroup keeps the GPU from power-gating but does + * not push the frequency governor; solo-vs-engine kernel gaps + * (~1.7x) suggest decode runs well below max clocks. More TGs raise + * apparent utilization without eating memory bandwidth. */ + uint32_t ka_tgs = ds4_gpu_tp_keepalive_tgs_from_env(); + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_tp_keepalive"); + if (!pipeline) { + fprintf(stderr, "ds4: TP keep-alive pipeline missing\n"); + return NULL; + } + while (!g_tp_shutdown) { + if (g_tp_keepalive_paused) { + usleep(200); + continue; + } + @autoreleasepool { + id cb = [g_tp_keepalive_queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:pipeline]; + [enc setBuffer:g_tp_keepalive_buffer offset:0 atIndex:0]; + [enc setBytes:&iters length:sizeof(iters) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)ka_tgs, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + } + } + return NULL; +} + +static void *ds4_gpu_tp_service_thread(void *arg) { + (void)arg; + const bool profile = getenv("DS4_TP_GATE_PROFILE") != NULL; + while (1) { + pthread_mutex_lock(&g_tp_mutex); + while (g_tp_queue_count == 0 && !g_tp_shutdown) + pthread_cond_wait(&g_tp_cond, &g_tp_mutex); + if (g_tp_shutdown && g_tp_queue_count == 0) { + pthread_mutex_unlock(&g_tp_mutex); + break; + } + ds4_gpu_tp_request req = g_tp_queue[g_tp_queue_head]; + g_tp_queue_head = (g_tp_queue_head + 1) % DS4_GPU_TP_QUEUE; + g_tp_queue_count--; + pthread_mutex_unlock(&g_tp_mutex); + + /* Wait for the GPU to reach this gate. Tight spin: gate arrival is + * on the decode critical path and normally tens of microseconds + * out; yielding here measurably delays the release wake-up. */ + const double t0 = profile ? ds4_gpu_now_ms() : 0.0; + uint32_t spins = 0; + if (req.big_bytes > 0) { + /* Big gates always signal arrival through the batch shared + * event (see ds4_gpu_tp_big_gate_kick): the event completion + * semantics are what guarantee the bounce payload is visible + * before the exchange reads it. */ + while (g_tp_batch_gpu_event.signaledValue < req.seq) { + if (g_tp_shutdown) break; + if (++spins > (1u << 16)) sched_yield(); + } + } else if (!req.event_arrival) { + const uint32_t slot = req.layer * 2u + req.gate; + uint32_t want = (uint32_t)req.seq; + if (req.rows > 0) + want = DS4_TP_BATCH_FLAG_TAG | (uint32_t)req.seq; + while (__atomic_load_n(&g_tp_gpu_flags[slot], __ATOMIC_ACQUIRE) != want) { + if (g_tp_shutdown) break; + if (++spins > (1u << 20)) { + sched_yield(); + spins = 0; + } + } + } else if (req.rows > 0) { + while (g_tp_batch_gpu_event.signaledValue < req.seq) { + if (g_tp_shutdown) break; + if (++spins > (1u << 16)) sched_yield(); + } + } else { + while (g_tp_gpu_event.signaledValue < req.seq) { + if (g_tp_shutdown) break; + if (++spins > (1u << 16)) sched_yield(); + } + } + const double t1 = profile ? ds4_gpu_now_ms() : 0.0; + int ok = 0; + if (!g_tp_shutdown && !g_tp_failed_flag) { + if (req.big_bytes > 0) { + if (g_tp_big_exchange_fn) + ok = g_tp_big_exchange_fn(g_tp_exchange_ud, req.layer, + req.seq, req.big_out, + req.big_in, req.big_bytes); + } else if (req.rows > 0) { + if (g_tp_batch_exchange_fn) + ok = g_tp_batch_exchange_fn(g_tp_exchange_ud, req.layer, + req.rows, req.seq); + } else if (g_tp_exchange_fn) { + ok = g_tp_exchange_fn(g_tp_exchange_ud, req.layer, req.gate, + req.seq); + } + } + if (!ok && !g_tp_shutdown) { + if (!g_tp_failed_flag) + fprintf(stderr, "ds4: TP gate exchange failed (layer %u gate %u seq %llu)\n", + req.layer, req.gate, (unsigned long long)req.seq); + g_tp_failed_flag = 1; + } + /* Release the GPU even on failure so end_commands can drain. */ + if (req.rows > 0) g_tp_batch_cpu_event.signaledValue = req.seq; + else g_tp_cpu_event.signaledValue = req.seq; + if (profile) { + g_tp_stat_gpu_wait_ms += t1 - t0; + g_tp_stat_exchange_ms += ds4_gpu_now_ms() - t1; + if (++g_tp_stat_gates % 860 == 0) { + fprintf(stderr, + "ds4: TP gates %llu: avg gpu-wait %.1f us, avg exchange %.1f us\n", + (unsigned long long)g_tp_stat_gates, + g_tp_stat_gpu_wait_ms / (double)g_tp_stat_gates * 1000.0, + g_tp_stat_exchange_ms / (double)g_tp_stat_gates * 1000.0); + } + } + } + return NULL; +} + +int ds4_gpu_tp_init(uint32_t rank, + ds4_gpu_tensor *slab, uint64_t gpu_flags_off, + ds4_gpu_tp_exchange_fn fn, void *ud) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (g_tp_thread_running || rank > 1) return 0; + g_tp_split_rank = (int32_t)rank; + g_tp_split_world = 2; + g_tp_slab_buffer = slab ? ds4_gpu_tensor_buffer(slab) : nil; + g_tp_slab_buffer_off = slab ? ds4_gpu_tensor_offset(slab) : 0; + g_tp_gpu_flags_off = gpu_flags_off; + g_tp_gpu_flags = slab ? + (volatile uint32_t *)((uint8_t *)ds4_gpu_tensor_contents(slab) + gpu_flags_off) : NULL; + /* Flag arrival is the default: the slab-word publish detects in ~1-3us + * where signaledValue polling costs 10-20, worth +2.3 t/s on the pair + * (A/B 2026-07-06, byte-identical output). DS4_TP_EVENT_GATES falls + * back to the shared-event arrival path. */ + g_tp_flag_gates = g_tp_gpu_flags != NULL && getenv("DS4_TP_EVENT_GATES") == NULL; + g_tp_gpu_event = [g_device newSharedEvent]; + g_tp_cpu_event = [g_device newSharedEvent]; + g_tp_batch_gpu_event = [g_device newSharedEvent]; + g_tp_batch_cpu_event = [g_device newSharedEvent]; + if (!g_tp_gpu_event || !g_tp_cpu_event || + !g_tp_batch_gpu_event || !g_tp_batch_cpu_event) { + fprintf(stderr, "ds4: failed to create TP shared events\n"); + return 0; + } + g_tp_exchange_fn = fn; + g_tp_exchange_ud = ud; + g_tp_seq = 0; + g_tp_batch_seq = 0; + g_tp_shutdown = 0; + g_tp_failed_flag = 0; + g_tp_queue_head = 0; + g_tp_queue_count = 0; + pthread_attr_t attr; + pthread_attr_init(&attr); + pthread_attr_set_qos_class_np(&attr, QOS_CLASS_USER_INTERACTIVE, 0); + if (pthread_create(&g_tp_thread, &attr, ds4_gpu_tp_service_thread, NULL) != 0) { + pthread_attr_destroy(&attr); + fprintf(stderr, "ds4: failed to start TP gate service thread\n"); + return 0; + } + pthread_attr_destroy(&attr); + g_tp_thread_running = 1; + if (getenv("DS4_TP_NO_KEEPALIVE") == NULL) { + uint32_t ka_tgs = ds4_gpu_tp_keepalive_tgs_from_env(); + g_tp_keepalive_queue = [g_device newCommandQueue]; + g_tp_keepalive_buffer = [g_device newBufferWithLength:(NSUInteger)ka_tgs * 256u * sizeof(float) + options:MTLResourceStorageModeShared]; + if (g_tp_keepalive_queue && g_tp_keepalive_buffer && + pthread_create(&g_tp_keepalive_thread, NULL, + ds4_gpu_tp_keepalive_thread, NULL) == 0) { + g_tp_keepalive_running = 1; + } else { + fprintf(stderr, "ds4: TP keep-alive setup failed (continuing without)\n"); + } + } + return 1; +} + +void ds4_gpu_tp_shutdown(void) { + if (!g_tp_thread_running) return; + pthread_mutex_lock(&g_tp_mutex); + g_tp_shutdown = 1; + pthread_cond_broadcast(&g_tp_cond); + pthread_mutex_unlock(&g_tp_mutex); + pthread_join(g_tp_thread, NULL); + g_tp_thread_running = 0; + if (g_tp_keepalive_running) { + pthread_join(g_tp_keepalive_thread, NULL); + g_tp_keepalive_running = 0; + g_tp_keepalive_queue = nil; + g_tp_keepalive_buffer = nil; + } + g_tp_exchange_fn = NULL; + g_tp_batch_exchange_fn = NULL; + g_tp_exchange_ud = NULL; + g_tp_split_rank = 0; + g_tp_split_world = 1; + g_tp_session_batch_mode = 0; +} + +void ds4_gpu_tp_suspend_expert_sharding(int suspend) { + if (!g_tp_thread_running) return; + g_tp_split_world = suspend ? 1 : 2; +} + +int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate) { + if (!g_batch_cb) { + fprintf(stderr, "ds4: TP gate encode without an open command batch (layer %u gate %u)\n", + layer, gate); + return 0; + } + if (!g_tp_thread_running) { + fprintf(stderr, "ds4: TP gate encode without the gate service (layer %u)\n", layer); + return 0; + } + const uint64_t seq = ++g_tp_seq; + const bool event_arrival = g_tp_session_batch_mode || !g_tp_flag_gates; + if (!event_arrival) { + /* Publish arrival through the slab word; the buffer hazard against + * the partial-output kernels orders the store after the payload. */ + const uint32_t slot = layer * 2u + gate; + const uint32_t value = (uint32_t)seq; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) return 0; + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_tp_flag_set"); + if (!pipeline) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBuffer:g_tp_slab_buffer + offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u) + atIndex:0]; + [enc setBytes:&value length:sizeof(value) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + ds4_gpu_close_batch_encoder(); + } else { + ds4_gpu_close_batch_encoder(); + [g_batch_cb encodeSignalEvent:g_tp_gpu_event value:seq]; + } + [g_batch_cb encodeWaitForEvent:g_tp_cpu_event value:seq]; + pthread_mutex_lock(&g_tp_mutex); + if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) { + pthread_mutex_unlock(&g_tp_mutex); + fprintf(stderr, "ds4: TP gate queue overflow\n"); + return 0; + } + uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE; + g_tp_queue[tail].layer = layer; + g_tp_queue[tail].gate = gate; + g_tp_queue[tail].rows = 0; + g_tp_queue[tail].event_arrival = event_arrival ? 1u : 0u; + g_tp_queue[tail].seq = seq; + g_tp_queue[tail].big_out = NULL; + g_tp_queue[tail].big_in = NULL; + g_tp_queue[tail].big_bytes = 0; + g_tp_queue_count++; + pthread_cond_signal(&g_tp_cond); + pthread_mutex_unlock(&g_tp_mutex); + return 1; +} + +void ds4_gpu_tp_set_batch_exchange(ds4_gpu_tp_batch_exchange_fn fn) { + g_tp_batch_exchange_fn = fn; +} + +/* Verify-block batch gate: same arrival/release machinery as the row gate + * (the FFN flag word and event pair are reused — a decode gate and a batch + * gate are never in flight together, and seq values stay globally unique), + * but the service thread runs the multi-row exchange callback. */ +int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows) { + if (!g_batch_cb) return 0; + if (!g_tp_thread_running || rows == 0) return 0; + const uint64_t seq = ++g_tp_batch_seq; + const bool event_arrival = g_tp_session_batch_mode || !g_tp_flag_gates; + if (!event_arrival) { + const uint32_t slot = layer * 2u + 1u; /* FFN gate slot */ + const uint32_t value = DS4_TP_BATCH_FLAG_TAG | (uint32_t)seq; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) return 0; + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_tp_flag_set"); + if (!pipeline) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBuffer:g_tp_slab_buffer + offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u) + atIndex:0]; + [enc setBytes:&value length:sizeof(value) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + ds4_gpu_close_batch_encoder(); + } else { + ds4_gpu_close_batch_encoder(); + [g_batch_cb encodeSignalEvent:g_tp_batch_gpu_event value:seq]; + } + [g_batch_cb encodeWaitForEvent:g_tp_batch_cpu_event value:seq]; + pthread_mutex_lock(&g_tp_mutex); + if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) { + pthread_mutex_unlock(&g_tp_mutex); + fprintf(stderr, "ds4: TP gate queue overflow\n"); + return 0; + } + uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE; + g_tp_queue[tail].layer = layer; + g_tp_queue[tail].gate = 1u; /* FFN */ + g_tp_queue[tail].rows = rows; + g_tp_queue[tail].event_arrival = event_arrival ? 1u : 0u; + g_tp_queue[tail].seq = seq; + g_tp_queue[tail].big_out = NULL; + g_tp_queue[tail].big_in = NULL; + g_tp_queue[tail].big_bytes = 0; + g_tp_queue_count++; + pthread_cond_signal(&g_tp_cond); + pthread_mutex_unlock(&g_tp_mutex); + return 1; +} + +/* Prefill batch gate kick: same seq space and release event as the verify + * batch gate, but the service thread exchanges big_bytes directly between + * the two shared bounce buffers instead of slab slots. The kick only + * publishes the GPU arrival marker and queues the exchange; the caller + * encodes the release wait later through ds4_gpu_tp_big_gate_wait, which + * lets it interleave more GPU work with the wire exchange. Arrival always + * uses the batch shared event, NOT the flag word: a flag write carries no + * memory-visibility guarantee for the payload buffer, and once the GPU + * keeps running past the kick (no event wait right behind it) the service + * thread can observe the flag before the producing kernels' stores reach + * CPU-visible memory (measured: stale rows in the first sub-kick). The + * shared-event signal only fires after every preceding command completes, + * which is exactly the payload ordering the exchange needs; the ~10 us + * slower arrival detection is noise against a multi-ms exchange. */ +uint64_t ds4_gpu_tp_big_gate_kick(uint32_t layer, uint32_t rows, + const ds4_gpu_tensor *out_t, + ds4_gpu_tensor *in_t, + uint64_t bytes) { + if (!g_batch_cb) return 0; + if (!g_tp_thread_running || rows == 0 || bytes == 0) return 0; + const void *out_ptr = ds4_gpu_tensor_contents((ds4_gpu_tensor *)out_t); + void *in_ptr = ds4_gpu_tensor_contents(in_t); + if (!out_ptr || !in_ptr) { + fprintf(stderr, "ds4: TP big gate needs CPU-visible bounce buffers\n"); + return 0; + } + const uint64_t seq = ++g_tp_batch_seq; + ds4_gpu_close_batch_encoder(); + [g_batch_cb encodeSignalEvent:g_tp_batch_gpu_event value:seq]; + pthread_mutex_lock(&g_tp_mutex); + if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) { + pthread_mutex_unlock(&g_tp_mutex); + fprintf(stderr, "ds4: TP gate queue overflow\n"); + return 0; + } + uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE; + g_tp_queue[tail].layer = layer; + g_tp_queue[tail].gate = 1u; + g_tp_queue[tail].rows = rows; + g_tp_queue[tail].event_arrival = 1u; + g_tp_queue[tail].seq = seq; + g_tp_queue[tail].big_out = out_ptr; + g_tp_queue[tail].big_in = in_ptr; + g_tp_queue[tail].big_bytes = bytes; + g_tp_queue_count++; + pthread_cond_signal(&g_tp_cond); + pthread_mutex_unlock(&g_tp_mutex); + return seq; +} + +/* Encode the GPU-side release wait for a previously kicked big gate. The + * batch release event is monotonic and the service thread completes queued + * exchanges in kick order, so waiting on the LAST kicked seq of a stage + * also covers every earlier kick. */ +int ds4_gpu_tp_big_gate_wait(uint64_t seq) { + if (!g_batch_cb || seq == 0) return 0; + ds4_gpu_close_batch_encoder(); + [g_batch_cb encodeWaitForEvent:g_tp_batch_cpu_event value:seq]; + return 1; +} + +int ds4_gpu_tp_big_gate_encode(uint32_t layer, uint32_t rows, + const ds4_gpu_tensor *out_t, + ds4_gpu_tensor *in_t, + uint64_t bytes) { + const uint64_t seq = ds4_gpu_tp_big_gate_kick(layer, rows, out_t, in_t, bytes); + if (seq == 0) return 0; + return ds4_gpu_tp_big_gate_wait(seq); +} + +int ds4_gpu_tp_failed(void) { + return g_tp_failed_flag; +} + +int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (event_value == 0) return 0; + + if (@available(macOS 12.0, *)) { + if (!g_selected_readback_event) return 0; + + const char *what = label ? label : "selected-id readback"; + const BOOL signaled = + [g_selected_readback_event waitUntilSignaledValue:event_value timeoutMS:60000]; + if (!signaled) { + fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", what); + return 0; + } + return 1; + } + + fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n"); + return 0; +} + +static int ds4_gpu_signal_batch_and_wait_event(const char *label) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!g_batch_cb) return 0; + + if (@available(macOS 12.0, *)) { + if (!g_selected_readback_event) { + g_selected_readback_event = [g_device newSharedEvent]; + if (!g_selected_readback_event) { + fprintf(stderr, "ds4: failed to create Metal shared event for %s\n", + label ? label : "selected-id readback"); + return 0; + } + } + + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + g_batch_cb = nil; + const uint64_t value = ++g_selected_readback_event_value; + [cb encodeSignalEvent:g_selected_readback_event value:value]; + g_batch_has_work = YES; + [cb commit]; + ds4_gpu_stream_expert_cache_note_batch_committed(); + + const BOOL signaled = [g_selected_readback_event waitUntilSignaledValue:value timeoutMS:60000]; + [g_pending_cbs addObject:cb]; + if (!signaled) { + fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", + label ? label : "selected-id readback"); + (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback"); + [g_transient_buffers removeAllObjects]; + return 0; + } + if (cb.status == MTLCommandBufferStatusError) { + fprintf(stderr, "ds4: Metal %s failed: %s\n", + label ? label : "selected-id readback", + [[cb.error localizedDescription] UTF8String]); + (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback"); + [g_transient_buffers removeAllObjects]; + return 0; + } + + g_batch_cb = ds4_gpu_new_command_buffer(); + g_batch_has_work = NO; + if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created(); + if (!g_batch_cb) { + (void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback"); + [g_transient_buffers removeAllObjects]; + return 0; + } + return 1; + } else { + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + g_batch_cb = nil; + g_batch_has_work = NO; + if (ds4_gpu_finish_command_buffer(cb, 1, label ? label : "selected-id readback") == 0) { + return 0; + } + return ds4_gpu_begin_commands(); + } +} + +int ds4_gpu_end_commands(void) { + if (!g_batch_cb) return 0; + ds4_gpu_close_batch_encoder(); + id cb = g_batch_cb; + g_batch_cb = nil; + g_batch_has_work = NO; + g_stream_expert_cache_owned_seq = g_stream_expert_cache_batch_seq; + g_stream_expert_cache_batch_seq = 0; + return ds4_gpu_finish_command_buffer(cb, 1, "command batch"); +} + +static int ds4_gpu_flash_attn_stage_profile_boundary( + id __strong *cbp, + const char *mode, + const char *stage, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t n_keys, + uint32_t n_head, + uint32_t head_dim, + uint32_t window, + uint32_t ratio, + double *stage_t0) { + if (!cbp || !*cbp || !stage_t0 || !stage) return 0; + if (ds4_gpu_end_commands() == 0) return 0; + + const double now_ms = ds4_gpu_now_ms(); + const char *filter = getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE_FILTER"); + const int print_stage = + !filter || !filter[0] || + strstr(stage, filter) != NULL || + (mode && strstr(mode, filter) != NULL); + if (print_stage) { + fprintf(stderr, + "ds4: Metal FlashAttention prefill stage mode=%s tokens=%u comp=%u " + "keys=%u heads=%u dim=%u window=%u ratio=%u %s=%.3f ms\n", + mode ? mode : "unknown", + n_tokens, + n_comp, + n_keys, + n_head, + head_dim, + window, + ratio, + stage, + now_ms - *stage_t0); + } + *stage_t0 = now_ms; + + if (ds4_gpu_begin_commands() == 0) return 0; + int owned = 0; + *cbp = ds4_gpu_command_buffer(&owned); + return *cbp != nil && owned == 0; +} + +int ds4_gpu_synchronize(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (g_batch_cb) return ds4_gpu_end_commands(); + if ([g_pending_cbs count] != 0) { + int ok = ds4_gpu_wait_pending_command_buffers("synchronize"); + [g_transient_buffers removeAllObjects]; + ds4_gpu_model_buffer_cache_maybe_evict("synchronize"); + return ok; + } + + id cb = ds4_gpu_new_command_buffer(); + if (!cb) return 0; + return ds4_gpu_finish_command_buffer(cb, 1, "synchronize"); +} + +void ds4_gpu_cleanup(void) { + if (!g_initialized) return; + + @autoreleasepool { + if (g_batch_cb) { + ds4_gpu_close_batch_encoder(); + [g_batch_cb commit]; + [g_batch_cb waitUntilCompleted]; + g_batch_cb = nil; + if (g_stream_expert_cache_batch_seq > g_stream_expert_cache_done_seq) { + g_stream_expert_cache_done_seq = g_stream_expert_cache_batch_seq; + } + g_stream_expert_cache_batch_seq = 0; + } + (void)ds4_gpu_wait_pending_command_buffers("cleanup"); + if (ds4_gpu_stream_expert_timing_summary_enabled() && + getenv("DS4_METAL_MEMORY_REPORT") == NULL) { + ds4_gpu_print_memory_report("at cleanup"); + } + g_selected_readback_event = nil; + g_selected_readback_event_value = 0; + [g_transient_buffers removeAllObjects]; + ds4_gpu_stream_expert_pread_pool_shutdown(); + ds4_gpu_stream_expert_cache_clear_all(1); + for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) { + g_stream_expert_cache_gate_addr_buffers[layer] = nil; + g_stream_expert_cache_up_addr_buffers[layer] = nil; + g_stream_expert_cache_down_addr_buffers[layer] = nil; + g_stream_compact_gate_addr_buffers[layer] = nil; + g_stream_compact_up_addr_buffers[layer] = nil; + g_stream_compact_down_addr_buffers[layer] = nil; + g_stream_compact_selected_buffers[layer] = nil; + g_stream_selected_id_buffers[layer] = nil; + } + g_set_rows_f32_i32_pipeline = nil; + g_get_rows_f32_pipeline = nil; + g_get_rows_f16_pipeline = nil; + g_get_rows_i32_pipeline = nil; + g_get_rows_q8_0_pipeline = nil; + g_get_rows_q4_0_pipeline = nil; + g_get_rows_q4_K_pipeline = nil; + g_repeat_f32_pipeline = nil; + g_concat_pipeline = nil; + g_cpy_f32_f32_pipeline = nil; + g_cpy_f32_f16_pipeline = nil; + g_cpy_contig_f32_f16_pipeline = nil; + g_cpy_f16_f32_pipeline = nil; + g_cpy_f16_f16_pipeline = nil; + g_cpy_contig_f16_f32_pipeline = nil; + g_cpy_contig_f16_f16_pipeline = nil; + g_flash_kv_stage_f16_pipeline = nil; + g_swiglu_pipeline = nil; + g_swiglu_flat_pipeline = nil; + g_add_pipeline = nil; + g_add2_pipeline = nil; + g_add3_pipeline = nil; + g_moe_sum6_pipeline = nil; + g_moe_sum8_pipeline = nil; + g_mul_pipeline = nil; + g_bin_mul_scalar_pipeline = nil; + g_bin_div_row_pipeline = nil; + g_unary_sigmoid_pipeline = nil; + g_unary_silu_pipeline = nil; + g_unary_softplus_pipeline = nil; + g_unary_sqrt_pipeline = nil; + g_unary_clamp_pipeline = nil; + g_unary_scale_pipeline = nil; + g_unary_fill_pipeline = nil; + g_unary_fill_f16_pipeline = nil; + g_rms_norm_pipeline = nil; + g_rms_norm_plain_pipeline = nil; + g_add_rms_norm_pipeline = nil; + g_rms_norm_scale_pipeline = nil; + g_dsv4_qkv_rms_norm_pipeline = nil; + g_hc_split_sinkhorn_pipeline = nil; + g_hc_split_weighted_sum_pipeline = nil; + g_hc_split_weighted_sum_norm_pipeline = nil; + g_hc_weighted_sum_pipeline = nil; + g_hc_weighted_sum_norm_pipeline = nil; + g_output_hc_weights4_pipeline = nil; + g_hc_expand_pipeline = nil; + g_moe_mul_mv_id_iq2_xxs_pipeline = nil; + g_moe_mul_mv_id_iq2_xxs_pair_pipeline = nil; + g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = nil; + g_moe_mul_mv_id_q2_k_pipeline = nil; + g_moe_mul_mv_id_q2_k_sum6_pipeline = nil; + g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = nil; + g_moe_mul_mv_id_q4_k_pipeline = nil; + g_moe_mul_mv_id_q4_k_pair_pipeline = nil; + g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_id_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_group_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_group6_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_group8_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_group24_q4_k_id_pipeline = nil; + g_moe_mul_mv_group24_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline = nil; + g_moe_mul_mv_slots6_q2_k_sum6_pipeline = nil; + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_slots6_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = nil; + g_moe_mul_mv_addr_iq2_xxs_pipeline = nil; + g_moe_mul_mv_addr_q2_k_sum6_pipeline = nil; + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline = nil; + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline = nil; + g_moe_stream_expert_cache_validate_pipeline = nil; + g_moe_q4_gather_slots6_pipeline = nil; + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_table_q4_k_sum6_pipeline = nil; + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline = nil; + g_moe_mul_mv_addr_q4_k_sum6_pipeline = nil; + g_moe_table_q4_pair_gate_encoder = nil; + g_moe_table_q4_pair_up_encoder = nil; + g_moe_table_q4_sum_down_encoder = nil; + g_rope_tail_batch_pipeline = nil; + g_rope_tail_inplace_pair_pipeline = nil; + g_rope_tail_inplace_pair_shared4_pipeline = nil; + g_rope_tail_inplace_pair_affine_pipeline = nil; + g_dsv4_fp8_kv_quantize_pipeline = nil; + g_dsv4_indexer_qat_pipeline = nil; + g_dsv4_kv_fp8_store_pipeline = nil; + g_dsv4_ratio4_shift_pipeline = nil; + g_dsv4_compressor_pack_ratio4_pipeline = nil; + g_dsv4_softmax_pool_ratio4_direct_pipeline = nil; + g_dsv4_softmax_pool_pipeline = nil; + g_soft_max_f32_pipeline = nil; + g_soft_max_f32_4_pipeline = nil; + g_argsort_f32_i32_desc_pipeline = nil; + g_argsort_merge_f32_i32_desc_pipeline = nil; + g_sum_rows_f32_f32_pipeline = nil; + g_dsv4_topk_mask_pipeline = nil; + g_dsv4_topk_mask_scatter_pipeline = nil; + g_dsv4_indexer_weighted_sum_pipeline = nil; + g_dsv4_indexer_score_one_direct_pipeline = nil; + g_dsv4_compressor_store_one_pipeline = nil; + g_dsv4_sort_i32_rows_asc_pipeline = nil; + g_dsv4_indexed_attention_heads8_pipeline = nil; + g_dsv4_indexed_attention_heads8_rb16_pipeline = nil; + g_dsv4_softplus_sqrt_pipeline = nil; + g_dsv4_router_finalize_one_pipeline = nil; + g_dsv4_router_finalize_one_simd_pipeline = nil; + g_dsv4_router_finalize_weights_one_simd_pipeline = nil; + g_dsv4_router_weights_one_pipeline = nil; + g_glm_router_select_one_pipeline = nil; + g_glm_kv_lora_rms_norm_pipeline = nil; + g_glm_k_b_project_pipeline = nil; + g_glm_store_compact_kv_pipeline = nil; + g_glm_qkv_norm_store_compact_kv_pipeline = nil; + g_glm_store_indexer_k_pipeline = nil; + g_glm_build_kv_cache_pipeline = nil; + g_glm_build_kv_cache_decode_group4_pipeline = nil; + g_glm_build_kv_cache_flash_pipeline = nil; + g_glm_attention_full_pipeline = nil; + g_glm_fill_selected_range_pipeline = nil; + g_glm_fill_selected_range_batch_pipeline = nil; + g_glm_indexer_rope_tail_pipeline = nil; + g_glm_indexer_score_one_pipeline = nil; + g_glm_indexer_score_one_direct_pipeline = nil; + g_glm_indexer_scores_batch_pipeline = nil; + g_glm_indexer_scores_tiled_pipeline = nil; + g_glm_indexer_scores_tiled_f32_pipeline = nil; + g_glm_qk_lowrank_pipeline = nil; + g_glm_qk_lowrank_glm52_pipeline = nil; + g_glm_qk_lowrank_glm52_sg_pipeline = nil; + g_glm_qk_lowrank_batch_pipeline = nil; + g_glm_qk_lowrank_batch_glm52_t4_pipeline = nil; + g_glm_value_project_q8_0_pipeline = nil; + g_glm_value_project_q8_0_batch_heads_pipeline = nil; + g_glm_value_project_q8_0_batch_heads_mma_pipeline = nil; + g_glm_attention_indexed_decode_pipeline = nil; + g_glm_attention_indexed_decode_split_group8_partial_pipeline = nil; + g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline = nil; + g_glm_attention_indexed_decode_split_group8_reduce_pipeline = nil; + g_glm_attention_indexed_decode_split_group8_reduce16_pipeline = nil; + g_glm_attention_indexed_batch_pipeline = nil; + g_glm_attention_indexed_batch_group2_pipeline = nil; + g_glm_attention_indexed_batch_q2_group4_pipeline = nil; + g_glm_attention_indexed_batch_group8_pipeline = nil; + g_glm_attention_indexed_batch_lora_group8_vec_pipeline = nil; + g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline = nil; + g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline = nil; + g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline = nil; + g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline = nil; + g_glm_q4_k_pair_swiglu_f32_pipeline = nil; + g_glm_q4_k_pair_swiglu2_f32_pipeline = nil; + g_glm_q4_k_pair_swiglu4_f32_pipeline = nil; + g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline = nil; + g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline = nil; + g_glm_q2_k_pair_swiglu_f32_pipeline = nil; + g_glm_q2_k_addr_pair_swiglu2_f32_pipeline = nil; + g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline = nil; + g_glm_q4_k_addr_pair_swiglu_f32_pipeline = nil; + g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline = nil; + g_glm_q2_k_down_f32_pipeline = nil; + g_glm_q4_k_down_f32_pipeline = nil; + g_glm_q2_k_addr_down_f32_pipeline = nil; + g_glm_q4_k_addr_down_f32_pipeline = nil; + g_glm_q5_k_pair_swiglu_f32_pipeline = nil; + g_glm_q5_k_pair_swiglu_mapped_f32_pipeline = nil; + g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline = nil; + g_glm_q5_k_down_f32_pipeline = nil; + g_glm_q6_k_down_f32_pipeline = nil; + g_dsv4_router_weights_batch_pipeline = nil; + g_dsv4_hc_expand4_pipeline = nil; + g_flash_attn_mask_buffer = nil; + g_flash_attn_zero_mask_buffer = nil; + g_flash_attn_pad_buffer = nil; + g_flash_attn_tmp_buffer = nil; + g_flash_attn_blk_buffer = nil; + ds4_gpu_clear_zero_prefix_prefill_mask_cache(); + g_flash_attn_ring_buffer = nil; + g_flash_attn_kv_buffer = nil; + g_glm_flash_attn_mask_buffer = nil; + g_compressor_pool_kv_buffer = nil; + g_compressor_pool_score_buffer = nil; + g_compressor_pool_score_cont_buffer = nil; + g_compressor_pool_softmax_buffer = nil; + g_compressor_pool_product_buffer = nil; + g_compressor_store_ape_buffer = nil; + g_compressor_store_score_buffer = nil; + g_embed_rows_buffer = nil; + g_router_selection_buffer = nil; + g_router_weight_sum_buffer = nil; + g_indexer_head_scores_buffer = nil; + g_indexer_topk_buffer = nil; + g_indexed_topk_buffer = nil; + g_stream_expert_validate_status_buffer = nil; + g_f16_round_scratch_buffer = nil; + g_raw_store_round_buffer = nil; + g_moe_gate_scratch_buffer = nil; + g_moe_down_scratch_buffer = nil; + g_moe_id_map_buffer = nil; + g_moe_q4_gate_slots_buffer = nil; + g_moe_q4_up_slots_buffer = nil; + g_moe_q4_down_slots_buffer = nil; + g_attn_out_group_ids_buffer = nil; + g_model_fd = -1; + g_model_map_ptr = NULL; + g_model_map_size = 0; + g_model_mapped_offset = 0; + g_model_mapped_size = 0; + g_model_mapped_max_tensor_bytes = 0; + ds4_gpu_tensor_tracking_reset(); + g_flash_attn_mask_bytes = 0; + g_flash_attn_zero_mask_bytes = 0; + g_flash_attn_pad_bytes = 0; + g_flash_attn_tmp_bytes = 0; + g_flash_attn_blk_bytes = 0; + g_flash_attn_ring_bytes = 0; + g_flash_attn_kv_bytes = 0; + g_glm_flash_attn_mask_bytes = 0; + g_glm_flash_attn_mask_valid = 0; + g_glm_flash_attn_mask_pos0 = 0; + g_glm_flash_attn_mask_tokens = 0; + g_glm_flash_attn_mask_cache_len = 0; + g_compressor_pool_kv_bytes = 0; + g_compressor_pool_score_bytes = 0; + g_compressor_pool_score_cont_bytes = 0; + g_compressor_pool_softmax_bytes = 0; + g_compressor_pool_product_bytes = 0; + g_compressor_store_ape_bytes = 0; + g_compressor_store_score_bytes = 0; + g_embed_rows_bytes = 0; + g_router_selection_bytes = 0; + g_router_weight_sum_bytes = 0; + g_indexer_head_scores_bytes = 0; + g_indexer_topk_bytes = 0; + g_indexed_topk_bytes = 0; + g_f16_round_scratch_bytes = 0; + g_raw_store_round_bytes = 0; + g_moe_gate_scratch_bytes = 0; + g_moe_down_scratch_bytes = 0; + g_moe_id_map_bytes = 0; + g_moe_q4_gate_slots_bytes = 0; + g_moe_q4_up_slots_bytes = 0; + g_moe_q4_down_slots_bytes = 0; + g_attn_out_group_ids_bytes = 0; + g_model_wrap_count = 0; + g_model_wrap_bytes = 0; + g_model_wrap_max_bytes = 0; + g_model_buffer_cache_bytes = 0; + g_model_buffer_cache_evictions = 0; + g_model_buffer_cache_over_limit = 0; + ds4_gpu_model_residency_clear(); + ds4_gpu_model_views_clear(); + [g_pipeline_cache removeAllObjects]; + g_pipeline_cache = nil; + [g_q4_expert_layer_residency_cache removeAllObjects]; + g_q4_expert_layer_residency_cache = nil; + [g_q4_expert_table_cache removeAllObjects]; + g_q4_expert_table_cache = nil; + [g_model_buffer_cache removeAllObjects]; + g_model_buffer_cache = nil; + g_transient_buffers = nil; + g_pending_cbs = nil; + g_library = nil; + g_queue = nil; + g_device = nil; + g_initialized = 0; + } +} + +static int ds4_gpu_encode_get_rows_f16( + id cb, + id weight, + NSUInteger weight_offset, + id tokens, + NSUInteger tokens_offset, + id out, + NSUInteger out_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd) { + if (!cb || !weight || !tokens || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { + return 0; + } + + const uint64_t src_row_bytes = (uint64_t)n_embd * sizeof(uint16_t); + const uint64_t dst_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); + ds4_gpu_get_rows_args args = { + .ne00t = (int32_t)n_embd, + .ne00 = (int32_t)n_embd, + .nb01 = src_row_bytes, + .nb02 = (uint64_t)n_vocab * src_row_bytes, + .nb03 = (uint64_t)n_vocab * src_row_bytes, + .ne10 = (int32_t)n_tokens, + .nb10 = sizeof(int32_t), + .nb11 = token_bytes, + .nb12 = token_bytes, + .nb1 = dst_row_bytes, + .nb2 = (uint64_t)n_tokens * dst_row_bytes, + .nb3 = (uint64_t)n_tokens * dst_row_bytes, + }; + + NSUInteger nth = (NSUInteger)n_embd; + const NSUInteger max_threads = g_get_rows_f16_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + const NSUInteger nw0 = ((NSUInteger)n_embd + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_get_rows_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weight offset:weight_offset atIndex:1]; + [enc setBuffer:tokens offset:tokens_offset atIndex:2]; + [enc setBuffer:out offset:out_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(nw0 * n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static uint64_t ds4_gpu_q8_0_row_bytes(uint32_t n_embd) { + return (((uint64_t)n_embd + 31u) / 32u) * 34u; +} + +static int ds4_gpu_quant_row_bytes( + uint32_t type, + uint32_t n_embd, + uint64_t *row_bytes_out) { + if (!row_bytes_out || n_embd == 0) return 0; + switch (type) { + case DS4_METAL_TENSOR_Q8_0: + *row_bytes_out = (((uint64_t)n_embd + 31u) / 32u) * 34u; + return 1; + case DS4_METAL_TENSOR_Q4_0: + *row_bytes_out = (((uint64_t)n_embd + 31u) / 32u) * 18u; + return 1; + case DS4_METAL_TENSOR_Q4_K: + if ((n_embd % 256u) != 0) return 0; + *row_bytes_out = ((uint64_t)n_embd / 256u) * 144u; + return 1; + default: + return 0; + } +} + +static int ds4_gpu_q8_0_table_bytes( + uint32_t n_vocab, + uint32_t n_embd, + uint64_t *bytes_out) { + if (!bytes_out || n_vocab == 0 || n_embd == 0) return 0; + const uint64_t row_bytes = ds4_gpu_q8_0_row_bytes(n_embd); + if (row_bytes != 0 && (uint64_t)n_vocab > UINT64_MAX / row_bytes) return 0; + *bytes_out = (uint64_t)n_vocab * row_bytes; + return 1; +} + +static int ds4_gpu_quant_table_bytes( + uint32_t type, + uint32_t n_vocab, + uint32_t n_embd, + uint64_t *bytes_out) { + if (!bytes_out || n_vocab == 0 || n_embd == 0) return 0; + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(type, n_embd, &row_bytes)) return 0; + if (row_bytes != 0 && (uint64_t)n_vocab > UINT64_MAX / row_bytes) return 0; + *bytes_out = (uint64_t)n_vocab * row_bytes; + return 1; +} + +static int ds4_gpu_encode_get_rows_q8_0( + id cb, + id weight, + NSUInteger weight_offset, + id tokens, + NSUInteger tokens_offset, + const int32_t *single_token, + id out, + NSUInteger out_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd) { + if (!cb || !weight || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { + return 0; + } + if (!tokens && (!single_token || n_tokens != 1)) { + return 0; + } + + ds4_gpu_get_rows_q8_0_args args = { + .n_embd = (int32_t)n_embd, + .n_vocab = (int32_t)n_vocab, + .n_tokens = (int32_t)n_tokens, + .src_row_bytes = ds4_gpu_q8_0_row_bytes(n_embd), + .dst_row_bytes = (uint64_t)n_embd * sizeof(float), + .token_stride = sizeof(int32_t), + }; + + NSUInteger nth = 32u; + const NSUInteger max_threads = g_get_rows_q8_0_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + const NSUInteger nblocks = ((NSUInteger)n_embd + 31u) / 32u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_get_rows_q8_0_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weight offset:weight_offset atIndex:1]; + if (tokens) { + [enc setBuffer:tokens offset:tokens_offset atIndex:2]; + } else { + [enc setBytes:single_token length:sizeof(*single_token) atIndex:2]; + } + [enc setBuffer:out offset:out_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(nblocks, n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_get_rows_quant( + id cb, + id weight, + NSUInteger weight_offset, + uint32_t weight_type, + id tokens, + NSUInteger tokens_offset, + const int32_t *single_token, + id out, + NSUInteger out_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd) { + if (weight_type == DS4_METAL_TENSOR_Q8_0) { + return ds4_gpu_encode_get_rows_q8_0(cb, + weight, + weight_offset, + tokens, + tokens_offset, + single_token, + out, + out_offset, + n_vocab, + n_tokens, + n_embd); + } + if (!cb || !weight || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { + return 0; + } + if (!tokens && (!single_token || n_tokens != 1)) { + return 0; + } + + uint64_t src_row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, n_embd, &src_row_bytes)) return 0; + ds4_gpu_get_rows_q8_0_args args = { + .n_embd = (int32_t)n_embd, + .n_vocab = (int32_t)n_vocab, + .n_tokens = (int32_t)n_tokens, + .src_row_bytes = src_row_bytes, + .dst_row_bytes = (uint64_t)n_embd * sizeof(float), + .token_stride = sizeof(int32_t), + }; + + id pipeline = nil; + NSUInteger block_width = 0; + if (weight_type == DS4_METAL_TENSOR_Q4_0) { + pipeline = g_get_rows_q4_0_pipeline; + block_width = 32u; + } else if (weight_type == DS4_METAL_TENSOR_Q4_K) { + pipeline = g_get_rows_q4_K_pipeline; + block_width = 256u; + } + if (!pipeline || block_width == 0) return 0; + + NSUInteger nth = block_width; + const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + const NSUInteger nblocks = ((NSUInteger)n_embd + block_width - 1u) / block_width; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weight offset:weight_offset atIndex:1]; + if (tokens) { + [enc setBuffer:tokens offset:tokens_offset atIndex:2]; + } else { + [enc setBytes:single_token length:sizeof(*single_token) atIndex:2]; + } + [enc setBuffer:out offset:out_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(nblocks, n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_repeat_hc_embedding( + id cb, + id rows, + NSUInteger rows_offset, + id out, + NSUInteger out_offset, + uint32_t n_tokens, + uint32_t n_embd, + uint32_t n_hc) { + if (!cb || !rows || !out || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0; + + const uint64_t embd_bytes = (uint64_t)n_embd * sizeof(float); + ds4_gpu_repeat_args args = { + .ne00 = (int32_t)n_embd, + .ne01 = 1, + .ne02 = (int32_t)n_tokens, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = embd_bytes, + .nb02 = embd_bytes, + .nb03 = (uint64_t)n_tokens * embd_bytes, + .ne0 = (int32_t)n_embd, + .ne1 = (int32_t)n_hc, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = embd_bytes, + .nb2 = (uint64_t)n_hc * embd_bytes, + .nb3 = (uint64_t)n_tokens * n_hc * embd_bytes, + }; + + NSUInteger nth = (NSUInteger)n_embd; + const NSUInteger max_threads = g_repeat_f32_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_repeat_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rows offset:rows_offset atIndex:1]; + [enc setBuffer:out offset:out_offset atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(n_hc, n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +int ds4_gpu_embed_token_q8_0_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t token, + uint32_t n_embd) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float); + if (!outbuf || ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal Q8_0 embedding received undersized output buffer\n"); + return 0; + } + + uint64_t weight_bytes = 0; + if (!ds4_gpu_q8_0_table_bytes(n_vocab, n_embd, &weight_bytes) || + weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8_0 embedding range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + uint32_t token_for_kernel = token; + id wbuf = nil; + const bool exact_token_row = + getenv("DS4_METAL_DISABLE_TOKEN_EMBED_EXACT_VIEW") == NULL; + if (exact_token_row) { + const uint64_t row_bytes = ds4_gpu_q8_0_row_bytes(n_embd); + const uint64_t token_rel = (uint64_t)token * row_bytes; + if (token_rel > weight_bytes || row_bytes > weight_bytes - token_rel) { + fprintf(stderr, "ds4: Metal Q8_0 embedding token row is outside the mapped table\n"); + return 0; + } + wbuf = ds4_gpu_wrap_model_exact_range(model_map, + model_size, + weight_offset + token_rel, + row_bytes, + &inner_offset); + token_for_kernel = 0; + } else { + wbuf = ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + } + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const int32_t token_i32 = (int32_t)token_for_kernel; + if (!ds4_gpu_encode_get_rows_q8_0(cb, + wbuf, + (NSUInteger)inner_offset, + nil, + 0, + &token_i32, + outbuf, + ds4_gpu_tensor_offset(out), + n_vocab, + 1, + n_embd)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "q8_0 embed token")) return 0; + } + + return 1; +} + +int ds4_gpu_embed_tokens_q8_0_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *tokens, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out); + id tokbuf = ds4_gpu_tensor_buffer(tokens); + const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * sizeof(float); + const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); + if (!outbuf || !tokbuf || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(tokens) < token_bytes) { + fprintf(stderr, "ds4: Metal Q8_0 batched embedding received undersized buffers\n"); + return 0; + } + + uint64_t weight_bytes = 0; + if (!ds4_gpu_q8_0_table_bytes(n_vocab, n_embd, &weight_bytes) || + weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8_0 batched embedding range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_get_rows_q8_0(cb, + wbuf, + (NSUInteger)inner_offset, + tokbuf, + ds4_gpu_tensor_offset(tokens), + NULL, + outbuf, + ds4_gpu_tensor_offset(out), + n_vocab, + n_tokens, + n_embd)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "q8_0 embed tokens")) return 0; + } + + return 1; +} + +int ds4_gpu_embed_token_quant_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_vocab, + uint32_t token, + uint32_t n_embd) { + if (weight_type == DS4_METAL_TENSOR_Q8_0) { + return ds4_gpu_embed_token_q8_0_tensor(out, + model_map, + model_size, + weight_offset, + n_vocab, + token, + n_embd); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float); + if (!outbuf || ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal quant embedding received undersized output buffer\n"); + return 0; + } + + uint64_t weight_bytes = 0; + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_table_bytes(weight_type, n_vocab, n_embd, &weight_bytes) || + !ds4_gpu_quant_row_bytes(weight_type, n_embd, &row_bytes) || + weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal quant embedding range is outside the mapped model\n"); + return 0; + } + + const uint64_t token_rel = (uint64_t)token * row_bytes; + if (token_rel > weight_bytes || row_bytes > weight_bytes - token_rel) { + fprintf(stderr, "ds4: Metal quant embedding token row is outside the mapped table\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + weight_offset + token_rel, + row_bytes, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const int32_t token_i32 = 0; + if (!ds4_gpu_encode_get_rows_quant(cb, + wbuf, + (NSUInteger)inner_offset, + weight_type, + nil, + 0, + &token_i32, + outbuf, + ds4_gpu_tensor_offset(out), + n_vocab, + 1, + n_embd)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "quant embed token")) return 0; + } + + return 1; +} + +int ds4_gpu_embed_tokens_quant_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *tokens, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd) { + if (weight_type == DS4_METAL_TENSOR_Q8_0) { + return ds4_gpu_embed_tokens_q8_0_tensor(out, + tokens, + model_map, + model_size, + weight_offset, + n_vocab, + n_tokens, + n_embd); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out); + id tokbuf = ds4_gpu_tensor_buffer(tokens); + const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * sizeof(float); + const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); + if (!outbuf || !tokbuf || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(tokens) < token_bytes) { + fprintf(stderr, "ds4: Metal quant batched embedding received undersized buffers\n"); + return 0; + } + + uint64_t weight_bytes = 0; + if (!ds4_gpu_quant_table_bytes(weight_type, n_vocab, n_embd, &weight_bytes) || + weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal quant batched embedding range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_get_rows_quant(cb, + wbuf, + (NSUInteger)inner_offset, + weight_type, + tokbuf, + ds4_gpu_tensor_offset(tokens), + NULL, + outbuf, + ds4_gpu_tensor_offset(out), + n_vocab, + n_tokens, + n_embd)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "quant embed tokens")) return 0; + } + + return 1; +} + +int ds4_gpu_embed_token_hc_tensor( + ds4_gpu_tensor *out_hc, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t token, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0 || n_hc == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out_hc); + const uint64_t out_bytes = (uint64_t)n_embd * n_hc * sizeof(float); + if (!outbuf || ds4_gpu_tensor_bytes(out_hc) < out_bytes) { + fprintf(stderr, "ds4: Metal graph embedding received undersized HC output buffer\n"); + return 0; + } + + const uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t); + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal graph embedding range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + const NSUInteger row_bytes = (NSUInteger)n_embd * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_embed_rows_buffer, + &g_embed_rows_bytes, + row_bytes, + "ds4_embed_rows")) { + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const int32_t token_i32 = (int32_t)token; + const uint64_t src_row_bytes = (uint64_t)n_embd * sizeof(uint16_t); + const uint64_t dst_row_bytes = (uint64_t)n_embd * sizeof(float); + ds4_gpu_get_rows_args args = { + .ne00t = (int32_t)n_embd, + .ne00 = (int32_t)n_embd, + .nb01 = src_row_bytes, + .nb02 = (uint64_t)n_vocab * src_row_bytes, + .nb03 = (uint64_t)n_vocab * src_row_bytes, + .ne10 = 1, + .nb10 = sizeof(int32_t), + .nb11 = sizeof(int32_t), + .nb12 = sizeof(int32_t), + .nb1 = dst_row_bytes, + .nb2 = dst_row_bytes, + .nb3 = dst_row_bytes, + }; + NSUInteger nth = (NSUInteger)n_embd; + const NSUInteger max_threads = g_get_rows_f16_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + const NSUInteger nw0 = ((NSUInteger)n_embd + nth - 1u) / nth; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_get_rows_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBytes:&token_i32 length:sizeof(token_i32) atIndex:2]; + [enc setBuffer:g_embed_rows_buffer offset:0 atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(nw0, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_encode_repeat_hc_embedding(cb, + g_embed_rows_buffer, + 0, + outbuf, + ds4_gpu_tensor_offset(out_hc), + 1, + n_embd, + n_hc)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph embed token")) return 0; + } + + return 1; +} + +int ds4_gpu_embed_tokens_hc_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *tokens, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0 || n_hc == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out_hc); + id tokbuf = ds4_gpu_tensor_buffer(tokens); + const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * n_hc * sizeof(float); + const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); + if (!outbuf || !tokbuf || + ds4_gpu_tensor_bytes(out_hc) < out_bytes || + ds4_gpu_tensor_bytes(tokens) < token_bytes) { + fprintf(stderr, "ds4: Metal graph batched embedding received undersized buffers\n"); + return 0; + } + + const uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t); + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal graph batched embedding range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + const NSUInteger rows_bytes = (NSUInteger)n_tokens * n_embd * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_embed_rows_buffer, + &g_embed_rows_bytes, + rows_bytes, + "ds4_embed_rows")) { + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_get_rows_f16(cb, + wbuf, + (NSUInteger)inner_offset, + tokbuf, + ds4_gpu_tensor_offset(tokens), + g_embed_rows_buffer, + 0, + n_vocab, + n_tokens, + n_embd) || + !ds4_gpu_encode_repeat_hc_embedding(cb, + g_embed_rows_buffer, + 0, + outbuf, + ds4_gpu_tensor_offset(out_hc), + n_tokens, + n_embd, + n_hc)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph embed tokens")) return 0; + } + + return 1; +} + +int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!model_map || model_size == 0) return 0; + if (map_offset > model_size || map_size == 0 || map_size > model_size - map_offset) return 0; + max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(map_size, max_tensor_bytes); + + @autoreleasepool { + if (g_model_map_ptr == model_map && + g_model_map_size == model_size && + g_model_mapped_offset == map_offset && + g_model_mapped_size == map_size && + g_model_mapped_max_tensor_bytes == max_tensor_bytes) { + return 1; + } + + for (uint32_t i = 0; i < g_model_view_count; i++) { + if (g_model_views[i].model_map == model_map && + g_model_views[i].model_size == model_size && + map_offset >= g_model_views[i].model_offset && + map_offset + map_size <= g_model_views[i].model_offset + g_model_views[i].bytes) { + return 1; + } + } + + ds4_gpu_model_residency_clear(); + if (!ds4_gpu_map_model_views(model_map, model_size, map_offset, map_size, max_tensor_bytes)) { + ds4_gpu_model_residency_clear(); + return 0; + } + g_model_map_ptr = model_map; + g_model_map_size = model_size; + g_model_mapped_offset = map_offset; + g_model_mapped_size = map_size; + g_model_mapped_max_tensor_bytes = max_tensor_bytes; + if (ds4_gpu_model_map_log_enabled()) { + fprintf(stderr, + "ds4: Metal mapped mmaped model as %u overlapping shared buffers\n", + g_model_view_count); + } + return 1; + } +} + +static int ds4_gpu_model_views_cover_spans( + const void *model_map, + uint64_t model_size, + const uint64_t *offsets, + const uint64_t *sizes, + uint32_t count); + +int ds4_gpu_set_model_map_spans( + const void *model_map, + uint64_t model_size, + const uint64_t *offsets, + const uint64_t *sizes, + uint32_t count, + uint64_t max_tensor_bytes) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!model_map || model_size == 0 || !offsets || !sizes || count == 0) return 0; + if (count == 1) { + return ds4_gpu_set_model_map_range(model_map, + model_size, + offsets[0], + sizes[0], + max_tensor_bytes); + } + if (ds4_gpu_model_views_cover_spans(model_map, model_size, offsets, sizes, count)) { + return 1; + } + + @autoreleasepool { + const double t0 = ds4_gpu_now_ms(); + max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(model_size, max_tensor_bytes); + + ds4_gpu_model_residency_clear(); + ds4_gpu_model_views_clear(); + + uint64_t mapped_total = 0; + uint64_t first_offset = UINT64_MAX; + for (uint32_t i = 0; i < count; i++) { + if (offsets[i] > model_size || sizes[i] == 0 || sizes[i] > model_size - offsets[i]) { + fprintf(stderr, "ds4: Metal model span %u is outside the GGUF mapping\n", i); + ds4_gpu_model_residency_clear(); + ds4_gpu_model_views_clear(); + return 0; + } + if (offsets[i] < first_offset) first_offset = offsets[i]; + uint64_t effective_max = max_tensor_bytes; + if (effective_max > sizes[i]) effective_max = sizes[i]; + if (!ds4_gpu_add_model_view_range(model_map, + model_size, + offsets[i], + sizes[i], + effective_max, + true, + &mapped_total)) { + ds4_gpu_model_residency_clear(); + ds4_gpu_model_views_clear(); + return 0; + } + } + if (!ds4_gpu_finish_model_views(t0, mapped_total, first_offset)) { + ds4_gpu_model_residency_clear(); + ds4_gpu_model_views_clear(); + return 0; + } + g_model_map_ptr = model_map; + g_model_map_size = model_size; + g_model_mapped_offset = first_offset == UINT64_MAX ? 0 : first_offset; + g_model_mapped_size = mapped_total; + g_model_mapped_max_tensor_bytes = max_tensor_bytes; + if (ds4_gpu_model_map_log_enabled()) { + fprintf(stderr, + "ds4: Metal mapped mmaped model as %u disjoint shared buffers across %u tensor spans\n", + g_model_view_count, + count); + } + return 1; + } +} + +int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size) { + return ds4_gpu_set_model_map_range(model_map, model_size, 0, model_size, 0); +} + +int ds4_gpu_set_model_fd(int fd) { + g_model_fd = fd; + return 1; +} + +static int ds4_gpu_model_views_cover_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t size) { + if (!model_map || model_size == 0 || size == 0 || + offset > model_size || size > model_size - offset) { + return 0; + } + const uint64_t end = offset + size; + for (uint32_t i = 0; i < g_model_view_count; i++) { + if (g_model_views[i].model_map != model_map || + g_model_views[i].model_size != model_size) { + continue; + } + const uint64_t view_start = g_model_views[i].model_offset; + const uint64_t view_end = view_start + g_model_views[i].bytes; + if (offset >= view_start && end <= view_end) return 1; + } + return 0; +} + +static int ds4_gpu_model_views_cover_spans( + const void *model_map, + uint64_t model_size, + const uint64_t *offsets, + const uint64_t *sizes, + uint32_t count) { + if (!offsets || !sizes || count == 0) return 0; + for (uint32_t i = 0; i < count; i++) { + if (!ds4_gpu_model_views_cover_range(model_map, + model_size, + offsets[i], + sizes[i])) { + return 0; + } + } + return 1; +} + +int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) { + (void)fd; + (void)model_map; + return 1; +} + +static id ds4_gpu_wrap_model_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset) { + (void)model_map; + if (model_size == 0 || offset > model_size || len > model_size - offset) { + fprintf(stderr, "ds4: Metal model range is outside the mapped model\n"); + return nil; + } + + const uint64_t end = offset + len; + for (uint32_t i = 0; i < g_model_view_count; i++) { + if (g_model_views[i].model_map != model_map || + g_model_views[i].model_size != model_size) { + continue; + } + const uint64_t view_start = g_model_views[i].model_offset; + const uint64_t view_end = view_start + g_model_views[i].bytes; + if (offset >= view_start && end <= view_end) { + *inner_offset = offset - view_start; + return g_model_views[i].buffer; + } + } + + fprintf(stderr, + "ds4: Metal model range %.2f..%.2f GiB is not covered by mapped model views\n", + ds4_gpu_gib(offset), + ds4_gpu_gib(end)); + return nil; +} + +typedef enum { + DS4_GPU_EXACT_VIEW_CACHED, + DS4_GPU_EXACT_VIEW_TRANSIENT, + DS4_GPU_EXACT_VIEW_OWNED, +} ds4_gpu_exact_view_lifetime; + +static id ds4_gpu_wrap_model_exact_range_impl( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset, + ds4_gpu_exact_view_lifetime lifetime) { + const bool cache_view = lifetime == DS4_GPU_EXACT_VIEW_CACHED; + const bool transient_view = lifetime == DS4_GPU_EXACT_VIEW_TRANSIENT; + if (!model_map || !g_device || + (cache_view && !g_model_buffer_cache) || + (transient_view && !g_transient_buffers) || + model_size == 0 || offset > model_size || len > model_size - offset) { + fprintf(stderr, "ds4: Metal exact model range is outside the mapped model\n"); + return nil; + } + + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t page_offset = offset & ~(page - 1); + const uint64_t leading = offset - page_offset; + if (len > UINT64_MAX - leading || + leading + len > UINT64_MAX - (page - 1)) { + fprintf(stderr, "ds4: Metal exact model range overflows page alignment\n"); + return nil; + } + uint64_t view_bytes = round_up_u64(leading + len, page); + if (view_bytes > model_size - page_offset) view_bytes = model_size - page_offset; + if (leading + len > view_bytes) { + fprintf(stderr, "ds4: Metal exact model range alignment exceeds mapped model\n"); + return nil; + } + if (view_bytes > (uint64_t)[g_device maxBufferLength]) { + fprintf(stderr, + "ds4: Metal exact model range %.2f GiB exceeds maxBufferLength %.2f GiB\n", + ds4_gpu_gib(view_bytes), + ds4_gpu_gib((uint64_t)[g_device maxBufferLength])); + return nil; + } + + NSString *key = nil; + id buffer = nil; + if (cache_view) { + key = [NSString stringWithFormat:@"%p:%llu:%llu:%llu", + model_map, + (unsigned long long)model_size, + (unsigned long long)page_offset, + (unsigned long long)view_bytes]; + buffer = [g_model_buffer_cache objectForKey:key]; + } + if (!buffer) { + const uintptr_t base = (uintptr_t)model_map; + buffer = [g_device newBufferWithBytesNoCopy:(void *)(base + page_offset) + length:(NSUInteger)view_bytes + options:ds4_gpu_model_resource_options() + deallocator:nil]; + if (!buffer) { + fprintf(stderr, + "ds4: Metal could not wrap exact mmaped model range at %.2f GiB, size %.2f MiB\n", + ds4_gpu_gib(page_offset), + ds4_gpu_mib(view_bytes)); + return nil; + } + if (cache_view) { + buffer.label = @"ds4_model_exact_view"; + } else if (transient_view) { + buffer.label = @"ds4_model_exact_transient_view"; + } else { + buffer.label = @"ds4_model_exact_owned_view"; + } + if (cache_view) { + [g_model_buffer_cache setObject:buffer forKey:key]; + ds4_gpu_model_buffer_cache_note_insert(view_bytes); + } else if (transient_view) { + [g_transient_buffers addObject:buffer]; + } + } + + if (inner_offset) *inner_offset = leading; + return buffer; +} + +static id ds4_gpu_wrap_model_exact_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset) { + return ds4_gpu_wrap_model_exact_range_impl(model_map, + model_size, + offset, + len, + inner_offset, + DS4_GPU_EXACT_VIEW_CACHED); +} + +static id ds4_gpu_wrap_model_exact_range_transient( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset) { + return ds4_gpu_wrap_model_exact_range_impl(model_map, + model_size, + offset, + len, + inner_offset, + DS4_GPU_EXACT_VIEW_TRANSIENT); +} + +static id ds4_gpu_wrap_model_exact_range_owned( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t *inner_offset) { + return ds4_gpu_wrap_model_exact_range_impl(model_map, + model_size, + offset, + len, + inner_offset, + DS4_GPU_EXACT_VIEW_OWNED); +} + +static id ds4_gpu_wrap_q8_decode_model_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t n_tokens, + uint64_t *inner_offset) { + const uint64_t exact_decode_max_mib = + ds4_gpu_env_u64("DS4_METAL_Q8_DECODE_EXACT_VIEW_MAX_MIB", + 1024u, + 1u, + 4096u); + const uint64_t exact_decode_max_bytes = + exact_decode_max_mib * 1024ull * 1024ull; + const bool exact_decode_weight_view = + n_tokens == 1u && + len <= exact_decode_max_bytes && + getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_DISABLE_Q8_DECODE_EXACT_VIEWS") == NULL; + return exact_decode_weight_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + offset, + len, + inner_offset) : + ds4_gpu_wrap_model_range(model_map, + model_size, + offset, + len, + inner_offset); +} + +static id ds4_gpu_wrap_f32_decode_model_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len, + uint64_t n_tokens, + uint64_t *inner_offset) { + const uint64_t exact_decode_max_mib = + ds4_gpu_env_u64("DS4_METAL_F32_DECODE_EXACT_VIEW_MAX_MIB", + 64u, + 1u, + 4096u); + const uint64_t exact_decode_max_bytes = + exact_decode_max_mib * 1024ull * 1024ull; + const bool exact_decode_weight_view = + n_tokens == 1u && + len <= exact_decode_max_bytes && + getenv("DS4_METAL_ENABLE_F32_DECODE_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_DISABLE_F32_DECODE_EXACT_VIEWS") == NULL; + return exact_decode_weight_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + offset, + len, + inner_offset) : + ds4_gpu_wrap_model_range(model_map, + model_size, + offset, + len, + inner_offset); +} + +uint32_t ds4_gpu_stream_expert_cache_configured_count(void) { + uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { + budget = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES; + } + return budget; +} + +uint32_t ds4_gpu_stream_expert_cache_current_count(void) { + return g_stream_expert_cache_entry_count; +} + +uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes)) { + return 0; + } + return ds4_gpu_stream_expert_cache_configured_budget(); +} + +static int ds4_gpu_stream_expert_cache_note_expert_size( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (gate_expert_bytes == 0 || down_expert_bytes == 0) return 0; + if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { + fprintf(stderr, "ds4: Metal streaming expert cache byte size overflow\n"); + return 0; + } + /* + * The cache is a single-size-class slab allocator: the expert byte size is + * frozen on first sight (or pre-seeded at startup from the model's slab + * class) and off-size layers are REJECTED rather than adopted. A rejected + * layer (mixed-precision boost: Q4_K experts among IQ2 layers) falls back + * to the mapped-model per-expert path; last-writer-wins here would instead + * poison the slab size class and deadlock slab reuse. + */ + const uint64_t bytes = gate_expert_bytes * 2ull + down_expert_bytes; + if (g_stream_expert_cache_expert_bytes == 0) { + g_stream_expert_cache_expert_bytes = bytes; + return 1; + } + return bytes == g_stream_expert_cache_expert_bytes; +} + +static uint32_t ds4_gpu_stream_expert_cache_requested_budget(void) { + if (!g_ssd_streaming_mode) return 0; + if (g_stream_expert_cache_budget_override != 0) { + return g_stream_expert_cache_budget_override; + } + return 0; +} + +static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void) { + return ds4_gpu_stream_expert_cache_requested_budget(); +} + +static uint32_t ds4_gpu_stream_expert_cache_effective_cap( + uint32_t layer, + uint32_t n_total_expert, + uint32_t n_selected) { + (void)layer; + if (ds4_gpu_stream_expert_cache_configured_budget() == 0) return 0; + + /* + * The residency policy is global: every layer can use any expert slot it + * routes to, and global pruning decides which existing entry is least + * valuable. A per-layer cap made cache size depend on model depth rather + * than the actual byte budget. + */ + uint32_t cap = n_total_expert; + if (cap < n_selected) cap = n_selected; + if (cap > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + cap = DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + } + return cap; +} + +static int ds4_gpu_stream_expert_timing_summary_enabled(void) { + static int checked = 0; + static int enabled = 0; + if (!checked) { + enabled = + (getenv("DS4_METAL_STREAMING_EXPERT_TIMING_SUMMARY") != NULL || + getenv("DS4_METAL_STREAMING_EXPERT_PROFILE_SUMMARY") != NULL) && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_TIMING_SUMMARY") == NULL; + checked = 1; + } + return enabled; +} + +static uint32_t ds4_gpu_stream_expert_popcount(uint32_t mask) { + return (uint32_t)__builtin_popcount(mask); +} + +static int ds4_gpu_stream_expert_split_worthwhile( + uint32_t resident_mask, + uint32_t missing_mask) { + if (resident_mask == 0 || missing_mask == 0) return 0; + /* + * The split path pays an extra command stage and a second routed-expert + * bind. It is worthwhile when several experts are missing and their SSD + * reads can be hidden by resident expert work. With one or two misses, + * especially in large caches, a single unsplit routed pass is faster. + */ + return ds4_gpu_stream_expert_popcount(missing_mask) >= 3u; +} + +static void ds4_gpu_stream_expert_timing_note_selected( + double sync_ms, + double copy_ms, + double bind_ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_selected_calls++; + g_stream_expert_timing_selected_sync_ms += sync_ms; + g_stream_expert_timing_selected_copy_ms += copy_ms; + g_stream_expert_timing_selected_read_ms += sync_ms + copy_ms; + g_stream_expert_timing_selected_bind_ms += bind_ms; +} + +static void ds4_gpu_stream_expert_timing_note_split( + uint32_t resident_mask, + uint32_t missing_mask, + double resident_ms, + double missing_ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_split_layers++; + g_stream_expert_timing_split_resident_experts += + ds4_gpu_stream_expert_popcount(resident_mask); + g_stream_expert_timing_split_missing_experts += + ds4_gpu_stream_expert_popcount(missing_mask); + g_stream_expert_timing_split_resident_ms += resident_ms; + g_stream_expert_timing_split_missing_ms += missing_ms; +} + +static void ds4_gpu_stream_expert_timing_note_split_missing_detail( + double load_ms, + double slot_ms, + double prune_ms, + double addr_ms, + double wait_ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_split_missing_load_ms += load_ms; + g_stream_expert_timing_split_missing_slot_ms += slot_ms; + g_stream_expert_timing_split_missing_prune_ms += prune_ms; + g_stream_expert_timing_split_missing_addr_ms += addr_ms; + g_stream_expert_timing_split_missing_wait_ms += wait_ms; +} + +static void ds4_gpu_stream_expert_timing_note_load_detail( + double prepare_ms, + double pread_ms, + double modify_ms, + double install_ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_load_calls++; + g_stream_expert_timing_load_prepare_ms += prepare_ms; + g_stream_expert_timing_load_pread_ms += pread_ms; + g_stream_expert_timing_load_modify_ms += modify_ms; + g_stream_expert_timing_load_install_ms += install_ms; +} + +static void ds4_gpu_stream_expert_timing_note_prepare_batch_reuse(double ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_prepare_batch_reuse_calls++; + g_stream_expert_timing_prepare_batch_reuse_ms += ms; +} + +static void ds4_gpu_stream_expert_timing_note_prepare_buffer(double ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_prepare_buffer_calls++; + g_stream_expert_timing_prepare_buffer_ms += ms; +} + +static void ds4_gpu_stream_expert_timing_note_prepare_task( + uint32_t experts, + double ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_prepare_task_experts += experts; + g_stream_expert_timing_prepare_task_ms += ms; +} + +static void ds4_gpu_stream_expert_timing_note_reuse_scan( + uint64_t entries, + double ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_reuse_scan_calls++; + g_stream_expert_timing_reuse_scan_entries += entries; + g_stream_expert_timing_reuse_scan_ms += ms; +} + +static void ds4_gpu_stream_expert_timing_note_reuse_clear(double ms) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + g_stream_expert_timing_reuse_clear_ms += ms; +} + +static void ds4_gpu_stream_expert_timing_note_cache_class( + uint32_t resident_mask, + uint32_t missing_mask) { + if (!ds4_gpu_stream_expert_timing_summary_enabled()) return; + const uint32_t resident = ds4_gpu_stream_expert_popcount(resident_mask); + const uint32_t missing = ds4_gpu_stream_expert_popcount(missing_mask); + if (missing == 0) { + g_stream_expert_timing_cache_all_resident_layers++; + } else if (resident == 0) { + g_stream_expert_timing_cache_all_missing_layers++; + } else { + g_stream_expert_timing_cache_mixed_layers++; + } + g_stream_expert_timing_cache_resident_experts += resident; + g_stream_expert_timing_cache_missing_experts += missing; +} + +static int ds4_gpu_stream_expert_readahead_enabled(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_READAHEAD") == NULL; +} + +static void ds4_gpu_stream_expert_readahead_range(uint64_t offset, uint64_t len) { + if (!ds4_gpu_stream_expert_readahead_enabled() || g_model_fd < 0 || len == 0) { + return; + } + +#if defined(F_RDADVISE) + const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); + const double t0 = timing ? ds4_gpu_now_ms() : 0.0; + uint64_t pos = offset; + uint64_t rem = len; + while (rem > 0) { + const uint64_t chunk64 = + rem > (uint64_t)INT_MAX ? (uint64_t)INT_MAX : rem; + if (pos > (uint64_t)LLONG_MAX) break; + + struct radvisory ra; + ra.ra_offset = (off_t)pos; + ra.ra_count = (int)chunk64; + (void)fcntl(g_model_fd, F_RDADVISE, &ra); + + pos += chunk64; + rem -= chunk64; + } + if (timing) { + g_stream_expert_timing_readahead_calls++; + g_stream_expert_timing_readahead_bytes += len; + g_stream_expert_timing_readahead_ms += ds4_gpu_now_ms() - t0; + } +#else + (void)offset; + (void)len; +#endif +} + +typedef struct { + uint64_t offset; + uint64_t len; + uint8_t *dst; + uint64_t read_bytes; + double ms; + int ok; +} ds4_gpu_stream_expert_pread_task; + +typedef struct { + int active; + const void *model_map; + uint64_t model_size; + uint32_t layer; + uint32_t n_total_expert; + uint32_t n_selected; + uint32_t missing_mask; + uint32_t load_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t source_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t n_loads; + uint32_t n_tasks; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint64_t gate_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint64_t up_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint64_t down_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + ds4_gpu_stream_expert_pread_task tasks[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED * 3u]; + double start_ms; + double prepare_ms; +} ds4_gpu_stream_expert_pending_load; + +static ds4_gpu_stream_expert_pending_load g_stream_expert_pending_load; + +typedef struct { + int active; + const void *model_map; + uint64_t model_size; + uint32_t layer; + uint32_t n_total_expert; + uint32_t n_selected; + uint64_t gate_offset; + uint64_t up_offset; + uint64_t down_offset; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; +} ds4_gpu_glm_stream_selected_prefetch; + +static ds4_gpu_glm_stream_selected_prefetch g_glm_stream_selected_prefetch; + +typedef struct { + ds4_gpu_stream_expert_pread_task *tasks; + uint32_t n_tasks; + uint32_t worker_index; + uint32_t n_workers; +} ds4_gpu_stream_expert_pread_worker_args; + +static void ds4_gpu_stream_expert_cache_note_pread( + uint32_t layer, + uint64_t bytes, + double ms) { + if (g_stream_expert_cache_pread_bytes > UINT64_MAX - bytes) { + g_stream_expert_cache_pread_bytes = UINT64_MAX; + } else { + g_stream_expert_cache_pread_bytes += bytes; + } + g_stream_expert_cache_pread_ms += ms; + if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) { + if (g_stream_expert_cache_layer_pread_bytes[layer] > UINT64_MAX - bytes) { + g_stream_expert_cache_layer_pread_bytes[layer] = UINT64_MAX; + } else { + g_stream_expert_cache_layer_pread_bytes[layer] += bytes; + } + g_stream_expert_cache_layer_pread_ms[layer] += ms; + } +} + +static uint32_t ds4_gpu_stream_expert_pread_thread_limit(void) { + uint32_t threads = 9; + const char *env = getenv("DS4_METAL_STREAMING_EXPERT_PREAD_THREADS"); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end != env && *end == '\0') { + threads = v > UINT32_MAX ? UINT32_MAX : (uint32_t)v; + } + } + if (threads == 0) threads = 1; + if (threads > 18) threads = 18; + return threads; +} + +static uint32_t ds4_gpu_stream_expert_pread_thread_count(uint32_t n_tasks) { + if (n_tasks <= 1) return n_tasks; + uint32_t threads = ds4_gpu_stream_expert_pread_thread_limit(); + if (threads > n_tasks) threads = n_tasks; + return threads; +} + +static int ds4_gpu_stream_expert_pread_into( + uint64_t offset, + uint64_t len, + uint8_t *dst, + uint64_t *read_bytes, + double *ms_out) { + if (read_bytes) *read_bytes = 0; + if (ms_out) *ms_out = 0.0; + if (g_model_fd < 0 || + !dst || + len == 0 || + offset > (uint64_t)LLONG_MAX || + len > (uint64_t)LLONG_MAX - offset) { + return 0; + } + + const double t0 = ds4_gpu_now_ms(); + uint64_t pos = 0; + int ok = 1; + while (pos < len) { + const uint64_t rem = len - pos; + const size_t want = rem > (uint64_t)SSIZE_MAX ? (size_t)SSIZE_MAX : (size_t)rem; + ssize_t nread; + do { + nread = pread(g_model_fd, dst + pos, want, (off_t)(offset + pos)); + } while (nread < 0 && errno == EINTR); + if (nread <= 0) { + ok = 0; + break; + } + pos += (uint64_t)nread; + } + const double dt = ds4_gpu_now_ms() - t0; + if (read_bytes) *read_bytes = pos; + if (ms_out) *ms_out = dt; + if (!ok || pos != len) { + fprintf(stderr, + "ds4: Metal streaming expert explicit pread failed offset=%.2f GiB len=%.2f MiB read=%.2f MiB\n", + ds4_gpu_gib(offset), + ds4_gpu_mib(len), + ds4_gpu_mib(pos)); + return 0; + } + return 1; +} + +static void *ds4_gpu_stream_expert_pread_worker(void *arg) { + ds4_gpu_stream_expert_pread_worker_args *wa = + (ds4_gpu_stream_expert_pread_worker_args *)arg; + for (uint32_t i = wa->worker_index; i < wa->n_tasks; i += wa->n_workers) { + ds4_gpu_stream_expert_pread_task *task = &wa->tasks[i]; + task->ok = ds4_gpu_stream_expert_pread_into(task->offset, + task->len, + task->dst, + &task->read_bytes, + &task->ms); + } + return NULL; +} + +static pthread_mutex_t g_stream_expert_pread_pool_mutex = PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t g_stream_expert_pread_pool_start_cond = PTHREAD_COND_INITIALIZER; +static pthread_cond_t g_stream_expert_pread_pool_done_cond = PTHREAD_COND_INITIALIZER; +static pthread_t g_stream_expert_pread_pool_threads[18]; +static uint32_t g_stream_expert_pread_pool_thread_count; +static uint32_t g_stream_expert_pread_pool_active_workers; +static uint32_t g_stream_expert_pread_pool_remaining_workers; +static uint32_t g_stream_expert_pread_pool_n_tasks; +static uint32_t g_stream_expert_pread_pool_next_task; +static uint64_t g_stream_expert_pread_pool_generation; +static ds4_gpu_stream_expert_pread_task *g_stream_expert_pread_pool_tasks; +static int g_stream_expert_pread_pool_initialized; +static int g_stream_expert_pread_pool_stopping; + +static int ds4_gpu_stream_expert_pread_pool_enabled(void) { + const char *env = getenv("DS4_METAL_STREAMING_EXPERT_PREAD_POOL"); + return !(env && strcmp(env, "0") == 0); +} + +static void *ds4_gpu_stream_expert_pread_pool_worker(void *arg) { + const uint32_t worker_index = (uint32_t)(uintptr_t)arg; + uint64_t seen_generation = 0; + + for (;;) { + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + while (!g_stream_expert_pread_pool_stopping && + g_stream_expert_pread_pool_generation == seen_generation) { + pthread_cond_wait(&g_stream_expert_pread_pool_start_cond, + &g_stream_expert_pread_pool_mutex); + } + if (g_stream_expert_pread_pool_stopping) { + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + break; + } + + seen_generation = g_stream_expert_pread_pool_generation; + if (worker_index >= g_stream_expert_pread_pool_active_workers) { + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + continue; + } + + for (;;) { + const uint32_t task_index = + g_stream_expert_pread_pool_next_task++; + if (task_index >= g_stream_expert_pread_pool_n_tasks) break; + + ds4_gpu_stream_expert_pread_task *task = + &g_stream_expert_pread_pool_tasks[task_index]; + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + + task->ok = ds4_gpu_stream_expert_pread_into(task->offset, + task->len, + task->dst, + &task->read_bytes, + &task->ms); + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + } + + if (g_stream_expert_pread_pool_remaining_workers > 0 && + --g_stream_expert_pread_pool_remaining_workers == 0) { + g_stream_expert_pread_pool_tasks = NULL; + g_stream_expert_pread_pool_n_tasks = 0; + g_stream_expert_pread_pool_active_workers = 0; + pthread_cond_signal(&g_stream_expert_pread_pool_done_cond); + } + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + } + + return NULL; +} + +static int ds4_gpu_stream_expert_pread_pool_init(uint32_t n_threads) { + if (g_stream_expert_pread_pool_initialized) return 1; + if (!ds4_gpu_stream_expert_pread_pool_enabled() || n_threads <= 1) return 0; + if (n_threads > 18) n_threads = 18; + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + g_stream_expert_pread_pool_thread_count = n_threads; + g_stream_expert_pread_pool_stopping = 0; + g_stream_expert_pread_pool_generation = 0; + g_stream_expert_pread_pool_tasks = NULL; + g_stream_expert_pread_pool_n_tasks = 0; + g_stream_expert_pread_pool_next_task = 0; + g_stream_expert_pread_pool_active_workers = 0; + g_stream_expert_pread_pool_remaining_workers = 0; + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + + uint32_t started = 0; + for (uint32_t i = 0; i < n_threads; i++) { + const int rc = pthread_create(&g_stream_expert_pread_pool_threads[i], + NULL, + ds4_gpu_stream_expert_pread_pool_worker, + (void *)(uintptr_t)i); + if (rc != 0) { + fprintf(stderr, + "ds4: Metal streaming expert pread pool thread creation failed: %s\n", + strerror(rc)); + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + g_stream_expert_pread_pool_stopping = 1; + g_stream_expert_pread_pool_generation++; + pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond); + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + for (uint32_t j = 0; j < started; j++) { + (void)pthread_join(g_stream_expert_pread_pool_threads[j], NULL); + } + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + g_stream_expert_pread_pool_thread_count = 0; + g_stream_expert_pread_pool_stopping = 0; + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 0; + } + started++; + } + + g_stream_expert_pread_pool_initialized = 1; + return 1; +} + +static int ds4_gpu_stream_expert_pread_pool_begin( + ds4_gpu_stream_expert_pread_task *tasks, + uint32_t n_tasks, + uint32_t n_workers) { + if (n_workers <= 1) return 0; + const uint32_t limit = ds4_gpu_stream_expert_pread_thread_limit(); + if (!ds4_gpu_stream_expert_pread_pool_init(limit)) return 0; + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + if (!g_stream_expert_pread_pool_initialized || + g_stream_expert_pread_pool_stopping || + g_stream_expert_pread_pool_thread_count == 0) { + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 0; + } + if (n_workers > g_stream_expert_pread_pool_thread_count) { + n_workers = g_stream_expert_pread_pool_thread_count; + } + if (n_workers == 0) { + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 0; + } + + if (g_stream_expert_pread_pool_remaining_workers != 0 || + g_stream_expert_pread_pool_tasks != NULL) { + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 0; + } + + g_stream_expert_pread_pool_tasks = tasks; + g_stream_expert_pread_pool_n_tasks = n_tasks; + g_stream_expert_pread_pool_next_task = 0; + g_stream_expert_pread_pool_active_workers = n_workers; + g_stream_expert_pread_pool_remaining_workers = n_workers; + g_stream_expert_pread_pool_generation++; + pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond); + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 1; +} + +static int ds4_gpu_stream_expert_pread_pool_wait(void) { + if (!g_stream_expert_pread_pool_initialized) return 0; + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + while (g_stream_expert_pread_pool_remaining_workers != 0) { + pthread_cond_wait(&g_stream_expert_pread_pool_done_cond, + &g_stream_expert_pread_pool_mutex); + } + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + return 1; +} + +static int ds4_gpu_stream_expert_pread_pool_dispatch( + ds4_gpu_stream_expert_pread_task *tasks, + uint32_t n_tasks, + uint32_t n_workers) { + if (!ds4_gpu_stream_expert_pread_pool_begin(tasks, n_tasks, n_workers)) { + return 0; + } + return ds4_gpu_stream_expert_pread_pool_wait(); +} + +static void ds4_gpu_stream_expert_pread_pool_shutdown(void) { + if (!g_stream_expert_pread_pool_initialized) return; + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + g_stream_expert_pread_pool_stopping = 1; + g_stream_expert_pread_pool_generation++; + pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond); + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); + + const uint32_t n_threads = g_stream_expert_pread_pool_thread_count; + for (uint32_t i = 0; i < n_threads; i++) { + (void)pthread_join(g_stream_expert_pread_pool_threads[i], NULL); + } + + pthread_mutex_lock(&g_stream_expert_pread_pool_mutex); + g_stream_expert_pread_pool_thread_count = 0; + g_stream_expert_pread_pool_active_workers = 0; + g_stream_expert_pread_pool_remaining_workers = 0; + g_stream_expert_pread_pool_n_tasks = 0; + g_stream_expert_pread_pool_next_task = 0; + g_stream_expert_pread_pool_tasks = NULL; + g_stream_expert_pread_pool_initialized = 0; + g_stream_expert_pread_pool_stopping = 0; + pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex); +} + +static int ds4_gpu_stream_expert_pread_tasks( + ds4_gpu_stream_expert_pread_task *tasks, + uint32_t n_tasks, + uint64_t *total_bytes, + double *wall_ms) { + if (total_bytes) *total_bytes = 0; + if (wall_ms) *wall_ms = 0.0; + if (!tasks || n_tasks == 0) return 1; + + const uint32_t n_workers = + ds4_gpu_stream_expert_pread_thread_count(n_tasks); + const double t0 = ds4_gpu_now_ms(); + int ok = 1; + if (n_workers <= 1) { + ds4_gpu_stream_expert_pread_worker_args wa = { + .tasks = tasks, + .n_tasks = n_tasks, + .worker_index = 0, + .n_workers = 1, + }; + (void)ds4_gpu_stream_expert_pread_worker(&wa); + } else if (!ds4_gpu_stream_expert_pread_pool_dispatch(tasks, + n_tasks, + n_workers)) { + pthread_t threads[18]; + ds4_gpu_stream_expert_pread_worker_args args[18]; + uint32_t started = 0; + for (uint32_t i = 0; i < n_workers; i++) { + args[i].tasks = tasks; + args[i].n_tasks = n_tasks; + args[i].worker_index = i; + args[i].n_workers = n_workers; + const int rc = pthread_create(&threads[i], + NULL, + ds4_gpu_stream_expert_pread_worker, + &args[i]); + if (rc != 0) { + fprintf(stderr, + "ds4: Metal streaming expert pread thread creation failed: %s\n", + strerror(rc)); + ok = 0; + break; + } + started++; + } + for (uint32_t i = 0; i < started; i++) { + if (pthread_join(threads[i], NULL) != 0) ok = 0; + } + } + const double dt = ds4_gpu_now_ms() - t0; + + uint64_t bytes = 0; + for (uint32_t i = 0; i < n_tasks; i++) { + if (!tasks[i].ok) ok = 0; + if (bytes > UINT64_MAX - tasks[i].read_bytes) { + bytes = UINT64_MAX; + } else { + bytes += tasks[i].read_bytes; + } + } + if (total_bytes) *total_bytes = bytes; + if (wall_ms) *wall_ms = dt; + return ok; +} + +static id ds4_gpu_stream_expert_alloc_buffer( + uint64_t len, + NSString *label) { + if (!g_device || + len == 0 || + len > (uint64_t)NSUIntegerMax) { + return nil; + } + + id buffer = [g_device newBufferWithLength:(NSUInteger)len + options:MTLResourceStorageModeShared]; + if (!buffer) { + fprintf(stderr, + "ds4: Metal streaming expert explicit buffer allocation failed (%.2f MiB)\n", + ds4_gpu_mib(len)); + return nil; + } + buffer.label = label; + g_stream_expert_cache_buffer_allocs++; + return buffer; +} + +static int ds4_gpu_stream_expert_combined_buffer_enabled(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_COMBINED_BUFFER") == NULL; +} + +static int ds4_gpu_stream_expert_slab_enabled(void) { + return ds4_gpu_stream_expert_combined_buffer_enabled() && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_SLABS") == NULL; +} + +/* + * Large PRO caches otherwise create thousands of small shared Metal buffers. + * Slabs keep the buffer object set small while locking pages only for slots + * that actually hold a streamed expert. + */ +static uint64_t ds4_gpu_stream_expert_slab_target_bytes(void) { + const uint64_t mib = 1024ull * 1024ull; + uint64_t target = 4096ull * mib; + const char *env = getenv("DS4_METAL_STREAMING_EXPERT_SLAB_MB"); + if (env && env[0]) { + char *end = NULL; + unsigned long long v = strtoull(env, &end, 10); + if (end != env && *end == '\0' && v != 0) { + target = v > UINT64_MAX / mib ? UINT64_MAX : (uint64_t)v * mib; + } + } + return target; +} + +static id ds4_gpu_stream_expert_alloc_slab_buffer( + uint64_t len, + NSString *label) { + if (!g_device || + len == 0 || + len > (uint64_t)NSUIntegerMax) { + return nil; + } + + id buffer = [g_device newBufferWithLength:(NSUInteger)len + options:MTLResourceStorageModeShared]; + if (!buffer) { + fprintf(stderr, + "ds4: Metal streaming expert slab allocation failed (%.2f MiB)\n", + ds4_gpu_mib(len)); + return nil; + } + buffer.label = label; + g_stream_expert_cache_buffer_allocs++; + return buffer; +} + +static int ds4_gpu_stream_expert_slab_slot_range( + uint32_t slot, + uint32_t *slab_index, + uint64_t *slot_base) { + for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) { + const uint32_t start = g_stream_expert_cache_slab_start_slot[i]; + const uint32_t count = g_stream_expert_cache_slab_slot_count[i]; + if (slot < start || slot >= start + count) continue; + if (slab_index) *slab_index = i; + if (slot_base) { + *slot_base = + (uint64_t)(slot - start) * g_stream_expert_cache_slab_slot_bytes; + } + return 1; + } + return 0; +} + +static int ds4_gpu_stream_expert_slab_slot_for_buffer( + id buffer, + NSUInteger gate_inner, + uint32_t *slot_out) { + if (!buffer || g_stream_expert_cache_slab_slot_bytes == 0 || !slot_out) { + return 0; + } + for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) { + if (g_stream_expert_cache_slabs[i] != buffer) continue; + const uint64_t inner = (uint64_t)gate_inner; + const uint64_t slot_bytes = g_stream_expert_cache_slab_slot_bytes; + if (slot_bytes == 0 || inner % slot_bytes != 0) return 0; + const uint64_t local_slot = inner / slot_bytes; + if (local_slot >= g_stream_expert_cache_slab_slot_count[i]) return 0; + const uint64_t slot = + (uint64_t)g_stream_expert_cache_slab_start_slot[i] + local_slot; + if (slot > UINT32_MAX) return 0; + *slot_out = (uint32_t)slot; + return 1; + } + return 0; +} + +static void ds4_gpu_stream_expert_slab_push_free_slot(uint32_t slot) { + if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES || + g_stream_expert_cache_free_slot_count >= + DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) { + return; + } + g_stream_expert_cache_free_slots[g_stream_expert_cache_free_slot_count++] = + slot; +} + +static int ds4_gpu_stream_expert_slab_slot_buffers( + uint32_t slot, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + __strong id *gate_buf, + __strong id *up_buf, + __strong id *down_buf, + NSUInteger *gate_inner, + NSUInteger *up_inner, + NSUInteger *down_inner) { + uint32_t slab = UINT32_MAX; + uint64_t base = 0; + if (!ds4_gpu_stream_expert_slab_slot_range(slot, &slab, &base) || + slab >= g_stream_expert_cache_slab_count || + !g_stream_expert_cache_slabs[slab]) { + return 0; + } + if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull || + base > UINT64_MAX - (gate_expert_bytes * 2ull + down_expert_bytes) || + base + gate_expert_bytes * 2ull + down_expert_bytes > + (uint64_t)NSUIntegerMax) { + return 0; + } + id b = g_stream_expert_cache_slabs[slab]; + *gate_buf = b; + *up_buf = b; + *down_buf = b; + *gate_inner = (NSUInteger)base; + *up_inner = (NSUInteger)(base + gate_expert_bytes); + *down_inner = (NSUInteger)(base + gate_expert_bytes * 2ull); + return 1; +} + +static int ds4_gpu_stream_expert_alloc_slab_slot( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + __strong id *gate_buf, + __strong id *up_buf, + __strong id *down_buf, + NSUInteger *gate_inner, + NSUInteger *up_inner, + NSUInteger *down_inner) { + if (!ds4_gpu_stream_expert_slab_enabled() || + !gate_buf || !up_buf || !down_buf || + !gate_inner || !up_inner || !down_inner || + gate_expert_bytes == 0 || down_expert_bytes == 0 || + gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { + return 0; + } + + uint64_t slot_bytes = gate_expert_bytes * 2ull + down_expert_bytes; + if (slot_bytes == 0 || slot_bytes > (uint64_t)NSUIntegerMax) return 0; + const uint64_t page = (uint64_t)getpagesize(); + if (page != 0) { + slot_bytes = round_up_u64(slot_bytes, page); + if (slot_bytes == 0 || slot_bytes > (uint64_t)NSUIntegerMax) return 0; + } + if (g_stream_expert_cache_slab_slot_bytes != 0 && + g_stream_expert_cache_slab_slot_bytes != slot_bytes) { + return 0; + } + g_stream_expert_cache_slab_slot_bytes = slot_bytes; + + if (g_stream_expert_cache_free_slot_count != 0) { + const uint32_t slot = + g_stream_expert_cache_free_slots[--g_stream_expert_cache_free_slot_count]; + return ds4_gpu_stream_expert_slab_slot_buffers(slot, + gate_expert_bytes, + down_expert_bytes, + gate_buf, + up_buf, + down_buf, + gate_inner, + up_inner, + down_inner); + } + + uint32_t slab = g_stream_expert_cache_slab_count; + if (slab != 0 && + g_stream_expert_cache_slab_slots_used[slab - 1] < + g_stream_expert_cache_slab_slot_count[slab - 1]) { + slab--; + } else { + if (g_stream_expert_cache_slab_count >= + DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS) { + return 0; + } + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget != 0 && g_stream_expert_cache_slab_total_slots >= budget) { + return 0; + } + uint64_t target = ds4_gpu_stream_expert_slab_target_bytes(); + uint64_t slots64 = target / slot_bytes; + if (slots64 == 0) slots64 = 1; + if (slots64 > UINT32_MAX) slots64 = UINT32_MAX; + uint32_t slots = (uint32_t)slots64; + if (budget != 0) { + const uint32_t remaining = + budget - g_stream_expert_cache_slab_total_slots; + if (slots > remaining) slots = remaining; + } + if (slots == 0) return 0; + id slab_buffer = nil; + while (slots != 0) { + if ((uint64_t)slots <= UINT64_MAX / slot_bytes && + (uint64_t)slots * slot_bytes <= (uint64_t)NSUIntegerMax) { + slab_buffer = + ds4_gpu_stream_expert_alloc_slab_buffer( + (uint64_t)slots * slot_bytes, + @"ds4_stream_expert_slab"); + if (slab_buffer) break; + } + slots /= 2u; + } + if (!slab_buffer || slots == 0) return 0; + + slab = g_stream_expert_cache_slab_count++; + g_stream_expert_cache_slabs[slab] = slab_buffer; + g_stream_expert_cache_slab_start_slot[slab] = + g_stream_expert_cache_slab_total_slots; + g_stream_expert_cache_slab_slot_count[slab] = slots; + g_stream_expert_cache_slab_slots_used[slab] = 0; + g_stream_expert_cache_slab_total_slots += slots; + } + + const uint32_t local_slot = g_stream_expert_cache_slab_slots_used[slab]++; + const uint32_t slot = + g_stream_expert_cache_slab_start_slot[slab] + local_slot; + return ds4_gpu_stream_expert_slab_slot_buffers(slot, + gate_expert_bytes, + down_expert_bytes, + gate_buf, + up_buf, + down_buf, + gate_inner, + up_inner, + down_inner); +} + +static uint64_t ds4_gpu_stream_expert_buffer_object_count( + id gate, + id up, + id down) { + if (!gate || !up || !down) return 0; + if (gate == up && gate == down) return 1; + if (gate == up || gate == down || up == down) return 2; + return 3; +} + +static int ds4_gpu_stream_expert_evict_dontneed_enabled(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_EVICT_DONTNEED") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_EVICT_DONTNEED") == NULL; +} + +static void ds4_gpu_stream_expert_evict_dontneed_range( + const void *model_map, + uint64_t model_size, + uint64_t offset, + uint64_t len) { + if (!ds4_gpu_stream_expert_evict_dontneed_enabled() || + !model_map || + model_size == 0 || + offset > model_size || + len == 0 || + len > model_size - offset) { + return; + } + +#if defined(POSIX_MADV_DONTNEED) + const uint64_t page = (uint64_t)getpagesize(); + const uint64_t page_offset = offset & ~(page - 1); + const uint64_t leading = offset - page_offset; + if (len > UINT64_MAX - leading || + leading + len > UINT64_MAX - (page - 1)) { + return; + } + uint64_t advise_bytes = round_up_u64(leading + len, page); + if (advise_bytes > model_size - page_offset) { + advise_bytes = model_size - page_offset; + } + if (advise_bytes == 0 || advise_bytes > (uint64_t)SIZE_MAX) return; + + const uintptr_t base = (uintptr_t)model_map; + if (page_offset > (uint64_t)(UINTPTR_MAX - base)) return; + void *addr = (void *)(base + (uintptr_t)page_offset); + const int rc = posix_madvise(addr, (size_t)advise_bytes, POSIX_MADV_DONTNEED); + if (rc == 0) { + if (g_stream_expert_cache_evict_advise_bytes > UINT64_MAX - advise_bytes) { + g_stream_expert_cache_evict_advise_bytes = UINT64_MAX; + } else { + g_stream_expert_cache_evict_advise_bytes += advise_bytes; + } + } else if (getenv("DS4_METAL_STREAMING_EXPERT_EVICT_DONTNEED_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal streaming expert evict DONTNEED failed offset=%.2f GiB len=%.2f MiB: %s\n", + ds4_gpu_gib(offset), + ds4_gpu_mib(len), + strerror(rc)); + } +#else + (void)model_map; + (void)model_size; + (void)offset; + (void)len; +#endif +} + +static int ds4_gpu_stream_expert_split_requested(void) { + return g_ssd_streaming_mode; +} + +static uint32_t ds4_gpu_stream_expert_split_min_decode_tokens(void) { + return 4; +} + +static uint32_t ds4_gpu_stream_expert_split_min_cached(void) { + uint32_t min_cached = 1024; + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget != 0 && budget < min_cached * 2u) { + min_cached = budget / 2u; + } + return min_cached; +} + +static int ds4_gpu_stream_expert_split_ready(void) { + if (!ds4_gpu_stream_expert_split_requested()) return 0; + if (g_stream_expert_cache_decode_tokens < + ds4_gpu_stream_expert_split_min_decode_tokens()) { + return 0; + } + return g_stream_expert_cache_entry_count >= + ds4_gpu_stream_expert_split_min_cached(); +} + +static void ds4_gpu_stream_expert_cache_decay_route_hotness(void) { + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + g_stream_expert_cache_route_hotness[layer][expert] >>= 1; + } + } +} + +void ds4_gpu_stream_expert_cache_reset_route_hotness(void) { + memset(g_stream_expert_cache_route_hotness, + 0, + sizeof(g_stream_expert_cache_route_hotness)); + g_stream_expert_cache_hotness_decay_token = + g_stream_expert_cache_decode_tokens; +} + +static void ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(void) { + if (g_stream_expert_cache_decode_tokens == 0) return; + if (g_stream_expert_cache_hotness_decay_token == 0) { + g_stream_expert_cache_hotness_decay_token = + g_stream_expert_cache_decode_tokens; + return; + } + while (g_stream_expert_cache_decode_tokens - + g_stream_expert_cache_hotness_decay_token >= + DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS) { + ds4_gpu_stream_expert_cache_decay_route_hotness(); + g_stream_expert_cache_hotness_decay_token += + DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS; + } +} + +static void ds4_gpu_stream_expert_cache_note_route_hotness( + uint32_t layer, + uint32_t expert, + uint32_t amount) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + amount == 0) { + return; + } + uint32_t *hotness = &g_stream_expert_cache_route_hotness[layer][expert]; + if (*hotness > UINT32_MAX - amount) { + *hotness = UINT32_MAX; + } else { + *hotness += amount; + } +} + +static void ds4_gpu_stream_expert_cache_note_selected_hotness( + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_selected) { + if (!selected_ids || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_selected == 0) { + return; + } + ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); + for (uint32_t i = 0; i < n_selected; i++) { + if (selected_ids[i] < 0 || + selected_ids[i] >= + (int32_t)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + continue; + } + ds4_gpu_stream_expert_cache_note_route_hotness( + layer, + (uint32_t)selected_ids[i], + 1); + } +} + +static void ds4_gpu_stream_expert_cache_note_frequency_hotness( + uint32_t layer, + const uint32_t *frequency, + uint32_t n_total_expert) { + if (!frequency || layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) { + return; + } + if (n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + n_total_expert = DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + } + ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); + for (uint32_t expert = 0; expert < n_total_expert; expert++) { + ds4_gpu_stream_expert_cache_note_route_hotness(layer, + expert, + frequency[expert]); + } +} + +static void ds4_gpu_stream_expert_cache_note_token(uint32_t layer_index) { + if (!g_ssd_streaming_mode || layer_index != 0 || + g_stream_expert_cache_decode_tokens == UINT64_MAX) { + return; + } + g_stream_expert_cache_decode_tokens++; + ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); +} + +static void ds4_gpu_stream_expert_cache_note_decode_token(void) { + if (!g_ssd_streaming_mode || + g_stream_expert_cache_decode_tokens == UINT64_MAX) { + return; + } + g_stream_expert_cache_decode_tokens++; + ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); +} + +static int ds4_gpu_stream_compact_addr_requested(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_ENABLE_STREAMING_COMPACT_ADDR") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_COMPACT_ADDR") == NULL && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL; +} + +static int ds4_gpu_stream_expert_addr_table_requested(void) { + return g_ssd_streaming_mode && + (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL || + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL || + g_stream_prefill_batch_selected_addr_building || + g_glm_stream_expert_addr_table_building || + (getenv("DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") == NULL) || + ds4_gpu_stream_expert_split_requested()) && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL; +} + +static int ds4_gpu_stream_expert_addr_table_kernel_requested(void) { + return g_ssd_streaming_mode && + (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL || + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL || + ds4_gpu_stream_expert_split_ready()) && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL; +} + +static int ds4_gpu_stream_expert_masked_addr_requested(void) { + return g_ssd_streaming_mode && + (getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL || + ds4_gpu_stream_expert_split_ready()) && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_MASKED_ADDR") == NULL; +} + +static int ds4_gpu_stream_expert_hit_validator_requested(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_HIT_VALIDATOR") == NULL; +} + +static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_max( + uint32_t n_total_expert) { + const char *env = getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MAX"); + if (env && env[0]) { + char *end = NULL; + const long v = strtol(env, &end, 10); + if (end != env) { + if (v <= 0) return 0; + if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; + return (uint32_t)v; + } + } + if (n_total_expert == 384) return 800u; + if (n_total_expert == 256) return 760u; + return 0; +} + +static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_min( + uint32_t n_total_expert) { + const char *env = getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MIN"); + if (env && env[0]) { + char *end = NULL; + const long v = strtol(env, &end, 10); + if (end != env) { + if (v <= 0) return 0; + if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX; + return (uint32_t)v; + } + } + if (n_total_expert == 384 || n_total_expert == 256) return 2u; + return 0; +} + +static int ds4_gpu_stream_prefill_batch_selected_addr_enabled( + uint32_t n_tokens, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t gate_type, + uint32_t down_type) { + if (!g_ssd_streaming_mode || + n_tokens <= 1 || + n_total_expert == 0 || + n_expert != 6 || + gate_type != DS4_METAL_TENSOR_IQ2_XXS || + down_type != DS4_METAL_TENSOR_Q2_K || + g_quality_mode || + getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL || + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL || + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL || + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") != NULL) { + return 0; + } + /* All unique experts for one layer must fit simultaneously because the + * address-table kernels consume them in one dispatch. Once the global + * cache fills, preparation reuses entries owned by other layers. */ + if (ds4_gpu_stream_expert_cache_configured_count() < n_total_expert) { + return 0; + } + if (getenv("DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL) { + return 1; + } + const uint32_t max_tokens = + ds4_gpu_stream_prefill_batch_selected_addr_auto_max(n_total_expert); + const uint32_t min_tokens = + ds4_gpu_stream_prefill_batch_selected_addr_auto_min(n_total_expert); + return max_tokens != 0 && n_tokens >= min_tokens && n_tokens <= max_tokens; +} + +static int ds4_gpu_glm_streaming_prefill_full_layer_active(void) { + return g_glm_streaming_prefill_full_layer_runtime || + getenv("DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER") != NULL; +} + +static int ds4_gpu_stream_full_expert_addr_table_requested(void) { + return g_ssd_streaming_mode && + getenv("DS4_METAL_ENABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") == NULL; +} + +static uint64_t ds4_gpu_buffer_address(id buffer, NSUInteger inner) { + if (!buffer) return 0; +#if TARGET_OS_OSX + if (@available(macOS 13.0, *)) { + return (uint64_t)[buffer gpuAddress] + (uint64_t)inner; + } +#endif + return 0; +} + +static int ds4_gpu_stream_compact_addr_ensure_buffers(uint32_t layer) { + if (!g_device || layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return 0; + + const NSUInteger addr_bytes = 6u * sizeof(uint64_t); + const NSUInteger ids_bytes = 6u * sizeof(int32_t); + for (uint32_t i = 0; i < 4; i++) { + id current = nil; + NSUInteger bytes = addr_bytes; + NSString *label = @"ds4_stream_compact_gate_addresses"; + switch (i) { + case 0: + current = g_stream_compact_gate_addr_buffers[layer]; + label = @"ds4_stream_compact_gate_addresses"; + break; + case 1: + current = g_stream_compact_up_addr_buffers[layer]; + label = @"ds4_stream_compact_up_addresses"; + break; + case 2: + current = g_stream_compact_down_addr_buffers[layer]; + label = @"ds4_stream_compact_down_addresses"; + break; + default: + current = g_stream_compact_selected_buffers[layer]; + bytes = ids_bytes; + label = @"ds4_stream_compact_selected_ids"; + break; + } + if (current) continue; + id b = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!b) { + fprintf(stderr, "ds4: Metal streaming compact address buffer allocation failed\n"); + return 0; + } + b.label = label; + memset([b contents], 0, bytes); + [b didModifyRange:NSMakeRange(0, bytes)]; + switch (i) { + case 0: + g_stream_compact_gate_addr_buffers[layer] = b; + break; + case 1: + g_stream_compact_up_addr_buffers[layer] = b; + break; + case 2: + g_stream_compact_down_addr_buffers[layer] = b; + break; + default: + g_stream_compact_selected_buffers[layer] = b; + break; + } + } + return 1; +} + +static int ds4_gpu_stream_compact_addr_prepare( + uint32_t layer, + ds4_gpu_stream_expert_cache_entry * const entries[6], + uint32_t n_entries, + id *gate_addrs, + id *up_addrs, + id *down_addrs, + id *selected_ids) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + !entries || !gate_addrs || !up_addrs || !down_addrs || !selected_ids || + n_entries == 0 || n_entries > 6) { + return 0; + } + if (!ds4_gpu_stream_compact_addr_ensure_buffers(layer)) return 0; + + uint64_t gate_values[6] = {0, 0, 0, 0, 0, 0}; + uint64_t up_values[6] = {0, 0, 0, 0, 0, 0}; + uint64_t down_values[6] = {0, 0, 0, 0, 0, 0}; + int32_t slot_ids[6] = {0, 1, 2, 3, 4, 5}; + + for (uint32_t i = 0; i < n_entries; i++) { + ds4_gpu_stream_expert_cache_entry *e = entries[i]; + if (!e || !e->gate_buffer || !e->up_buffer || !e->down_buffer) { + return 0; + } + gate_values[i] = ds4_gpu_buffer_address(e->gate_buffer, e->gate_inner); + up_values[i] = ds4_gpu_buffer_address(e->up_buffer, e->up_inner); + down_values[i] = ds4_gpu_buffer_address(e->down_buffer, e->down_inner); + if (gate_values[i] == 0 || up_values[i] == 0 || down_values[i] == 0) { + fprintf(stderr, "ds4: Metal streaming compact address path requires GPU addresses\n"); + return 0; + } + } + + const NSUInteger addr_bytes = 6u * sizeof(uint64_t); + const NSUInteger ids_bytes = 6u * sizeof(int32_t); + id gb = g_stream_compact_gate_addr_buffers[layer]; + id ub = g_stream_compact_up_addr_buffers[layer]; + id db = g_stream_compact_down_addr_buffers[layer]; + id ib = g_stream_compact_selected_buffers[layer]; + memcpy([gb contents], gate_values, addr_bytes); + memcpy([ub contents], up_values, addr_bytes); + memcpy([db contents], down_values, addr_bytes); + memcpy([ib contents], slot_ids, ids_bytes); + [gb didModifyRange:NSMakeRange(0, addr_bytes)]; + [ub didModifyRange:NSMakeRange(0, addr_bytes)]; + [db didModifyRange:NSMakeRange(0, addr_bytes)]; + [ib didModifyRange:NSMakeRange(0, ids_bytes)]; + + *gate_addrs = gb; + *up_addrs = ub; + *down_addrs = db; + *selected_ids = ib; + return 1; +} + +static int ds4_gpu_stream_selected_ids_prepare( + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_selected, + id *selected_buf, + NSUInteger *selected_off) { + if (!g_device || + !selected_ids || + !selected_buf || + !selected_off || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_selected == 0 || + n_selected > DS4_METAL_MAX_ROUTED_EXPERT_USED) { + return 0; + } + + const NSUInteger bytes = + (NSUInteger)DS4_METAL_MAX_ROUTED_EXPERT_USED * sizeof(int32_t); + id b = g_stream_selected_id_buffers[layer]; + if (!b) { + b = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!b) { + fprintf(stderr, "ds4: Metal streaming selected-id buffer allocation failed\n"); + return 0; + } + b.label = @"ds4_stream_selected_ids"; + g_stream_selected_id_buffers[layer] = b; + } + + int32_t ids[DS4_METAL_MAX_ROUTED_EXPERT_USED] = {0}; + memcpy(ids, selected_ids, (size_t)n_selected * sizeof(ids[0])); + memcpy([b contents], ids, bytes); + [b didModifyRange:NSMakeRange(0, bytes)]; + + *selected_buf = b; + *selected_off = 0; + return 1; +} + +static int ds4_gpu_stream_expert_cache_ensure_addr_buffers(uint32_t layer) { + if (!g_device || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) { + return 0; + } + + const NSUInteger bytes = + (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t); + id buffers[3] = { + g_stream_expert_cache_gate_addr_buffers[layer], + g_stream_expert_cache_up_addr_buffers[layer], + g_stream_expert_cache_down_addr_buffers[layer], + }; + for (uint32_t i = 0; i < 3; i++) { + if (buffers[i]) continue; + id b = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!b) { + fprintf(stderr, "ds4: Metal streaming expert address table allocation failed\n"); + return 0; + } + b.label = + i == 0 ? @"ds4_stream_expert_gate_addresses" : + (i == 1 ? @"ds4_stream_expert_up_addresses" : + @"ds4_stream_expert_down_addresses"); + memset([b contents], 0, bytes); + [b didModifyRange:NSMakeRange(0, bytes)]; + if (i == 0) { + g_stream_expert_cache_gate_addr_buffers[layer] = b; + } else if (i == 1) { + g_stream_expert_cache_up_addr_buffers[layer] = b; + } else { + g_stream_expert_cache_down_addr_buffers[layer] = b; + } + } + return 1; +} + +static void ds4_gpu_stream_expert_cache_zero_addr_slot(uint32_t layer, uint32_t expert) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + return; + } + + id buffers[3] = { + g_stream_expert_cache_gate_addr_buffers[layer], + g_stream_expert_cache_up_addr_buffers[layer], + g_stream_expert_cache_down_addr_buffers[layer], + }; + const NSUInteger off = (NSUInteger)expert * sizeof(uint64_t); + for (uint32_t i = 0; i < 3; i++) { + if (!buffers[i]) continue; + uint64_t *addr = (uint64_t *)((uint8_t *)[buffers[i] contents] + off); + *addr = 0; + [buffers[i] didModifyRange:NSMakeRange(off, sizeof(uint64_t))]; + } +} + +static int ds4_gpu_stream_expert_cache_set_addr_slot_raw( + uint32_t layer, + uint32_t expert, + id gate_buffer, + NSUInteger gate_inner, + id up_buffer, + NSUInteger up_inner, + id down_buffer, + NSUInteger down_inner) { + if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) return 0; + + const uint64_t values[3] = { + ds4_gpu_buffer_address(gate_buffer, gate_inner), + ds4_gpu_buffer_address(up_buffer, up_inner), + ds4_gpu_buffer_address(down_buffer, down_inner), + }; + if (values[0] == 0 || values[1] == 0 || values[2] == 0) { + fprintf(stderr, "ds4: Metal streaming expert address table requires GPU addresses\n"); + return 0; + } + + id buffers[3] = { + g_stream_expert_cache_gate_addr_buffers[layer], + g_stream_expert_cache_up_addr_buffers[layer], + g_stream_expert_cache_down_addr_buffers[layer], + }; + const NSUInteger off = (NSUInteger)expert * sizeof(uint64_t); + for (uint32_t i = 0; i < 3; i++) { + uint64_t *addr = (uint64_t *)((uint8_t *)[buffers[i] contents] + off); + *addr = values[i]; + [buffers[i] didModifyRange:NSMakeRange(off, sizeof(uint64_t))]; + } + return 1; +} + +static int ds4_gpu_stream_expert_cache_set_addr_slot( + uint32_t layer, + uint32_t expert, + id gate_buffer, + NSUInteger gate_inner, + id up_buffer, + NSUInteger up_inner, + id down_buffer, + NSUInteger down_inner) { + if (!ds4_gpu_stream_expert_addr_table_requested()) return 1; + return ds4_gpu_stream_expert_cache_set_addr_slot_raw(layer, + expert, + gate_buffer, + gate_inner, + up_buffer, + up_inner, + down_buffer, + down_inner); +} + +static int ds4_gpu_stream_expert_cache_addr_buffers( + uint32_t layer, + id *gate, + id *up, + id *down) { + if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) return 0; + if (gate) *gate = g_stream_expert_cache_gate_addr_buffers[layer]; + if (up) *up = g_stream_expert_cache_up_addr_buffers[layer]; + if (down) *down = g_stream_expert_cache_down_addr_buffers[layer]; + return g_stream_expert_cache_gate_addr_buffers[layer] && + g_stream_expert_cache_up_addr_buffers[layer] && + g_stream_expert_cache_down_addr_buffers[layer]; +} + +static void ds4_gpu_stream_full_expert_addr_clear_layer(uint32_t layer) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return; + ds4_gpu_stream_expert_cache_entry *e = &g_stream_full_expert_addr_entry[layer]; + e->gate_buffer = nil; + e->up_buffer = nil; + e->down_buffer = nil; + e->model_map = NULL; + e->model_size = 0; + e->gate_abs_offset = 0; + e->up_abs_offset = 0; + e->down_abs_offset = 0; + e->gate_expert_bytes = 0; + e->down_expert_bytes = 0; + e->logical_bytes = 0; + e->last_used = 0; + e->use_count = 0; + e->gate_inner = 0; + e->up_inner = 0; + e->down_inner = 0; + e->slab_slot = 0; + e->valid = 0; + e->slab_backed = 0; +} + +static int ds4_gpu_stream_full_expert_addr_table_prepare( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t n_total_expert, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + id *gate_addrs, + id *up_addrs, + id *down_addrs, + ds4_gpu_stream_expert_cache_entry **entry_out) { + if (!ds4_gpu_stream_full_expert_addr_table_requested()) return 0; + if (!model_map || model_size == 0 || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + gate_expert_bytes == 0 || + down_expert_bytes == 0 || + n_total_expert > UINT64_MAX / gate_expert_bytes || + n_total_expert > UINT64_MAX / down_expert_bytes) { + return 0; + } + + const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes; + if (gate_abs_offset > model_size || + up_abs_offset > model_size || + down_abs_offset > model_size || + gate_tensor_bytes > model_size - gate_abs_offset || + gate_tensor_bytes > model_size - up_abs_offset || + down_tensor_bytes > model_size - down_abs_offset) { + return 0; + } + + ds4_gpu_stream_expert_cache_entry *entry = + &g_stream_full_expert_addr_entry[layer]; + if (!entry->valid || + entry->model_map != model_map || + entry->model_size != model_size || + entry->gate_abs_offset != gate_abs_offset || + entry->up_abs_offset != up_abs_offset || + entry->down_abs_offset != down_abs_offset || + entry->gate_expert_bytes != gate_expert_bytes || + entry->down_expert_bytes != down_expert_bytes || + !entry->gate_buffer || + !entry->up_buffer || + !entry->down_buffer) { + ds4_gpu_stream_full_expert_addr_clear_layer(layer); + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gate_buf = + ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + gate_abs_offset, + gate_tensor_bytes, + &gate_inner); + id up_buf = + ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + up_abs_offset, + gate_tensor_bytes, + &up_inner); + id down_buf = + ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + down_abs_offset, + down_tensor_bytes, + &down_inner); + if (!gate_buf || !up_buf || !down_buf) return 0; + + entry->gate_buffer = gate_buf; + entry->up_buffer = up_buf; + entry->down_buffer = down_buf; + entry->model_map = model_map; + entry->model_size = model_size; + entry->gate_abs_offset = gate_abs_offset; + entry->up_abs_offset = up_abs_offset; + entry->down_abs_offset = down_abs_offset; + entry->gate_expert_bytes = gate_expert_bytes; + entry->down_expert_bytes = down_expert_bytes; + entry->logical_bytes = gate_tensor_bytes * 2ull + down_tensor_bytes; + entry->gate_inner = (NSUInteger)gate_inner; + entry->up_inner = (NSUInteger)up_inner; + entry->down_inner = (NSUInteger)down_inner; + entry->valid = 1; + + if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) { + ds4_gpu_stream_full_expert_addr_clear_layer(layer); + return 0; + } + + const uint64_t gate_base = ds4_gpu_buffer_address(entry->gate_buffer, + entry->gate_inner); + const uint64_t up_base = ds4_gpu_buffer_address(entry->up_buffer, + entry->up_inner); + const uint64_t down_base = ds4_gpu_buffer_address(entry->down_buffer, + entry->down_inner); + if (gate_base == 0 || up_base == 0 || down_base == 0) { + fprintf(stderr, "ds4: Metal full streaming expert address table requires GPU addresses\n"); + ds4_gpu_stream_full_expert_addr_clear_layer(layer); + return 0; + } + + id buffers[3] = { + g_stream_expert_cache_gate_addr_buffers[layer], + g_stream_expert_cache_up_addr_buffers[layer], + g_stream_expert_cache_down_addr_buffers[layer], + }; + uint64_t *gate_addr = (uint64_t *)[buffers[0] contents]; + uint64_t *up_addr = (uint64_t *)[buffers[1] contents]; + uint64_t *down_addr = (uint64_t *)[buffers[2] contents]; + for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) { + if (expert < n_total_expert) { + gate_addr[expert] = gate_base + (uint64_t)expert * gate_expert_bytes; + up_addr[expert] = up_base + (uint64_t)expert * gate_expert_bytes; + down_addr[expert] = down_base + (uint64_t)expert * down_expert_bytes; + } else { + gate_addr[expert] = 0; + up_addr[expert] = 0; + down_addr[expert] = 0; + } + } + const NSUInteger bytes = + (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t); + for (uint32_t i = 0; i < 3; i++) { + [buffers[i] didModifyRange:NSMakeRange(0, bytes)]; + } + } + + if (gate_addrs) *gate_addrs = g_stream_expert_cache_gate_addr_buffers[layer]; + if (up_addrs) *up_addrs = g_stream_expert_cache_up_addr_buffers[layer]; + if (down_addrs) *down_addrs = g_stream_expert_cache_down_addr_buffers[layer]; + if (entry_out) *entry_out = entry; + return entry->valid && + g_stream_expert_cache_gate_addr_buffers[layer] && + g_stream_expert_cache_up_addr_buffers[layer] && + g_stream_expert_cache_down_addr_buffers[layer]; +} + +static id ds4_gpu_stream_expert_validate_status_buffer(void) { + if (!g_device) return nil; + if (g_stream_expert_validate_status_buffer) { + return g_stream_expert_validate_status_buffer; + } + + const NSUInteger bytes = + (NSUInteger)DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS * sizeof(uint32_t); + id b = [g_device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!b) { + fprintf(stderr, "ds4: Metal streaming expert validator allocation failed\n"); + return nil; + } + b.label = @"ds4_stream_expert_validate_status"; + memset([b contents], 0, bytes); + [b didModifyRange:NSMakeRange(0, bytes)]; + g_stream_expert_validate_status_buffer = b; + return b; +} + +static int ds4_gpu_encode_stream_expert_cache_validate( + id cb, + const ds4_gpu_stream_expert_validate_args *args, + id selected, + NSUInteger selected_off, + id gate_addrs, + id up_addrs, + id down_addrs, + id status) { + if (!cb || !args || !selected || !gate_addrs || !up_addrs || !down_addrs || + !status || !g_moe_stream_expert_cache_validate_pipeline || + args->n_total_expert == 0 || args->n_total_expert > 384 || + args->n_expert == 0 || args->n_expert > 6) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_moe_stream_expert_cache_validate_pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:selected offset:selected_off atIndex:1]; + [enc setBuffer:gate_addrs offset:0 atIndex:2]; + [enc setBuffer:up_addrs offset:0 atIndex:3]; + [enc setBuffer:down_addrs offset:0 atIndex:4]; + [enc setBuffer:status offset:0 atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_stream_expert_cache_validate_selected( + const ds4_gpu_tensor *selected, + id gate_addrs, + id up_addrs, + id down_addrs, + uint32_t n_total_expert, + uint32_t n_expert, + int32_t selected_ids[6], + uint32_t *all_cached, + uint32_t *miss_mask, + uint32_t *invalid_mask) { + if (!selected || !selected_ids || !all_cached || !miss_mask || !invalid_mask) { + return 0; + } + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id status = ds4_gpu_stream_expert_validate_status_buffer(); + if (!selectedbuf || !status) return 0; + + const NSUInteger status_bytes = + (NSUInteger)DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS * sizeof(uint32_t); + memset([status contents], 0, status_bytes); + [status didModifyRange:NSMakeRange(0, status_bytes)]; + + ds4_gpu_stream_expert_validate_args args = { + .n_total_expert = n_total_expert, + .n_expert = n_expert, + }; + + const int had_batch = g_batch_cb != nil; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_stream_expert_cache_validate(cb, + &args, + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_addrs, + up_addrs, + down_addrs, + status)) { + return 0; + } + + if (had_batch) { + if (ds4_gpu_end_commands() == 0) return 0; + } else if (!ds4_gpu_finish_command_buffer(cb, owned, + "streaming expert cache validator")) { + return 0; + } + + const uint32_t *words = (const uint32_t *)[status contents]; + *all_cached = words[0]; + *miss_mask = words[1]; + *invalid_mask = words[2]; + for (uint32_t i = 0; i < 6; i++) { + selected_ids[i] = (int32_t)words[4 + i]; + } + + if (had_batch && ds4_gpu_begin_commands() == 0) return 0; + return 1; +} + +static void ds4_gpu_stream_expert_cache_clear_entry_internal( + uint32_t layer, + uint32_t expert, + int count_eviction, + int recycle_slab_slot, + ds4_gpu_stream_expert_reusable_buffers *reuse) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + return; + } + if (reuse) { + reuse->gate_buffer = nil; + reuse->up_buffer = nil; + reuse->down_buffer = nil; + reuse->gate_inner = 0; + reuse->up_inner = 0; + reuse->down_inner = 0; + } + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!e->valid) return; + if (ds4_gpu_stream_expert_cache_entry_inflight(e)) { + return; + } + + const uint64_t bytes = e->logical_bytes; + ds4_gpu_stream_expert_evict_dontneed_range(e->model_map, + e->model_size, + e->gate_abs_offset, + e->gate_expert_bytes); + ds4_gpu_stream_expert_evict_dontneed_range(e->model_map, + e->model_size, + e->up_abs_offset, + e->gate_expert_bytes); + ds4_gpu_stream_expert_evict_dontneed_range(e->model_map, + e->model_size, + e->down_abs_offset, + e->down_expert_bytes); + ds4_gpu_stream_expert_cache_zero_addr_slot(layer, expert); + if (reuse) { + reuse->gate_buffer = e->gate_buffer; + reuse->up_buffer = e->up_buffer; + reuse->down_buffer = e->down_buffer; + reuse->gate_inner = e->gate_inner; + reuse->up_inner = e->up_inner; + reuse->down_inner = e->down_inner; + } else if (e->slab_backed && recycle_slab_slot) { + ds4_gpu_stream_expert_slab_push_free_slot(e->slab_slot); + } + e->gate_buffer = nil; + e->up_buffer = nil; + e->down_buffer = nil; + e->model_map = NULL; + e->model_size = 0; + e->gate_abs_offset = 0; + e->up_abs_offset = 0; + e->down_abs_offset = 0; + e->gate_expert_bytes = 0; + e->down_expert_bytes = 0; + e->logical_bytes = 0; + e->last_used = 0; + e->use_count = 0; + e->gate_inner = 0; + e->up_inner = 0; + e->down_inner = 0; + e->inflight_seq = 0; + e->slab_slot = 0; + e->valid = 0; + e->slab_backed = 0; + + if (g_stream_expert_cache_layer_count[layer] > 0) { + g_stream_expert_cache_layer_count[layer]--; + } + if (g_stream_expert_cache_entry_count > 0) { + g_stream_expert_cache_entry_count--; + } + if (g_stream_expert_cache_bytes >= bytes) { + g_stream_expert_cache_bytes -= bytes; + } else { + g_stream_expert_cache_bytes = 0; + } + if (count_eviction) { + g_stream_expert_cache_evictions++; + g_stream_expert_cache_layer_evictions[layer]++; + } +} + +static void ds4_gpu_stream_expert_cache_clear_entry( + uint32_t layer, + uint32_t expert, + int count_eviction) { + ds4_gpu_stream_expert_cache_clear_entry_internal(layer, + expert, + count_eviction, + 1, + NULL); +} + +static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats) { + ds4_gpu_stream_expert_pending_load_clear(); + g_stream_expert_cache_done_seq = g_stream_expert_cache_cb_seq; + g_stream_expert_cache_batch_seq = 0; + g_stream_expert_cache_owned_seq = 0; + g_stream_expert_cache_pending_max_seq = 0; + for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) { + for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) { + g_stream_expert_cache[layer][expert].inflight_seq = 0; + ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); + } + ds4_gpu_stream_full_expert_addr_clear_layer(layer); + g_stream_expert_cache_layer_count[layer] = 0; + id buffers[3] = { + g_stream_expert_cache_gate_addr_buffers[layer], + g_stream_expert_cache_up_addr_buffers[layer], + g_stream_expert_cache_down_addr_buffers[layer], + }; + for (uint32_t i = 0; i < 3; i++) { + if (!buffers[i]) continue; + const NSUInteger bytes = + (NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t); + memset([buffers[i] contents], 0, bytes); + [buffers[i] didModifyRange:NSMakeRange(0, bytes)]; + } + } + g_stream_expert_cache_bytes = 0; + g_stream_expert_cache_entry_count = 0; + for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) { + g_stream_expert_cache_slabs[i] = nil; + g_stream_expert_cache_slab_start_slot[i] = 0; + g_stream_expert_cache_slab_slot_count[i] = 0; + g_stream_expert_cache_slab_slots_used[i] = 0; + } + g_stream_expert_cache_slab_count = 0; + g_stream_expert_cache_slab_total_slots = 0; + g_stream_expert_cache_free_slot_count = 0; + g_stream_expert_cache_slab_slot_bytes = 0; + if (reset_stats) { + g_stream_expert_cache_hits = 0; + g_stream_expert_cache_misses = 0; + g_stream_expert_cache_evictions = 0; + g_stream_expert_cache_wraps = 0; + g_stream_expert_cache_clock = 0; + g_stream_expert_cache_evict_advise_bytes = 0; + g_stream_expert_cache_willneed_advise_bytes = 0; + g_stream_expert_cache_pread_bytes = 0; + g_stream_expert_cache_pread_ms = 0.0; + g_stream_expert_cache_buffer_allocs = 0; + g_stream_expert_cache_buffer_reuses = 0; + g_stream_expert_cache_decode_tokens = 0; + g_stream_expert_cache_hotness_decay_token = 0; + memset(g_stream_expert_cache_route_hotness, + 0, + sizeof(g_stream_expert_cache_route_hotness)); + g_stream_expert_timing_selected_calls = 0; + g_stream_expert_timing_selected_read_ms = 0.0; + g_stream_expert_timing_selected_sync_ms = 0.0; + g_stream_expert_timing_selected_copy_ms = 0.0; + g_stream_expert_timing_selected_bind_ms = 0.0; + g_stream_expert_timing_split_layers = 0; + g_stream_expert_timing_split_resident_experts = 0; + g_stream_expert_timing_split_missing_experts = 0; + g_stream_expert_timing_split_resident_ms = 0.0; + g_stream_expert_timing_split_missing_ms = 0.0; + g_stream_expert_timing_split_missing_load_ms = 0.0; + g_stream_expert_timing_split_missing_slot_ms = 0.0; + g_stream_expert_timing_split_missing_prune_ms = 0.0; + g_stream_expert_timing_split_missing_addr_ms = 0.0; + g_stream_expert_timing_split_missing_wait_ms = 0.0; + g_stream_expert_timing_load_calls = 0; + g_stream_expert_timing_load_prepare_ms = 0.0; + g_stream_expert_timing_load_pread_ms = 0.0; + g_stream_expert_timing_load_modify_ms = 0.0; + g_stream_expert_timing_load_install_ms = 0.0; + g_stream_expert_timing_prepare_batch_reuse_calls = 0; + g_stream_expert_timing_prepare_batch_reuse_ms = 0.0; + g_stream_expert_timing_prepare_buffer_calls = 0; + g_stream_expert_timing_prepare_buffer_ms = 0.0; + g_stream_expert_timing_prepare_task_experts = 0; + g_stream_expert_timing_prepare_task_ms = 0.0; + g_stream_expert_timing_reuse_scan_calls = 0; + g_stream_expert_timing_reuse_scan_entries = 0; + g_stream_expert_timing_reuse_scan_ms = 0.0; + g_stream_expert_timing_reuse_clear_ms = 0.0; + g_stream_expert_timing_readahead_calls = 0; + g_stream_expert_timing_readahead_bytes = 0; + g_stream_expert_timing_readahead_ms = 0.0; + g_stream_expert_timing_cache_all_resident_layers = 0; + g_stream_expert_timing_cache_all_missing_layers = 0; + g_stream_expert_timing_cache_mixed_layers = 0; + g_stream_expert_timing_cache_resident_experts = 0; + g_stream_expert_timing_cache_missing_experts = 0; + g_stream_expert_timing_last_report = + (ds4_gpu_stream_expert_timing_snapshot){0}; + memset(g_stream_expert_cache_layer_hits, + 0, + sizeof(g_stream_expert_cache_layer_hits)); + memset(g_stream_expert_cache_layer_misses, + 0, + sizeof(g_stream_expert_cache_layer_misses)); + memset(g_stream_expert_cache_layer_evictions, + 0, + sizeof(g_stream_expert_cache_layer_evictions)); + memset(g_stream_expert_cache_layer_pread_bytes, + 0, + sizeof(g_stream_expert_cache_layer_pread_bytes)); + memset(g_stream_expert_cache_layer_pread_ms, + 0, + sizeof(g_stream_expert_cache_layer_pread_ms)); + } +} + +static int ds4_gpu_stream_expert_cache_is_protected( + uint32_t expert, + const int32_t *protect_ids, + uint32_t n_protect) { + if (!protect_ids) return 0; + for (uint32_t i = 0; i < n_protect; i++) { + if (protect_ids[i] >= 0 && (uint32_t)protect_ids[i] == expert) { + return 1; + } + } + return 0; +} + +static void ds4_gpu_stream_expert_cache_prune_layer( + uint32_t layer, + uint32_t n_total_expert, + uint32_t n_selected, + const int32_t *protect_ids, + uint32_t n_protect) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return; + uint32_t cap = ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected); + if (cap == 0) return; + + /* + * Route hotness counts selected experts even when they miss. Hit-count + * LFU penalizes experts that are repeatedly selected but evicted before a + * second hit, which keeps too many decode layers in the mixed-cache path. + */ + while (g_stream_expert_cache_layer_count[layer] > cap) { + uint32_t victim = UINT32_MAX; + uint32_t lowest_hotness = UINT32_MAX; + uint64_t oldest = UINT64_MAX; + for (uint32_t expert = 0; expert < n_total_expert; expert++) { + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!e->valid || + ds4_gpu_stream_expert_cache_entry_inflight(e) || + ds4_gpu_stream_expert_cache_is_protected(expert, protect_ids, n_protect)) { + continue; + } + const uint32_t hotness = + g_stream_expert_cache_route_hotness[layer][expert]; + if (hotness < lowest_hotness || + (hotness == lowest_hotness && e->last_used < oldest)) { + lowest_hotness = hotness; + oldest = e->last_used; + victim = expert; + } + } + if (victim == UINT32_MAX) break; + ds4_gpu_stream_expert_cache_clear_entry(layer, victim, 1); + } +} + +static int ds4_gpu_stream_expert_cache_entry_protected( + uint32_t layer, + uint32_t expert, + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect) { + if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && + ds4_gpu_stream_expert_cache_entry_inflight( + &g_stream_expert_cache[layer][expert])) { + return 1; + } + return layer == protect_layer && + ds4_gpu_stream_expert_cache_is_protected(expert, + protect_ids, + n_protect); +} + +static int ds4_gpu_stream_expert_cache_entry_reusable( + const ds4_gpu_stream_expert_cache_entry *e, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + return e && + e->valid && + e->gate_buffer && + e->up_buffer && + e->down_buffer && + e->gate_expert_bytes == gate_expert_bytes && + e->down_expert_bytes == down_expert_bytes; +} + +static int ds4_gpu_stream_expert_cache_take_reusable( + int force_reuse, + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + ds4_gpu_stream_expert_reusable_buffers *reuse) { + if (!reuse) return 0; + reuse->gate_buffer = nil; + reuse->up_buffer = nil; + reuse->down_buffer = nil; + reuse->gate_inner = 0; + reuse->up_inner = 0; + reuse->down_inner = 0; + + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget == 0 || + (!force_reuse && g_stream_expert_cache_entry_count < budget)) { + return 0; + } + + int waited_inflight = 0; +retry: + ; + uint32_t victim_layer = UINT32_MAX; + uint32_t victim_expert = UINT32_MAX; + uint32_t lowest_hotness = UINT32_MAX; + uint64_t oldest = UINT64_MAX; + int skipped_inflight = 0; + const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); + const double scan_t0 = timing ? ds4_gpu_now_ms() : 0.0; + uint64_t scan_entries = 0; + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + scan_entries++; + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!ds4_gpu_stream_expert_cache_entry_reusable(e, + gate_expert_bytes, + down_expert_bytes)) { + continue; + } + if (ds4_gpu_stream_expert_cache_entry_inflight(e)) { + skipped_inflight = 1; + continue; + } + if (ds4_gpu_stream_expert_cache_entry_protected(layer, + expert, + protect_layer, + protect_ids, + n_protect)) { + continue; + } + const uint32_t hotness = + g_stream_expert_cache_route_hotness[layer][expert]; + if (hotness < lowest_hotness || + (hotness == lowest_hotness && e->last_used < oldest)) { + lowest_hotness = hotness; + oldest = e->last_used; + victim_layer = layer; + victim_expert = expert; + } + } + } + if (timing) { + ds4_gpu_stream_expert_timing_note_reuse_scan(scan_entries, + ds4_gpu_now_ms() - scan_t0); + } + + if (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) { + if (skipped_inflight && !waited_inflight && + !ds4_gpu_stream_expert_cache_on_service_thread()) { + waited_inflight = 1; + if (!ds4_gpu_stream_expert_cache_wait_inflight( + "streaming expert cache reuse")) { + return 0; + } + goto retry; + } + return 0; + } + const double clear_t0 = timing ? ds4_gpu_now_ms() : 0.0; + ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layer, + victim_expert, + 1, + 1, + reuse); + if (timing) { + ds4_gpu_stream_expert_timing_note_reuse_clear(ds4_gpu_now_ms() - + clear_t0); + } + if (!reuse->gate_buffer || !reuse->up_buffer || !reuse->down_buffer) { + reuse->gate_buffer = nil; + reuse->up_buffer = nil; + reuse->down_buffer = nil; + reuse->gate_inner = 0; + reuse->up_inner = 0; + reuse->down_inner = 0; + return 0; + } + g_stream_expert_cache_buffer_reuses += + ds4_gpu_stream_expert_buffer_object_count(reuse->gate_buffer, + reuse->up_buffer, + reuse->down_buffer); + return 1; +} + +static uint32_t ds4_gpu_stream_expert_cache_take_reusable_batch( + uint32_t n_needed, + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + ds4_gpu_stream_expert_reusable_buffers *reuses) { + if (!reuses || n_needed == 0) return 0; + if (n_needed > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { + n_needed = DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; + } + for (uint32_t i = 0; i < n_needed; i++) { + reuses[i].gate_buffer = nil; + reuses[i].up_buffer = nil; + reuses[i].down_buffer = nil; + reuses[i].gate_inner = 0; + reuses[i].up_inner = 0; + reuses[i].down_inner = 0; + } + + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget == 0 || g_stream_expert_cache_entry_count < budget) { + return 0; + } + + int waited_inflight = 0; +retry: + ; + uint32_t victim_layers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t victim_experts[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t victim_hotness[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint64_t victim_last_used[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t victim_count = 0; + int skipped_inflight = 0; + const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); + const double scan_t0 = timing ? ds4_gpu_now_ms() : 0.0; + uint64_t scan_entries = 0; + for (uint32_t i = 0; i < n_needed; i++) { + victim_layers[i] = UINT32_MAX; + victim_experts[i] = UINT32_MAX; + victim_hotness[i] = UINT32_MAX; + victim_last_used[i] = UINT64_MAX; + } + + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + scan_entries++; + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!ds4_gpu_stream_expert_cache_entry_reusable(e, + gate_expert_bytes, + down_expert_bytes)) { + continue; + } + if (ds4_gpu_stream_expert_cache_entry_inflight(e)) { + skipped_inflight = 1; + continue; + } + if (ds4_gpu_stream_expert_cache_entry_protected(layer, + expert, + protect_layer, + protect_ids, + n_protect)) { + continue; + } + + const uint32_t hotness = + g_stream_expert_cache_route_hotness[layer][expert]; + const uint64_t last_used = e->last_used; + if (victim_count < n_needed) { + victim_layers[victim_count] = layer; + victim_experts[victim_count] = expert; + victim_hotness[victim_count] = hotness; + victim_last_used[victim_count] = last_used; + victim_count++; + continue; + } + + uint32_t worst = 0; + for (uint32_t i = 1; i < victim_count; i++) { + if (victim_hotness[i] > victim_hotness[worst] || + (victim_hotness[i] == victim_hotness[worst] && + victim_last_used[i] > victim_last_used[worst])) { + worst = i; + } + } + if (hotness < victim_hotness[worst] || + (hotness == victim_hotness[worst] && + last_used < victim_last_used[worst])) { + victim_layers[worst] = layer; + victim_experts[worst] = expert; + victim_hotness[worst] = hotness; + victim_last_used[worst] = last_used; + } + } + } + if (timing) { + ds4_gpu_stream_expert_timing_note_reuse_scan(scan_entries, + ds4_gpu_now_ms() - scan_t0); + } + + if (victim_count == 0) { + if (skipped_inflight && !waited_inflight) { + waited_inflight = 1; + if (!ds4_gpu_stream_expert_cache_wait_inflight( + "streaming expert cache batch reuse")) { + return 0; + } + goto retry; + } + return 0; + } + + uint32_t reuse_count = 0; + const double clear_t0 = timing ? ds4_gpu_now_ms() : 0.0; + for (uint32_t i = 0; i < victim_count; i++) { + if (victim_layers[i] == UINT32_MAX || + victim_experts[i] == UINT32_MAX) { + continue; + } + ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layers[i], + victim_experts[i], + 1, + 1, + &reuses[reuse_count]); + if (!reuses[reuse_count].gate_buffer || + !reuses[reuse_count].up_buffer || + !reuses[reuse_count].down_buffer) { + reuses[reuse_count].gate_buffer = nil; + reuses[reuse_count].up_buffer = nil; + reuses[reuse_count].down_buffer = nil; + reuses[reuse_count].gate_inner = 0; + reuses[reuse_count].up_inner = 0; + reuses[reuse_count].down_inner = 0; + continue; + } + g_stream_expert_cache_buffer_reuses += + ds4_gpu_stream_expert_buffer_object_count( + reuses[reuse_count].gate_buffer, + reuses[reuse_count].up_buffer, + reuses[reuse_count].down_buffer); + reuse_count++; + } + if (timing) { + ds4_gpu_stream_expert_timing_note_reuse_clear(ds4_gpu_now_ms() - + clear_t0); + } + return reuse_count; +} + +static int ds4_gpu_stream_expert_batch_reuse_enabled( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { + return 0; + } + const uint64_t slot_bytes = gate_expert_bytes * 2ull + down_expert_bytes; + /* + * One global victim scan per selected miss is measurable for GLM Q2-size + * slots, but batching larger Q4-size slots regressed short decode on M5. + * Keep larger slots on the older single-victim path until profiling says + * otherwise. + */ + return slot_bytes <= 16ull * 1024ull * 1024ull; +} + +static int ds4_gpu_stream_expert_cache_prepare_load_buffers( + uint32_t layer, + uint32_t expert, + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + int force_reuse, + __strong id *gate_buf, + __strong id *up_buf, + __strong id *down_buf, + NSUInteger *gate_inner, + NSUInteger *up_inner, + NSUInteger *down_inner) { + if (!gate_buf || !up_buf || !down_buf || + !gate_inner || !up_inner || !down_inner || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + return 0; + } + + *gate_buf = nil; + *up_buf = nil; + *down_buf = nil; + *gate_inner = 0; + *up_inner = 0; + *down_inner = 0; + + ds4_gpu_stream_expert_reusable_buffers reuse = { nil, nil, nil, 0, 0, 0 }; + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (e->valid && ds4_gpu_stream_expert_cache_entry_inflight(e)) { + if (ds4_gpu_stream_expert_cache_on_service_thread()) return 0; + if (!ds4_gpu_stream_expert_cache_wait_inflight( + "streaming expert cache replacement")) { + return 0; + } + e = &g_stream_expert_cache[layer][expert]; + } + if (ds4_gpu_stream_expert_cache_entry_reusable(e, + gate_expert_bytes, + down_expert_bytes)) { + ds4_gpu_stream_expert_cache_clear_entry_internal(layer, + expert, + 0, + 1, + &reuse); + if (reuse.gate_buffer && reuse.up_buffer && reuse.down_buffer) { + g_stream_expert_cache_buffer_reuses += + ds4_gpu_stream_expert_buffer_object_count(reuse.gate_buffer, + reuse.up_buffer, + reuse.down_buffer); + } + } else if (e->valid) { + ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); + } + + if (!reuse.gate_buffer || !reuse.up_buffer || !reuse.down_buffer) { + if (!ds4_gpu_stream_expert_cache_take_reusable(force_reuse, + protect_layer, + protect_ids, + n_protect, + gate_expert_bytes, + down_expert_bytes, + &reuse)) { + reuse.gate_buffer = nil; + reuse.up_buffer = nil; + reuse.down_buffer = nil; + } + } + + if (reuse.gate_buffer && reuse.up_buffer && reuse.down_buffer) { + *gate_buf = reuse.gate_buffer; + *up_buf = reuse.up_buffer; + *down_buf = reuse.down_buffer; + *gate_inner = reuse.gate_inner; + *up_inner = reuse.up_inner; + *down_inner = reuse.down_inner; + return 1; + } + + if (ds4_gpu_stream_expert_combined_buffer_enabled()) { + if (gate_expert_bytes > UINT64_MAX - gate_expert_bytes || + gate_expert_bytes * 2ull > UINT64_MAX - down_expert_bytes || + gate_expert_bytes > (uint64_t)NSUIntegerMax || + gate_expert_bytes * 2ull > (uint64_t)NSUIntegerMax || + gate_expert_bytes * 2ull + down_expert_bytes > + (uint64_t)NSUIntegerMax) { + return 0; + } + if (ds4_gpu_stream_expert_alloc_slab_slot(gate_expert_bytes, + down_expert_bytes, + gate_buf, + up_buf, + down_buf, + gate_inner, + up_inner, + down_inner)) { + return 1; + } + const uint64_t up_off = gate_expert_bytes; + const uint64_t down_off = gate_expert_bytes * 2ull; + const uint64_t combined_bytes = down_off + down_expert_bytes; + id combined = + ds4_gpu_stream_expert_alloc_buffer(combined_bytes, + @"ds4_stream_expert_combined"); + if (!combined) return 0; + *gate_buf = combined; + *up_buf = combined; + *down_buf = combined; + *gate_inner = 0; + *up_inner = (NSUInteger)up_off; + *down_inner = (NSUInteger)down_off; + return 1; + } + + *gate_buf = ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes, + @"ds4_stream_expert_gate"); + *up_buf = ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes, + @"ds4_stream_expert_up"); + *down_buf = ds4_gpu_stream_expert_alloc_buffer(down_expert_bytes, + @"ds4_stream_expert_down"); + return *gate_buf && *up_buf && *down_buf; +} + +static void ds4_gpu_stream_expert_cache_prune_global( + uint32_t protect_layer, + const int32_t *protect_ids, + uint32_t n_protect) { + const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget(); + if (budget == 0 || g_stream_expert_cache_entry_count <= budget) return; + + while (g_stream_expert_cache_entry_count > budget) { + uint32_t victim_layer = UINT32_MAX; + uint32_t victim_expert = UINT32_MAX; + uint32_t lowest_hotness = UINT32_MAX; + uint64_t oldest = UINT64_MAX; + for (uint32_t layer = 0; + layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; + layer++) { + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!e->valid || + ds4_gpu_stream_expert_cache_entry_protected(layer, + expert, + protect_layer, + protect_ids, + n_protect)) { + continue; + } + const uint32_t hotness = + g_stream_expert_cache_route_hotness[layer][expert]; + if (hotness < lowest_hotness || + (hotness == lowest_hotness && e->last_used < oldest)) { + lowest_hotness = hotness; + oldest = e->last_used; + victim_layer = layer; + victim_expert = expert; + } + } + } + if (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) break; + ds4_gpu_stream_expert_cache_clear_entry(victim_layer, victim_expert, 1); + } +} + +static int ds4_gpu_stream_expert_cache_entry_matches( + const ds4_gpu_stream_expert_cache_entry *e, + const void *model_map, + uint64_t model_size, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + return e && + e->valid && + e->model_map == model_map && + e->model_size == model_size && + e->gate_abs_offset == gate_abs_offset && + e->up_abs_offset == up_abs_offset && + e->down_abs_offset == down_abs_offset && + e->gate_expert_bytes == gate_expert_bytes && + e->down_expert_bytes == down_expert_bytes && + e->gate_buffer && e->up_buffer && e->down_buffer; +} + +static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_peek( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t expert, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (!g_ssd_streaming_mode || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + expert >= n_total_expert || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected) == 0) { + return NULL; + } + + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (!ds4_gpu_stream_expert_cache_entry_matches(e, + model_map, + model_size, + gate_abs_offset, + up_abs_offset, + down_abs_offset, + gate_expert_bytes, + down_expert_bytes)) { + return NULL; + } + + e->last_used = ++g_stream_expert_cache_clock; + e->use_count++; + g_stream_expert_cache_hits++; + g_stream_expert_cache_layer_hits[layer]++; + return e; +} + +static ds4_gpu_stream_expert_cache_entry * +ds4_gpu_stream_expert_cache_install_loaded( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t expert, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + id gate_buf, + id up_buf, + id down_buf, + NSUInteger gate_inner, + NSUInteger up_inner, + NSUInteger down_inner) { + if (!gate_buf || !up_buf || !down_buf || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) { + return NULL; + } + if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { + fprintf(stderr, "ds4: Metal streaming expert cache byte size overflow\n"); + return NULL; + } + const uint64_t logical_bytes = gate_expert_bytes * 2ull + down_expert_bytes; + + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (e->valid) { + if (ds4_gpu_stream_expert_cache_entry_inflight(e) && + ds4_gpu_stream_expert_cache_on_service_thread()) { + return NULL; + } + if (ds4_gpu_stream_expert_cache_entry_inflight(e) && + !ds4_gpu_stream_expert_cache_wait_inflight( + "streaming expert cache install")) { + return NULL; + } + if (ds4_gpu_stream_expert_cache_entry_inflight(e)) return NULL; + ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); + if (e->valid) return NULL; + } + + if (!ds4_gpu_stream_expert_cache_set_addr_slot(layer, + expert, + gate_buf, + gate_inner, + up_buf, + up_inner, + down_buf, + down_inner)) { + return NULL; + } + + e->gate_buffer = gate_buf; + e->up_buffer = up_buf; + e->down_buffer = down_buf; + e->model_map = model_map; + e->model_size = model_size; + e->gate_abs_offset = gate_abs_offset; + e->up_abs_offset = up_abs_offset; + e->down_abs_offset = down_abs_offset; + e->gate_expert_bytes = gate_expert_bytes; + e->down_expert_bytes = down_expert_bytes; + e->logical_bytes = logical_bytes; + e->last_used = ++g_stream_expert_cache_clock; + e->use_count = 1; + e->gate_inner = gate_inner; + e->up_inner = up_inner; + e->down_inner = down_inner; + e->inflight_seq = 0; + uint32_t slab_slot = 0; + if (gate_buf == up_buf && + gate_buf == down_buf && + ds4_gpu_stream_expert_slab_slot_for_buffer(gate_buf, + gate_inner, + &slab_slot)) { + e->slab_backed = 1; + e->slab_slot = slab_slot; + } else { + e->slab_backed = 0; + e->slab_slot = 0; + } + e->valid = 1; + g_stream_expert_cache_layer_count[layer]++; + if (g_stream_expert_cache_entry_count < UINT32_MAX) { + g_stream_expert_cache_entry_count++; + } + if (g_stream_expert_cache_bytes > UINT64_MAX - logical_bytes) { + g_stream_expert_cache_bytes = UINT64_MAX; + } else { + g_stream_expert_cache_bytes += logical_bytes; + } + g_stream_expert_cache_misses++; + g_stream_expert_cache_layer_misses[layer]++; + g_stream_expert_cache_wraps += 3; + return e; +} + +static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_get_protected( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t expert, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + const int32_t *protect_ids, + uint32_t n_protect) { + if (!g_ssd_streaming_mode || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + expert >= n_total_expert || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected) == 0) { + return NULL; + } + + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][expert]; + if (ds4_gpu_stream_expert_cache_entry_matches(e, + model_map, + model_size, + gate_abs_offset, + up_abs_offset, + down_abs_offset, + gate_expert_bytes, + down_expert_bytes)) { + e->last_used = ++g_stream_expert_cache_clock; + e->use_count++; + g_stream_expert_cache_hits++; + g_stream_expert_cache_layer_hits[layer]++; + return e; + } + + ds4_gpu_stream_expert_readahead_range(gate_abs_offset, gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(up_abs_offset, gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(down_abs_offset, down_expert_bytes); + + id gate_buf = nil; + id up_buf = nil; + id down_buf = nil; + NSUInteger gate_inner = 0; + NSUInteger up_inner = 0; + NSUInteger down_inner = 0; + const int32_t protect_one = (int32_t)expert; + if (!protect_ids || n_protect == 0) { + protect_ids = &protect_one; + n_protect = 1; + } + const uint32_t cache_budget = + ds4_gpu_stream_expert_cache_configured_budget(); + const int force_reuse = + cache_budget != 0 && g_stream_expert_cache_entry_count >= cache_budget; + if (!ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, + expert, + layer, + protect_ids, + n_protect, + gate_expert_bytes, + down_expert_bytes, + force_reuse, + &gate_buf, + &up_buf, + &down_buf, + &gate_inner, + &up_inner, + &down_inner)) { + return NULL; + } + if (!gate_buf || !up_buf || !down_buf) return NULL; + + uint8_t *gate_dst = (uint8_t *)[gate_buf contents] + gate_inner; + uint8_t *up_dst = (uint8_t *)[up_buf contents] + up_inner; + uint8_t *down_dst = (uint8_t *)[down_buf contents] + down_inner; + if (!gate_dst || !up_dst || !down_dst) return NULL; + + ds4_gpu_stream_expert_pread_task tasks[3] = { + { + .offset = gate_abs_offset, + .len = gate_expert_bytes, + .dst = gate_dst, + }, + { + .offset = up_abs_offset, + .len = gate_expert_bytes, + .dst = up_dst, + }, + { + .offset = down_abs_offset, + .len = down_expert_bytes, + .dst = down_dst, + }, + }; + uint64_t read_bytes = 0; + double read_ms = 0.0; + if (!ds4_gpu_stream_expert_pread_tasks(tasks, 3, &read_bytes, &read_ms)) { + return NULL; + } + ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); + + [gate_buf didModifyRange:NSMakeRange(gate_inner, (NSUInteger)gate_expert_bytes)]; + [up_buf didModifyRange:NSMakeRange(up_inner, (NSUInteger)gate_expert_bytes)]; + [down_buf didModifyRange:NSMakeRange(down_inner, (NSUInteger)down_expert_bytes)]; + if (getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal streaming expert parallel pread layer=%u experts=1 tensors=3 " + "threads=%u bytes=%.2f GiB wall=%.3f ms\n", + layer, + ds4_gpu_stream_expert_pread_thread_count(3), + ds4_gpu_gib(read_bytes), + read_ms); + } + return ds4_gpu_stream_expert_cache_install_loaded(model_map, + model_size, + layer, + expert, + gate_abs_offset, + up_abs_offset, + down_abs_offset, + gate_expert_bytes, + down_expert_bytes, + gate_buf, + up_buf, + down_buf, + gate_inner, + up_inner, + down_inner); +} + +static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_get( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t expert, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_abs_offset, + uint64_t up_abs_offset, + uint64_t down_abs_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + return ds4_gpu_stream_expert_cache_get_protected(model_map, + model_size, + layer, + expert, + n_total_expert, + n_selected, + gate_abs_offset, + up_abs_offset, + down_abs_offset, + gate_expert_bytes, + down_expert_bytes, + NULL, + 0); +} + +static int ds4_gpu_stream_expert_pending_load_profile_enabled(void) { + return getenv("DS4_METAL_STREAMING_EXPERT_EARLY_LOAD_PROFILE") != NULL; +} + +static void ds4_gpu_stream_expert_pending_load_release_buffers( + ds4_gpu_stream_expert_pending_load *p) { + if (!p) return; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + p->gate_bufs[i] = nil; + p->up_bufs[i] = nil; + p->down_bufs[i] = nil; + p->gate_inners[i] = 0; + p->up_inners[i] = 0; + p->down_inners[i] = 0; + } +} + +static int ds4_gpu_stream_expert_pending_load_install( + ds4_gpu_stream_expert_pending_load *p, + ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED], + double elapsed_ms) { + if (!p || p->n_loads == 0) return 1; + + uint64_t read_bytes = 0; + int ok = 1; + for (uint32_t i = 0; i < p->n_tasks; i++) { + if (!p->tasks[i].ok) ok = 0; + if (read_bytes > UINT64_MAX - p->tasks[i].read_bytes) { + read_bytes = UINT64_MAX; + } else { + read_bytes += p->tasks[i].read_bytes; + } + } + if (!ok) return 0; + + ds4_gpu_stream_expert_cache_note_pread(p->layer, read_bytes, elapsed_ms); + const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled(); + double load_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + double load_modify_ms = 0.0; + double load_install_ms = 0.0; + for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) { + [p->gate_bufs[load_i] didModifyRange:NSMakeRange(p->gate_inners[load_i], (NSUInteger)p->gate_expert_bytes)]; + [p->up_bufs[load_i] didModifyRange:NSMakeRange(p->up_inners[load_i], (NSUInteger)p->gate_expert_bytes)]; + [p->down_bufs[load_i] didModifyRange:NSMakeRange(p->down_inners[load_i], (NSUInteger)p->down_expert_bytes)]; + } + if (load_timing) { + const double now_ms = ds4_gpu_now_ms(); + load_modify_ms = now_ms - load_t0; + load_t0 = now_ms; + } + + ds4_gpu_stream_expert_cache_entry + *loaded_entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + loaded_entries[i] = NULL; + } + for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) { + const uint32_t slot = p->load_slots[load_i]; + const uint32_t expert = (uint32_t)p->selected_ids[slot]; + ds4_gpu_stream_expert_cache_entry *entry = + ds4_gpu_stream_expert_cache_install_loaded(p->model_map, + p->model_size, + p->layer, + expert, + p->gate_abs_offsets[slot], + p->up_abs_offsets[slot], + p->down_abs_offsets[slot], + p->gate_expert_bytes, + p->down_expert_bytes, + p->gate_bufs[load_i], + p->up_bufs[load_i], + p->down_bufs[load_i], + p->gate_inners[load_i], + p->up_inners[load_i], + p->down_inners[load_i]); + if (!entry) return 0; + loaded_entries[slot] = entry; + if (entries) entries[slot] = entry; + } + for (uint32_t i = 0; i < p->n_selected; i++) { + if ((p->missing_mask & (1u << i)) == 0) continue; + if (entries && entries[i]) continue; + const uint32_t source = p->source_slots[i]; + if (source >= p->n_selected) return 0; + ds4_gpu_stream_expert_cache_entry *entry = entries && entries[source] ? + entries[source] : loaded_entries[source]; + if (!entry) return 0; + entry->use_count++; + if (entries) entries[i] = entry; + } + if (load_timing) { + load_install_ms = ds4_gpu_now_ms() - load_t0; + ds4_gpu_stream_expert_timing_note_load_detail(p->prepare_ms, + elapsed_ms, + load_modify_ms, + load_install_ms); + } + if (ds4_gpu_stream_expert_pending_load_profile_enabled()) { + fprintf(stderr, + "ds4: Metal streaming expert early-load finish layer=%u experts=%u tensors=%u bytes=%.2f GiB wall=%.3f ms\n", + p->layer, + p->n_loads, + p->n_tasks, + ds4_gpu_gib(read_bytes), + elapsed_ms); + } + return 1; +} + +static int ds4_gpu_stream_expert_pending_load_finish( + ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]) { + ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load; + if (!p->active) return 1; + + const double start_ms = p->start_ms; + if (!ds4_gpu_stream_expert_pread_pool_wait()) { + ds4_gpu_stream_expert_pending_load_release_buffers(p); + p->active = 0; + return 0; + } + const double elapsed_ms = ds4_gpu_now_ms() - start_ms; + p->active = 0; + const int ok = ds4_gpu_stream_expert_pending_load_install(p, + entries, + elapsed_ms); + ds4_gpu_stream_expert_pending_load_release_buffers(p); + p->n_tasks = 0; + p->n_loads = 0; + p->prepare_ms = 0.0; + return ok; +} + +static void ds4_gpu_stream_expert_pending_load_clear(void) { + if (!g_stream_expert_pending_load.active) { + ds4_gpu_stream_expert_pending_load_release_buffers( + &g_stream_expert_pending_load); + return; + } + (void)ds4_gpu_stream_expert_pending_load_finish(NULL); +} + +static int ds4_gpu_stream_expert_pending_load_matches( + const void *model_map, + uint64_t model_size, + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load; + if (!p->active || + p->model_map != model_map || + p->model_size != model_size || + p->layer != layer || + p->n_total_expert != n_total_expert || + p->n_selected != n_selected || + p->gate_expert_bytes != gate_expert_bytes || + p->down_expert_bytes != down_expert_bytes) { + return 0; + } + for (uint32_t i = 0; i < n_selected; i++) { + if (p->selected_ids[i] != selected_ids[i]) return 0; + } + return 1; +} + +int ds4_gpu_stream_expert_cache_begin_selected_load( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_selected) { + if (!g_ssd_streaming_mode || + getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_EARLY_LOAD") != NULL) { + return 1; + } + if (!table) return 0; + const void *model_map = table->model_map; + const uint64_t model_size = table->model_size; + const uint32_t layer = table->layer; + const uint32_t n_total_expert = table->n_total_expert; + const uint64_t gate_offset = table->gate_offset; + const uint64_t up_offset = table->up_offset; + const uint64_t down_offset = table->down_offset; + const uint64_t gate_expert_bytes = table->gate_expert_bytes; + const uint64_t down_expert_bytes = table->down_expert_bytes; + if (!model_map || !selected_ids || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected) == 0) { + return 1; + } + if (!g_initialized && !ds4_gpu_init()) return 0; + + if (ds4_gpu_stream_expert_pending_load_matches(model_map, + model_size, + layer, + selected_ids, + n_total_expert, + n_selected, + gate_expert_bytes, + down_expert_bytes)) { + return 1; + } + + ds4_gpu_stream_expert_pending_load_clear(); + ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load; + p->active = 0; + p->model_map = model_map; + p->model_size = model_size; + p->layer = layer; + p->n_total_expert = n_total_expert; + p->n_selected = n_selected; + p->missing_mask = 0; + p->n_loads = 0; + p->n_tasks = 0; + p->gate_expert_bytes = gate_expert_bytes; + p->down_expert_bytes = down_expert_bytes; + p->prepare_ms = 0.0; + const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled(); + const double load_prepare_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + p->selected_ids[i] = -1; + p->load_slots[i] = 0; + p->source_slots[i] = UINT32_MAX; + p->gate_abs_offsets[i] = 0; + p->up_abs_offsets[i] = 0; + p->down_abs_offsets[i] = 0; + p->gate_bufs[i] = nil; + p->up_bufs[i] = nil; + p->down_bufs[i] = nil; + p->gate_inners[i] = 0; + p->up_inners[i] = 0; + p->down_inners[i] = 0; + } + memset(p->tasks, 0, sizeof(p->tasks)); + + for (uint32_t i = 0; i < n_selected; i++) { + if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { + fprintf(stderr, + "ds4: Metal streaming early-load expert id %d is outside 0..%u\n", + selected_ids[i], + n_total_expert); + return 0; + } + p->selected_ids[i] = selected_ids[i]; + const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i]; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal streaming early-load offset overflow\n"); + return 0; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal streaming early-load offset overflow\n"); + return 0; + } + p->gate_abs_offsets[i] = gate_offset + gate_rel; + p->up_abs_offsets[i] = up_offset + gate_rel; + p->down_abs_offsets[i] = down_offset + down_rel; + + ds4_gpu_stream_expert_cache_entry *e = + &g_stream_expert_cache[layer][(uint32_t)selected_ids[i]]; + if (ds4_gpu_stream_expert_cache_entry_matches(e, + model_map, + model_size, + p->gate_abs_offsets[i], + p->up_abs_offsets[i], + p->down_abs_offsets[i], + gate_expert_bytes, + down_expert_bytes)) { + continue; + } + + uint32_t source = UINT32_MAX; + for (uint32_t prev = 0; prev < i; prev++) { + if (selected_ids[prev] == selected_ids[i] && + p->gate_abs_offsets[prev] == p->gate_abs_offsets[i] && + p->up_abs_offsets[prev] == p->up_abs_offsets[i] && + p->down_abs_offsets[prev] == p->down_abs_offsets[i] && + (p->missing_mask & (1u << prev)) != 0) { + source = p->source_slots[prev] != UINT32_MAX ? + p->source_slots[prev] : prev; + break; + } + } + if (source != UINT32_MAX) { + p->source_slots[i] = source; + p->missing_mask |= 1u << i; + continue; + } + p->source_slots[i] = i; + p->missing_mask |= 1u << i; + p->load_slots[p->n_loads++] = i; + } + if (p->n_loads == 0) return 1; + + const uint32_t cache_budget = + ds4_gpu_stream_expert_cache_configured_budget(); + uint32_t reserved_entries = g_stream_expert_cache_entry_count; + ds4_gpu_stream_expert_reusable_buffers + batch_reuse[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + batch_reuse[i] = + (ds4_gpu_stream_expert_reusable_buffers){ nil, nil, nil, 0, 0, 0 }; + } + uint32_t batch_reuse_count = 0; + if (cache_budget != 0 && + reserved_entries >= cache_budget && + p->n_loads > 1 && + ds4_gpu_stream_expert_batch_reuse_enabled(gate_expert_bytes, + down_expert_bytes)) { + const double reuse_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + batch_reuse_count = + ds4_gpu_stream_expert_cache_take_reusable_batch( + p->n_loads, + layer, + selected_ids, + n_selected, + gate_expert_bytes, + down_expert_bytes, + batch_reuse); + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_batch_reuse( + ds4_gpu_now_ms() - reuse_t0); + } + } + for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) { + const uint32_t slot = p->load_slots[load_i]; + const uint32_t expert = (uint32_t)p->selected_ids[slot]; + const int force_reuse = + cache_budget != 0 && reserved_entries >= cache_budget; + + ds4_gpu_stream_expert_readahead_range(p->gate_abs_offsets[slot], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(p->up_abs_offsets[slot], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(p->down_abs_offsets[slot], + down_expert_bytes); + + if (load_i < batch_reuse_count && + batch_reuse[load_i].gate_buffer && + batch_reuse[load_i].up_buffer && + batch_reuse[load_i].down_buffer) { + p->gate_bufs[load_i] = batch_reuse[load_i].gate_buffer; + p->up_bufs[load_i] = batch_reuse[load_i].up_buffer; + p->down_bufs[load_i] = batch_reuse[load_i].down_buffer; + p->gate_inners[load_i] = batch_reuse[load_i].gate_inner; + p->up_inners[load_i] = batch_reuse[load_i].up_inner; + p->down_inners[load_i] = batch_reuse[load_i].down_inner; + } else { + const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + const int prepared = + ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, + expert, + layer, + selected_ids, + n_selected, + gate_expert_bytes, + down_expert_bytes, + force_reuse, + &p->gate_bufs[load_i], + &p->up_bufs[load_i], + &p->down_bufs[load_i], + &p->gate_inners[load_i], + &p->up_inners[load_i], + &p->down_inners[load_i]); + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_buffer( + ds4_gpu_now_ms() - buffer_t0); + } + if (!prepared) { + ds4_gpu_stream_expert_pending_load_release_buffers(p); + return 0; + } + } + if (!force_reuse && reserved_entries < UINT32_MAX) { + reserved_entries++; + } + uint8_t *gate_dst = (uint8_t *)[p->gate_bufs[load_i] contents] + + p->gate_inners[load_i]; + uint8_t *up_dst = (uint8_t *)[p->up_bufs[load_i] contents] + + p->up_inners[load_i]; + uint8_t *down_dst = (uint8_t *)[p->down_bufs[load_i] contents] + + p->down_inners[load_i]; + if (!gate_dst || !up_dst || !down_dst) { + ds4_gpu_stream_expert_pending_load_release_buffers(p); + return 0; + } + const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = p->gate_abs_offsets[slot], + .len = gate_expert_bytes, + .dst = gate_dst, + }; + p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = p->up_abs_offsets[slot], + .len = gate_expert_bytes, + .dst = up_dst, + }; + p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = p->down_abs_offsets[slot], + .len = down_expert_bytes, + .dst = down_dst, + }; + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_task( + 1, + ds4_gpu_now_ms() - task_t0); + } + } + + const uint32_t n_workers = + ds4_gpu_stream_expert_pread_thread_count(p->n_tasks); + p->start_ms = ds4_gpu_now_ms(); + if (load_timing) { + p->prepare_ms = p->start_ms - load_prepare_t0; + } + if (ds4_gpu_stream_expert_pread_pool_begin(p->tasks, + p->n_tasks, + n_workers)) { + p->active = 1; + if (ds4_gpu_stream_expert_pending_load_profile_enabled()) { + fprintf(stderr, + "ds4: Metal streaming expert early-load begin layer=%u experts=%u tensors=%u threads=%u\n", + layer, + p->n_loads, + p->n_tasks, + n_workers); + } + return 1; + } + + uint64_t read_bytes = 0; + double read_ms = 0.0; + if (!ds4_gpu_stream_expert_pread_tasks(p->tasks, + p->n_tasks, + &read_bytes, + &read_ms)) { + ds4_gpu_stream_expert_pending_load_release_buffers(p); + return 0; + } + (void)read_bytes; + if (!ds4_gpu_stream_expert_pending_load_install(p, NULL, read_ms)) { + ds4_gpu_stream_expert_pending_load_release_buffers(p); + return 0; + } + ds4_gpu_stream_expert_pending_load_release_buffers(p); + p->n_tasks = 0; + p->n_loads = 0; + p->prepare_ms = 0.0; + return 1; +} + +static int ds4_gpu_stream_expert_cache_load_selected_missing( + const void *model_map, + uint64_t model_size, + uint32_t layer, + const int32_t *selected_ids, + uint32_t n_total_expert, + uint32_t n_selected, + const uint64_t *gate_abs_offsets, + const uint64_t *up_abs_offsets, + const uint64_t *down_abs_offsets, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + uint32_t missing_mask, + ds4_gpu_stream_expert_cache_entry **entries) { + if (!g_ssd_streaming_mode || + !model_map || + !selected_ids || + !gate_abs_offsets || + !up_abs_offsets || + !down_abs_offsets || + !entries || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected) == 0) { + return 0; + } + missing_mask &= (1u << n_selected) - 1u; + if (missing_mask == 0) return 1; + if (g_stream_expert_pending_load.active && + n_selected <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { + if (ds4_gpu_stream_expert_pending_load_matches(model_map, + model_size, + layer, + selected_ids, + n_total_expert, + n_selected, + gate_expert_bytes, + down_expert_bytes)) { + if (!ds4_gpu_stream_expert_pending_load_finish(entries)) return 0; + for (uint32_t i = 0; i < n_selected; i++) { + if ((missing_mask & (1u << i)) != 0 && entries[i]) { + missing_mask &= ~(1u << i); + } + } + if (missing_mask == 0) return 1; + } else { + ds4_gpu_stream_expert_pending_load_clear(); + } + } else if (g_stream_expert_pending_load.active) { + ds4_gpu_stream_expert_pending_load_clear(); + } + + const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled(); + double load_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + double load_prepare_ms = 0.0; + double load_modify_ms = 0.0; + double load_install_ms = 0.0; + + uint32_t load_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + uint32_t source_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + load_slots[i] = 0; + source_slots[i] = UINT32_MAX; + } + uint32_t n_loads = 0; + for (uint32_t i = 0; i < n_selected; i++) { + if ((missing_mask & (1u << i)) == 0) continue; + if (entries[i]) { + source_slots[i] = i; + continue; + } + if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { + fprintf(stderr, + "ds4: Metal streaming selected missing expert id %d is outside 0..%u\n", + selected_ids[i], + n_total_expert); + return 0; + } + uint32_t source = UINT32_MAX; + for (uint32_t prev = 0; prev < i; prev++) { + if (selected_ids[prev] == selected_ids[i] && + gate_abs_offsets[prev] == gate_abs_offsets[i] && + up_abs_offsets[prev] == up_abs_offsets[i] && + down_abs_offsets[prev] == down_abs_offsets[i] && + (entries[prev] || (missing_mask & (1u << prev)) != 0)) { + source = source_slots[prev] != UINT32_MAX ? + source_slots[prev] : prev; + break; + } + } + if (source != UINT32_MAX) { + source_slots[i] = source; + continue; + } + source_slots[i] = i; + load_slots[n_loads++] = i; + } + if (n_loads == 0) { + for (uint32_t i = 0; i < n_selected; i++) { + if ((missing_mask & (1u << i)) == 0 || entries[i]) continue; + const uint32_t source = source_slots[i]; + if (source >= n_selected || !entries[source]) return 0; + entries[i] = entries[source]; + entries[i]->use_count++; + } + return 1; + } + + __strong id gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + __strong id down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + gate_bufs[i] = nil; + up_bufs[i] = nil; + down_bufs[i] = nil; + gate_inners[i] = 0; + up_inners[i] = 0; + down_inners[i] = 0; + } + ds4_gpu_stream_expert_pread_task tasks[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED * 3u]; + memset(tasks, 0, sizeof(tasks)); + uint32_t n_tasks = 0; + const uint32_t cache_budget = + ds4_gpu_stream_expert_cache_configured_budget(); + uint32_t reserved_entries = g_stream_expert_cache_entry_count; + ds4_gpu_stream_expert_reusable_buffers + batch_reuse[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + batch_reuse[i] = + (ds4_gpu_stream_expert_reusable_buffers){ nil, nil, nil, 0, 0, 0 }; + } + uint32_t batch_reuse_count = 0; + if (cache_budget != 0 && + reserved_entries >= cache_budget && + n_loads > 1 && + ds4_gpu_stream_expert_batch_reuse_enabled(gate_expert_bytes, + down_expert_bytes)) { + const double reuse_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + batch_reuse_count = + ds4_gpu_stream_expert_cache_take_reusable_batch( + n_loads, + layer, + selected_ids, + n_selected, + gate_expert_bytes, + down_expert_bytes, + batch_reuse); + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_batch_reuse( + ds4_gpu_now_ms() - reuse_t0); + } + } + + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + const uint32_t slot = load_slots[load_i]; + const uint32_t expert = (uint32_t)selected_ids[slot]; + const int force_reuse = + cache_budget != 0 && reserved_entries >= cache_budget; + + ds4_gpu_stream_expert_readahead_range(gate_abs_offsets[slot], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(up_abs_offsets[slot], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(down_abs_offsets[slot], + down_expert_bytes); + + if (load_i < batch_reuse_count && + batch_reuse[load_i].gate_buffer && + batch_reuse[load_i].up_buffer && + batch_reuse[load_i].down_buffer) { + gate_bufs[load_i] = batch_reuse[load_i].gate_buffer; + up_bufs[load_i] = batch_reuse[load_i].up_buffer; + down_bufs[load_i] = batch_reuse[load_i].down_buffer; + gate_inners[load_i] = batch_reuse[load_i].gate_inner; + up_inners[load_i] = batch_reuse[load_i].up_inner; + down_inners[load_i] = batch_reuse[load_i].down_inner; + } else { + const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + const int prepared = + ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, + expert, + layer, + selected_ids, + n_selected, + gate_expert_bytes, + down_expert_bytes, + force_reuse, + &gate_bufs[load_i], + &up_bufs[load_i], + &down_bufs[load_i], + &gate_inners[load_i], + &up_inners[load_i], + &down_inners[load_i]); + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_buffer( + ds4_gpu_now_ms() - buffer_t0); + } + if (!prepared) { + return 0; + } + } + if (!force_reuse && reserved_entries < UINT32_MAX) { + reserved_entries++; + } + if (!gate_bufs[load_i] || !up_bufs[load_i] || !down_bufs[load_i]) { + return 0; + } + + uint8_t *gate_dst = (uint8_t *)[gate_bufs[load_i] contents] + + gate_inners[load_i]; + uint8_t *up_dst = (uint8_t *)[up_bufs[load_i] contents] + + up_inners[load_i]; + uint8_t *down_dst = (uint8_t *)[down_bufs[load_i] contents] + + down_inners[load_i]; + if (!gate_dst || !up_dst || !down_dst) return 0; + + const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = gate_abs_offsets[slot], + .len = gate_expert_bytes, + .dst = gate_dst, + }; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = up_abs_offsets[slot], + .len = gate_expert_bytes, + .dst = up_dst, + }; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = down_abs_offsets[slot], + .len = down_expert_bytes, + .dst = down_dst, + }; + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_task( + 1, + ds4_gpu_now_ms() - task_t0); + } + } + + if (load_timing) { + const double now_ms = ds4_gpu_now_ms(); + load_prepare_ms = now_ms - load_t0; + load_t0 = now_ms; + } + + uint64_t read_bytes = 0; + double read_ms = 0.0; + const int ok = ds4_gpu_stream_expert_pread_tasks(tasks, + n_tasks, + &read_bytes, + &read_ms); + if (!ok) return 0; + if (load_timing) { + load_t0 = ds4_gpu_now_ms(); + } + ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); + + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + [gate_bufs[load_i] didModifyRange:NSMakeRange(gate_inners[load_i], (NSUInteger)gate_expert_bytes)]; + [up_bufs[load_i] didModifyRange:NSMakeRange(up_inners[load_i], (NSUInteger)gate_expert_bytes)]; + [down_bufs[load_i] didModifyRange:NSMakeRange(down_inners[load_i], (NSUInteger)down_expert_bytes)]; + } + if (load_timing) { + const double now_ms = ds4_gpu_now_ms(); + load_modify_ms = now_ms - load_t0; + load_t0 = now_ms; + } + if (getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal streaming expert parallel pread layer=%u experts=%u tensors=%u " + "threads=%u bytes=%.2f GiB wall=%.3f ms\n", + layer, + n_loads, + n_tasks, + ds4_gpu_stream_expert_pread_thread_count(n_tasks), + ds4_gpu_gib(read_bytes), + read_ms); + } + + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + const uint32_t slot = load_slots[load_i]; + const uint32_t expert = (uint32_t)selected_ids[slot]; + ds4_gpu_stream_expert_cache_entry *entry = + ds4_gpu_stream_expert_cache_install_loaded(model_map, + model_size, + layer, + expert, + gate_abs_offsets[slot], + up_abs_offsets[slot], + down_abs_offsets[slot], + gate_expert_bytes, + down_expert_bytes, + gate_bufs[load_i], + up_bufs[load_i], + down_bufs[load_i], + gate_inners[load_i], + up_inners[load_i], + down_inners[load_i]); + if (!entry) return 0; + entries[slot] = entry; + } + + for (uint32_t i = 0; i < n_selected; i++) { + if ((missing_mask & (1u << i)) == 0 || entries[i]) continue; + const uint32_t source = source_slots[i]; + if (source >= n_selected || !entries[source]) return 0; + entries[i] = entries[source]; + entries[i]->use_count++; + } + for (uint32_t i = 0; i < n_selected; i++) { + if ((missing_mask & (1u << i)) != 0 && !entries[i]) return 0; + } + if (load_timing) { + load_install_ms = ds4_gpu_now_ms() - load_t0; + ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms, + read_ms, + load_modify_ms, + load_install_ms); + } + return 1; +} + +static void ds4_gpu_glm_stream_selected_prefetch_set( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_selected) { + g_glm_stream_selected_prefetch.active = 0; + if (!table || !selected_ids || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { + return; + } + g_glm_stream_selected_prefetch.model_map = table->model_map; + g_glm_stream_selected_prefetch.model_size = table->model_size; + g_glm_stream_selected_prefetch.layer = table->layer; + g_glm_stream_selected_prefetch.n_total_expert = table->n_total_expert; + g_glm_stream_selected_prefetch.n_selected = n_selected; + g_glm_stream_selected_prefetch.gate_offset = table->gate_offset; + g_glm_stream_selected_prefetch.up_offset = table->up_offset; + g_glm_stream_selected_prefetch.down_offset = table->down_offset; + g_glm_stream_selected_prefetch.gate_expert_bytes = table->gate_expert_bytes; + g_glm_stream_selected_prefetch.down_expert_bytes = table->down_expert_bytes; + for (uint32_t i = 0; i < n_selected; i++) { + g_glm_stream_selected_prefetch.selected_ids[i] = selected_ids[i]; + } + g_glm_stream_selected_prefetch.active = 1; +} + +static int ds4_gpu_glm_stream_selected_prefetch_take( + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + int32_t *selected_ids_out) { + if (!selected_ids_out || !g_glm_stream_selected_prefetch.active) return 0; + ds4_gpu_glm_stream_selected_prefetch *p = &g_glm_stream_selected_prefetch; + if (p->model_map != model_map || + p->model_size != model_size || + p->layer != layer || + p->n_total_expert != n_total_expert || + p->n_selected != n_selected || + p->gate_offset != gate_offset || + p->up_offset != up_offset || + p->down_offset != down_offset || + p->gate_expert_bytes != gate_expert_bytes || + p->down_expert_bytes != down_expert_bytes) { + return 0; + } + for (uint32_t i = 0; i < n_selected; i++) { + selected_ids_out[i] = p->selected_ids[i]; + } + p->active = 0; + return 1; +} + +int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( + const ds4_gpu_stream_expert_table *table, + const ds4_gpu_tensor *selected, + uint32_t n_selected) { + g_glm_stream_selected_prefetch.active = 0; + if (!g_ssd_streaming_mode || + getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_EARLY_LOAD") != NULL) { + return 1; + } + if (!table || !selected || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) { + return 1; + } + + const int timing = ds4_gpu_stream_expert_timing_summary_enabled(); + double t0 = timing ? ds4_gpu_now_ms() : 0.0; + double selected_sync_ms = 0.0; + double selected_copy_ms = 0.0; + double selected_bind_ms = 0.0; + const int had_batch = g_batch_cb != nil; + if (had_batch) { + if (ds4_gpu_end_commands() == 0) return 0; + if (timing) { + selected_sync_ms = ds4_gpu_now_ms() - t0; + t0 = ds4_gpu_now_ms(); + } + } + + int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) { + selected_ids[i] = -1; + } + int ok = ds4_gpu_tensor_read(selected, + 0, + selected_ids, + (uint64_t)n_selected * sizeof(selected_ids[0])) != 0; + if (timing) { + selected_copy_ms = ds4_gpu_now_ms() - t0; + t0 = ds4_gpu_now_ms(); + } + if (ok) { + ok = ds4_gpu_stream_expert_cache_begin_selected_load(table, + selected_ids, + n_selected) != 0; + } + if (timing) { + selected_bind_ms = ds4_gpu_now_ms() - t0; + ds4_gpu_stream_expert_timing_note_selected(selected_sync_ms, + selected_copy_ms, + selected_bind_ms); + } + if (ok) { + ds4_gpu_glm_stream_selected_prefetch_set(table, + selected_ids, + n_selected); + } + if (had_batch && ds4_gpu_begin_commands() == 0) ok = 0; + return ok; +} + +static void ds4_gpu_stream_expert_cache_clear_layer(uint32_t layer) { + if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return; + for (uint32_t expert = 0; + expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; + expert++) { + ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0); + } + g_stream_expert_cache_layer_count[layer] = 0; +} + +static int ds4_gpu_stream_expert_cache_prepare_selected_batch( + const void *model_map, + uint64_t model_size, + uint32_t layer, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_total_expert, + uint32_t n_selected, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + id *gate_addrs, + id *up_addrs, + id *down_addrs, + ds4_gpu_stream_expert_cache_entry **resources, + uint32_t *n_resources, + uint32_t *unique_out, + id *overflow_gate, + id *overflow_up, + id *overflow_down) { + if (overflow_gate) *overflow_gate = nil; + if (overflow_up) *overflow_up = nil; + if (overflow_down) *overflow_down = nil; + if (!g_ssd_streaming_mode || + !model_map || + !selected || + !gate_addrs || + !up_addrs || + !down_addrs || + !resources || + !n_resources || + !unique_out || + !overflow_gate || + !overflow_up || + !overflow_down || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_tokens == 0 || + n_selected == 0 || + n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + gate_expert_bytes == 0 || + down_expert_bytes == 0) { + return 0; + } + if (n_tokens > UINT32_MAX / n_selected) return 0; + + const uint64_t n_ids = (uint64_t)n_tokens * n_selected; + if (n_ids > SIZE_MAX / sizeof(int32_t)) return 0; + int32_t *ids = malloc((size_t)n_ids * sizeof(ids[0])); + if (!ids) return 0; + + const bool profile = + getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_PROFILE") != NULL; + const double t0 = profile ? ds4_gpu_now_ms() : 0.0; + int ok = ds4_gpu_tensor_read(selected, + 0, + ids, + n_ids * sizeof(ids[0])); + const double t_read = profile ? ds4_gpu_now_ms() : 0.0; + + bool seen[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { false }; + uint32_t frequency[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { 0 }; + int32_t unique_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + uint32_t unique_count = 0; + *n_resources = 0; + *unique_out = 0; + if (ok) { + if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) { + ok = 0; + } + } + if (ok) { + for (uint64_t i = 0; i < n_ids; i++) { + const int32_t selected_id = ids[i]; + if (selected_id < 0 || (uint32_t)selected_id >= n_total_expert) { + fprintf(stderr, + "ds4: Metal streaming batch selected expert id %d is outside 0..%u at layer %u\n", + selected_id, + n_total_expert, + layer); + ok = 0; + break; + } + frequency[(uint32_t)selected_id]++; + if (!seen[(uint32_t)selected_id]) { + seen[(uint32_t)selected_id] = true; + unique_ids[unique_count++] = selected_id; + } + } + } + if (ok) { + ds4_gpu_stream_expert_cache_note_frequency_hotness(layer, + frequency, + n_total_expert); + } + /* + * When the layer's unique selected set does not fit the cache budget, the + * extra experts are addressed straight into whole-tensor mapped model views + * instead of falling back to a different MoE kernel path. The address-table + * kernels read identical expert bytes either way, so the cache/view split + * does not change the computed logits. + */ + uint32_t view_served = 0; + uint64_t overflow_gate_inner = 0; + uint64_t overflow_up_inner = 0; + uint64_t overflow_down_inner = 0; + + uint64_t unique_gate_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + uint64_t unique_up_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + uint64_t unique_down_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + ds4_gpu_stream_expert_cache_entry + *unique_entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + uint32_t load_unique[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + __strong id + gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + __strong id + up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + __strong id + down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT]; + for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; i++) { + unique_gate_offsets[i] = 0; + unique_up_offsets[i] = 0; + unique_down_offsets[i] = 0; + unique_entries[i] = NULL; + load_unique[i] = 0; + gate_bufs[i] = nil; + up_bufs[i] = nil; + down_bufs[i] = nil; + gate_inners[i] = 0; + up_inners[i] = 0; + down_inners[i] = 0; + } + + ds4_gpu_stream_expert_pread_task *tasks = NULL; + uint32_t n_loads = 0; + uint32_t n_tasks = 0; + double load_prepare_ms = 0.0; + double load_modify_ms = 0.0; + double load_install_ms = 0.0; + double load_timing_t0 = ds4_gpu_stream_expert_timing_summary_enabled() ? + ds4_gpu_now_ms() : 0.0; + const int load_timing = load_timing_t0 != 0.0; + if (ok && unique_count != 0) { + tasks = calloc((size_t)unique_count * 3u, sizeof(tasks[0])); + if (!tasks) ok = 0; + } + if (ok) { + const uint32_t cache_budget = + ds4_gpu_stream_expert_cache_configured_budget(); + uint32_t reserved_entries = g_stream_expert_cache_entry_count; + + for (uint32_t u = 0; u < unique_count; u++) { + const uint32_t expert = (uint32_t)unique_ids[u]; + + if ((uint64_t)expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)expert > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal streaming batch selected expert offset overflow\n"); + ok = 0; + break; + } + const uint64_t gate_rel = (uint64_t)expert * gate_expert_bytes; + const uint64_t down_rel = (uint64_t)expert * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal streaming batch selected expert offset overflow\n"); + ok = 0; + break; + } + unique_gate_offsets[u] = gate_offset + gate_rel; + unique_up_offsets[u] = up_offset + gate_rel; + unique_down_offsets[u] = down_offset + down_rel; + + ds4_gpu_stream_expert_cache_entry *entry = + ds4_gpu_stream_expert_cache_peek(model_map, + model_size, + layer, + expert, + n_total_expert, + n_selected, + unique_gate_offsets[u], + unique_up_offsets[u], + unique_down_offsets[u], + gate_expert_bytes, + down_expert_bytes); + if (entry) { + unique_entries[u] = entry; + continue; + } + + const int force_reuse = + cache_budget != 0 && reserved_entries >= cache_budget; + ds4_gpu_stream_expert_readahead_range(unique_gate_offsets[u], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(unique_up_offsets[u], + gate_expert_bytes); + ds4_gpu_stream_expert_readahead_range(unique_down_offsets[u], + down_expert_bytes); + const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + const int prepared = + ds4_gpu_stream_expert_cache_prepare_load_buffers(layer, + expert, + layer, + unique_ids, + unique_count, + gate_expert_bytes, + down_expert_bytes, + force_reuse, + &gate_bufs[n_loads], + &up_bufs[n_loads], + &down_bufs[n_loads], + &gate_inners[n_loads], + &up_inners[n_loads], + &down_inners[n_loads]); + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_buffer( + ds4_gpu_now_ms() - buffer_t0); + } + if (!prepared) { + ok = 0; + break; + } + if (!force_reuse && reserved_entries < UINT32_MAX) { + reserved_entries++; + } + if (!gate_bufs[n_loads] || + !up_bufs[n_loads] || + !down_bufs[n_loads]) { + ok = 0; + break; + } + + uint8_t *gate_dst = (uint8_t *)[gate_bufs[n_loads] contents] + + gate_inners[n_loads]; + uint8_t *up_dst = (uint8_t *)[up_bufs[n_loads] contents] + + up_inners[n_loads]; + uint8_t *down_dst = (uint8_t *)[down_bufs[n_loads] contents] + + down_inners[n_loads]; + if (!gate_dst || !up_dst || !down_dst) { + ok = 0; + break; + } + + load_unique[n_loads] = u; + const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = unique_gate_offsets[u], + .len = gate_expert_bytes, + .dst = gate_dst, + }; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = unique_up_offsets[u], + .len = gate_expert_bytes, + .dst = up_dst, + }; + tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) { + .offset = unique_down_offsets[u], + .len = down_expert_bytes, + .dst = down_dst, + }; + if (load_timing) { + ds4_gpu_stream_expert_timing_note_prepare_task( + 1, + ds4_gpu_now_ms() - task_t0); + } + n_loads++; + } + } + if (ok && n_loads != 0) { + if (load_timing_t0 != 0.0) { + const double now_ms = ds4_gpu_now_ms(); + load_prepare_ms = now_ms - load_timing_t0; + load_timing_t0 = now_ms; + } + uint64_t read_bytes = 0; + double read_ms = 0.0; + ok = ds4_gpu_stream_expert_pread_tasks(tasks, + n_tasks, + &read_bytes, + &read_ms); + if (ok) { + ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms); + } + if (load_timing_t0 != 0.0) { + load_timing_t0 = ds4_gpu_now_ms(); + } + if (ok) { + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + [gate_bufs[load_i] didModifyRange:NSMakeRange(gate_inners[load_i], (NSUInteger)gate_expert_bytes)]; + [up_bufs[load_i] didModifyRange:NSMakeRange(up_inners[load_i], (NSUInteger)gate_expert_bytes)]; + [down_bufs[load_i] didModifyRange:NSMakeRange(down_inners[load_i], (NSUInteger)down_expert_bytes)]; + } + } + if (load_timing_t0 != 0.0) { + const double now_ms = ds4_gpu_now_ms(); + load_modify_ms = now_ms - load_timing_t0; + load_timing_t0 = now_ms; + } + if (ok && getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal streaming batch expert parallel pread layer=%u experts=%u tensors=%u " + "threads=%u bytes=%.2f GiB wall=%.3f ms\n", + layer, + n_loads, + n_tasks, + ds4_gpu_stream_expert_pread_thread_count(n_tasks), + ds4_gpu_gib(read_bytes), + read_ms); + } + if (ok) { + for (uint32_t load_i = 0; load_i < n_loads; load_i++) { + const uint32_t u = load_unique[load_i]; + const uint32_t expert = (uint32_t)unique_ids[u]; + ds4_gpu_stream_expert_cache_entry *entry = + ds4_gpu_stream_expert_cache_install_loaded(model_map, + model_size, + layer, + expert, + unique_gate_offsets[u], + unique_up_offsets[u], + unique_down_offsets[u], + gate_expert_bytes, + down_expert_bytes, + gate_bufs[load_i], + up_bufs[load_i], + down_bufs[load_i], + gate_inners[load_i], + up_inners[load_i], + down_inners[load_i]); + if (!entry) { + ok = 0; + break; + } + unique_entries[u] = entry; + } + } + if (load_timing_t0 != 0.0) { + load_install_ms = ds4_gpu_now_ms() - load_timing_t0; + ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms, + read_ms, + load_modify_ms, + load_install_ms); + } + } + if (tasks) free(tasks); + if (ok) { + for (uint32_t u = 0; u < unique_count; u++) { + ds4_gpu_stream_expert_cache_entry *entry = unique_entries[u]; + const uint32_t expert = (uint32_t)unique_ids[u]; + if (!entry) { + const uint64_t gate_rel = unique_gate_offsets[u] - gate_offset; + const uint64_t down_rel = unique_down_offsets[u] - down_offset; + if (!*overflow_gate) { + const uint64_t gate_tensor_bytes = + (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t down_tensor_bytes = + (uint64_t)n_total_expert * down_expert_bytes; + uint64_t gate_view_inner = 0; + uint64_t up_view_inner = 0; + uint64_t down_view_inner = 0; + id gv = ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset, + gate_tensor_bytes, + &gate_view_inner); + id uv = ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset, + gate_tensor_bytes, + &up_view_inner); + id dv = ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset, + down_tensor_bytes, + &down_view_inner); + if (!gv || !uv || !dv) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr " + "failed to map overflow expert views at layer %u\n", + layer); + ok = 0; + break; + } + *overflow_gate = gv; + *overflow_up = uv; + *overflow_down = dv; + overflow_gate_inner = gate_view_inner; + overflow_up_inner = up_view_inner; + overflow_down_inner = down_view_inner; + } + if (!ds4_gpu_stream_expert_cache_set_addr_slot( + layer, + expert, + *overflow_gate, + (NSUInteger)(overflow_gate_inner + gate_rel), + *overflow_up, + (NSUInteger)(overflow_up_inner + gate_rel), + *overflow_down, + (NSUInteger)(overflow_down_inner + down_rel))) { + ok = 0; + break; + } + view_served++; + continue; + } + const uint32_t extra_uses = + frequency[expert] > 0 ? frequency[expert] - 1u : 0; + if (extra_uses != 0) { + if (entry->use_count > UINT64_MAX - extra_uses) { + entry->use_count = UINT64_MAX; + } else { + entry->use_count += extra_uses; + } + if (g_stream_expert_cache_hits > UINT64_MAX - extra_uses) { + g_stream_expert_cache_hits = UINT64_MAX; + } else { + g_stream_expert_cache_hits += extra_uses; + } + if (g_stream_expert_cache_layer_hits[layer] > + UINT64_MAX - extra_uses) { + g_stream_expert_cache_layer_hits[layer] = UINT64_MAX; + } else { + g_stream_expert_cache_layer_hits[layer] += extra_uses; + } + } + resources[*n_resources] = entry; + (*n_resources)++; + } + } + free(ids); + + if (!ok || (*n_resources == 0 && view_served == 0)) { + ds4_gpu_stream_expert_cache_clear_layer(layer); + return 0; + } + if (!ds4_gpu_stream_expert_cache_addr_buffers(layer, + gate_addrs, + up_addrs, + down_addrs)) { + ds4_gpu_stream_expert_cache_clear_layer(layer); + return 0; + } + *unique_out = *n_resources + view_served; + if (view_served != 0 && + ds4_gpu_stream_expert_timing_summary_enabled()) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "served %u/%u unique experts via mapped views (cache budget %u)\n", + layer, + view_served, + unique_count, + ds4_gpu_stream_expert_cache_configured_budget()); + } + if (profile) { + const double t_done = ds4_gpu_now_ms(); + uint64_t per_expert_bytes = UINT64_MAX; + uint64_t logical_bytes = UINT64_MAX; + if (gate_expert_bytes <= (UINT64_MAX - down_expert_bytes) / 2ull) { + per_expert_bytes = gate_expert_bytes * 2ull + down_expert_bytes; + if (per_expert_bytes == 0 || + *n_resources <= UINT64_MAX / per_expert_bytes) { + logical_bytes = (uint64_t)(*n_resources) * per_expert_bytes; + } + } + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "tokens=%u unique=%u read=%.3f ms wrap=%.3f ms bytes=%.2f GiB\n", + layer, + n_tokens, + *n_resources, + t_read - t0, + t_done - t_read, + ds4_gpu_gib(logical_bytes)); + } + return 1; +} + +int ds4_gpu_stream_expert_cache_seed_selected( + const ds4_gpu_stream_expert_table *table, + const int32_t *selected_ids, + uint32_t n_selected) { + if (!g_ssd_streaming_mode) return 1; + if (!table) return 0; + const void *model_map = table->model_map; + const uint64_t model_size = table->model_size; + const uint32_t layer = table->layer; + const uint32_t n_total_expert = table->n_total_expert; + const uint64_t gate_offset = table->gate_offset; + const uint64_t up_offset = table->up_offset; + const uint64_t down_offset = table->down_offset; + const uint64_t gate_expert_bytes = table->gate_expert_bytes; + const uint64_t down_expert_bytes = table->down_expert_bytes; + if (!model_map || !selected_ids || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_selected == 0 || + n_selected > 6 || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + n_selected) == 0) { + return 1; + } + if (!g_initialized && !ds4_gpu_init()) return 0; + + ds4_gpu_stream_expert_cache_note_selected_hotness(layer, + selected_ids, + n_selected); + for (uint32_t i = 0; i < n_selected; i++) { + if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { + fprintf(stderr, + "ds4: Metal prefill expert-cache seed selected expert id %d is outside 0..%u\n", + selected_ids[i], + n_total_expert); + return 0; + } + const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i]; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal prefill expert-cache seed offset overflow\n"); + return 0; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal prefill expert-cache seed offset overflow\n"); + return 0; + } + + if (!ds4_gpu_stream_expert_cache_get(model_map, + model_size, + layer, + (uint32_t)selected_ids[i], + n_total_expert, + n_selected, + gate_offset + gate_rel, + up_offset + gate_rel, + down_offset + down_rel, + gate_expert_bytes, + down_expert_bytes)) { + return 0; + } + } + ds4_gpu_stream_expert_cache_prune_layer(layer, + n_total_expert, + n_selected, + selected_ids, + n_selected); + ds4_gpu_stream_expert_cache_prune_global(layer, + selected_ids, + n_selected); + return 1; +} + +int ds4_gpu_stream_expert_cache_seed_experts( + const ds4_gpu_stream_expert_table *table, + const int32_t *expert_ids, + const uint32_t *expert_priorities, + uint32_t n_experts) { + if (!g_ssd_streaming_mode) return 1; + if (!table) return 0; + const void *model_map = table->model_map; + const uint64_t model_size = table->model_size; + const uint32_t layer = table->layer; + const uint32_t n_total_expert = table->n_total_expert; + const uint64_t gate_offset = table->gate_offset; + const uint64_t up_offset = table->up_offset; + const uint64_t down_offset = table->down_offset; + const uint64_t gate_expert_bytes = table->gate_expert_bytes; + const uint64_t down_expert_bytes = table->down_expert_bytes; + if (!model_map || !expert_ids || + layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER || + n_experts == 0 || + n_total_expert == 0 || + n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || + !ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) || + ds4_gpu_stream_expert_cache_effective_cap(layer, + n_total_expert, + 1) == 0) { + return 1; + } + if (!g_initialized && !ds4_gpu_init()) return 0; + + ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(); + uint32_t remaining = n_experts; + while (remaining != 0) { + const uint32_t batch = + remaining > 6u ? 6u : remaining; + remaining -= batch; + + int32_t selected_ids[6] = { -1, -1, -1, -1, -1, -1 }; + uint64_t gate_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + uint64_t up_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + uint64_t down_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + ds4_gpu_stream_expert_cache_entry *entries[6] = { + NULL, NULL, NULL, NULL, NULL, NULL + }; + uint32_t missing_mask = 0; + + for (uint32_t i = 0; i < batch; i++) { + const int32_t expert = expert_ids[remaining + i]; + const uint32_t priority = + expert_priorities ? expert_priorities[remaining + i] : 0; + if (expert < 0 || (uint32_t)expert >= n_total_expert) { + fprintf(stderr, + "ds4: Metal streaming hotlist seed expert id %d is outside 0..%u\n", + expert, + n_total_expert); + return 0; + } + const uint64_t expert_id = (uint64_t)(uint32_t)expert; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal streaming hotlist seed offset overflow\n"); + return 0; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal streaming hotlist seed offset overflow\n"); + return 0; + } + + ds4_gpu_stream_expert_cache_note_route_hotness( + layer, + (uint32_t)expert, + priority != 0 ? priority : 1u); + selected_ids[i] = expert; + gate_abs_offsets[i] = gate_offset + gate_rel; + up_abs_offsets[i] = up_offset + gate_rel; + down_abs_offsets[i] = down_offset + down_rel; + entries[i] = ds4_gpu_stream_expert_cache_peek(model_map, + model_size, + layer, + (uint32_t)expert, + n_total_expert, + batch, + gate_abs_offsets[i], + up_abs_offsets[i], + down_abs_offsets[i], + gate_expert_bytes, + down_expert_bytes); + if (!entries[i]) missing_mask |= 1u << i; + if (entries[i] && priority != 0 && + entries[i]->use_count < (uint64_t)priority) { + entries[i]->use_count = (uint64_t)priority; + } + } + + if (missing_mask != 0 && + !ds4_gpu_stream_expert_cache_load_selected_missing(model_map, + model_size, + layer, + selected_ids, + n_total_expert, + batch, + gate_abs_offsets, + up_abs_offsets, + down_abs_offsets, + gate_expert_bytes, + down_expert_bytes, + missing_mask, + entries)) { + return 0; + } + if (expert_priorities) { + for (uint32_t i = 0; i < batch; i++) { + const uint32_t priority = expert_priorities[remaining + i]; + if (entries[i] && priority != 0 && + entries[i]->use_count < (uint64_t)priority) { + entries[i]->use_count = (uint64_t)priority; + } + } + } + } + + const uint32_t protect_n = n_experts < 6u ? n_experts : 6u; + ds4_gpu_stream_expert_cache_prune_layer(layer, + n_total_expert, + 1, + expert_ids, + protect_n); + ds4_gpu_stream_expert_cache_prune_global(layer, + expert_ids, + protect_n); + return 1; +} + +static uint32_t ds4_gpu_q4_expert_table_group_size(uint32_t n_total_expert) { + const char *env = getenv("DS4_METAL_Q4_EXPERT_TABLE_GROUP_SIZE"); + if (!env || !env[0]) return 1; + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end == env || *end != '\0' || v < 2 || v > n_total_expert) { + return 1; + } + return (uint32_t)v; +} + +static bool ds4_gpu_q4_table_queue_residency_requested(void) { + return getenv("DS4_METAL_Q4_TABLE_QUEUE_RESIDENCY_SET") != NULL; +} + +static bool ds4_gpu_q4_table_queue_residency_available(void) { +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + return g_queue && [g_queue respondsToSelector:@selector(addResidencySet:)]; + } +#endif + return false; +} + +static bool ds4_gpu_q4_non_streaming_opt_in_enabled(void) { + return getenv("DS4_METAL_ENABLE_Q4_SELECTED_EXPERT_VIEWS") != NULL || + getenv("DS4_METAL_ENABLE_PRO_Q4_SELECTED_EXPERT_VIEWS") != NULL || + getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL || + getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL; +} + +static bool ds4_gpu_q4_selected_paths_allowed(void) { + if (g_ssd_streaming_mode) return true; + if (g_glm_model_mode) return false; + return ds4_gpu_q4_non_streaming_opt_in_enabled(); +} + +int ds4_gpu_pro_q4_expert_table_auto_available(void) { + if (!g_initialized && !ds4_gpu_init()) return 0; + return (g_ssd_streaming_mode || + getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL) && + getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL && + ds4_gpu_q4_table_queue_residency_available(); +} + +static bool ds4_gpu_q4_table_queue_residency_enabled(bool auto_queue_residency) { + return (auto_queue_residency || ds4_gpu_q4_table_queue_residency_requested()) && + ds4_gpu_q4_table_queue_residency_available(); +} + +static bool ds4_gpu_q4_table_model_residency_enabled(void) { + return getenv("DS4_METAL_Q4_TABLE_MODEL_RESIDENCY_SET") != NULL; +} + +static bool ds4_gpu_pro_q4_expert_indirect_shape_supported( + uint32_t n_total_expert, + uint32_t n_expert, + uint64_t gate_tensor_bytes, + uint64_t down_tensor_bytes) { + const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; + return n_total_expert == 384 && + n_expert == 6 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes; +} + +static bool ds4_gpu_pro_q4_expert_table_auto_enabled( + uint32_t n_total_expert, + uint32_t n_expert, + uint64_t gate_tensor_bytes, + uint64_t down_tensor_bytes) { + /* + * This path lets the shader choose among exact per-expert resources. + * It is only automatic when a Metal queue residency set can make every + * indirect expert resource visible up front; otherwise the selected + * exact-slice path remains the fallback. + */ + return ds4_gpu_pro_q4_expert_indirect_shape_supported(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes) && + ds4_gpu_pro_q4_expert_table_auto_available(); +} + +static bool ds4_gpu_pro_q4_expert_address_auto_enabled( + uint32_t n_total_expert, + uint32_t n_expert, + uint64_t gate_tensor_bytes, + uint64_t down_tensor_bytes) { + /* + * GPU-address expert tables are useful for experiments, but they are not + * safe as an automatic path unless every indirect expert resource is made + * visible to Metal. Keep this behind an explicit opt-in while the selected + * active-slice path remains the correctness/performance baseline. + */ + return getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL && + ds4_gpu_q4_selected_paths_allowed() && + ds4_gpu_pro_q4_expert_indirect_shape_supported(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes) && + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil && + getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL && + getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") == NULL && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL && + ds4_gpu_q4_table_queue_residency_available(); +} + +static bool ds4_gpu_q4_table_bind_anchors_enabled(void) { + return getenv("DS4_METAL_Q4_TABLE_BIND_ANCHORS") != NULL; +} + +static int ds4_gpu_use_model_residency_set(id cb) { + if (!ds4_gpu_q4_table_model_residency_enabled()) return 1; +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + if (cb && g_model_residency_set && [cb respondsToSelector:@selector(useResidencySet:)]) { + [cb useResidencySet:g_model_residency_set]; + return 1; + } + } +#endif + fprintf(stderr, "ds4: Metal Q4 table model residency set is not available\n"); + return 0; +} + +static int ds4_gpu_bind_q4_expert_table_anchors( + id enc, + DS4MetalQ4ExpertTable *table, + NSUInteger first_index, + NSUInteger max_count) { + if (!ds4_gpu_q4_table_bind_anchors_enabled()) return 1; + if (!enc || !table || !table.expertBuffers) return 0; + + const NSUInteger count = [table.expertBuffers count]; + if (count > max_count) { + fprintf(stderr, + "ds4: Metal Q4 table anchor count %lu exceeds available slots %lu\n", + (unsigned long)count, + (unsigned long)max_count); + return 0; + } + for (NSUInteger i = 0; i < count; i++) { + [enc setBuffer:[table.expertBuffers objectAtIndex:i] + offset:0 + atIndex:first_index + i]; + } + return 1; +} + +static id ds4_gpu_q4_expert_table_residency_set(NSMutableArray> *buffers) { + if (!buffers || [buffers count] == 0 || + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") == NULL || + getenv("DS4_METAL_Q4_TABLE_PER_TENSOR_RESIDENCY_SET") == NULL) { + return nil; + } +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init]; + desc.label = @"ds4_q4_expert_table"; + desc.initialCapacity = [buffers count]; + NSError *error = nil; + id residency_set = [g_device newResidencySetWithDescriptor:desc error:&error]; + if (!residency_set) { + fprintf(stderr, "ds4: Metal Q4 expert table residency set creation failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + for (id buffer in buffers) { + [residency_set addAllocation:buffer]; + } + [residency_set commit]; + [residency_set requestResidency]; + return residency_set; + } +#endif + return nil; +} + +static void ds4_gpu_q4_residency_add_table(id residency_set, + DS4MetalQ4ExpertTable *table) { +#if TARGET_OS_OSX + if (!residency_set || !table) return; + if (@available(macOS 15.0, *)) { + if (table.argumentBuffer) { + [residency_set addAllocation:table.argumentBuffer]; + } + if (table.addressBuffer) { + [residency_set addAllocation:table.addressBuffer]; + } + for (id buffer in table.expertBuffers) { + [residency_set addAllocation:buffer]; + } + } +#else + (void)residency_set; + (void)table; +#endif +} + +static id ds4_gpu_q4_expert_layer_residency_set(DS4MetalQ4ExpertTable *gate_table, + DS4MetalQ4ExpertTable *up_table, + DS4MetalQ4ExpertTable *down_table, + bool auto_queue_residency) { + const bool queue_residency = + ds4_gpu_q4_table_queue_residency_enabled(auto_queue_residency); + if (!gate_table || !up_table || !down_table || + !g_device || !g_q4_expert_layer_residency_cache || + (getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") == NULL && + !queue_residency)) { + return nil; + } +#if TARGET_OS_OSX + if (@available(macOS 15.0, *)) { + NSString *key = [NSString stringWithFormat:@"%p:%p:%p", + gate_table, up_table, down_table]; + DS4MetalQ4LayerResidency *cached = + [g_q4_expert_layer_residency_cache objectForKey:key]; + if (cached) return cached.residencySet; + + const NSUInteger capacity = + [gate_table.expertBuffers count] + + [up_table.expertBuffers count] + + [down_table.expertBuffers count] + 3u; + MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init]; + desc.label = @"ds4_q4_expert_layer"; + desc.initialCapacity = capacity; + NSError *error = nil; + id residency_set = [g_device newResidencySetWithDescriptor:desc error:&error]; + if (!residency_set) { + fprintf(stderr, + "ds4: Metal Q4 expert layer residency set creation failed: %s\n", + [[error localizedDescription] UTF8String]); + return nil; + } + + ds4_gpu_q4_residency_add_table(residency_set, gate_table); + ds4_gpu_q4_residency_add_table(residency_set, up_table); + ds4_gpu_q4_residency_add_table(residency_set, down_table); + [residency_set commit]; + [residency_set requestResidency]; + + DS4MetalQ4LayerResidency *entry = [DS4MetalQ4LayerResidency new]; + entry.residencySet = residency_set; + if (queue_residency) { + [g_queue addResidencySet:residency_set]; + entry.addedToQueue = YES; + } + [g_q4_expert_layer_residency_cache setObject:entry forKey:key]; + return residency_set; + } +#endif + return nil; +} + +static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_table( + const void *model_map, + uint64_t model_size, + uint64_t tensor_offset, + uint64_t expert_bytes, + uint32_t n_total_expert, + id encoder); + +static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table( + const void *model_map, + uint64_t model_size, + uint64_t tensor_offset, + uint64_t expert_bytes, + uint32_t n_total_expert); + +int ds4_gpu_preload_q4_expert_tables(const void *model_map, uint64_t model_size, + uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, + uint64_t gate_expert_bytes, uint64_t down_expert_bytes, + uint32_t n_total_expert) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!model_map || model_size == 0 || gate_expert_bytes == 0 || + down_expert_bytes == 0 || n_total_expert == 0) { + return 0; + } + if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert table preload byte size overflow\n"); + return 0; + } + + const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes; + const bool address_auto = + ds4_gpu_pro_q4_expert_address_auto_enabled(n_total_expert, + 6, + gate_tensor_bytes, + down_tensor_bytes); + const bool table_auto = + ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, + 6, + gate_tensor_bytes, + down_tensor_bytes); + if (!address_auto && !table_auto) { + return 1; + } + + if (address_auto) { + @autoreleasepool { + DS4MetalQ4ExpertTable *gate_table = + ds4_gpu_q4_expert_address_table(model_map, + model_size, + gate_offset, + gate_expert_bytes, + n_total_expert); + DS4MetalQ4ExpertTable *up_table = + ds4_gpu_q4_expert_address_table(model_map, + model_size, + up_offset, + gate_expert_bytes, + n_total_expert); + DS4MetalQ4ExpertTable *down_table = + ds4_gpu_q4_expert_address_table(model_map, + model_size, + down_offset, + down_expert_bytes, + n_total_expert); + if (!gate_table || !up_table || !down_table) { + return 0; + } + id residency = ds4_gpu_q4_expert_layer_residency_set(gate_table, + up_table, + down_table, + true); + if (!residency) { + fprintf(stderr, "ds4: Metal Q4 expert address table preload failed to create queue residency set\n"); + return 0; + } + } + return 1; + } + + if (!g_moe_table_q4_pair_gate_encoder || !g_moe_table_q4_pair_up_encoder || + !g_moe_table_q4_sum_down_encoder) { + fprintf(stderr, "ds4: Metal Q4 expert table preload missing argument encoders\n"); + return 0; + } + + @autoreleasepool { + DS4MetalQ4ExpertTable *gate_table = + ds4_gpu_q4_expert_table(model_map, + model_size, + gate_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_gate_encoder); + DS4MetalQ4ExpertTable *up_table = + ds4_gpu_q4_expert_table(model_map, + model_size, + up_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_up_encoder); + DS4MetalQ4ExpertTable *down_table = + ds4_gpu_q4_expert_table(model_map, + model_size, + down_offset, + down_expert_bytes, + n_total_expert, + g_moe_table_q4_sum_down_encoder); + if (!gate_table || !up_table || !down_table) { + return 0; + } + id residency = ds4_gpu_q4_expert_layer_residency_set(gate_table, + up_table, + down_table, + true); + if (!residency) { + fprintf(stderr, "ds4: Metal Q4 expert table preload failed to create queue residency set\n"); + return 0; + } + } + return 1; +} + +static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_table( + const void *model_map, + uint64_t model_size, + uint64_t tensor_offset, + uint64_t expert_bytes, + uint32_t n_total_expert, + id encoder) { + if (!model_map || !g_device || !g_q4_expert_table_cache || !encoder || + model_size == 0 || expert_bytes == 0 || + n_total_expert == 0 || n_total_expert > 384) { + return nil; + } + if ((uint64_t)n_total_expert > UINT64_MAX / expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert table byte size overflow\n"); + return nil; + } + const uint64_t tensor_bytes = (uint64_t)n_total_expert * expert_bytes; + if (tensor_offset > model_size || tensor_bytes > model_size - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert table is outside the mapped model\n"); + return nil; + } + + const uint32_t table_group_size = + ds4_gpu_q4_expert_table_group_size(n_total_expert); + NSString *key = [NSString stringWithFormat:@"%p:%llu:%llu:%llu:%u:%llu:%u", + model_map, + (unsigned long long)model_size, + (unsigned long long)tensor_offset, + (unsigned long long)expert_bytes, + n_total_expert, + (unsigned long long)[encoder encodedLength], + table_group_size]; + DS4MetalQ4ExpertTable *cached = [g_q4_expert_table_cache objectForKey:key]; + if (cached) return cached; + + id arg_buffer = + [g_device newBufferWithLength:[encoder encodedLength] + options:MTLResourceStorageModeShared]; + if (!arg_buffer) { + fprintf(stderr, "ds4: Metal Q4 expert table argument buffer allocation failed\n"); + return nil; + } + arg_buffer.label = @"ds4_q4_expert_table"; + [encoder setArgumentBuffer:arg_buffer offset:0]; + + NSMutableArray> *expert_buffers = + [NSMutableArray arrayWithCapacity:table_group_size > 1 ? + (n_total_expert + table_group_size - 1u) / table_group_size : + n_total_expert]; + if (!expert_buffers) return nil; + + if (table_group_size > 1) { + for (uint32_t first = 0; first < n_total_expert; first += table_group_size) { + const uint32_t remaining = n_total_expert - first; + const uint32_t group_n = + remaining < table_group_size ? remaining : table_group_size; + if ((uint64_t)first > UINT64_MAX / expert_bytes || + (uint64_t)group_n > UINT64_MAX / expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert table group byte overflow\n"); + return nil; + } + const uint64_t rel = (uint64_t)first * expert_bytes; + const uint64_t group_bytes = (uint64_t)group_n * expert_bytes; + if (rel > UINT64_MAX - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert table group offset overflow\n"); + return nil; + } + uint64_t inner = 0; + id group_buf = + ds4_gpu_wrap_model_range(model_map, + model_size, + tensor_offset + rel, + group_bytes, + &inner); + if (!group_buf) return nil; + for (uint32_t j = 0; j < group_n; j++) { + const uint64_t expert_inner = inner + (uint64_t)j * expert_bytes; + [encoder setBuffer:group_buf offset:(NSUInteger)expert_inner atIndex:first + j]; + } + [expert_buffers addObject:group_buf]; + } + } else { + for (uint32_t i = 0; i < n_total_expert; i++) { + const uint64_t rel = (uint64_t)i * expert_bytes; + if (rel > UINT64_MAX - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert table offset overflow\n"); + return nil; + } + uint64_t inner = 0; + id expert_buf = + ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + tensor_offset + rel, + expert_bytes, + &inner); + if (!expert_buf) return nil; + [encoder setBuffer:expert_buf offset:(NSUInteger)inner atIndex:i]; + [expert_buffers addObject:expert_buf]; + } + } + [arg_buffer didModifyRange:NSMakeRange(0, [encoder encodedLength])]; + + if (getenv("DS4_METAL_Q4_EXPERT_TABLE_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal Q4 expert table: experts=%u group=%u buffers=%lu expert_bytes=%.2f MiB\n", + n_total_expert, + table_group_size, + (unsigned long)[expert_buffers count], + ds4_gpu_mib(expert_bytes)); + } + + DS4MetalQ4ExpertTable *table = [DS4MetalQ4ExpertTable new]; + table.argumentBuffer = arg_buffer; + table.expertBuffers = expert_buffers; + table.residencySet = ds4_gpu_q4_expert_table_residency_set(expert_buffers); + table.nExpert = n_total_expert; + table.expertBytes = expert_bytes; + [g_q4_expert_table_cache setObject:table forKey:key]; + return table; +} + +static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table( + const void *model_map, + uint64_t model_size, + uint64_t tensor_offset, + uint64_t expert_bytes, + uint32_t n_total_expert) { + if (!model_map || !g_device || !g_q4_expert_table_cache || + model_size == 0 || expert_bytes == 0 || + n_total_expert == 0 || n_total_expert > 384) { + return nil; + } + if ((uint64_t)n_total_expert > UINT64_MAX / expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert address table byte size overflow\n"); + return nil; + } + const uint64_t tensor_bytes = (uint64_t)n_total_expert * expert_bytes; + if (tensor_offset > model_size || tensor_bytes > model_size - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert address table is outside the mapped model\n"); + return nil; + } + + const uint32_t table_group_size = + ds4_gpu_q4_expert_table_group_size(n_total_expert); + NSString *key = [NSString stringWithFormat:@"addr:%p:%llu:%llu:%llu:%u:%u", + model_map, + (unsigned long long)model_size, + (unsigned long long)tensor_offset, + (unsigned long long)expert_bytes, + n_total_expert, + table_group_size]; + DS4MetalQ4ExpertTable *cached = [g_q4_expert_table_cache objectForKey:key]; + if (cached) return cached; + + id address_buffer = + [g_device newBufferWithLength:(NSUInteger)n_total_expert * sizeof(uint64_t) + options:MTLResourceStorageModeShared]; + if (!address_buffer) { + fprintf(stderr, "ds4: Metal Q4 expert address table allocation failed\n"); + return nil; + } + address_buffer.label = @"ds4_q4_expert_address_table"; + uint64_t *addresses = (uint64_t *)[address_buffer contents]; + if (!addresses) return nil; + + NSMutableArray> *expert_buffers = + [NSMutableArray arrayWithCapacity:table_group_size > 1 ? + (n_total_expert + table_group_size - 1u) / table_group_size : + n_total_expert]; + if (!expert_buffers) return nil; + +#if TARGET_OS_OSX + if (@available(macOS 13.0, *)) { + for (uint32_t first = 0; first < n_total_expert; first += table_group_size) { + const uint32_t remaining = n_total_expert - first; + const uint32_t group_n = + remaining < table_group_size ? remaining : table_group_size; + if ((uint64_t)first > UINT64_MAX / expert_bytes || + (uint64_t)group_n > UINT64_MAX / expert_bytes) { + fprintf(stderr, "ds4: Metal Q4 expert address table group byte overflow\n"); + return nil; + } + const uint64_t rel = (uint64_t)first * expert_bytes; + const uint64_t group_bytes = (uint64_t)group_n * expert_bytes; + if (rel > UINT64_MAX - tensor_offset) { + fprintf(stderr, "ds4: Metal Q4 expert address table group offset overflow\n"); + return nil; + } + uint64_t inner = 0; + id group_buf = nil; + if (table_group_size > 1) { + group_buf = ds4_gpu_wrap_model_range(model_map, + model_size, + tensor_offset + rel, + group_bytes, + &inner); + } else { + group_buf = ds4_gpu_wrap_model_exact_range_owned(model_map, + model_size, + tensor_offset + rel, + expert_bytes, + &inner); + } + if (!group_buf) return nil; + const uint64_t base_address = (uint64_t)[group_buf gpuAddress] + inner; + for (uint32_t j = 0; j < group_n; j++) { + addresses[first + j] = base_address + (uint64_t)j * expert_bytes; + } + [expert_buffers addObject:group_buf]; + } + } else +#endif + { + fprintf(stderr, "ds4: Metal GPU addresses require macOS 13 or newer\n"); + return nil; + } + + [address_buffer didModifyRange:NSMakeRange(0, + (NSUInteger)n_total_expert * sizeof(uint64_t))]; + + if (getenv("DS4_METAL_Q4_EXPERT_TABLE_PROFILE") != NULL) { + fprintf(stderr, + "ds4: Metal Q4 expert address table: experts=%u group=%u buffers=%lu expert_bytes=%.2f MiB\n", + n_total_expert, + table_group_size, + (unsigned long)[expert_buffers count], + ds4_gpu_mib(expert_bytes)); + } + + DS4MetalQ4ExpertTable *table = [DS4MetalQ4ExpertTable new]; + table.addressBuffer = address_buffer; + table.expertBuffers = expert_buffers; + table.nExpert = n_total_expert; + table.expertBytes = expert_bytes; + [g_q4_expert_table_cache setObject:table forKey:key]; + return table; +} + +static void ds4_gpu_use_q4_expert_table_resources( + id cb, + id enc, + DS4MetalQ4ExpertTable *table, + bool queue_residency) { + if (!enc || !table) return; + if (table.argumentBuffer) { + [enc useResource:table.argumentBuffer usage:MTLResourceUsageRead]; + } + if (table.addressBuffer) { + [enc useResource:table.addressBuffer usage:MTLResourceUsageRead]; + } + if (!queue_residency && + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL && + table.residencySet && + cb && + [cb respondsToSelector:@selector(useResidencySet:)]) { + [cb useResidencySet:table.residencySet]; + } + if (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") == NULL && + getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") == NULL) { + return; + } + const NSUInteger count = [table.expertBuffers count]; + for (NSUInteger base = 0; base < count; base += 64u) { + const NSUInteger n = count - base < 64u ? count - base : 64u; + __unsafe_unretained id resources[64]; + for (NSUInteger i = 0; i < n; i++) { + resources[i] = [table.expertBuffers objectAtIndex:base + i]; + } + [enc useResources:resources count:n usage:MTLResourceUsageRead]; + } +} + +int ds4_gpu_indexer_score_one_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_head, + uint32_t head_dim, + float scale) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!scores || !q || !weights || !index_comp || + n_comp == 0 || n_head == 0 || head_dim == 0) { + return 0; + } + + @autoreleasepool { + const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t weight_bytes = (uint64_t)n_head * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t score_bytes = (uint64_t)n_comp * sizeof(float); + id qbuf = ds4_gpu_tensor_buffer(q); + id wbuf = ds4_gpu_tensor_buffer(weights); + id compbuf = ds4_gpu_tensor_buffer(index_comp); + id scorebuf = ds4_gpu_tensor_buffer(scores); + if (!qbuf || !wbuf || !compbuf || !scorebuf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(weights) < weight_bytes || + ds4_gpu_tensor_bytes(index_comp) < comp_bytes || + ds4_gpu_tensor_bytes(scores) < score_bytes) { + fprintf(stderr, "ds4: Metal graph indexer score received undersized buffers\n"); + return 0; + } + + if (n_head == 64 && head_dim == 128) { + id direct_pipeline = + ds4_gpu_hot_pipeline(g_dsv4_indexer_score_one_direct_pipeline, + "kernel_dsv4_indexer_score_one_direct"); + if (!direct_pipeline) return 0; + + ds4_gpu_dsv4_indexer_scores_fused_args args = { + .n_comp = n_comp, + .n_tokens = 1, + .n_head = n_head, + .head_dim = head_dim, + .pos0 = 0, + .ratio = 4, + .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float), + .q_head_stride = (uint64_t)head_dim * sizeof(float), + .weights_token_stride = (uint64_t)n_head * sizeof(float), + .index_row_stride = (uint64_t)head_dim * sizeof(float), + .score_token_stride = (uint64_t)n_comp * sizeof(float), + .scale = scale, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:direct_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3]; + [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + [enc setThreadgroupMemoryLength:(128u + 4u) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_comp, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer direct score")) return 0; + return 1; + } + + const uint64_t head_score_bytes = (uint64_t)n_comp * n_head * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_indexer_head_scores_buffer, + &g_indexer_head_scores_bytes, + (NSUInteger)head_score_bytes, + "ds4_indexer_head_scores")) { + return 0; + } + + ds4_gpu_q8_0_matvec_args dot_args = + ds4_gpu_make_f32_mv_args(head_dim, n_comp, n_head); + ds4_gpu_mv_dispatch dot_dispatch = + ds4_gpu_make_plain_mv_dispatch(head_dim, 1); + dot_args.nr0 = dot_dispatch.nr0; + id dot_pipeline = + ds4_gpu_get_mul_mv_pipeline(dot_dispatch.function_name, dot_dispatch.nsg); + if (!dot_pipeline) return 0; + ds4_gpu_dsv4_indexer_weighted_sum_args sum_args = { + .ne00 = (int64_t)n_comp, + .ne01 = 1, + .ne02 = (int64_t)n_head, + .nb00 = sizeof(float), + .nb01 = (uint64_t)n_comp * sizeof(float), + .nb02 = (uint64_t)n_comp * sizeof(float), + .ne10 = (int64_t)n_head, + .ne11 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)n_head * sizeof(float), + .ne0 = (int64_t)n_comp, + .ne1 = 1, + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_comp * sizeof(float), + .scale = scale, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:dot_pipeline]; + [enc setBytes:&dot_args length:sizeof(dot_args) atIndex:0]; + [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:1]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2]; + [enc setBuffer:g_indexer_head_scores_buffer offset:0 atIndex:3]; + if (dot_dispatch.smem) { + [enc setThreadgroupMemoryLength:dot_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + (NSUInteger)dot_dispatch.nr0 - 1u) / (NSUInteger)dot_dispatch.nr0, + n_head, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dot_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_indexer_weighted_sum_pipeline]; + [enc setBytes:&sum_args length:sizeof(sum_args) atIndex:0]; + [enc setBuffer:g_indexer_head_scores_buffer offset:0 atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer score")) return 0; + } + + return 1; +} + +static int ds4_gpu_indexer_scores_batch_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!scores || !q || !weights || !index_comp || + n_comp == 0 || n_tokens == 0 || n_head == 0 || head_dim == 0 || ratio == 0) { + return 0; + } + + @autoreleasepool { + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t weight_bytes = (uint64_t)n_tokens * n_head * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t score_bytes = (uint64_t)n_comp * n_tokens * sizeof(float); + id qbuf = ds4_gpu_tensor_buffer(q); + id wbuf = ds4_gpu_tensor_buffer(weights); + id compbuf = ds4_gpu_tensor_buffer(index_comp); + id scorebuf = ds4_gpu_tensor_buffer(scores); + if (!qbuf || !wbuf || !compbuf || !scorebuf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(weights) < weight_bytes || + ds4_gpu_tensor_bytes(index_comp) < comp_bytes || + ds4_gpu_tensor_bytes(scores) < score_bytes) { + fprintf(stderr, "ds4: Metal graph indexer prefill scores received undersized buffers\n"); + return 0; + } + if (head_dim != 128) { + fprintf(stderr, "ds4: Metal fused DS4 indexer scores expect 128-wide rows\n"); + return 0; + } + /* + * The NAX/TensorOps score builder is a prefill-only win. At small + * batches and in one-token decode the setup cost is not amortized, so + * those paths keep the older direct/tiled score kernels. + */ + const bool use_nax = ds4_gpu_mpp_available() && n_tokens >= 16u; + id pipeline = ds4_gpu_get_pipeline( + use_nax ? "kernel_dsv4_indexer_scores_nax" : + (g_quality_mode ? "kernel_dsv4_indexer_scores_tiled_f32" + : "kernel_dsv4_indexer_scores_tiled")); + if (!pipeline) return 0; + + ds4_gpu_dsv4_indexer_scores_fused_args args = { + .n_comp = n_comp, + .n_tokens = n_tokens, + .n_head = n_head, + .head_dim = head_dim, + .pos0 = pos0, + .ratio = ratio, + .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float), + .q_head_stride = (uint64_t)head_dim * sizeof(float), + .weights_token_stride = (uint64_t)n_head * sizeof(float), + .index_row_stride = (uint64_t)head_dim * sizeof(float), + .score_token_stride = (uint64_t)n_comp * sizeof(float), + .scale = scale, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3]; + [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + if (use_nax) { + const NSUInteger q_shared = 2u * 32u * 32u; + const NSUInteger k_shared = 32u * 128u; + const NSUInteger dot_shared = 32u * 32u; + [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) + + dot_shared * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u, + ((NSUInteger)n_tokens + 15u) / 16u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + } else if (g_quality_mode) { + const NSUInteger q_shared = 8u * 128u; + const NSUInteger k_shared = 32u * 128u; + const NSUInteger dot_shared = 8u * 32u; + [enc setThreadgroupMemoryLength:(q_shared + k_shared + dot_shared) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u, + ((NSUInteger)n_tokens + 7u) / 8u, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + } else { + const NSUInteger q_shared = 8u * 128u; + const NSUInteger k_shared = 32u * 128u; + const NSUInteger dot_shared = 8u * 32u; + [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) + + dot_shared * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u, + ((NSUInteger)n_tokens + 7u) / 8u, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer prefill scores")) return 0; + } + + return 1; +} + +int ds4_gpu_indexer_scores_prefill_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale) { + return ds4_gpu_indexer_scores_batch_tensor(scores, + q, + weights, + index_comp, + n_comp, + n_tokens, + 0, + n_head, + head_dim, + ratio, + scale); +} + +int ds4_gpu_indexer_scores_decode_batch_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *index_comp, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t ratio, + float scale) { + return ds4_gpu_indexer_scores_batch_tensor(scores, + q, + weights, + index_comp, + n_comp, + n_tokens, + pos0, + n_head, + head_dim, + ratio, + scale); +} + +int ds4_gpu_indexer_topk_tensor( + ds4_gpu_tensor *selected, + const ds4_gpu_tensor *scores, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !scores || n_comp == 0 || n_tokens == 0 || top_k == 0 || top_k > n_comp) return 0; + + @autoreleasepool { + const uint64_t score_bytes = (uint64_t)n_comp * n_tokens * sizeof(float); + const uint64_t selected_bytes = (uint64_t)top_k * n_tokens * sizeof(uint32_t); + id scorebuf = ds4_gpu_tensor_buffer(scores); + id selbuf = ds4_gpu_tensor_buffer(selected); + if (!scorebuf || !selbuf || + ds4_gpu_tensor_bytes(scores) < score_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal graph indexer top-k received undersized buffers\n"); + return 0; + } + NSUInteger max_threads = g_argsort_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup; + if (max_threads == 0) max_threads = 256; + int32_t nth = 1; + while ((uint32_t)nth < n_comp && (uint64_t)2u * (uint64_t)nth <= (uint64_t)max_threads) { + nth *= 2; + } + const int32_t npr = (int32_t)((n_comp + (uint32_t)nth - 1u) / (uint32_t)nth); + const int32_t block_top_k = (int32_t)(top_k < (uint32_t)nth ? top_k : (uint32_t)nth); + int32_t work_width = (int32_t)top_k; + if (npr > 1) { + const int32_t last_block = (int32_t)n_comp - (npr - 1) * nth; + work_width = (npr - 1) * block_top_k + (last_block < block_top_k ? last_block : block_top_k); + } + const uint64_t scratch_row_bytes = (uint64_t)work_width * sizeof(uint32_t); + const bool one_pass = npr <= 1; + const uint64_t scratch_bytes = one_pass ? scratch_row_bytes * n_tokens : + 2u * scratch_row_bytes * n_tokens; + if (!ds4_gpu_ensure_scratch_buffer(&g_indexer_topk_buffer, + &g_indexer_topk_bytes, + (NSUInteger)scratch_bytes, + "ds4_indexer_topk")) { + return 0; + } + + ds4_gpu_kargs_argsort args = { + .ne00 = (int32_t)n_comp, + .ne01 = (int32_t)n_tokens, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = (uint64_t)n_comp * sizeof(float), + .nb02 = (uint64_t)n_comp * n_tokens * sizeof(float), + .nb03 = (uint64_t)n_comp * n_tokens * sizeof(float), + .ne0 = work_width, + .ne1 = (int32_t)n_tokens, + .ne2 = 1, + .ne3 = 1, + .top_k = block_top_k, + }; + // kernel_argsort_f32_i32_desc stages the block's scores behind the + // index array: nth int32 indices + nth float scores. + const NSUInteger smem = (((NSUInteger)nth * (sizeof(int32_t) + sizeof(float))) + 15u) & ~(NSUInteger)15u; + + NSUInteger cur_off = 0; + NSUInteger next_off = (NSUInteger)scratch_row_bytes * n_tokens; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_argsort_f32_i32_desc_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1]; + [enc setBuffer:one_pass ? selbuf : g_indexer_topk_buffer + offset:one_pass ? ds4_gpu_tensor_offset(selected) : cur_off + atIndex:2]; + [enc setThreadgroupMemoryLength:smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)npr * n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake((NSUInteger)nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + int32_t len = block_top_k; + while (len < work_width) { + const int32_t nm = (work_width + 2 * len - 1) / (2 * len); + const bool final_merge = nm == 1; + NSUInteger merge_threads = g_argsort_merge_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup; + if (merge_threads == 0 || merge_threads > 512u) merge_threads = 512u; + if (merge_threads > (NSUInteger)len) merge_threads = (NSUInteger)len; + if (merge_threads == 0) merge_threads = 1; + + ds4_gpu_kargs_argsort_merge merge_args = { + .ne00 = (int64_t)n_comp, + .ne01 = (int64_t)n_tokens, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = (uint64_t)n_comp * sizeof(float), + .nb02 = (uint64_t)n_comp * n_tokens * sizeof(float), + .nb03 = (uint64_t)n_comp * n_tokens * sizeof(float), + .ne0 = work_width, + .ne1 = (int32_t)n_tokens, + .ne2 = 1, + .ne3 = 1, + .top_k = nm == 1 ? (int32_t)top_k : work_width, + .len = len, + }; + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_argsort_merge_f32_i32_desc_pipeline]; + [enc setBytes:&merge_args length:sizeof(merge_args) atIndex:0]; + [enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1]; + [enc setBuffer:g_indexer_topk_buffer offset:cur_off atIndex:2]; + [enc setBuffer:final_merge ? selbuf : g_indexer_topk_buffer + offset:final_merge ? ds4_gpu_tensor_offset(selected) : next_off + atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nm * n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(merge_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + const NSUInteger tmp = cur_off; + cur_off = next_off; + next_off = tmp; + len <<= 1; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer top-k")) return 0; + } + + return 1; +} + +int ds4_gpu_argmax_tensor( + ds4_gpu_tensor *out_idx, + const ds4_gpu_tensor *logits, + uint32_t n_vocab) { + if (!out_idx || !logits || n_vocab == 0) return 0; + if (ds4_gpu_tensor_bytes(out_idx) < sizeof(int32_t) || + ds4_gpu_tensor_bytes(logits) < (uint64_t)n_vocab * sizeof(float)) { + fprintf(stderr, "ds4: Metal graph argmax received undersized buffers\n"); + return 0; + } + + return ds4_gpu_indexer_topk_tensor(out_idx, logits, n_vocab, 1, 1); +} + +int ds4_gpu_dsv4_topk_mask_tensor( + ds4_gpu_tensor *mask, + const ds4_gpu_tensor *topk, + uint32_t n_comp, + uint32_t n_tokens, + uint32_t top_k) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!mask || !topk || n_comp == 0 || n_tokens == 0 || top_k == 0) return 0; + + @autoreleasepool { + const uint64_t topk_bytes = (uint64_t)top_k * n_tokens * sizeof(int32_t); + const uint64_t mask_bytes = (uint64_t)n_comp * n_tokens * sizeof(float); + id topkbuf = ds4_gpu_tensor_buffer(topk); + id maskbuf = ds4_gpu_tensor_buffer(mask); + if (!topkbuf || !maskbuf || + ds4_gpu_tensor_bytes(topk) < topk_bytes || + ds4_gpu_tensor_bytes(mask) < mask_bytes) { + fprintf(stderr, "ds4: Metal dsv4 top-k mask received undersized buffers\n"); + return 0; + } + + ds4_gpu_dsv4_topk_mask_args args = { + .ne00 = (int64_t)top_k, + .ne01 = (int64_t)n_tokens, + .nb00 = sizeof(int32_t), + .nb01 = (uint64_t)top_k * sizeof(int32_t), + .ne0 = (int64_t)n_comp, + .ne1 = (int64_t)n_tokens, + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_comp * sizeof(float), + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_topk_mask_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1]; + [enc setBuffer:maskbuf offset:ds4_gpu_tensor_offset(mask) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake((((NSUInteger)n_comp * n_tokens) + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_topk_mask_scatter_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1]; + [enc setBuffer:maskbuf offset:ds4_gpu_tensor_offset(mask) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake((((NSUInteger)top_k * n_tokens) + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "dsv4 top-k mask")) return 0; + } + + return 1; +} + +static int ds4_gpu_matmul_q8_0_legacy_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + bool prefer_decode_mpp, + bool force_model_view) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if ((in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal Q8_0 tensor matmul received undersized activation buffers\n"); + return 0; + } + + const uint64_t blocks = in_dim / 32; + const uint64_t row_bytes = blocks * 34; + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8_0 tensor matmul range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = force_model_view ? + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset) : + ds4_gpu_wrap_q8_decode_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + n_tok, + &inner_offset); + if (!wbuf) { + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (n_tok == 1) { + if (ds4_gpu_mpp_available() && + (prefer_decode_mpp || getenv("DS4_METAL_Q8_DECODE_MPP") != NULL) && + getenv("DS4_METAL_DISABLE_Q8_DECODE_MPP") == NULL && + (in_dim % 64u) == 0) { + const char *nax_fn = "kernel_mul_mm_q8_0_f32_nax_direct_rhs"; + id mpp_pipeline = + ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false); + if (mpp_pipeline) { + ds4_gpu_mul_mm_args args = + ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:mpp_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:64u * 32u * sizeof(uint16_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1u, + ((NSUInteger)out_dim + 63u) / 64u, + 1u) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 decode MPP matmul")) { + return 0; + } + return 1; + } + ds4_gpu_warn_mpp_fallback(); + } + + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + if (out_dim > 65536u) mv_dispatch.nsg = 8; + const bool force_output_nr4 = + getenv("DS4_METAL_ENABLE_OUTPUT_Q8_NR4") != NULL; + const bool output_shape = + in_dim == 4096u && out_dim == 129280u; + const bool use_output_nr4 = + !g_quality_mode && (out_dim & 3u) == 0u && + (force_output_nr4 || + (output_shape && ds4_gpu_device_name_contains("M3"))) && + getenv("DS4_METAL_DISABLE_M3_OUTPUT_Q8_NR4") == NULL; + if (use_output_nr4) { + mv_dispatch.function_name = + "kernel_mul_mv_q8_0_f32_nr4"; + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor matvec")) { + return 0; + } + return 1; + } + + const uint64_t mv_ext_max_tokens = + ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u); + if (n_tok <= mv_ext_max_tokens && (in_dim % 128u) == 0) { + const int16_t nsg = 2; + const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); + const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); + const char *fn_name = ds4_gpu_mv_ext_name(1, r1ptg); + id pipeline = + fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; + if (!pipeline) return 0; + + const int16_t nypsg = 32 / nxpsg; + const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; + ds4_gpu_mul_mv_ext_args args = + ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, 34, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, + ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor mul_mv_ext")) { + return 0; + } + return 1; + } + + /* + * Dense Q8_0 prefill is the cleanest DS4 TensorOps shape: M/N/K are + * aligned and the RHS activation matrix is already dense. The retained + * kernel dequantizes each 64x32 weight tile to half in threadgroup + * memory, then uses direct-RHS MPP for the activation tile. This avoids + * staging RHS into threadgroup memory and was the direct replacement for + * the slower generic MPP prototype. + */ + if (ds4_gpu_mpp_available() && + n_tok >= 32u && + (in_dim % 64u) == 0 && + (out_dim % 64u) == 0 && + (n_tok % 32u) == 0) { + uint64_t nax_tile_n = 32u; + if ((n_tok % 128u) == 0) { + nax_tile_n = 128u; + } else if ((n_tok % 64u) == 0) { + nax_tile_n = 64u; + } + const char *nax_fn = nax_tile_n == 128u + ? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128" + : (nax_tile_n == 64u + ? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64" + : "kernel_mul_mm_q8_0_f32_nax_direct_rhs"); + id pipeline = + ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false); + if (pipeline) { + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), + (NSUInteger)out_dim / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 NAX tensor matmul")) { + return 0; + } + return 1; + } + ds4_gpu_warn_mpp_fallback(); + } + + const bool bc_inp = (in_dim % 32u) != 0; + const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0; + id pipeline = + ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_q8_0_f32", bc_inp, bc_out); + if (!pipeline) return 0; + + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u, + ((NSUInteger)out_dim + 63u) / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor matmul")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_matmul_q8_0_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if ((in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) { + return 0; + } + + const int profile_requested = + n_tok > 8u && ds4_gpu_env_bool("DS4_METAL_Q8_PREFILL_PROFILE") > 0; + int profile_prefill = 0; + int split_batch_for_profile = 0; + const char *profile_label = NULL; + char profile_label_buf[128]; + char profile_fallback[128]; + if (profile_requested) { + snprintf(profile_fallback, sizeof(profile_fallback), + "q8 weight_off=%llu in=%llu out=%llu tok=%llu", + (unsigned long long)weight_offset, + (unsigned long long)in_dim, + (unsigned long long)out_dim, + (unsigned long long)n_tok); + snprintf(profile_label_buf, sizeof(profile_label_buf), "%s", profile_fallback); + profile_label = profile_label_buf; + const char *profile_filter = getenv("DS4_METAL_Q8_PREFILL_PROFILE_FILTER"); + profile_prefill = + profile_requested && + (!profile_filter || !profile_filter[0] || + strstr(profile_label, profile_filter) != NULL); + } + if (profile_prefill) { + if (g_batch_cb) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + split_batch_for_profile = 1; + } + } + + const double profile_t0 = profile_prefill ? ds4_gpu_now_ms() : 0.0; + int ok = ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, + weight_offset, in_dim, out_dim, + x, n_tok, false, false); + if (profile_prefill) { + if (split_batch_for_profile && ds4_gpu_end_commands() == 0) { + ok = 0; + } + const double elapsed_ms = ds4_gpu_now_ms() - profile_t0; + fprintf(stderr, + "ds4: Metal Q8_0 prefill profile %s in=%llu out=%llu tok=%llu %.3f ms\n", + profile_label ? profile_label : profile_fallback, + (unsigned long long)in_dim, + (unsigned long long)out_dim, + (unsigned long long)n_tok, + elapsed_ms); + if (split_batch_for_profile && ds4_gpu_begin_commands() == 0) { + ok = 0; + } + } + return ok; +} + +int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_rows) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !x || !model_map || n_rows == 0 || + n_rows > INT32_MAX || in_dim == 0 || out_dim == 0 || + (in_dim & 31u) != 0 || in_dim > UINT32_MAX || + out_dim > UINT32_MAX || + in_dim > UINT64_MAX / n_rows / sizeof(float) || + out_dim > UINT64_MAX / n_rows / sizeof(float) || + ds4_gpu_tensor_bytes(x) < + (uint64_t)n_rows * in_dim * sizeof(float) || + ds4_gpu_tensor_bytes(out) < + (uint64_t)n_rows * out_dim * sizeof(float)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t blocks = in_dim / 32u; + const uint64_t row_bytes = blocks * 34u; + if (!xbuf || !outbuf || + out_dim > UINT64_MAX / row_bytes) { + return 0; + } + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = ds4_gpu_wrap_q8_decode_model_range( + model_map, model_size, weight_offset, weight_bytes, 1u, + &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_mv_dispatch dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + if (out_dim > 65536u) dispatch.nsg = 8; + ds4_gpu_q8_0_matvec_args args = + ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + args.ne11 = (int32_t)n_rows; + args.nb12 = (uint64_t)n_rows * in_dim * sizeof(float); + args.nb13 = args.nb12; + args.ne1 = (int32_t)n_rows; + args.nr0 = dispatch.nr0; + + id pipeline = + ds4_gpu_get_mul_mv_pipeline(dispatch.function_name, dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:dispatch.smem atIndex:0]; + [enc dispatchThreadgroups: + MTLSizeMake(((NSUInteger)out_dim + + (NSUInteger)dispatch.nr0 - 1u) / + (NSUInteger)dispatch.nr0, + (NSUInteger)n_rows, + 1) + threadsPerThreadgroup: + MTLSizeMake(32, (NSUInteger)dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return ds4_gpu_finish_command_buffer( + cb, owned, "Q8_0 exact decode-row matvec"); + } +} + +int ds4_gpu_matmul_q8_0_decode_mpp_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, + weight_offset, in_dim, out_dim, + x, n_tok, true, false); +} + +int ds4_gpu_matmul_q8_0_decode_mpp_model_view_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size, + weight_offset, in_dim, out_dim, + x, n_tok, true, true); +} + +static const char *ds4_gpu_q4_mv_ext_name(uint32_t weight_type, int16_t r1ptg) { + const char *prefix = NULL; + if (weight_type == DS4_METAL_TENSOR_Q4_K) { + prefix = "kernel_mul_mv_ext_q4_K_f32_r1_"; + } else if (weight_type == DS4_METAL_TENSOR_Q4_0) { + prefix = "kernel_mul_mv_ext_q4_0_f32_r1_"; + } else { + return NULL; + } + switch (r1ptg) { + case 1: return weight_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mv_ext_q4_K_f32_r1_1" : "kernel_mul_mv_ext_q4_0_f32_r1_1"; + case 2: return weight_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mv_ext_q4_K_f32_r1_2" : "kernel_mul_mv_ext_q4_0_f32_r1_2"; + case 3: return weight_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mv_ext_q4_K_f32_r1_3" : "kernel_mul_mv_ext_q4_0_f32_r1_3"; + case 4: return weight_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mv_ext_q4_K_f32_r1_4" : "kernel_mul_mv_ext_q4_0_f32_r1_4"; + case 5: return weight_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mv_ext_q4_K_f32_r1_5" : "kernel_mul_mv_ext_q4_0_f32_r1_5"; + default: + (void)prefix; + return NULL; + } +} + +static const char *ds4_gpu_q4_mm_name(uint32_t weight_type) { + switch (weight_type) { + case DS4_METAL_TENSOR_Q4_0: return "kernel_mul_mm_q4_0_f32"; + case DS4_METAL_TENSOR_Q4_K: return "kernel_mul_mm_q4_K_f32"; + default: return NULL; + } +} + +static const char *ds4_gpu_q4_nax_name(uint32_t weight_type, uint64_t tile_n) { + if (weight_type == DS4_METAL_TENSOR_Q4_0) { + return tile_n == 128u ? "kernel_mul_mm_q4_0_f32_nax_direct_rhs_n128" : + tile_n == 64u ? "kernel_mul_mm_q4_0_f32_nax_direct_rhs_n64" : + "kernel_mul_mm_q4_0_f32_nax_direct_rhs"; + } + if (weight_type == DS4_METAL_TENSOR_Q4_K) { + return tile_n == 128u ? "kernel_mul_mm_q4_K_f32_nax_direct_rhs_n128" : + tile_n == 64u ? "kernel_mul_mm_q4_K_f32_nax_direct_rhs_n64" : + "kernel_mul_mm_q4_K_f32_nax_direct_rhs"; + } + return NULL; +} + +static int ds4_gpu_matmul_quant_impl_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + bool prefer_decode_mpp, + bool force_model_view) { + if (weight_type == DS4_METAL_TENSOR_Q8_0) { + return ds4_gpu_matmul_q8_0_legacy_tensor(out, + model_map, + model_size, + weight_offset, + in_dim, + out_dim, + x, + n_tok, + prefer_decode_mpp, + force_model_view); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !x || !model_map || + in_dim == 0 || out_dim == 0 || n_tok == 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) { + return 0; + } + + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, (uint32_t)in_dim, &row_bytes)) { + fprintf(stderr, "ds4: Metal quant matmul received unsupported type/dim (%u, in=%llu)\n", + weight_type, + (unsigned long long)in_dim); + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal quant tensor matmul received undersized activation buffers\n"); + return 0; + } + + if (out_dim > UINT64_MAX / row_bytes) return 0; + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal quant tensor matmul range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = force_model_view ? + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset) : + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + /* + * Small-batch Q4_K goes to the classic (llama.cpp-style) matvec: + * the mul_mv_ext family tops out around 220 GB/s on M5 for the GLM + * DenseQ4 decode shapes while this impl streams 530-650 GB/s + * (misc/q4mv_bench.m). Falls through to ext when unavailable. + */ + if (weight_type == DS4_METAL_TENSOR_Q4_K && + n_tok <= 8u && + (in_dim % 256u) == 0 && + getenv("DS4_METAL_DISABLE_Q4_MV_CLASSIC") == NULL) { + const int16_t nsg = 2; + id pipeline = + ds4_gpu_get_mul_mv_ext_pipeline("kernel_mul_mv_q4_K_dense_f32", nsg, 8); + if (pipeline) { + ds4_gpu_q8_0_matvec_args args = { + .ne00 = (int32_t)in_dim, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = 1, + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)in_dim, + .ne11 = (int32_t)n_tok, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = in_dim * sizeof(float), + .nb12 = in_dim * n_tok * sizeof(float), + .nb13 = in_dim * n_tok * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = (int32_t)n_tok, + .nr0 = 2, + .r2 = 1, + .r3 = 1, + }; + const uint64_t rows_ptg = (uint64_t)nsg * 2u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:32 atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + rows_ptg - 1u) / rows_ptg, + (NSUInteger)n_tok, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4_K classic mul_mv")) return 0; + return 1; + } + } + + if (n_tok <= 8u && (in_dim % 128u) == 0) { + const int16_t nsg = 2; + const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); + const int16_t r1ptg = (n_tok == 1u) ? 1 : ds4_gpu_mv_ext_r1ptg(n_tok); + const char *fn_name = ds4_gpu_q4_mv_ext_name(weight_type, r1ptg); + id pipeline = + fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; + if (pipeline) { + const int16_t nypsg = 32 / nxpsg; + const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; + ds4_gpu_mul_mv_ext_args args = + ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, row_bytes, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, + ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 tensor mul_mv_ext")) return 0; + return 1; + } + } + + if (ds4_gpu_mpp_available() && + n_tok >= 32u && + (in_dim % 64u) == 0 && + (out_dim % 64u) == 0 && + (n_tok % 32u) == 0) { + uint64_t nax_tile_n = 32u; + if ((n_tok % 128u) == 0) { + nax_tile_n = 128u; + } else if ((n_tok % 64u) == 0) { + nax_tile_n = 64u; + } + const char *nax_fn = ds4_gpu_q4_nax_name(weight_type, nax_tile_n); + id pipeline = + nax_fn ? ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false) : nil; + if (pipeline) { + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:64u * 32u * sizeof(uint16_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), + (NSUInteger)out_dim / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 NAX tensor matmul")) return 0; + return 1; + } + ds4_gpu_warn_mpp_fallback(); + } + + const char *mm_fn = ds4_gpu_q4_mm_name(weight_type); + const bool bc_inp = (in_dim % 32u) != 0; + const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0; + id pipeline = + mm_fn ? ds4_gpu_get_mul_mm_pipeline(mm_fn, bc_inp, bc_out) : nil; + if (!pipeline) return 0; + + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u, + ((NSUInteger)out_dim + 63u) / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 tensor matmul")) return 0; + } + + return 1; +} + +int ds4_gpu_matmul_quant_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + return ds4_gpu_matmul_quant_impl_tensor(out, + model_map, + model_size, + weight_offset, + weight_type, + in_dim, + out_dim, + x, + n_tok, + false, + false); +} + +int ds4_gpu_matmul_quant_decode_mpp_model_view_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + /* GLM decode dense matvecs: the classic Q8_0 kernels stream 570-613 + * GB/s on the GLM shapes while the MPP/nax matrix path measures ~150 + * GB/s at n_tok=1 (TP decode 8.75 -> 16.64 t/s on the IQ2+Q8 gguf). + * MPP stays available as an opt-in via DS4_METAL_Q8_DECODE_MPP. */ + const bool prefer_mpp = getenv("DS4_METAL_Q8_DECODE_MPP") != NULL; + return ds4_gpu_matmul_quant_impl_tensor(out, + model_map, + model_size, + weight_offset, + weight_type, + in_dim, + out_dim, + x, + n_tok, + prefer_mpp, + true); +} + +int ds4_gpu_matmul_quant_rows_scalar_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + return ds4_gpu_matmul_quant_tensor(out, + model_map, + model_size, + weight_offset, + weight_type, + in_dim, + out_dim, + x, + n_tok); +} + +int ds4_gpu_matmul_q8_0_rows_scalar_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (n_tok == 1) { + return ds4_gpu_matmul_q8_0_tensor(out, + model_map, + model_size, + weight_offset, + in_dim, + out_dim, + x, + 1); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !x || !model_map || + n_tok == 0 || n_tok > INT32_MAX || + (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || + in_dim > UINT64_MAX / sizeof(float) || + out_dim > UINT64_MAX / sizeof(float)) { + return 0; + } + + const uint64_t x_row_bytes = in_dim * sizeof(float); + const uint64_t out_row_bytes = out_dim * sizeof(float); + if ((x_row_bytes != 0 && n_tok > UINT64_MAX / x_row_bytes) || + (out_row_bytes != 0 && n_tok > UINT64_MAX / out_row_bytes)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t x_bytes = n_tok * x_row_bytes; + const uint64_t out_bytes = n_tok * out_row_bytes; + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal Q8_0 scalar-row matmul received undersized activation buffers\n"); + return 0; + } + + const uint64_t blocks = in_dim / 32; + if (blocks > UINT64_MAX / 34u) return 0; + const uint64_t row_bytes = blocks * 34u; + if (out_dim > UINT64_MAX / row_bytes) return 0; + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8_0 scalar-row matmul range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + if (out_dim > 65536u) mv_dispatch.nsg = 8; + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + + const NSUInteger x_base = ds4_gpu_tensor_offset(x); + const NSUInteger out_base = ds4_gpu_tensor_offset(out); + for (uint64_t t = 0; t < n_tok; t++) { + [enc setBuffer:xbuf offset:x_base + (NSUInteger)(t * x_row_bytes) atIndex:2]; + [enc setBuffer:outbuf offset:out_base + (NSUInteger)(t * out_row_bytes) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 scalar-row matvecs")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_matmul_q8_0_pair_tensor( + ds4_gpu_tensor *out0, + ds4_gpu_tensor *out1, + const void *model_map, + uint64_t model_size, + uint64_t weight0_offset, + uint64_t weight1_offset, + uint64_t in_dim, + uint64_t out0_dim, + uint64_t out1_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out0 || !out1 || !model_map || !x || n_tok != 1 || + out0_dim == 0 || out1_dim == 0 || (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out0_dim > UINT32_MAX || out1_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id out0buf = ds4_gpu_tensor_buffer(out0); + id out1buf = ds4_gpu_tensor_buffer(out1); + const uint64_t x_bytes = in_dim * sizeof(float); + const uint64_t out0_bytes = out0_dim * sizeof(float); + const uint64_t out1_bytes = out1_dim * sizeof(float); + if (!xbuf || !out0buf || !out1buf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out0) < out0_bytes || + ds4_gpu_tensor_bytes(out1) < out1_bytes) { + fprintf(stderr, "ds4: Metal paired Q8_0 matvec received undersized activation buffers\n"); + return 0; + } + + const uint64_t row_bytes = (in_dim / 32u) * 34u; + const uint64_t weight0_bytes = out0_dim * row_bytes; + const uint64_t weight1_bytes = out1_dim * row_bytes; + if (weight0_offset > model_size || weight0_bytes > model_size - weight0_offset || + weight1_offset > model_size || weight1_bytes > model_size - weight1_offset) { + fprintf(stderr, "ds4: Metal paired Q8_0 matvec range is outside the mapped model\n"); + return 0; + } + + uint64_t inner0 = 0; + uint64_t inner1 = 0; + id weight0buf = + ds4_gpu_wrap_model_range(model_map, model_size, + weight0_offset, weight0_bytes, &inner0); + id weight1buf = + ds4_gpu_wrap_model_range(model_map, model_size, + weight1_offset, weight1_bytes, &inner1); + if (!weight0buf || !weight1buf) return 0; + + ds4_gpu_mv_dispatch dispatch0 = ds4_gpu_make_q8_0_mv_dispatch(); + ds4_gpu_mv_dispatch dispatch1 = ds4_gpu_make_q8_0_mv_dispatch(); + if (out0_dim > 65536u) dispatch0.nsg = 8; + if (out1_dim > 65536u) dispatch1.nsg = 8; + /* A common threadgroup shape is required to retain each standalone + * reduction tree. Mixed 4/8-simdgroup extents use the existing fallback. */ + if (dispatch0.nsg != dispatch1.nsg) return 0; + + ds4_gpu_q8_0_matvec_args args0 = ds4_gpu_make_q8_0_mv_args(in_dim, out0_dim); + ds4_gpu_q8_0_matvec_args args1 = ds4_gpu_make_q8_0_mv_args(in_dim, out1_dim); + args0.nr0 = dispatch0.nr0; + args1.nr0 = dispatch1.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_q8_0_f32_pair", dispatch0.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args0 length:sizeof(args0) atIndex:0]; + [enc setBytes:&args1 length:sizeof(args1) atIndex:1]; + [enc setBuffer:weight0buf offset:(NSUInteger)inner0 atIndex:2]; + [enc setBuffer:weight1buf offset:(NSUInteger)inner1 atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:out0buf offset:ds4_gpu_tensor_offset(out0) atIndex:5]; + [enc setBuffer:out1buf offset:ds4_gpu_tensor_offset(out1) atIndex:6]; + [enc setThreadgroupMemoryLength:2u * dispatch0.smem atIndex:0]; + const uint64_t max_out_dim = out0_dim > out1_dim ? out0_dim : out1_dim; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)max_out_dim + + (NSUInteger)dispatch0.nr0 - 1u) / + (NSUInteger)dispatch0.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dispatch0.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "paired Q8_0 matvec")) return 0; + } + + return 1; +} + +int ds4_gpu_matmul_q8_0_f16_out_tensor( + ds4_gpu_tensor *out_h, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + (void)out_h; (void)model_map; (void)model_size; (void)weight_offset; + (void)in_dim; (void)out_dim; (void)x; (void)n_tok; + return 0; +} + +static int ds4_gpu_shared_gate_up_swiglu_q8_0_impl( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp, + int store_gate_up, + bool force_model_view) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!mid || !x || !model_map || + (store_gate_up && (!gate || !up)) || + (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || + !isfinite(clamp) || clamp < 0.0f) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id midbuf = ds4_gpu_tensor_buffer(mid); + id gatebuf = store_gate_up ? + ds4_gpu_tensor_buffer(gate) : midbuf; + id upbuf = store_gate_up ? + ds4_gpu_tensor_buffer(up) : midbuf; + const uint64_t x_bytes = in_dim * sizeof(float); + const uint64_t out_bytes = out_dim * sizeof(float); + if (!xbuf || !gatebuf || !upbuf || !midbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + (store_gate_up && ds4_gpu_tensor_bytes(gate) < out_bytes) || + (store_gate_up && ds4_gpu_tensor_bytes(up) < out_bytes) || + ds4_gpu_tensor_bytes(mid) < out_bytes) { + fprintf(stderr, "ds4: Metal shared expert fused gate/up received undersized activation buffers\n"); + return 0; + } + + const uint64_t blocks = in_dim / 32; + const uint64_t row_bytes = blocks * 34; + const uint64_t weight_bytes = out_dim * row_bytes; + if (gate_offset > model_size || weight_bytes > model_size - gate_offset || + up_offset > model_size || weight_bytes > model_size - up_offset) { + fprintf(stderr, "ds4: Metal shared expert fused gate/up range is outside the mapped model\n"); + return 0; + } + + const bool exact_decode_views = + !force_model_view && + getenv("DS4_METAL_ENABLE_Q8_DECODE_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_DISABLE_Q8_DECODE_EXACT_VIEWS") == NULL; + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + id gate_wbuf = exact_decode_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + gate_offset, + weight_bytes, + &gate_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset, + weight_bytes, + &gate_inner); + id up_wbuf = exact_decode_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + up_offset, + weight_bytes, + &up_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset, + weight_bytes, + &up_inner); + if (!gate_wbuf || !up_wbuf) return 0; + + ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + args.nr0 = mv_dispatch.nr0; + const char *fn_name = store_gate_up ? + (mv_dispatch.nr0 >= 4 ? + "kernel_dsv4_shared_gate_up_swiglu_q8_0_r4" : + "kernel_dsv4_shared_gate_up_swiglu_q8_0") : + (mv_dispatch.nr0 >= 4 ? + "kernel_dsv4_shared_mid_swiglu_q8_0_r4" : + "kernel_dsv4_shared_mid_swiglu_q8_0"); + id pipeline = + ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; + [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:gatebuf offset:(store_gate_up ? + ds4_gpu_tensor_offset(gate) : + ds4_gpu_tensor_offset(mid)) atIndex:4]; + [enc setBuffer:upbuf offset:(store_gate_up ? + ds4_gpu_tensor_offset(up) : + ds4_gpu_tensor_offset(mid)) atIndex:5]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; + [enc setBytes:&clamp length:sizeof(clamp) atIndex:7]; + [enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, + owned, + store_gate_up ? + "shared expert fused gate/up" : + "shared expert fused mid")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(gate, + up, + mid, + model_map, + model_size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp, + 1, + false); +} + +int ds4_gpu_shared_mid_swiglu_q8_0_tensor( + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(NULL, + NULL, + mid, + model_map, + model_size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp, + 0, + true); +} + +int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + float clamp) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(gate, + up, + mid, + model_map, + model_size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp, + 1, + true); +} + +int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + float clamp) { + if (n_tok == 1) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(gate, + up, + mid, + model_map, + model_size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!gate || !up || !mid || !x || !model_map || + n_tok == 0 || + (in_dim & 31u) != 0 || (in_dim % 128u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX || + !isfinite(clamp) || clamp < 0.0f) { + return 0; + } + + const uint64_t mv_ext_max_tokens = + ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u); + if (n_tok > mv_ext_max_tokens) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id midbuf = ds4_gpu_tensor_buffer(mid); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (!xbuf || !gatebuf || !upbuf || !midbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(gate) < out_bytes || + ds4_gpu_tensor_bytes(up) < out_bytes || + ds4_gpu_tensor_bytes(mid) < out_bytes) { + fprintf(stderr, "ds4: Metal fused Q8_0 gate/up rows received undersized activation buffers\n"); + return 0; + } + + const uint64_t blocks = in_dim / 32; + const uint64_t row_bytes = blocks * 34; + const uint64_t weight_bytes = out_dim * row_bytes; + if (gate_offset > model_size || weight_bytes > model_size - gate_offset || + up_offset > model_size || weight_bytes > model_size - up_offset) { + fprintf(stderr, "ds4: Metal fused Q8_0 gate/up rows range is outside the mapped model\n"); + return 0; + } + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + id gate_wbuf = + ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner); + id up_wbuf = + ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner); + if (!gate_wbuf || !up_wbuf) return 0; + + const int16_t nsg = 2; + const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); + const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); + const char *fn_name = ds4_gpu_mv_ext_q8_pair_swiglu_name(r1ptg); + id pipeline = + fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; + if (!pipeline) return 0; + + const int16_t nypsg = 32 / nxpsg; + const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; + ds4_gpu_mul_mv_ext_args args = + ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, 34, row_bytes); + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; + [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:4]; + [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:5]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; + [enc setBytes:&clamp length:sizeof(clamp) atIndex:7]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / + (NSUInteger)r0ptg, + ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / + (NSUInteger)r1ptg, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "fused Q8_0 gate/up rows")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok, + float clamp) { + if (n_tok == 1) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(gate, + up, + mid, + model_map, + model_size, + gate_offset, + up_offset, + in_dim, + out_dim, + x, + clamp); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!gate || !up || !mid || !x || !model_map || + n_tok == 0 || + (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX || + in_dim > UINT64_MAX / sizeof(float) || + out_dim > UINT64_MAX / sizeof(float) || + !isfinite(clamp) || clamp < 0.0f) { + return 0; + } + + const uint64_t x_row_bytes = in_dim * sizeof(float); + const uint64_t out_row_bytes = out_dim * sizeof(float); + if ((x_row_bytes != 0 && n_tok > UINT64_MAX / x_row_bytes) || + (out_row_bytes != 0 && n_tok > UINT64_MAX / out_row_bytes)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id midbuf = ds4_gpu_tensor_buffer(mid); + const uint64_t x_bytes = n_tok * x_row_bytes; + const uint64_t out_bytes = n_tok * out_row_bytes; + if (!xbuf || !gatebuf || !upbuf || !midbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(gate) < out_bytes || + ds4_gpu_tensor_bytes(up) < out_bytes || + ds4_gpu_tensor_bytes(mid) < out_bytes) { + fprintf(stderr, "ds4: Metal shared expert scalar-row fused gate/up received undersized activation buffers\n"); + return 0; + } + + const uint64_t blocks = in_dim / 32; + if (blocks > UINT64_MAX / 34u) return 0; + const uint64_t row_bytes = blocks * 34u; + if (out_dim > UINT64_MAX / row_bytes) return 0; + const uint64_t weight_bytes = out_dim * row_bytes; + if (gate_offset > model_size || weight_bytes > model_size - gate_offset || + up_offset > model_size || weight_bytes > model_size - up_offset) { + fprintf(stderr, "ds4: Metal shared expert scalar-row fused gate/up range is outside the mapped model\n"); + return 0; + } + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + id gate_wbuf = + ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner); + id up_wbuf = + ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner); + if (!gate_wbuf || !up_wbuf) return 0; + + ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + args.nr0 = mv_dispatch.nr0; + args.ne11 = (int32_t)n_tok; + args.nb12 = n_tok * x_row_bytes; + args.nb13 = args.nb12; + args.ne1 = (int32_t)n_tok; + const char *fn_name = mv_dispatch.nr0 >= 4 ? + "kernel_dsv4_shared_gate_up_swiglu_q8_0_r4" : + "kernel_dsv4_shared_gate_up_swiglu_q8_0"; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; + [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:4]; + [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:5]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; + [enc setBytes:&clamp length:sizeof(clamp) atIndex:7]; + [enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + (NSUInteger)n_tok, + 1) + threadsPerThreadgroup: + MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "shared expert scalar-row fused gate/up")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_matmul_f16_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal F16 tensor matmul received undersized activation buffers\n"); + return 0; + } + + const uint64_t row_bytes = in_dim * sizeof(uint16_t); + const uint64_t weight_bytes = row_bytes * out_dim; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal F16 tensor matmul range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_f32_decode_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + n_tok, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (n_tok == 1) { + ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = + ds4_gpu_make_plain_mv_dispatch(in_dim, 0); + if (!g_quality_mode && (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) { + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor matvec")) return 0; + return 1; + } + + if (n_tok <= 8 && (in_dim % 128u) == 0) { + const int16_t nsg = 2; + const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); + const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); + const char *fn_name = ds4_gpu_mv_ext_name(0, r1ptg); + id pipeline = + fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; + if (!pipeline) return 0; + + const int16_t nypsg = 32 / nxpsg; + const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; + ds4_gpu_mul_mv_ext_args args = + ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, sizeof(uint16_t), row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, + ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor mul_mv_ext")) return 0; + return 1; + } + + /* + * Same direct-RHS TensorOps structure as Q8_0, but for F16 model + * matrices. The 128-token RHS tile is kept when the batch alignment + * allows it because the later tile_n=64 retest was neutral/slower. + */ + if (ds4_gpu_mpp_available() && + n_tok >= 32u && + (in_dim % 32u) == 0 && + (out_dim % 64u) == 0 && + (n_tok % 32u) == 0) { + uint64_t nax_tile_n = 32u; + if ((n_tok % 128u) == 0) { + nax_tile_n = 128u; + } else if ((n_tok % 64u) == 0) { + nax_tile_n = 64u; + } + const char *nax_fn = nax_tile_n == 128u + ? "kernel_mul_mm_f16_f32_mpp_direct_rhs_n128" + : (nax_tile_n == 64u + ? "kernel_mul_mm_f16_f32_mpp_direct_rhs_n64" + : "kernel_mul_mm_f16_f32_mpp_direct_rhs"); + id pipeline = + ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false); + if (pipeline) { + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n), + (NSUInteger)out_dim / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 NAX tensor matmul")) { + return 0; + } + return 1; + } + ds4_gpu_warn_mpp_fallback(); + } + + const bool bc_inp = (in_dim % 32u) != 0; + const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0; + id pipeline = + ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_f16_f32", bc_inp, bc_out); + if (!pipeline) return 0; + + ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u, + ((NSUInteger)out_dim + 63u) / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor matmul")) return 0; + } + + return 1; +} + +int ds4_gpu_matmul_f16_pair_tensor( + ds4_gpu_tensor *out_a, + ds4_gpu_tensor *out_b, + const void *model_map, + uint64_t model_size, + uint64_t weight_a_offset, + uint64_t weight_b_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok != 1 || (in_dim & 3u) != 0) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outabuf = ds4_gpu_tensor_buffer(out_a); + id outbbuf = ds4_gpu_tensor_buffer(out_b); + const uint64_t x_bytes = in_dim * sizeof(float); + const uint64_t out_bytes = out_dim * sizeof(float); + if (!xbuf || !outabuf || !outbbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out_a) < out_bytes || + ds4_gpu_tensor_bytes(out_b) < out_bytes) { + fprintf(stderr, "ds4: Metal F16 paired matvec received undersized activation buffers\n"); + return 0; + } + + const uint64_t row_bytes = in_dim * sizeof(uint16_t); + const uint64_t weight_bytes = row_bytes * out_dim; + if (weight_a_offset > model_size || weight_bytes > model_size - weight_a_offset || + weight_b_offset > model_size || weight_bytes > model_size - weight_b_offset) { + fprintf(stderr, "ds4: Metal F16 paired matvec range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_a = 0; + uint64_t inner_b = 0; + id wabuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_a_offset, weight_bytes, + &inner_a); + id wbbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_b_offset, weight_bytes, + &inner_b); + if (!wabuf || !wbbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 0); + if (ds4_gpu_use_compressor_pair_nr4() && + (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) { + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_f16_f32_pair_4", mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wabuf offset:(NSUInteger)inner_a atIndex:1]; + [enc setBuffer:wbbuf offset:(NSUInteger)inner_b atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:outabuf offset:ds4_gpu_tensor_offset(out_a) atIndex:4]; + [enc setBuffer:outbbuf offset:ds4_gpu_tensor_offset(out_b) atIndex:5]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 paired matvec")) return 0; + } + + return 1; +} + +int ds4_gpu_matmul_f16_pair_compressor_store_tensor( + ds4_gpu_tensor *out_kv, + ds4_gpu_tensor *out_score, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const void *model_map, + uint64_t model_size, + uint64_t weight_kv_offset, + uint64_t weight_score_offset, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t in_dim, + uint32_t width, + const ds4_gpu_tensor *x, + uint32_t ratio, + uint32_t pos) { + if (!g_initialized && !ds4_gpu_init()) return -1; + const bool force = + getenv("DS4_METAL_ENABLE_COMPRESSOR_PAIR_STATE_STORE") != NULL; + if ((g_quality_mode || + (!ds4_gpu_device_name_contains("M3") && + !ds4_gpu_device_name_contains("M5") && !force)) || + getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_PAIR_STATE_STORE") != NULL || + getenv("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ") != NULL || + getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") != NULL) { + return 0; + } + if (!out_kv || !out_score || !state_kv || !state_score || !model_map || !x || + in_dim == 0 || width == 0 || ratio == 0 || + (ape_type != 0u && ape_type != 1u)) { + return -1; + } + if (in_dim != 4096u || + (width != 256u && width != 512u && width != 1024u) || + (ratio != 4u && ratio != 128u)) { + return 0; + } + + @autoreleasepool { + const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio; + const uint64_t row_bytes = in_dim * sizeof(uint16_t); + const uint64_t weight_bytes = row_bytes * width; + const uint64_t out_bytes = (uint64_t)width * sizeof(float); + const uint64_t state_bytes = + (uint64_t)state_rows * width * sizeof(float); + const uint64_t ape_elem = ape_type == 1u ? sizeof(uint16_t) : sizeof(float); + const uint64_t ape_bytes = (uint64_t)ratio * width * ape_elem; + if (weight_kv_offset > model_size || + weight_bytes > model_size - weight_kv_offset || + weight_score_offset > model_size || + weight_bytes > model_size - weight_score_offset || + ape_offset > model_size || ape_bytes > model_size - ape_offset) { + return -1; + } + + id xbuf = ds4_gpu_tensor_buffer(x); + id outkvbuf = ds4_gpu_tensor_buffer(out_kv); + id outscorebuf = ds4_gpu_tensor_buffer(out_score); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!xbuf || !outkvbuf || !outscorebuf || !statekvbuf || !statescbuf || + ds4_gpu_tensor_bytes(x) < in_dim * sizeof(float) || + ds4_gpu_tensor_bytes(out_kv) < out_bytes || + ds4_gpu_tensor_bytes(out_score) < out_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + return -1; + } + + uint64_t weight_kv_inner = 0; + uint64_t weight_score_inner = 0; + uint64_t ape_inner = 0; + id weightkvbuf = ds4_gpu_wrap_model_range( + model_map, model_size, weight_kv_offset, weight_bytes, + &weight_kv_inner); + id weightscorebuf = ds4_gpu_wrap_model_range( + model_map, model_size, weight_score_offset, weight_bytes, + &weight_score_inner); + id apebuf = ds4_gpu_wrap_model_range( + model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!weightkvbuf || !weightscorebuf || !apebuf) return -1; + + ds4_gpu_f16_matvec_args mv_args = + ds4_gpu_make_f16_mv_args(in_dim, width); + ds4_gpu_mv_dispatch mv_dispatch = + ds4_gpu_make_plain_mv_dispatch(in_dim, 0); + if (ds4_gpu_use_compressor_pair_nr4() && + (width == 512u || width == 1024u) && in_dim >= 4096u) { + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + mv_args.nr0 = mv_dispatch.nr0; + ds4_gpu_dsv4_compressor_store_one_args store_args = { + .width = width, + .ratio = ratio, + .pos = pos, + .ape_type = ape_type, + }; + id pipeline = ds4_gpu_get_mul_mv_pipeline( + "kernel_mul_mv_f16_f32_pair_compressor_store_4", + mv_dispatch.nsg); + if (!pipeline) return -1; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return -1; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBytes:&store_args length:sizeof(store_args) atIndex:1]; + [enc setBuffer:weightkvbuf offset:(NSUInteger)weight_kv_inner atIndex:2]; + [enc setBuffer:weightscorebuf offset:(NSUInteger)weight_score_inner atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:outkvbuf offset:ds4_gpu_tensor_offset(out_kv) atIndex:5]; + [enc setBuffer:outscorebuf offset:ds4_gpu_tensor_offset(out_score) atIndex:6]; + [enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:7]; + [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:8]; + [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:9]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)width + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake( + 32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "F16 paired matvec compressor state store")) { + return -1; + } + } + + return 1; +} + +int ds4_gpu_matmul_f32_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t n_tok) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok == 0 || n_tok > UINT32_MAX) return 0; + if (in_dim > UINT64_MAX / n_tok / sizeof(float) || + out_dim > UINT64_MAX / n_tok / sizeof(float)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t x_bytes = n_tok * in_dim * sizeof(float); + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal F32 tensor matmul received undersized activation buffers\n"); + return 0; + } + + const uint64_t row_bytes = in_dim * sizeof(float); + const uint64_t weight_bytes = row_bytes * out_dim; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal F32 tensor matmul range is outside the mapped model\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = + ds4_gpu_wrap_f32_decode_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + n_tok, + &inner_offset); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (n_tok == 1) { + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, 1); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1); + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor matvec")) return 0; + return 1; + } + + if (n_tok <= 8 && (in_dim % 128u) == 0) { + const int16_t nsg = 2; + const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok); + const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok); + const char *fn_name = ds4_gpu_mv_ext_f32_name(r1ptg); + id pipeline = + fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil; + if (!pipeline) return 0; + + const int16_t nypsg = 32 / nxpsg; + const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg; + ds4_gpu_mul_mv_ext_args args = + ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, sizeof(float), row_bytes); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg, + ((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor mul_mv_ext")) return 0; + return 1; + } + + /* Generic multi-row path (GLM prefill shapes: one grid row per + * token through the plain matvec pipeline). */ + { + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, n_tok); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1); + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + if (mv_dispatch.smem) { + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + (NSUInteger)n_tok, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor matmul")) return 0; + } + } + + return 1; +} + +int ds4_gpu_repeat_hc_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *row, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !row || n_embd == 0 || n_hc == 0) return 0; + + @autoreleasepool { + id rowbuf = ds4_gpu_tensor_buffer(row); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t out_bytes = row_bytes * n_hc; + if (!rowbuf || !outbuf || + ds4_gpu_tensor_bytes(row) < row_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal HC repeat received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_repeat_hc_embedding(cb, + rowbuf, + ds4_gpu_tensor_offset(row), + outbuf, + ds4_gpu_tensor_offset(out), + 1, + n_embd, + n_hc)) { + return 0; + } + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC repeat")) return 0; + } + + return 1; +} + +int ds4_gpu_repeat_hc_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *rows, + uint32_t n_tokens, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !rows || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0; + + @autoreleasepool { + id rowsbuf = ds4_gpu_tensor_buffer(rows); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t rows_bytes = (uint64_t)n_tokens * row_bytes; + const uint64_t out_bytes = rows_bytes * n_hc; + if (!rowsbuf || !outbuf || + ds4_gpu_tensor_bytes(rows) < rows_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal HC row repeat received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_repeat_hc_embedding(cb, + rowsbuf, + ds4_gpu_tensor_offset(rows), + outbuf, + ds4_gpu_tensor_offset(out), + n_tokens, + n_embd, + n_hc)) { + return 0; + } + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC row repeat")) return 0; + } + + return 1; +} + +int ds4_gpu_rms_norm_plain_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + uint32_t n, + float eps) { + return ds4_gpu_rms_norm_plain_rows_tensor(out, x, n, 1, eps); +} + +int ds4_gpu_rms_norm_plain_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + uint32_t n, + uint32_t rows, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (n == 0 || rows == 0 || (n & 3u) != 0) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t bytes = (uint64_t)n * rows * sizeof(float); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal plain RMS norm received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps); + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_rms_norm_plain_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "plain RMS norm")) return 0; + } + + return 1; +} + +static int ds4_gpu_hc_rms_scale_project_mode( + uint32_t in_dim, + uint32_t out_dim, + uint32_t n_rows) { + const bool hard_shape = n_rows > 8u && + (in_dim == 16384u || in_dim == 28672u) && out_dim == 24u; + if (!hard_shape) return 0; + + const bool force = + getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL; + if (getenv("DS4_METAL_DISABLE_M3_HC_RMS_SCALE_PROJ") != NULL || + (!ds4_gpu_device_name_contains("M3") && !force)) { + return 0; + } + if (g_rms_norm_scale_pipeline == nil) { + return force ? -1 : 0; + } + return 1; +} + +int ds4_gpu_hc_rms_scale_project_f16_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *scale_scratch, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t in_dim, + uint32_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_rows, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !scale_scratch || !model_map || !x || + in_dim == 0u || out_dim == 0u || n_rows == 0u || + n_rows > INT32_MAX) { + return 0; + } + + const int mode = + ds4_gpu_hc_rms_scale_project_mode(in_dim, out_dim, n_rows); + if (mode < 0) return 0; + if (mode == 0) { + return ds4_gpu_rms_norm_plain_rows_tensor( + scale_scratch, x, in_dim, n_rows, eps) != 0 && + ds4_gpu_matmul_f16_tensor( + out, model_map, model_size, weight_offset, + in_dim, out_dim, scale_scratch, n_rows) != 0; + } + + @autoreleasepool { + const uint64_t x_elems = (uint64_t)in_dim * n_rows; + const uint64_t out_elems = (uint64_t)out_dim * n_rows; + if (x_elems > UINT64_MAX / sizeof(float) || + out_elems > UINT64_MAX / sizeof(float)) { + return 0; + } + const uint64_t x_bytes = x_elems * sizeof(float); + const uint64_t out_bytes = out_elems * sizeof(float); + const uint64_t scale_bytes = (uint64_t)n_rows * sizeof(float); + + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + id scalebuf = ds4_gpu_tensor_buffer(scale_scratch); + if (!xbuf || !outbuf || !scalebuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(scale_scratch) < scale_bytes) { + fprintf(stderr, "ds4: Metal HC RMS scale projection received undersized activation buffers\n"); + return 0; + } + + const uint64_t weight_row_bytes = (uint64_t)in_dim * sizeof(uint16_t); + if ((uint64_t)out_dim > UINT64_MAX / weight_row_bytes) return 0; + const uint64_t weight_bytes = (uint64_t)out_dim * weight_row_bytes; + if (weight_offset > model_size || + weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal HC RMS scale projection weight range is outside the mapped model\n"); + return 0; + } + + const bool bc_inp = (in_dim % 32u) != 0u; + const bool bc_out = (out_dim % 64u) != 0u || + (n_rows % 32u) != 0u; + id mm_pipeline = + ds4_gpu_get_mul_mm_pipeline( + "kernel_mul_mm_f16_f32_scaled", bc_inp, bc_out); + if (!mm_pipeline) { + if (getenv("DS4_METAL_ENABLE_HC_RMS_SCALE_PROJ") != NULL) { + return 0; + } + return ds4_gpu_rms_norm_plain_rows_tensor( + scale_scratch, x, in_dim, n_rows, eps) != 0 && + ds4_gpu_matmul_f16_tensor( + out, model_map, model_size, weight_offset, + in_dim, out_dim, scale_scratch, n_rows) != 0; + } + + id scale_pipeline = ds4_gpu_hot_pipeline( + g_rms_norm_scale_pipeline, "kernel_rms_norm_scale_f32_4"); + if (!scale_pipeline) return 0; + + uint64_t weight_inner = 0; + id weightbuf = ds4_gpu_wrap_model_range( + model_map, model_size, weight_offset, weight_bytes, &weight_inner); + if (!weightbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_rms_norm_args norm_args = + ds4_gpu_make_rms_norm_args(in_dim, n_rows, eps); + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:scale_pipeline]; + [enc setBytes:&norm_args length:sizeof(norm_args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:scalebuf + offset:ds4_gpu_tensor_offset(scale_scratch) + atIndex:2]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake( + ds4_gpu_rms_norm_threads(in_dim), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_mul_mm_args mm_args = ds4_gpu_make_mm_args( + in_dim, out_dim, n_rows, weight_row_bytes); + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:mm_pipeline]; + [enc setBytes:&mm_args length:sizeof(mm_args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setBuffer:scalebuf + offset:ds4_gpu_tensor_offset(scale_scratch) + atIndex:4]; + [enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)n_rows + 31u) / 32u, + ((NSUInteger)out_dim + 63u) / 64u, + 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "HC RMS scale F16 projection")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_rms_norm_weight_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + float eps) { + return ds4_gpu_rms_norm_weight_rows_tensor(out, x, model_map, model_size, weight_offset, n, 1, eps); +} + +int ds4_gpu_rms_norm_weight_rows_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *x, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + uint32_t rows, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (n == 0 || rows == 0 || (n & 3u) != 0) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + const uint64_t bytes = row_bytes * rows; + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal weighted RMS norm received undersized activation buffers\n"); + return 0; + } + if (weight_offset > model_size || row_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal weighted RMS norm range is outside the mapped model\n"); + return 0; + } + + const bool exact_decode_weight_view = + rows == 1u && + row_bytes <= (1ull << 20) && + getenv("DS4_METAL_ENABLE_DECODE_NORM_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; + uint64_t inner_offset = 0; + id wbuf = exact_decode_weight_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + weight_offset, + row_bytes, + &inner_offset) : + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + row_bytes, + &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps); + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_rms_norm_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "weighted RMS norm")) return 0; + } + + return 1; +} + +int ds4_gpu_add_rms_norm_weight_tensor( + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *sum_out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!norm_out || !sum_out || !a || !b || n == 0 || (n & 3u) != 0) return 0; + + @autoreleasepool { + id abuf = ds4_gpu_tensor_buffer(a); + id bbuf = ds4_gpu_tensor_buffer(b); + id sumbuf = ds4_gpu_tensor_buffer(sum_out); + id normbuf = ds4_gpu_tensor_buffer(norm_out); + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + if (!abuf || !bbuf || !sumbuf || !normbuf || + ds4_gpu_tensor_bytes(a) < row_bytes || + ds4_gpu_tensor_bytes(b) < row_bytes || + ds4_gpu_tensor_bytes(sum_out) < row_bytes || + ds4_gpu_tensor_bytes(norm_out) < row_bytes) { + fprintf(stderr, "ds4: Metal add+RMS norm received undersized activation buffers\n"); + return 0; + } + if (weight_offset > model_size || row_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal add+RMS norm range is outside the mapped model\n"); + return 0; + } + + const bool exact_decode_weight_view = + row_bytes <= (1ull << 20) && + getenv("DS4_METAL_ENABLE_DECODE_NORM_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; + uint64_t inner_offset = 0; + id wbuf = exact_decode_weight_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + weight_offset, + row_bytes, + &inner_offset) : + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + row_bytes, + &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, 1, eps); + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_add_rms_norm_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1]; + [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:3]; + [enc setBuffer:sumbuf offset:ds4_gpu_tensor_offset(sum_out) atIndex:4]; + [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "add+RMS norm")) return 0; + } + + return 1; +} + +int ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( + ds4_gpu_tensor *q_out, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t q_weight_offset, + uint32_t q_n, + ds4_gpu_tensor *kv_out, + const ds4_gpu_tensor *kv, + uint64_t kv_weight_offset, + uint32_t kv_n, + uint32_t rows, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!q_out || !q || !kv_out || !kv || q_n == 0 || kv_n == 0 || rows == 0 || + (q_n & 3u) != 0 || (kv_n & 3u) != 0) { + return 0; + } + + @autoreleasepool { + id qbuf = ds4_gpu_tensor_buffer(q); + id qoutbuf = ds4_gpu_tensor_buffer(q_out); + id kvbuf = ds4_gpu_tensor_buffer(kv); + id kvoutbuf = ds4_gpu_tensor_buffer(kv_out); + + const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float); + const uint64_t kv_row_bytes = (uint64_t)kv_n * sizeof(float); + if (!qbuf || !qoutbuf || !kvbuf || !kvoutbuf || + ds4_gpu_tensor_bytes(q) < q_row_bytes * rows || + ds4_gpu_tensor_bytes(q_out) < q_row_bytes * rows || + ds4_gpu_tensor_bytes(kv) < kv_row_bytes * rows || + ds4_gpu_tensor_bytes(kv_out) < kv_row_bytes * rows) { + fprintf(stderr, "ds4: Metal fused q/kv RMS norm received undersized activation buffers\n"); + return 0; + } + if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset || + kv_weight_offset > model_size || kv_row_bytes > model_size - kv_weight_offset) { + fprintf(stderr, "ds4: Metal fused q/kv RMS norm weight range is outside the mapped model\n"); + return 0; + } + + uint64_t q_inner_offset = 0; + uint64_t kv_inner_offset = 0; + id q_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + q_weight_offset, q_row_bytes, + &q_inner_offset); + if (!q_wbuf) return 0; + id kv_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + kv_weight_offset, kv_row_bytes, + &kv_inner_offset); + if (!kv_wbuf) return 0; + + ds4_gpu_qkv_rms_norm_args args = { + .q_n = (int32_t)q_n, + .q_n4 = (int32_t)(q_n / 4u), + .kv_n = (int32_t)kv_n, + .kv_n4 = (int32_t)(kv_n / 4u), + .q_row_stride = q_row_bytes, + .kv_row_stride = kv_row_bytes, + .eps = eps, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_qkv_rms_norm_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:q_wbuf offset:(NSUInteger)q_inner_offset atIndex:2]; + [enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:3]; + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:4]; + [enc setBuffer:kv_wbuf offset:(NSUInteger)kv_inner_offset atIndex:5]; + [enc setBuffer:kvoutbuf offset:ds4_gpu_tensor_offset(kv_out) atIndex:6]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 2, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "fused q/kv RMS norm")) return 0; + } + + return 1; +} + +int ds4_gpu_head_rms_norm_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || (head_dim & 3u) != 0) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = (uint64_t)n_tok * n_head * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, "ds4: Metal head RMS norm received undersized activation buffer\n"); + return 0; + } + + ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_3d_args(head_dim, n_head, n_tok, eps); + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_rms_norm_plain_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_head, n_tok, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_pipeline_threads(head_dim, g_rms_norm_plain_pipeline), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "head RMS norm")) return 0; + } + + return 1; +} + +int ds4_gpu_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || n_rot > head_dim || (n_rot & 1u) != 0) { + return 0; + } + if (n_rot == 0) return 1; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = (uint64_t)n_tok * n_head * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, "ds4: Metal RoPE received undersized activation buffer\n"); + return 0; + } + + ds4_gpu_rope_tail_batch_args args = ds4_gpu_make_rope_tail_args( + n_tok, n_head, head_dim, n_rot, n_ctx_orig, inverse, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_rope_tail_inplace(cb, + xbuf, + ds4_gpu_tensor_offset(x), + &args, + n_tok, + n_head, + head_dim, + pos0, + 1)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "RoPE tail")) return 0; + } + + return 1; +} + +int ds4_gpu_head_rms_norm_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps) { + (void)x; (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; + (void)pos0; (void)n_ctx_orig; (void)inverse; (void)freq_base; + (void)freq_scale; (void)ext_factor; (void)attn_factor; + (void)beta_fast; (void)beta_slow; (void)eps; + return 0; +} + +int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *q_half, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t n_rot, + uint32_t pos0, + uint32_t n_ctx_orig, + bool inverse, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float eps) { + (void)out; (void)q_half; (void)model_map; (void)model_size; + (void)weight_offset; (void)in_dim; (void)out_dim; (void)x; + (void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; (void)pos0; + (void)n_ctx_orig; (void)inverse; (void)freq_base; (void)freq_scale; + (void)ext_factor; (void)attn_factor; (void)beta_fast; (void)beta_slow; + (void)eps; + return 0; +} + +int ds4_gpu_dsv4_fp8_kv_quantize_tensor( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t head_dim, + uint32_t n_rot) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_tok == 0 || head_dim == 0 || n_rot > head_dim) return 0; + if (n_rot == head_dim) return 1; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = (uint64_t)n_tok * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, "ds4: Metal DSV4 FP8 KV quantize received undersized activation buffer\n"); + return 0; + } + + ds4_gpu_dsv4_fp8_kv_quantize_args args = { + .ne00 = head_dim, + .ne01 = n_tok, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = (uint64_t)head_dim * sizeof(float), + .nb02 = (uint64_t)n_tok * head_dim * sizeof(float), + .nb03 = (uint64_t)n_tok * head_dim * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)head_dim * sizeof(float), + .nb2 = (uint64_t)n_tok * head_dim * sizeof(float), + .nb3 = (uint64_t)n_tok * head_dim * sizeof(float), + .n_rot = (int32_t)n_rot, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_fp8_kv_quantize_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tok, 1, 1) + threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "DSV4 FP8 KV quantize")) return 0; + } + + return 1; +} + +int ds4_gpu_dsv4_indexer_qat_tensor( + ds4_gpu_tensor *x, + uint32_t n_rows, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_rows == 0 || head_dim != 128u) return 0; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = (uint64_t)n_rows * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, "ds4: Metal DSV4 indexer QAT received undersized activation buffer\n"); + return 0; + } + + ds4_gpu_dsv4_indexer_qat_args args = { + .n_rows = n_rows, + .head_dim = head_dim, + .row_stride = (uint64_t)head_dim * sizeof(float), + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_indexer_qat_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setThreadgroupMemoryLength:256u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "DSV4 indexer Hadamard+FP4")) return 0; + } + + return 1; +} + +static void ds4_gpu_set_rows_thread_shape( + uint32_t width, + NSUInteger *nth_out, + NSUInteger *nrptg_out) { + const NSUInteger nk0 = width ? (NSUInteger)width : 1u; + const NSUInteger max_threads = g_set_rows_f32_i32_pipeline + ? (NSUInteger)g_set_rows_f32_i32_pipeline.maxTotalThreadsPerThreadgroup + : 1024u; + + NSUInteger nth = 32u; + while (nth < nk0 && nth < max_threads) { + nth *= 2u; + } + + NSUInteger nrptg = 1u; + if (nth > nk0) { + nrptg = (nth + nk0 - 1u) / nk0; + nth = nk0; + if (nrptg * nth > max_threads) { + nrptg--; + } + } + + if (nth > nk0) nth = nk0; + if (nth == 0u) nth = 1u; + if (nrptg == 0u) nrptg = 1u; + + *nth_out = nth; + *nrptg_out = nrptg; +} + +static int ds4_gpu_encode_f16_round_copy_for_raw_store( + id cb, + const ds4_gpu_tensor *src, + uint32_t n) { + id srcbuf = ds4_gpu_tensor_buffer(src); + const uint64_t src_bytes = (uint64_t)n * sizeof(float); + if (!srcbuf || ds4_gpu_tensor_bytes(src) < src_bytes) { + fprintf(stderr, "ds4: Metal raw KV store received undersized source buffer\n"); + return 0; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_f16_round_scratch_buffer, + &g_f16_round_scratch_bytes, + (NSUInteger)n * sizeof(uint16_t), + "ds4_f16_round_scratch") || + !ds4_gpu_ensure_scratch_buffer(&g_raw_store_round_buffer, + &g_raw_store_round_bytes, + (NSUInteger)n * sizeof(float), + "ds4_raw_store_round")) { + return 0; + } + + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + srcbuf, + ds4_gpu_tensor_offset(src), + g_f16_round_scratch_buffer, + 0, + n)) { + return 0; + } + return ds4_gpu_encode_cpy_f16_f32_1d(cb, + g_f16_round_scratch_buffer, + 0, + g_raw_store_round_buffer, + 0, + n); +} + +static int ds4_gpu_encode_set_rows_f32_i32( + id cb, + ds4_gpu_tensor *dst, + id srcbuf, + NSUInteger src_off, + const int32_t *rows, + uint32_t n_rows, + uint32_t dst_rows, + uint32_t width) { + id dstbuf = ds4_gpu_tensor_buffer(dst); + const uint64_t dst_bytes = (uint64_t)dst_rows * width * sizeof(float); + const uint64_t src_bytes = (uint64_t)n_rows * width * sizeof(float); + if (!dstbuf || !srcbuf || !rows || n_rows == 0 || width == 0 || + ds4_gpu_tensor_bytes(dst) < dst_bytes || + src_bytes > NSUIntegerMax - src_off) { + fprintf(stderr, "ds4: Metal DS4 set_rows received invalid buffers\n"); + return 0; + } + + const uint64_t row_bytes = (uint64_t)width * sizeof(float); + const uint64_t rows_bytes = (uint64_t)n_rows * sizeof(int32_t); + ds4_gpu_set_rows_args args = { + .nk0 = (int32_t)width, + .ne01 = (int32_t)n_rows, + .nb01 = row_bytes, + .nb02 = (uint64_t)n_rows * row_bytes, + .nb03 = (uint64_t)n_rows * row_bytes, + .ne11 = 1, + .ne12 = 1, + .nb10 = sizeof(int32_t), + .nb11 = rows_bytes, + .nb12 = rows_bytes, + .nb1 = row_bytes, + .nb2 = (uint64_t)dst_rows * row_bytes, + .nb3 = (uint64_t)dst_rows * row_bytes, + }; + + NSUInteger nth; + NSUInteger nrptg; + ds4_gpu_set_rows_thread_shape(width, &nth, &nrptg); + + id rowsbuf = nil; + if (rows_bytes > 4096u) { + rowsbuf = ds4_gpu_new_transient_buffer((NSUInteger)rows_bytes, "ds4_set_rows_indices"); + if (!rowsbuf) return 0; + memcpy([rowsbuf contents], rows, (NSUInteger)rows_bytes); + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_set_rows_f32_i32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:srcbuf offset:src_off atIndex:1]; + if (rowsbuf) { + [enc setBuffer:rowsbuf offset:0 atIndex:2]; + } else { + [enc setBytes:rows length:(NSUInteger)rows_bytes atIndex:2]; + } + [enc setBuffer:dstbuf offset:ds4_gpu_tensor_offset(dst) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_rows + nrptg - 1u) / nrptg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, nrptg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_add_f32_1d( + id cb, + id a, + NSUInteger a_off, + id b, + NSUInteger b_off, + id out, + NSUInteger out_off, + uint32_t n) { + if (!cb || !a || !b || !out || n == 0) return 0; + + ds4_gpu_add_flat_args args = { .n = n }; + NSUInteger nth = g_add2_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + if (nth == 0u) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_add2_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:a offset:a_off atIndex:1]; + [enc setBuffer:b offset:b_off atIndex:2]; + [enc setBuffer:out offset:out_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +int ds4_gpu_store_raw_kv_tensor( + ds4_gpu_tensor *raw_cache, + const ds4_gpu_tensor *kv, + uint32_t raw_cap, + uint32_t row, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!raw_cache || !kv || raw_cap == 0 || row >= raw_cap || head_dim == 0 || raw_cap > INT32_MAX) return 0; + + @autoreleasepool { + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { + fprintf(stderr, "ds4: Metal raw KV store received undersized destination buffer\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const int32_t row_i32 = (int32_t)row; + if (!ds4_gpu_encode_f16_round_copy_for_raw_store(cb, kv, head_dim) || + !ds4_gpu_encode_set_rows_f32_i32(cb, raw_cache, + g_raw_store_round_buffer, + 0, + &row_i32, + 1, + raw_cap, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "raw KV DS4 set_rows store")) return 0; + } + + return 1; +} + +/* Release decode fused KV finalizer. Reference paths are selected by the C + * graph driver; this Objective-C entry point always means "use the fused + * Metal kernel." */ +int ds4_gpu_kv_fp8_store_raw_tensor( + ds4_gpu_tensor *kv, + ds4_gpu_tensor *raw_cache, + uint32_t raw_cap, + uint32_t row, + uint32_t head_dim, + uint32_t n_rot) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!kv || !raw_cache || raw_cap == 0 || row >= raw_cap || head_dim == 0 || + n_rot > head_dim || raw_cap > INT32_MAX) { + return 0; + } + + @autoreleasepool { + id kvbuf = ds4_gpu_tensor_buffer(kv); + id rawbuf = ds4_gpu_tensor_buffer(raw_cache); + const uint64_t kv_bytes = (uint64_t)head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + if (!kvbuf || !rawbuf || + ds4_gpu_tensor_bytes(kv) < kv_bytes || + ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { + fprintf(stderr, "ds4: Metal fused KV FP8/raw-store received undersized buffers\n"); + return 0; + } + + ds4_gpu_dsv4_kv_fp8_store_args args = { + .head_dim = (int32_t)head_dim, + .n_rot = (int32_t)n_rot, + .raw_row = (int32_t)row, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_kv_fp8_store_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:1]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_cache) atIndex:2]; + [enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "KV FP8/raw-store fused")) return 0; + } + + return 1; +} + +int ds4_gpu_store_raw_kv_batch_tensor( + ds4_gpu_tensor *raw_cache, + const ds4_gpu_tensor *kv, + uint32_t raw_cap, + uint32_t pos0, + uint32_t n_tokens, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!raw_cache || !kv || raw_cap == 0 || n_tokens == 0 || head_dim == 0 || raw_cap > INT32_MAX) return 0; + + @autoreleasepool { + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) { + fprintf(stderr, "ds4: Metal raw KV batch store received undersized destination buffer\n"); + return 0; + } + + int32_t rows_stack[512]; + int32_t *rows = rows_stack; + if (n_tokens > (uint32_t)(sizeof(rows_stack) / sizeof(rows_stack[0]))) { + rows = malloc((size_t)n_tokens * sizeof(*rows)); + if (!rows) { + fprintf(stderr, "ds4: failed to allocate raw KV set_rows index list\n"); + return 0; + } + } + for (uint32_t t = 0; t < n_tokens; t++) { + rows[t] = (int32_t)((pos0 + t) % raw_cap); + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) { + if (rows != rows_stack) free(rows); + return 0; + } + + const uint64_t n = (uint64_t)n_tokens * head_dim; + const int ok = n <= UINT32_MAX && + ds4_gpu_encode_f16_round_copy_for_raw_store(cb, kv, (uint32_t)n) && + ds4_gpu_encode_set_rows_f32_i32(cb, raw_cache, + g_raw_store_round_buffer, + 0, + rows, + n_tokens, + raw_cap, + head_dim); + if (rows != rows_stack) free(rows); + if (!ok) return 0; + + if (!ds4_gpu_finish_command_buffer(cb, owned, "raw KV batch DS4 set_rows store")) return 0; + } + + return 1; +} + +static int ds4_gpu_encode_compressor_score_with_ape( + id cb, + id score_src, + NSUInteger score_src_offset, + id score_dst, + NSUInteger score_dst_offset, + id apebuf, + NSUInteger ape_offset, + uint32_t ape_type, + uint32_t width, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens) { + if (!cb || !score_src || !score_dst || !apebuf || + width == 0 || ratio == 0 || n_tokens == 0 || + (ape_type != 0u && ape_type != 1u)) { + return 0; + } + + const uint64_t total_elems64 = (uint64_t)n_tokens * width; + if (total_elems64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal compressor APE add received too many elements\n"); + return 0; + } + const uint32_t total_elems = (uint32_t)total_elems64; + + const bool force_fused = + getenv("DS4_METAL_ENABLE_COMPRESSOR_APE_ADD") != NULL; + const bool use_fused = + (ds4_gpu_device_name_contains("M3") || force_fused) && + getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_APE_ADD") == NULL; + if (use_fused) { + id pipeline = ds4_gpu_get_pipeline( + ape_type == 1u ? "kernel_dsv4_compressor_score_ape_f16" + : "kernel_dsv4_compressor_score_ape_f32"); + if (pipeline) { + ds4_gpu_dsv4_compressor_score_ape_args args = { + .width = width, + .ratio = ratio, + .pos0 = pos0, + .n_tokens = n_tokens, + }; + NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth == 0) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:score_src offset:score_src_offset atIndex:1]; + [enc setBuffer:apebuf offset:ape_offset atIndex:2]; + [enc setBuffer:score_dst offset:score_dst_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake( + ((NSUInteger)total_elems + nth - 1u) / nth, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + if (force_fused) return 0; + } + + const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer, + &g_compressor_store_ape_bytes, + scratch_bytes, + "ds4_compressor_store_ape")) { + return 0; + } + + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + uint32_t copied_rows = 0; + uint32_t pos_mod = pos0 % ratio; + while (copied_rows < n_tokens) { + uint32_t seg_rows = ratio - pos_mod; + if (seg_rows > n_tokens - copied_rows) seg_rows = n_tokens - copied_rows; + const uint32_t seg_elems = seg_rows * width; + const NSUInteger src_off = ape_offset + (NSUInteger)pos_mod * width * elem_ape; + const NSUInteger dst_off = (NSUInteger)copied_rows * width * sizeof(float); + int ok; + if (ape_type == 1u) { + ok = ds4_gpu_encode_cpy_f16_f32_1d(cb, + apebuf, + src_off, + g_compressor_store_ape_buffer, + dst_off, + seg_elems); + } else { + ok = ds4_gpu_encode_cpy_f32_f32_1d(cb, + apebuf, + src_off, + g_compressor_store_ape_buffer, + dst_off, + seg_elems); + } + if (!ok) return 0; + copied_rows += seg_rows; + pos_mod = 0; + } + + return ds4_gpu_encode_add_f32_1d(cb, + score_src, + score_src_offset, + g_compressor_store_ape_buffer, + 0, + score_dst, + score_dst_offset, + total_elems); +} + +static int ds4_gpu_encode_compressor_set_rows_projected( + id cb, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + id kvbuf, + NSUInteger kv_offset, + id scorebuf, + NSUInteger score_offset, + id apebuf, + NSUInteger ape_offset, + uint32_t ape_type, + uint32_t width, + uint32_t ratio, + uint32_t pos0, + const int32_t *rows, + uint32_t n_rows, + uint32_t state_rows) { + if (!cb || !state_kv || !state_score || !kvbuf || !scorebuf || + !apebuf || !rows || width == 0 || n_rows == 0 || state_rows == 0) { + return 0; + } + + const NSUInteger score_scratch_bytes = (NSUInteger)n_rows * width * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, + &g_compressor_store_score_bytes, + score_scratch_bytes, + "ds4_compressor_store_score")) { + return 0; + } + + return ds4_gpu_encode_compressor_score_with_ape(cb, + scorebuf, + score_offset, + g_compressor_store_score_buffer, + 0, + apebuf, + ape_offset, + ape_type, + width, + ratio, + pos0, + n_rows) && + ds4_gpu_encode_set_rows_f32_i32(cb, + state_kv, + kvbuf, + kv_offset, + rows, + n_rows, + state_rows, + width) && + ds4_gpu_encode_set_rows_f32_i32(cb, + state_score, + g_compressor_store_score_buffer, + 0, + rows, + n_rows, + state_rows, + width); +} + +static int ds4_gpu_compressor_store_one_tensor( + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t width, + uint32_t ratio, + uint32_t pos) { + if (!kv || !sc || !state_kv || !state_score || !model_map || + width == 0 || ratio == 0 || (ape_type != 0u && ape_type != 1u)) { + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_dsv4_compressor_store_one_pipeline, + "kernel_dsv4_compressor_store_one"); + if (!pipeline) return 0; + + const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t row_bytes = (uint64_t)width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * row_bytes; + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + ds4_gpu_tensor_bytes(kv) < row_bytes || + ds4_gpu_tensor_bytes(sc) < row_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + return 0; + } + + uint64_t ape_inner = 0; + id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, + ape_offset, ape_bytes, + &ape_inner); + id kvbuf = ds4_gpu_tensor_buffer(kv); + id scbuf = ds4_gpu_tensor_buffer(sc); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!apebuf || !kvbuf || !scbuf || !statekvbuf || !statescbuf) return 0; + + ds4_gpu_dsv4_compressor_store_one_args args = { + .width = width, + .ratio = ratio, + .pos = pos, + .ape_type = ape_type, + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const NSUInteger nth = 256u; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:1]; + [enc setBuffer:scbuf offset:ds4_gpu_tensor_offset(sc) atIndex:2]; + [enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:3]; + [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:4]; + [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)width + nth - 1u) / nth, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return ds4_gpu_finish_command_buffer(cb, owned, "compressor one-row store"); +} + +int ds4_gpu_compressor_store_batch_tensor( + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!kv || !sc || !state_kv || !state_score || !model_map || + head_dim == 0 || ratio == 0 || n_tokens == 0 || + (ape_type != 0u && ape_type != 1u)) { + return 0; + } + + @autoreleasepool { + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t state_rows = coff * ratio; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset) { + fprintf(stderr, "ds4: Metal compressor batch APE range is outside the mapped model\n"); + return 0; + } + + id kvbuf = ds4_gpu_tensor_buffer(kv); + id scbuf = ds4_gpu_tensor_buffer(sc); + if (!kvbuf || !scbuf || + ds4_gpu_tensor_bytes(kv) < kv_bytes || + ds4_gpu_tensor_bytes(sc) < kv_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + fprintf(stderr, "ds4: Metal compressor batch store received undersized buffers\n"); + return 0; + } + + uint64_t ape_inner = 0; + id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!apebuf) return 0; + + const uint64_t total_elems64 = (uint64_t)n_tokens * width; + if (total_elems64 > UINT32_MAX || state_rows > INT32_MAX) { + fprintf(stderr, "ds4: Metal compressor batch store received too many elements\n"); + return 0; + } + const uint32_t total_elems = (uint32_t)total_elems64; + const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer, + &g_compressor_store_ape_bytes, + scratch_bytes, + "ds4_compressor_store_ape") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, + &g_compressor_store_score_bytes, + scratch_bytes, + "ds4_compressor_store_score")) { + return 0; + } + + int32_t rows_stack[16]; + int32_t *rows = rows_stack; + if (n_tokens > (uint32_t)(sizeof(rows_stack) / sizeof(rows_stack[0]))) { + rows = malloc((size_t)n_tokens * sizeof(*rows)); + if (!rows) { + fprintf(stderr, "ds4: failed to allocate compressor set_rows index list\n"); + return 0; + } + } + for (uint32_t t = 0; t < n_tokens; t++) { + const uint32_t pos_mod = (pos0 + t) % ratio; + rows[t] = (int32_t)(ratio == 4u ? ratio + pos_mod : pos_mod); + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) { + if (rows != rows_stack) free(rows); + return 0; + } + + int ok = 1; + uint32_t copied_rows = 0; + uint32_t pos_mod = pos0 % ratio; + while (ok && copied_rows < n_tokens) { + uint32_t seg_rows = ratio - pos_mod; + if (seg_rows > n_tokens - copied_rows) seg_rows = n_tokens - copied_rows; + const uint32_t seg_elems = seg_rows * width; + const NSUInteger src_off = (NSUInteger)ape_inner + + (NSUInteger)pos_mod * width * elem_ape; + const NSUInteger dst_off = (NSUInteger)copied_rows * width * sizeof(float); + if (ape_type == 1u) { + ok = ds4_gpu_encode_cpy_f16_f32_1d(cb, + apebuf, + src_off, + g_compressor_store_ape_buffer, + dst_off, + seg_elems); + } else { + ok = ds4_gpu_encode_cpy_f32_f32_1d(cb, + apebuf, + src_off, + g_compressor_store_ape_buffer, + dst_off, + seg_elems); + } + copied_rows += seg_rows; + pos_mod = 0; + } + + if (ok) { + ok = ds4_gpu_encode_add_f32_1d(cb, + scbuf, + ds4_gpu_tensor_offset(sc), + g_compressor_store_ape_buffer, + 0, + g_compressor_store_score_buffer, + 0, + total_elems); + } + if (ok) { + ok = ds4_gpu_encode_set_rows_f32_i32(cb, + state_kv, + kvbuf, + ds4_gpu_tensor_offset(kv), + rows, + n_tokens, + state_rows, + width); + } + if (ok) { + ok = ds4_gpu_encode_set_rows_f32_i32(cb, + state_score, + g_compressor_store_score_buffer, + 0, + rows, + n_tokens, + state_rows, + width); + } + if (rows != rows_stack) free(rows); + if (!ok) return 0; + + if (!ds4_gpu_finish_command_buffer(cb, owned, "compressor batch DS4 store")) return 0; + } + + return 1; +} + +static ds4_gpu_bin_args ds4_gpu_make_bin_contiguous_3d_args( + uint32_t cols, + uint32_t rows, + uint32_t planes) { + const uint64_t row_bytes = (uint64_t)cols * sizeof(float); + const uint64_t plane_bytes = (uint64_t)rows * row_bytes; + return (ds4_gpu_bin_args) { + .ne00 = (int32_t)cols, + .ne01 = (int32_t)rows, + .ne02 = (int32_t)planes, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = row_bytes, + .nb02 = plane_bytes, + .nb03 = (uint64_t)planes * plane_bytes, + .ne10 = (int32_t)cols, + .ne11 = (int32_t)rows, + .ne12 = (int32_t)planes, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = row_bytes, + .nb12 = plane_bytes, + .nb13 = (uint64_t)planes * plane_bytes, + .ne0 = (int32_t)cols, + .ne1 = (int32_t)rows, + .ne2 = (int32_t)planes, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = row_bytes, + .nb2 = plane_bytes, + .nb3 = (uint64_t)planes * plane_bytes, + .offs = 0, + .o1 = { 0 }, + }; +} + +static int ds4_gpu_encode_softmax_f32_contiguous( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t width, + uint32_t rows, + uint32_t planes) { + if (!cb || !src || !dst || width == 0 || rows == 0 || planes == 0) return 0; + + const uint64_t row_bytes = (uint64_t)width * sizeof(float); + const uint64_t plane_bytes = (uint64_t)rows * row_bytes; + ds4_gpu_softmax_args args = { + .ne00 = (int32_t)width, + .ne01 = (int32_t)rows, + .ne02 = (int32_t)planes, + .nb01 = row_bytes, + .nb02 = plane_bytes, + .nb03 = (uint64_t)planes * plane_bytes, + .ne11 = (int32_t)width, + .ne12 = (int32_t)rows, + .ne13 = (int32_t)planes, + .nb11 = row_bytes, + .nb12 = plane_bytes, + .nb13 = (uint64_t)planes * plane_bytes, + .nb1 = row_bytes, + .nb2 = plane_bytes, + .nb3 = (uint64_t)planes * plane_bytes, + .scale = 1.0f, + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 1, + }; + + id pipeline = + (width % 4u) == 0 ? g_soft_max_f32_4_pipeline : g_soft_max_f32_pipeline; + if (!pipeline) return 0; + + NSUInteger nth = 32u; + if ((width % 4u) == 0) { + while (nth < (NSUInteger)(width / 4u) && + nth * (NSUInteger)rows * (NSUInteger)planes < 256u) { + nth *= 2u; + } + } else { + while (nth < (NSUInteger)width && + nth * (NSUInteger)rows * (NSUInteger)planes < 256u) { + nth *= 2u; + } + } + const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:src offset:src_off atIndex:2]; + [enc setBuffer:src offset:src_off atIndex:3]; + [enc setBuffer:dst offset:dst_off atIndex:4]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, planes, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml( + id cb, + ds4_gpu_tensor *out, + id kvbuf, + NSUInteger kv_offset, + uint64_t kv_nb0, + uint64_t kv_nb1, + uint64_t kv_nb2, + id scorebuf, + NSUInteger score_offset, + uint64_t score_nb0, + uint64_t score_nb1, + uint64_t score_nb2, + uint32_t n_rows, + uint32_t head_dim) { + id outbuf = ds4_gpu_tensor_buffer(out); + if (!cb || !outbuf || !kvbuf || !scorebuf || n_rows == 0 || head_dim == 0 || + ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * sizeof(float)) { + return 0; + } + + const NSUInteger pack_bytes = (NSUInteger)n_rows * head_dim * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_product_buffer, + &g_compressor_pool_product_bytes, + pack_bytes, + "ds4_compressor_pool_product") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_cont_buffer, + &g_compressor_pool_score_cont_bytes, + pack_bytes, + "ds4_compressor_pool_score_cont") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_softmax_buffer, + &g_compressor_pool_softmax_bytes, + pack_bytes, + "ds4_compressor_pool_softmax")) { + return 0; + } + + const uint64_t cont_row_stride = (uint64_t)n_rows * sizeof(float); + const uint64_t cont_plane_stride = (uint64_t)head_dim * cont_row_stride; + + /* + * Keep the n_comp == 1 compressor path as the unfused graph sequence: + * + * score = soft_max(contiguous(score)) + * pooled = sum_rows(contiguous(kv) * score) + * + * The fused DS4 pool kernel is mathematically equivalent, but it reduces in + * a different order. That is enough to create ~1e-6 compressor differences + * and later FP8/routing flips, so this path intentionally keeps the same + * operation boundary and memory layout as the graph. + */ + ds4_gpu_bin_args mul_args = + ds4_gpu_make_bin_contiguous_3d_args(n_rows, head_dim, 1); + + return + ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb, + kvbuf, + kv_offset, + g_compressor_pool_product_buffer, + 0, + n_rows, + head_dim, + 1, + kv_nb0, + kv_nb1, + kv_nb2, + cont_row_stride, + cont_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb, + scorebuf, + score_offset, + g_compressor_pool_score_cont_buffer, + 0, + n_rows, + head_dim, + 1, + score_nb0, + score_nb1, + score_nb2, + cont_row_stride, + cont_plane_stride) && + ds4_gpu_encode_softmax_f32_contiguous(cb, + g_compressor_pool_score_cont_buffer, + 0, + g_compressor_pool_softmax_buffer, + 0, + n_rows, + head_dim, + 1) && + ds4_gpu_encode_bin_f32_rows(cb, + g_mul_pipeline, + &mul_args, + g_compressor_pool_product_buffer, + 0, + g_compressor_pool_softmax_buffer, + 0, + g_compressor_pool_product_buffer, + 0) && + ds4_gpu_encode_sum_rows_f32(cb, + g_compressor_pool_product_buffer, + 0, + outbuf, + ds4_gpu_tensor_offset(out), + n_rows, + head_dim); +} + +static int ds4_gpu_encode_dsv4_softmax_pool( + id cb, + ds4_gpu_tensor *out, + id kvbuf, + NSUInteger kv_offset, + uint64_t kv_nb0, + uint64_t kv_nb1, + uint64_t kv_nb2, + id scorebuf, + NSUInteger score_offset, + uint64_t score_nb0, + uint64_t score_nb1, + uint64_t score_nb2, + uint32_t n_rows, + uint32_t head_dim, + uint32_t n_comp) { + id outbuf = ds4_gpu_tensor_buffer(out); + if (!cb || !outbuf || !kvbuf || !scorebuf || + n_rows == 0 || head_dim == 0 || n_comp == 0 || + ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * n_comp * sizeof(float)) { + return 0; + } + + if (n_comp == 1) { + return ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml(cb, + out, + kvbuf, + kv_offset, + kv_nb0, + kv_nb1, + kv_nb2, + scorebuf, + score_offset, + score_nb0, + score_nb1, + score_nb2, + n_rows, + head_dim); + } + + ds4_gpu_dsv4_softmax_pool_args args = { + .ne00 = (int64_t)n_rows, + .ne01 = (int64_t)head_dim, + .ne02 = (int64_t)n_comp, + .nb00 = kv_nb0, + .nb01 = kv_nb1, + .nb02 = kv_nb2, + .nb10 = score_nb0, + .nb11 = score_nb1, + .nb12 = score_nb2, + .ne0 = (int64_t)head_dim, + .ne1 = (int64_t)n_comp, + .nb0 = sizeof(float), + .nb1 = (uint64_t)head_dim * sizeof(float), + }; + const uint64_t n = (uint64_t)head_dim * n_comp; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_softmax_pool_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; + [enc setBuffer:scorebuf offset:score_offset atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_concat_f32_dim1( + id cb, + id src0, + NSUInteger src0_offset, + uint32_t src0_rows, + uint64_t src0_row_stride, + id src1, + NSUInteger src1_offset, + uint32_t src1_rows, + uint64_t src1_row_stride, + id dst, + NSUInteger dst_offset, + uint32_t cols, + uint64_t dst_row_stride) { + if (!cb || !src0 || !src1 || !dst || cols == 0 || src0_rows == 0 || src1_rows == 0) { + return 0; + } + + const uint32_t rows = src0_rows + src1_rows; + const uint64_t src0_plane = (uint64_t)src0_rows * src0_row_stride; + const uint64_t src1_plane = (uint64_t)src1_rows * src1_row_stride; + const uint64_t dst_plane = (uint64_t)rows * dst_row_stride; + ds4_gpu_concat_args args = { + .ne00 = (int32_t)cols, + .ne01 = (int32_t)src0_rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src0_row_stride, + .nb02 = src0_plane, + .nb03 = src0_plane, + .ne10 = (int32_t)cols, + .ne11 = (int32_t)src1_rows, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = src1_row_stride, + .nb12 = src1_plane, + .nb13 = src1_plane, + .ne0 = (int32_t)cols, + .ne1 = (int32_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = dst_row_stride, + .nb2 = dst_plane, + .nb3 = dst_plane, + .dim = 1, + }; + + NSUInteger nth = cols < 1024u ? (NSUInteger)cols : 1024u; + const NSUInteger max_threads = g_concat_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_concat_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src0 offset:src0_offset atIndex:1]; + [enc setBuffer:src1 offset:src1_offset atIndex:2]; + [enc setBuffer:dst offset:dst_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_compressor_pack_ratio4_fusion_mode(uint32_t head_dim) { + const bool force = + getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_PACK_FUSION") != NULL; + const bool default_shape = head_dim == 128u || head_dim == 512u; + if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_PACK_FUSION") != NULL || + (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) { + return 0; + } + if (g_dsv4_compressor_pack_ratio4_pipeline == nil) { + return force ? -1 : 0; + } + return 1; +} + +static bool ds4_gpu_buffer_ranges_overlap( + id a, + NSUInteger a_offset, + uint64_t a_bytes, + id b, + NSUInteger b_offset, + uint64_t b_bytes) { + if (a != b || a_bytes == 0u || b_bytes == 0u) return false; + if (a_offset <= b_offset) { + return a_bytes > (uint64_t)(b_offset - a_offset); + } + return b_bytes > (uint64_t)(a_offset - b_offset); +} + +static int ds4_gpu_compressor_ratio4_direct_pool_mode( + uint32_t head_dim, + uint32_t n_comp) { + // One compressed row intentionally uses the legacy GGML reduction graph. + if (n_comp <= 1u) return 0; + + const bool force = + getenv("DS4_METAL_ENABLE_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL; + const bool default_shape = head_dim == 128u || head_dim == 512u; + if (getenv("DS4_METAL_DISABLE_M3_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL || + (!(ds4_gpu_device_name_contains("M3") && default_shape) && !force)) { + return 0; + } + if (g_dsv4_softmax_pool_ratio4_direct_pipeline == nil) { + return force ? -1 : 0; + } + return 1; +} + +static int ds4_gpu_encode_compressor_ratio4_direct_pool( + id cb, + ds4_gpu_tensor *out, + id kvbuf, + NSUInteger kv_offset, + id scorebuf, + NSUInteger score_offset, + id statekvbuf, + NSUInteger state_kv_offset, + id statescbuf, + NSUInteger state_score_offset, + uint32_t head_dim, + uint32_t n_comp, + bool replay) { + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t n = (uint64_t)head_dim * n_comp; + if (!cb || !outbuf || !kvbuf || !scorebuf || !statekvbuf || !statescbuf || + head_dim == 0u || n_comp <= 1u || n > UINT32_MAX || + ds4_gpu_tensor_bytes(out) < n * sizeof(float)) { + return 0; + } + + id pipeline = ds4_gpu_hot_pipeline( + g_dsv4_softmax_pool_ratio4_direct_pipeline, + "kernel_dsv4_softmax_pool_ratio4_direct"); + if (!pipeline) return 0; + + ds4_gpu_dsv4_softmax_pool_ratio4_direct_args args = { + .n_rows = 8, + .head_dim = head_dim, + .n_comp = n_comp, + .replay = replay ? 1u : 0u, + .pad = 0u, + }; + + NSUInteger nth = 256u; + if (nth > pipeline.maxTotalThreadsPerThreadgroup) { + nth = pipeline.maxTotalThreadsPerThreadgroup; + } + if (nth == 0u) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; + [enc setBuffer:scorebuf offset:score_offset atIndex:2]; + [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3]; + [enc setBuffer:statescbuf offset:state_score_offset atIndex:4]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + nth - 1u) / nth, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_compressor_pack_ratio4( + id cb, + id kvbuf, + NSUInteger kv_offset, + id scorebuf, + NSUInteger score_offset, + id statekvbuf, + NSUInteger state_kv_offset, + id statescbuf, + NSUInteger state_score_offset, + uint32_t head_dim, + uint32_t n_comp, + bool replay) { + if (!cb || !kvbuf || !scorebuf || !statekvbuf || !statescbuf || + !g_compressor_pool_kv_buffer || !g_compressor_pool_score_buffer || + head_dim == 0 || n_comp == 0 || head_dim > UINT32_MAX / 2u) { + return 0; + } + + const uint64_t total_elems64 = (uint64_t)n_comp * 8u * head_dim; + if (total_elems64 > UINT32_MAX || n_comp > UINT32_MAX / 8u) return 0; + id pipeline = ds4_gpu_hot_pipeline( + g_dsv4_compressor_pack_ratio4_pipeline, + "kernel_dsv4_compressor_pack_ratio4"); + if (!pipeline) return 0; + + NSUInteger nth = head_dim; + if (nth > 256u) nth = 256u; + if (nth > pipeline.maxTotalThreadsPerThreadgroup) { + nth = pipeline.maxTotalThreadsPerThreadgroup; + } + if (nth == 0) return 0; + + ds4_gpu_dsv4_compressor_pack_ratio4_args args = { + .head_dim = head_dim, + .n_comp = n_comp, + .replay = replay ? 1u : 0u, + .n_threads = (uint32_t)nth, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvbuf offset:kv_offset atIndex:1]; + [enc setBuffer:scorebuf offset:score_offset atIndex:2]; + [enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3]; + [enc setBuffer:statescbuf offset:state_score_offset atIndex:4]; + [enc setBuffer:g_compressor_pool_kv_buffer offset:0 atIndex:5]; + [enc setBuffer:g_compressor_pool_score_buffer offset:0 atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_comp, 8u, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_compressor_pool( + id cb, + ds4_gpu_tensor *out, + const ds4_gpu_tensor *state_kv, + const ds4_gpu_tensor *state_score, + uint32_t head_dim, + uint32_t ratio) { + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!cb || !out || !statekvbuf || !statescbuf || head_dim == 0 || ratio == 0) return 0; + + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t rows = coff * ratio; + const uint64_t state_bytes = (uint64_t)width * rows * sizeof(float); + if (ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + return 0; + } + + if (ratio != 4u) { + const uint64_t row_stride = (uint64_t)width * sizeof(float); + return ds4_gpu_encode_dsv4_softmax_pool(cb, + out, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + row_stride, + sizeof(float), + (uint64_t)rows * row_stride, + statescbuf, + ds4_gpu_tensor_offset(state_score), + row_stride, + sizeof(float), + (uint64_t)rows * row_stride, + ratio, + head_dim, + 1); + } + + const NSUInteger packed_bytes = (NSUInteger)8u * head_dim * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, + &g_compressor_pool_kv_bytes, + packed_bytes, + "ds4_compressor_pool_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, + &g_compressor_pool_score_bytes, + packed_bytes, + "ds4_compressor_pool_score")) { + return 0; + } + + const uint64_t state_row_stride = (uint64_t)width * sizeof(float); + const uint64_t pool_row_stride = (uint64_t)head_dim * sizeof(float); + const NSUInteger curr_offset = (NSUInteger)4u * state_row_stride + + (NSUInteger)head_dim * sizeof(float); + if (!ds4_gpu_encode_concat_f32_dim1(cb, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + 4, + state_row_stride, + statekvbuf, + ds4_gpu_tensor_offset(state_kv) + curr_offset, + 4, + state_row_stride, + g_compressor_pool_kv_buffer, + 0, + head_dim, + pool_row_stride) || + !ds4_gpu_encode_concat_f32_dim1(cb, + statescbuf, + ds4_gpu_tensor_offset(state_score), + 4, + state_row_stride, + statescbuf, + ds4_gpu_tensor_offset(state_score) + curr_offset, + 4, + state_row_stride, + g_compressor_pool_score_buffer, + 0, + head_dim, + pool_row_stride)) { + return 0; + } + + return ds4_gpu_encode_dsv4_softmax_pool(cb, + out, + g_compressor_pool_kv_buffer, + 0, + pool_row_stride, + sizeof(float), + packed_bytes, + g_compressor_pool_score_buffer, + 0, + pool_row_stride, + sizeof(float), + packed_bytes, + 8, + head_dim, + 1); +} + +static int ds4_gpu_encode_compressor_shift_ratio4( + id cb, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + uint32_t width) { + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!cb || !statekvbuf || !statescbuf || !g_dsv4_ratio4_shift_pipeline || width == 0) return 0; + + ds4_gpu_dsv4_ratio4_shift_args args = { .width = width }; + const uint32_t n = 4u * width; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_ratio4_shift_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:1]; + [enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +int ds4_gpu_compressor_prefill_tensor( + ds4_gpu_tensor *comp_cache, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos0, + uint32_t n_tokens, + uint32_t n_rot, + uint32_t n_ctx_orig, + bool quantize_fp8, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map || + head_dim == 0 || ratio == 0 || n_tokens == 0 || + n_rot > head_dim || (n_rot & 1u) != 0 || + (ape_type != 0u && ape_type != 1u) || + norm_type != 0u) { + return 0; + } + + @autoreleasepool { + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t state_rows = coff * ratio; + const uint32_t n_comp = n_tokens / ratio; + const uint32_t cutoff = n_comp * ratio; + const uint32_t rem = n_tokens - cutoff; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + norm_offset > model_size || norm_bytes > model_size - norm_offset) { + fprintf(stderr, "ds4: Metal compressor prefill tensor range is outside the mapped model\n"); + return 0; + } + + id kvbuf = ds4_gpu_tensor_buffer(kv); + id scbuf = ds4_gpu_tensor_buffer(sc); + id compbuf = ds4_gpu_tensor_buffer(comp_cache); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!kvbuf || !scbuf || !compbuf || !statekvbuf || !statescbuf || + ds4_gpu_tensor_bytes(kv) < kv_bytes || + ds4_gpu_tensor_bytes(sc) < kv_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes || + (n_comp && ds4_gpu_tensor_bytes(comp_cache) < comp_bytes)) { + fprintf(stderr, "ds4: Metal compressor prefill received undersized buffers\n"); + return 0; + } + + uint64_t ape_inner = 0; + id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!apebuf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + int ok = 1; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) ok = 0; + + if (ok) { + ok = ds4_gpu_encode_fill_f32_rows(cb, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + width, + state_rows, + 0.0f) && + ds4_gpu_encode_fill_f32_rows(cb, + statescbuf, + ds4_gpu_tensor_offset(state_score), + width, + state_rows, + ds4_gpu_negative_infinity()); + } + + if (ok && ratio == 4u) { + int32_t rows_prev[4] = { 0, 1, 2, 3 }; + const int have_prev = cutoff >= ratio ? 1 : 0; + const uint32_t prev_start = rem == 0 ? cutoff - ratio : cutoff - ratio; + if (have_prev) { + ok = ds4_gpu_encode_compressor_set_rows_projected(cb, + state_kv, + state_score, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)prev_start * width * sizeof(float), + scbuf, + ds4_gpu_tensor_offset(sc) + + (NSUInteger)prev_start * width * sizeof(float), + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0 + prev_start, + rows_prev, + 4, + state_rows); + } + if (ok && rem != 0) { + int32_t rows_cur[4]; + for (uint32_t i = 0; i < rem; i++) rows_cur[i] = (int32_t)(ratio + i); + ok = ds4_gpu_encode_compressor_set_rows_projected(cb, + state_kv, + state_score, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)cutoff * width * sizeof(float), + scbuf, + ds4_gpu_tensor_offset(sc) + + (NSUInteger)cutoff * width * sizeof(float), + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0 + cutoff, + rows_cur, + rem, + state_rows); + } + } else if (ok && rem != 0) { + int32_t rows[128]; + if (rem > (uint32_t)(sizeof(rows) / sizeof(rows[0]))) { + fprintf(stderr, "ds4: Metal compressor prefill remainder exceeds local row list\n"); + ok = 0; + } else { + for (uint32_t i = 0; i < rem; i++) rows[i] = (int32_t)i; + ok = ds4_gpu_encode_compressor_set_rows_projected(cb, + state_kv, + state_score, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)cutoff * width * sizeof(float), + scbuf, + ds4_gpu_tensor_offset(sc) + + (NSUInteger)cutoff * width * sizeof(float), + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0 + cutoff, + rows, + rem, + state_rows); + } + } + + if (ok && n_comp != 0) { + const NSUInteger score_bytes = (NSUInteger)cutoff * width * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, + &g_compressor_store_score_bytes, + score_bytes, + "ds4_compressor_store_score")) { + ok = 0; + } + if (ok) { + ok = ds4_gpu_encode_compressor_score_with_ape(cb, + scbuf, + ds4_gpu_tensor_offset(sc), + g_compressor_store_score_buffer, + 0, + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0, + cutoff); + } + + if (ok && ratio == 4u) { + const int direct_pool_mode = + ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp); + if (ok && direct_pool_mode < 0) ok = 0; + const uint64_t direct_output_bytes = + (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t direct_input_bytes = + (uint64_t)n_comp * 4u * width * sizeof(float); + const bool direct_pool_overlap = direct_pool_mode > 0 && + ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + kvbuf, + ds4_gpu_tensor_offset(kv), + direct_input_bytes); + const bool use_direct_pool = + direct_pool_mode > 0 && !direct_pool_overlap; + const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float); + if (ok && !use_direct_pool && + (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, + &g_compressor_pool_kv_bytes, + pack_bytes, + "ds4_compressor_pool_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, + &g_compressor_pool_score_bytes, + pack_bytes, + "ds4_compressor_pool_score"))) { + ok = 0; + } + const int pack_fusion_mode = use_direct_pool ? 0 : + ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim); + if (ok && pack_fusion_mode < 0) ok = 0; + const bool use_pack_fusion = pack_fusion_mode > 0; + if (ok && use_direct_pool) { + ok = ds4_gpu_encode_compressor_ratio4_direct_pool( + cb, + comp_cache, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + statescbuf, + ds4_gpu_tensor_offset(state_score), + head_dim, + n_comp, + false); + } + if (ok && use_pack_fusion) { + ok = ds4_gpu_encode_compressor_pack_ratio4( + cb, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + statescbuf, + ds4_gpu_tensor_offset(state_score), + head_dim, + n_comp, + false); + } + if (ok && !use_direct_pool && !use_pack_fusion) { + ok = ds4_gpu_encode_fill_f32_rows(cb, + g_compressor_pool_kv_buffer, + 0, + head_dim, + 8u * n_comp, + 0.0f) && + ds4_gpu_encode_fill_f32_rows(cb, + g_compressor_pool_score_buffer, + 0, + head_dim, + 8u * n_comp, + ds4_gpu_negative_infinity()); + } + if (ok && !use_direct_pool && !use_pack_fusion) { + const uint64_t src_row_stride = (uint64_t)width * sizeof(float); + const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; + const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); + const uint64_t dst_plane_stride = 8ull * dst_row_stride; + ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)head_dim * sizeof(float), + g_compressor_pool_kv_buffer, + (NSUInteger)4u * head_dim * sizeof(float), + head_dim, + ratio, + n_comp, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d(cb, + g_compressor_store_score_buffer, + (NSUInteger)head_dim * sizeof(float), + g_compressor_pool_score_buffer, + (NSUInteger)4u * head_dim * sizeof(float), + head_dim, + ratio, + n_comp, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride); + } + if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) { + const uint64_t src_row_stride = (uint64_t)width * sizeof(float); + const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; + const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); + const uint64_t dst_plane_stride = 8ull * dst_row_stride; + ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_pool_kv_buffer, + dst_plane_stride, + head_dim, + ratio, + n_comp - 1u, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d(cb, + g_compressor_store_score_buffer, + 0, + g_compressor_pool_score_buffer, + dst_plane_stride, + head_dim, + ratio, + n_comp - 1u, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride); + } + if (ok && !use_direct_pool) { + ok = ds4_gpu_encode_dsv4_softmax_pool(cb, + comp_cache, + g_compressor_pool_kv_buffer, + 0, + (uint64_t)head_dim * sizeof(float), + sizeof(float), + 8ull * head_dim * sizeof(float), + g_compressor_pool_score_buffer, + 0, + (uint64_t)head_dim * sizeof(float), + sizeof(float), + 8ull * head_dim * sizeof(float), + 8, + head_dim, + n_comp); + } + } else if (ok) { + const uint64_t row_stride = (uint64_t)width * sizeof(float); + ok = ds4_gpu_encode_dsv4_softmax_pool(cb, + comp_cache, + kvbuf, + ds4_gpu_tensor_offset(kv), + row_stride, + sizeof(float), + (uint64_t)ratio * row_stride, + g_compressor_store_score_buffer, + 0, + row_stride, + sizeof(float), + (uint64_t)ratio * row_stride, + ratio, + head_dim, + n_comp); + } + } + + if (ok && n_comp != 0) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, + comp_cache, + model_map, + model_size, + norm_offset, + head_dim, + n_comp, + rms_eps) != 0; + } + if (ok && n_comp != 0 && n_rot != 0) { + ds4_gpu_rope_tail_batch_args rope_args = ds4_gpu_make_rope_tail_args( + n_comp, 1, head_dim, n_rot, n_ctx_orig, false, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + cb = ds4_gpu_command_buffer(&owned); + ok = cb && !owned && + ds4_gpu_encode_rope_tail_inplace(cb, + compbuf, + ds4_gpu_tensor_offset(comp_cache), + &rope_args, + n_comp, + 1, + head_dim, + pos0, + ratio); + } + if (ok && n_comp != 0 && quantize_fp8) { + ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot) != 0; + } + + if (!had_batch) { + const int end_ok = ds4_gpu_end_commands(); + ok = end_ok && ok; + } + return ok ? 1 : 0; + } +} + +int ds4_gpu_compressor_prefill_ratio4_replay_tensor( + ds4_gpu_tensor *comp_cache, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv, + const ds4_gpu_tensor *sc, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t pos0, + uint32_t n_tokens, + uint32_t n_rot, + uint32_t n_ctx_orig, + bool quantize_fp8, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map || + head_dim == 0 || n_tokens == 0 || (n_tokens & 3u) != 0 || (pos0 & 3u) != 0 || + n_rot > head_dim || (n_rot & 1u) != 0 || + (ape_type != 0u && ape_type != 1u) || + norm_type != 0u) { + return 0; + } + + @autoreleasepool { + const uint32_t ratio = 4u; + const uint32_t width = 2u * head_dim; + const uint32_t state_rows = 8u; + const uint32_t n_comp = n_tokens / ratio; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + norm_offset > model_size || norm_bytes > model_size - norm_offset) { + fprintf(stderr, "ds4: Metal compressor replay tensor range is outside the mapped model\n"); + return 0; + } + + id kvbuf = ds4_gpu_tensor_buffer(kv); + id scbuf = ds4_gpu_tensor_buffer(sc); + id compbuf = ds4_gpu_tensor_buffer(comp_cache); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!kvbuf || !scbuf || !compbuf || !statekvbuf || !statescbuf || + ds4_gpu_tensor_bytes(kv) < kv_bytes || + ds4_gpu_tensor_bytes(sc) < kv_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes || + ds4_gpu_tensor_bytes(comp_cache) < comp_bytes) { + fprintf(stderr, "ds4: Metal compressor replay received undersized buffers\n"); + return 0; + } + + uint64_t ape_inner = 0; + id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!apebuf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + int ok = 1; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) ok = 0; + + const NSUInteger score_bytes = (NSUInteger)n_tokens * width * sizeof(float); + const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float); + const int direct_pool_mode = + ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp); + if (ok && direct_pool_mode < 0) ok = 0; + const uint64_t direct_output_bytes = + (uint64_t)n_comp * head_dim * sizeof(float); + const uint64_t direct_input_bytes = + (uint64_t)n_tokens * width * sizeof(float); + const uint64_t direct_state_bytes = + (uint64_t)4u * width * sizeof(float); + const bool direct_pool_overlap = direct_pool_mode > 0 && + (ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + kvbuf, + ds4_gpu_tensor_offset(kv), + direct_input_bytes) || + ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + direct_state_bytes) || + ds4_gpu_buffer_ranges_overlap( + compbuf, + ds4_gpu_tensor_offset(comp_cache), + direct_output_bytes, + statescbuf, + ds4_gpu_tensor_offset(state_score), + direct_state_bytes)); + const bool use_direct_pool = + direct_pool_mode > 0 && !direct_pool_overlap; + if (ok && (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer, + &g_compressor_store_score_bytes, + score_bytes, + "ds4_compressor_store_score") || + (!use_direct_pool && + (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer, + &g_compressor_pool_kv_bytes, + pack_bytes, + "ds4_compressor_pool_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer, + &g_compressor_pool_score_bytes, + pack_bytes, + "ds4_compressor_pool_score"))))) { + ok = 0; + } + + if (ok) { + ok = ds4_gpu_encode_compressor_score_with_ape(cb, + scbuf, + ds4_gpu_tensor_offset(sc), + g_compressor_store_score_buffer, + 0, + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0, + n_tokens); + } + + const uint64_t src_row_stride = (uint64_t)width * sizeof(float); + const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride; + const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float); + const uint64_t dst_plane_stride = 8ull * dst_row_stride; + const NSUInteger state_off = ds4_gpu_tensor_offset(state_kv); + const NSUInteger state_score_off = ds4_gpu_tensor_offset(state_score); + const int pack_fusion_mode = use_direct_pool ? 0 : + ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim); + if (ok && pack_fusion_mode < 0) ok = 0; + const bool use_pack_fusion = pack_fusion_mode > 0; + + if (ok && use_direct_pool) { + ok = ds4_gpu_encode_compressor_ratio4_direct_pool( + cb, + comp_cache, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + state_off, + statescbuf, + state_score_off, + head_dim, + n_comp, + true); + } + + if (ok && use_pack_fusion) { + ok = ds4_gpu_encode_compressor_pack_ratio4( + cb, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_store_score_buffer, + 0, + statekvbuf, + state_off, + statescbuf, + state_score_off, + head_dim, + n_comp, + true); + } + + if (ok && !use_direct_pool && !use_pack_fusion) { + ok = ds4_gpu_encode_fill_f32_rows(cb, + g_compressor_pool_kv_buffer, + 0, + head_dim, + 8u * n_comp, + 0.0f) && + ds4_gpu_encode_fill_f32_rows(cb, + g_compressor_pool_score_buffer, + 0, + head_dim, + 8u * n_comp, + ds4_gpu_negative_infinity()); + } + + if (ok && !use_direct_pool && !use_pack_fusion) { + /* + * The aligned nonzero ratio-4 path replays the current ubatch + * compressor, but seeds the first compressed row with the previous + * compressor state. Rows 0..3 are the previous half, rows 4..7 are + * the current half. + */ + ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, + statekvbuf, + state_off, + g_compressor_pool_kv_buffer, + 0, + head_dim, + ratio, + 1, + src_row_stride, + (uint64_t)ratio * src_row_stride, + dst_row_stride, + dst_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d(cb, + statescbuf, + state_score_off, + g_compressor_pool_score_buffer, + 0, + head_dim, + ratio, + 1, + src_row_stride, + (uint64_t)ratio * src_row_stride, + dst_row_stride, + dst_plane_stride); + } + if (ok && !use_direct_pool && !use_pack_fusion) { + ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)head_dim * sizeof(float), + g_compressor_pool_kv_buffer, + (NSUInteger)4u * head_dim * sizeof(float), + head_dim, + ratio, + n_comp, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d(cb, + g_compressor_store_score_buffer, + (NSUInteger)head_dim * sizeof(float), + g_compressor_pool_score_buffer, + (NSUInteger)4u * head_dim * sizeof(float), + head_dim, + ratio, + n_comp, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride); + } + if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) { + ok = ds4_gpu_encode_cpy_f32_f32_3d(cb, + kvbuf, + ds4_gpu_tensor_offset(kv), + g_compressor_pool_kv_buffer, + dst_plane_stride, + head_dim, + ratio, + n_comp - 1u, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride) && + ds4_gpu_encode_cpy_f32_f32_3d(cb, + g_compressor_store_score_buffer, + 0, + g_compressor_pool_score_buffer, + dst_plane_stride, + head_dim, + ratio, + n_comp - 1u, + src_row_stride, + src_plane_stride, + dst_row_stride, + dst_plane_stride); + } + if (ok && !use_direct_pool) { + ok = ds4_gpu_encode_dsv4_softmax_pool(cb, + comp_cache, + g_compressor_pool_kv_buffer, + 0, + dst_row_stride, + sizeof(float), + dst_plane_stride, + g_compressor_pool_score_buffer, + 0, + dst_row_stride, + sizeof(float), + dst_plane_stride, + 8, + head_dim, + n_comp); + } + if (ok) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, + comp_cache, + model_map, + model_size, + norm_offset, + head_dim, + n_comp, + rms_eps) != 0; + } + if (ok && n_rot != 0) { + ds4_gpu_rope_tail_batch_args rope_args = ds4_gpu_make_rope_tail_args( + n_comp, 1, head_dim, n_rot, n_ctx_orig, false, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + cb = ds4_gpu_command_buffer(&owned); + ok = cb && !owned && + ds4_gpu_encode_rope_tail_inplace(cb, + compbuf, + ds4_gpu_tensor_offset(comp_cache), + &rope_args, + n_comp, + 1, + head_dim, + pos0, + ratio); + } + if (ok && quantize_fp8) { + ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot) != 0; + } + + if (ok) { + ok = ds4_gpu_encode_fill_f32_rows(cb, + statekvbuf, + state_off, + width, + state_rows, + 0.0f) && + ds4_gpu_encode_fill_f32_rows(cb, + statescbuf, + state_score_off, + width, + state_rows, + ds4_gpu_negative_infinity()); + } + if (ok) { + int32_t rows_prev[4] = { 0, 1, 2, 3 }; + const uint32_t prev_start = n_tokens - ratio; + ok = ds4_gpu_encode_compressor_set_rows_projected(cb, + state_kv, + state_score, + kvbuf, + ds4_gpu_tensor_offset(kv) + + (NSUInteger)prev_start * width * sizeof(float), + scbuf, + ds4_gpu_tensor_offset(sc) + + (NSUInteger)prev_start * width * sizeof(float), + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0 + prev_start, + rows_prev, + ratio, + state_rows); + } + + if (!had_batch) { + const int end_ok = ds4_gpu_end_commands(); + ok = end_ok && ok; + } + return ok ? 1 : 0; + } +} + +int ds4_gpu_compressor_prefill_state_ratio4_tensor( + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + const ds4_gpu_tensor *kv_tail, + const ds4_gpu_tensor *sc_tail, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint32_t head_dim, + uint32_t pos0) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!state_kv || !state_score || !kv_tail || !sc_tail || !model_map || + head_dim == 0 || (ape_type != 0u && ape_type != 1u)) { + return 0; + } + + @autoreleasepool { + const uint32_t ratio = 4u; + const uint32_t width = 2u * head_dim; + const uint32_t state_rows = 8u; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t tail_bytes = (uint64_t)ratio * width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t ape_bytes = (uint64_t)ratio * width * elem_ape; + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset) { + fprintf(stderr, "ds4: Metal compressor prefill-state APE range is outside the mapped model\n"); + return 0; + } + + id kvbuf = ds4_gpu_tensor_buffer(kv_tail); + id scbuf = ds4_gpu_tensor_buffer(sc_tail); + id statekvbuf = ds4_gpu_tensor_buffer(state_kv); + id statescbuf = ds4_gpu_tensor_buffer(state_score); + if (!kvbuf || !scbuf || !statekvbuf || !statescbuf || + ds4_gpu_tensor_bytes(kv_tail) < tail_bytes || + ds4_gpu_tensor_bytes(sc_tail) < tail_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes) { + fprintf(stderr, "ds4: Metal compressor prefill-state received undersized buffers\n"); + return 0; + } + + uint64_t ape_inner = 0; + id apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner); + if (!apebuf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + int ok = 1; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) ok = 0; + + if (ok) { + ok = ds4_gpu_encode_fill_f32_rows(cb, + statekvbuf, + ds4_gpu_tensor_offset(state_kv), + width, + state_rows, + 0.0f) && + ds4_gpu_encode_fill_f32_rows(cb, + statescbuf, + ds4_gpu_tensor_offset(state_score), + width, + state_rows, + ds4_gpu_negative_infinity()); + } + if (ok) { + int32_t rows[4] = { 0, 1, 2, 3 }; + ok = ds4_gpu_encode_compressor_set_rows_projected(cb, + state_kv, + state_score, + kvbuf, + ds4_gpu_tensor_offset(kv_tail), + scbuf, + ds4_gpu_tensor_offset(sc_tail), + apebuf, + (NSUInteger)ape_inner, + ape_type, + width, + ratio, + pos0, + rows, + ratio, + state_rows); + } + + if (!had_batch) { + const int end_ok = ds4_gpu_end_commands(); + ok = end_ok && ok; + } + return ok ? 1 : 0; + } +} + +int ds4_gpu_compressor_update_tensor( + const ds4_gpu_tensor *kv_cur, + const ds4_gpu_tensor *sc_cur, + ds4_gpu_tensor *state_kv, + ds4_gpu_tensor *state_score, + ds4_gpu_tensor *comp_cache, + const void *model_map, + uint64_t model_size, + uint64_t ape_offset, + uint32_t ape_type, + uint64_t norm_offset, + uint32_t norm_type, + uint32_t head_dim, + uint32_t ratio, + uint32_t pos, + uint32_t comp_row, + uint32_t n_rot, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + float rms_eps, + bool state_already_stored) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache || + !model_map || head_dim == 0 || ratio == 0 || + n_rot > head_dim || (n_rot & 1u) != 0 || + (ape_type != 0u && ape_type != 1u) || + norm_type != 0u) { + return 0; + } + + @autoreleasepool { + const uint32_t coff = ratio == 4u ? 2u : 1u; + const uint32_t width = coff * head_dim; + const uint32_t state_rows = coff * ratio; + const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u; + const uint64_t elem_ape = ape_type == 1u ? 2u : 4u; + const uint64_t kv_bytes = (uint64_t)width * sizeof(float); + const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float); + const uint64_t comp_bytes = (uint64_t)(comp_row + (emit ? 1u : 0u)) * head_dim * sizeof(float); + const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape; + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + + if (ape_offset > model_size || ape_bytes > model_size - ape_offset || + norm_offset > model_size || norm_bytes > model_size - norm_offset) { + fprintf(stderr, "ds4: Metal compressor tensor range is outside the mapped model\n"); + return 0; + } + + id kvbuf = ds4_gpu_tensor_buffer(kv_cur); + id scbuf = ds4_gpu_tensor_buffer(sc_cur); + id compbuf = ds4_gpu_tensor_buffer(comp_cache); + if (!kvbuf || !scbuf || !compbuf || + ds4_gpu_tensor_bytes(kv_cur) < kv_bytes || + ds4_gpu_tensor_bytes(sc_cur) < kv_bytes || + ds4_gpu_tensor_bytes(state_kv) < state_bytes || + ds4_gpu_tensor_bytes(state_score) < state_bytes || + (emit && ds4_gpu_tensor_bytes(comp_cache) < comp_bytes)) { + fprintf(stderr, "ds4: Metal compressor update received undersized buffers\n"); + return 0; + } + + if (!state_already_stored) { + const bool use_store_one = + getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") == NULL; + const int store_ok = use_store_one + ? ds4_gpu_compressor_store_one_tensor(kv_cur, + sc_cur, + state_kv, + state_score, + model_map, + model_size, + ape_offset, + ape_type, + width, + ratio, + pos) + : ds4_gpu_compressor_store_batch_tensor(kv_cur, + sc_cur, + state_kv, + state_score, + model_map, + model_size, + ape_offset, + ape_type, + head_dim, + ratio, + pos, + 1); + if (!store_ok) { + return 0; + } + } + if (!emit) return 1; + + ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view( + comp_cache, + (uint64_t)comp_row * head_dim * sizeof(float), + (uint64_t)head_dim * sizeof(float)); + if (!comp_row_view) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + int ok = cb && + ds4_gpu_encode_compressor_pool(cb, + comp_row_view, + state_kv, + state_score, + head_dim, + ratio); + if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor DS4 softmax pool"); + if (ok) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view, + comp_row_view, + model_map, + model_size, + norm_offset, + head_dim, + 1, + rms_eps) != 0; + } + if (ok) { + const uint32_t comp_pos = pos + 1u - ratio; + ok = ds4_gpu_rope_tail_tensor(comp_row_view, + 1, + 1, + head_dim, + n_rot, + comp_pos, + n_ctx_orig, + false, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow) != 0; + } + if (ok && ratio == 4u) { + cb = ds4_gpu_command_buffer(&owned); + ok = cb && + ds4_gpu_encode_compressor_shift_ratio4(cb, + state_kv, + state_score, + width); + if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor ratio4 state shift"); + } + ds4_gpu_tensor_free(comp_row_view); + if (!ok) return 0; + } + + return 1; +} + +static int ds4_gpu_encode_fill_f32_rows( + id cb, + id buf, + NSUInteger offset, + uint32_t width, + uint32_t rows, + float value) { + if (!cb || !buf || width == 0 || rows == 0 || (width & 3u) != 0) return 0; + + ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(width, rows, 1, 0.0f, 0.0f); + args.val = value; + + NSUInteger nth_max = g_unary_fill_pipeline.maxTotalThreadsPerThreadgroup; + if (nth_max > 256u) nth_max = 256u; + NSUInteger nth = (NSUInteger)args.ne00; + if (nth > nth_max) nth = nth_max; + if (nth == 0) nth = 1u; + const NSUInteger nk0 = ((NSUInteger)args.ne00 + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_unary_fill_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:buf offset:offset atIndex:1]; + [enc setBuffer:buf offset:offset atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nk0 * (NSUInteger)args.ne01, + (NSUInteger)args.ne02, + (NSUInteger)args.ne03) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +int ds4_gpu_attention_output_q8_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, + ds4_gpu_tensor *low_tmp, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !low || !group_tmp || !low_tmp || !heads || !model_map || + group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 || + group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + const uint64_t low_dim = (uint64_t)n_groups * rank; + if ((group_dim % 32u) != 0 || (low_dim % 32u) != 0 || low_dim > UINT32_MAX) { + fprintf(stderr, "ds4: Metal attention output batch received invalid q8 dimensions\n"); + return 0; + } + const uint64_t row_a_bytes = (group_dim / 32u) * 34u; + const uint64_t row_b_bytes = (low_dim / 32u) * 34u; + const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; + const uint64_t out_b_bytes = out_dim * row_b_bytes; + if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset || + out_b_offset > model_size || out_b_bytes > model_size - out_b_offset) { + fprintf(stderr, "ds4: Metal attention output batch weights are outside the mapped model\n"); + return 0; + } + + const uint64_t heads_bytes = (uint64_t)n_tokens * n_groups * group_dim * sizeof(float); + const uint64_t low_bytes = (uint64_t)n_tokens * low_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(low) < low_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal attention output batch received undersized buffers\n"); + return 0; + } + (void)group_tmp; + (void)low_tmp; + + const bool use_direct_low = + n_tokens < 32u && getenv("DS4_METAL_DISABLE_ATTN_OUT_LOW_DIRECT") == NULL; + /* The exported TensorOps attention-output kernel is a 64-token tile. + * Keep this on full tiles only; smaller multiples of 32 use the legacy + * path instead of relying on cooperative tensor partial RHS bounds. */ + const bool use_mpp_low = + n_tokens >= 32u && + (n_tokens % DS4_METAL_ATTN_OUT_MPP_TILE_N) == 0 && + ds4_gpu_use_mpp_attn_out_low_matmul(); + const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t); + id group_ids_buffer = nil; + if (!use_direct_low && !use_mpp_low) { + if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { + group_ids_buffer = + ds4_gpu_new_transient_buffer(ids_bytes, "attention output group ids"); + if (!group_ids_buffer) { + return 0; + } + } else { + if (!ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, + &g_attn_out_group_ids_bytes, + ids_bytes, + "ds4_attention_output_group_ids")) { + return 0; + } + group_ids_buffer = g_attn_out_group_ids_buffer; + } + int32_t *ids = (int32_t *)[group_ids_buffer contents]; + for (uint32_t t = 0; t < n_tokens; t++) { + for (uint32_t group = 0; group < n_groups; group++) { + ids[(uint64_t)t * n_groups + group] = (int32_t)group; + } + } + } + + uint64_t out_a_inner = 0; + id out_a_buf = + ds4_gpu_wrap_model_range(model_map, model_size, + out_a_offset, out_a_bytes, + &out_a_inner); + if (!out_a_buf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + bool ok = true; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) { + ok = false; + } + const bool attn_out_profile = + getenv("DS4_METAL_ATTN_OUT_STAGE_PROFILE") != NULL && g_batch_cb != nil; + if (ok && attn_out_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + ok = false; + } else { + cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) ok = false; + } + } + double attn_out_t0 = attn_out_profile ? ds4_gpu_now_ms() : 0.0; +#define DS4_METAL_PROFILE_ATTN_OUT_STAGE(name) do { \ + if (ok && attn_out_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = false; \ + } else { \ + const double now_ms = ds4_gpu_now_ms(); \ + fprintf(stderr, \ + "ds4: Metal attention output stage tokens=%u %s=%.3f ms\n", \ + n_tokens, (name), now_ms - attn_out_t0); \ + attn_out_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = false; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb || owned) ok = false; \ + } \ + } \ + } \ + } while (0) + + if (ok) { + /* + * Batched attention-output projections switch from the vector + * kernel to the SIMD matrix kernel once the batch has at least 32 + * tokens. This preserves the single-token generation path while + * keeping prefill accumulation stable. + */ + if (use_mpp_low) { + ds4_gpu_mul_mm_id_args mm_args = + ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim, + (uint32_t)rank, + n_groups, + row_a_bytes, + (uint64_t)rank * row_a_bytes, + n_groups, + n_groups, + n_tokens); + /* + * Direct RHS lets MPP read the dense low-rank activation tile + * directly from device memory instead of staging a second + * threadgroup tile. The retained attention-output path is the + * 64-token direct-RHS kernel; the older staged-RHS and 32-token + * variants were not kept as alternate runtime modes. + */ + const char *attn_out_pipeline_name = + "kernel_attn_out_low_q8_0_mpp_direct_rhs_n64"; + id mm_pipeline = + ds4_gpu_get_mul_mm_id_pipeline(attn_out_pipeline_name, false); + ok = ds4_gpu_encode_attn_out_low_q8_mpp(cb, + mm_pipeline, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low)) != 0; + if (!ok) { + ds4_gpu_warn_mpp_fallback(); + if (ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) { + if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { + group_ids_buffer = + ds4_gpu_new_transient_buffer(ids_bytes, "attention output group ids"); + } else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, + &g_attn_out_group_ids_bytes, + ids_bytes, + "ds4_attention_output_group_ids")) { + group_ids_buffer = g_attn_out_group_ids_buffer; + } + if (group_ids_buffer) { + int32_t *ids = (int32_t *)[group_ids_buffer contents]; + for (uint32_t t = 0; t < n_tokens; t++) { + for (uint32_t group = 0; group < n_groups; group++) { + ids[(uint64_t)t * n_groups + group] = (int32_t)group; + } + } + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, + n_groups, + n_groups, + n_groups, + n_tokens); + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); + id fallback_pipeline = + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false); + ok = ds4_gpu_encode_mul_mm_id(cb, + map_pipeline, + fallback_pipeline, + &map_args, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + group_ids_buffer, + 0) != 0; + } + } + } + } else if (n_tokens >= 32u && ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) { + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, + n_groups, + n_groups, + n_groups, + n_tokens); + ds4_gpu_mul_mm_id_args mm_args = + ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim, + (uint32_t)rank, + n_groups, + row_a_bytes, + (uint64_t)rank * row_a_bytes, + n_groups, + n_groups, + n_tokens); + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); + id mm_pipeline = + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false); + ok = ds4_gpu_encode_mul_mm_id(cb, + map_pipeline, + mm_pipeline, + &map_args, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + group_ids_buffer, + 0) != 0; + } else if (use_direct_low) { + ds4_gpu_mul_mv_id_args args = { + .nei0 = (int32_t)n_groups, + .nei1 = (int32_t)n_tokens, + .nbi1 = 0, + .ne00 = (int32_t)group_dim, + .ne01 = (int32_t)rank, + .ne02 = (int32_t)n_groups, + .nb00 = 34, + .nb01 = row_a_bytes, + .nb02 = (uint64_t)rank * row_a_bytes, + .ne10 = (int32_t)group_dim, + .ne11 = (int32_t)n_groups, + .ne12 = (int32_t)n_tokens, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)group_dim * sizeof(float), + .nb12 = (uint64_t)n_groups * group_dim * sizeof(float), + .ne0 = (int32_t)rank, + .ne1 = (int32_t)n_groups, + .nb1 = (uint64_t)rank * sizeof(float), + .nr0 = 2, + }; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4); + ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, + pipeline, + &args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + 32u * 2u * sizeof(float), + 4, + true) != 0; + } else { + ds4_gpu_mul_mv_id_args args = { + .nei0 = (int32_t)n_groups, + .nei1 = (int32_t)n_tokens, + .nbi1 = (uint64_t)n_groups * sizeof(int32_t), + .ne00 = (int32_t)group_dim, + .ne01 = (int32_t)rank, + .ne02 = (int32_t)n_groups, + .nb00 = 34, + .nb01 = row_a_bytes, + .nb02 = (uint64_t)rank * row_a_bytes, + .ne10 = (int32_t)group_dim, + .ne11 = (int32_t)n_groups, + .ne12 = (int32_t)n_tokens, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)group_dim * sizeof(float), + .nb12 = (uint64_t)n_groups * group_dim * sizeof(float), + .ne0 = (int32_t)rank, + .ne1 = (int32_t)n_groups, + .nb1 = (uint64_t)rank * sizeof(float), + .nr0 = 2, + }; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_0_f32", 4); + ok = ds4_gpu_encode_mul_mv_id(cb, + pipeline, + &args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + group_ids_buffer, + 0, + 32u * 2u * sizeof(float), + 4, + true) != 0; + } + } + DS4_METAL_PROFILE_ATTN_OUT_STAGE("low_proj"); + + if (ok) { + ok = ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size, + out_b_offset, + low_dim, out_dim, low, n_tokens) != 0; + } + DS4_METAL_PROFILE_ATTN_OUT_STAGE("out_proj"); + + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } +#undef DS4_METAL_PROFILE_ATTN_OUT_STAGE + return ok ? 1 : 0; + } +} + +int ds4_gpu_attention_output_q4_K_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + ds4_gpu_tensor *group_tmp, + ds4_gpu_tensor *low_tmp, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint32_t out_b_type, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !low || !heads || !model_map || + group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 || + group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + if (n_tokens < 32u) return 0; + + @autoreleasepool { + const uint64_t low_dim = (uint64_t)n_groups * rank; + if ((group_dim % 256u) != 0 || (low_dim % 256u) != 0 || low_dim > UINT32_MAX) { + return 0; + } + + uint64_t row_a_bytes = 0; + uint64_t row_b_bytes = 0; + if (!ds4_gpu_quant_row_bytes(DS4_METAL_TENSOR_Q4_K, (uint32_t)group_dim, &row_a_bytes) || + !ds4_gpu_quant_row_bytes(out_b_type, (uint32_t)low_dim, &row_b_bytes)) { + return 0; + } + + const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; + const uint64_t out_b_bytes = out_dim * row_b_bytes; + if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset || + out_b_offset > model_size || out_b_bytes > model_size - out_b_offset) { + fprintf(stderr, "ds4: Metal Q4 attention output batch weights are outside the mapped model\n"); + return 0; + } + + const uint64_t heads_bytes = (uint64_t)n_tokens * n_groups * group_dim * sizeof(float); + const uint64_t low_bytes = (uint64_t)n_tokens * low_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(low) < low_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal Q4 attention output batch received undersized buffers\n"); + return 0; + } + (void)group_tmp; + (void)low_tmp; + + const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t); + id group_ids_buffer = nil; + if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) { + group_ids_buffer = + ds4_gpu_new_transient_buffer(ids_bytes, "attention output Q4 group ids"); + } else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer, + &g_attn_out_group_ids_bytes, + ids_bytes, + "ds4_attention_output_group_ids")) { + group_ids_buffer = g_attn_out_group_ids_buffer; + } + if (!group_ids_buffer) return 0; + + int32_t *ids = (int32_t *)[group_ids_buffer contents]; + for (uint32_t t = 0; t < n_tokens; t++) { + for (uint32_t group = 0; group < n_groups; group++) { + ids[(uint64_t)t * n_groups + group] = (int32_t)group; + } + } + + uint64_t out_a_inner = 0; + id out_a_buf = + ds4_gpu_wrap_model_range(model_map, model_size, + out_a_offset, out_a_bytes, + &out_a_inner); + if (!out_a_buf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + bool ok = true; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) ok = false; + + if (ok) { + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim, + n_groups, + n_groups, + n_groups, + n_tokens); + ds4_gpu_mul_mm_id_args mm_args = + ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim, + (uint32_t)rank, + n_groups, + row_a_bytes, + (uint64_t)rank * row_a_bytes, + n_groups, + n_groups, + n_tokens); + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups)); + id mm_pipeline = + ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false); + ok = ds4_gpu_encode_mul_mm_id(cb, + map_pipeline, + mm_pipeline, + &map_args, + &mm_args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + group_ids_buffer, + 0) != 0; + } + + if (ok) { + ok = ds4_gpu_matmul_quant_tensor(out, + model_map, + model_size, + out_b_offset, + out_b_type, + low_dim, + out_dim, + low, + n_tokens) != 0; + } + + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + return ok ? 1 : 0; + } +} + +int ds4_gpu_attention_output_q8_batch_f16_tensor( + ds4_gpu_tensor *out_h, + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + uint64_t out_dim, + const ds4_gpu_tensor *heads, + uint32_t n_tokens) { + (void)out_h; (void)low; (void)model_map; (void)model_size; + (void)out_a_offset; (void)out_b_offset; (void)group_dim; (void)rank; + (void)n_groups; (void)out_dim; (void)heads; (void)n_tokens; + return 0; +} + +int ds4_gpu_matmul_q8_0_kslice_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t full_in_dim, + uint64_t k_off, + uint64_t k_cnt, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t x_elem_off) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if ((full_in_dim & 31u) != 0 || (k_off & 31u) != 0 || (k_cnt & 31u) != 0 || + k_cnt == 0 || k_off + k_cnt > full_in_dim || + full_in_dim > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < (x_elem_off + k_cnt) * sizeof(float) || + ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) { + fprintf(stderr, "ds4: Metal Q8_0 kslice matmul received undersized buffers\n"); + return 0; + } + const uint64_t row_bytes = (full_in_dim / 32u) * 34u; + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8_0 kslice weights are outside the mapped model\n"); + return 0; + } + uint64_t inner = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, weight_bytes, &inner); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + /* Same matvec kernel as the full projection: ne00 bounds the k loop + * while nb01/nb02 keep the full-row stride, so each row reads only + * the owned k window. */ + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(full_in_dim, out_dim); + mv_args.ne00 = (int32_t)k_cnt; + mv_args.ne10 = (int32_t)k_cnt; + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + if (out_dim > 65536u) mv_dispatch.nsg = 8; + mv_args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg); + if (!pipeline) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)(inner + (k_off / 32u) * 34u) atIndex:1]; + [enc setBuffer:xbuf + offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + x_elem_off * sizeof(float)) + atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 kslice matvec")) { + return 0; + } + return 1; + } +} + +int ds4_gpu_matmul_quant_kslice_tensor( + ds4_gpu_tensor *out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint64_t full_in_dim, + uint64_t k_off, + uint64_t k_cnt, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint64_t x_elem_off) { + if (weight_type == DS4_METAL_TENSOR_Q8_0) { + return ds4_gpu_matmul_q8_0_kslice_tensor(out, + model_map, + model_size, + weight_offset, + full_in_dim, + k_off, + k_cnt, + out_dim, + x, + x_elem_off); + } + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !x || !model_map || + full_in_dim == 0 || k_cnt == 0 || out_dim == 0 || + k_off + k_cnt > full_in_dim || + full_in_dim > UINT32_MAX || k_cnt > UINT32_MAX || + out_dim > UINT32_MAX) { + return 0; + } + + uint64_t block_elems = 0; + uint64_t block_bytes = 0; + if (weight_type == DS4_METAL_TENSOR_Q4_K) { + block_elems = 256u; + block_bytes = 144u; + } else if (weight_type == DS4_METAL_TENSOR_Q4_0) { + block_elems = 32u; + block_bytes = 18u; + } else { + fprintf(stderr, "ds4: Metal quant kslice received unsupported type %u\n", weight_type); + return 0; + } + if ((full_in_dim % block_elems) != 0 || + (k_off % block_elems) != 0 || + (k_cnt % block_elems) != 0) { + fprintf(stderr, "ds4: Metal quant kslice dimensions are not block aligned\n"); + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id outbuf = ds4_gpu_tensor_buffer(out); + if (!xbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < (x_elem_off + k_cnt) * sizeof(float) || + ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) { + fprintf(stderr, "ds4: Metal quant kslice matmul received undersized buffers\n"); + return 0; + } + + const uint64_t row_bytes = (full_in_dim / block_elems) * block_bytes; + const uint64_t weight_bytes = out_dim * row_bytes; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal quant kslice weights are outside the mapped model\n"); + return 0; + } + uint64_t inner = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, weight_bytes, &inner); + if (!wbuf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (weight_type == DS4_METAL_TENSOR_Q4_K && + (k_cnt % 256u) == 0 && + getenv("DS4_METAL_DISABLE_Q4_MV_CLASSIC") == NULL) { + const int16_t nsg = 2; + id pipeline = + ds4_gpu_get_mul_mv_ext_pipeline("kernel_mul_mv_q4_K_dense_f32", nsg, 8); + if (pipeline) { + ds4_gpu_q8_0_matvec_args args = { + .ne00 = (int32_t)k_cnt, + .ne01 = (int32_t)out_dim, + .ne02 = 1, + .nb00 = 1, + .nb01 = row_bytes, + .nb02 = row_bytes * out_dim, + .nb03 = row_bytes * out_dim, + .ne10 = (int32_t)k_cnt, + .ne11 = 1, + .ne12 = 1, + .nb10 = sizeof(float), + .nb11 = k_cnt * sizeof(float), + .nb12 = k_cnt * sizeof(float), + .nb13 = k_cnt * sizeof(float), + .ne0 = (int32_t)out_dim, + .ne1 = 1, + .nr0 = 2, + .r2 = 1, + .r3 = 1, + }; + const uint64_t rows_ptg = (uint64_t)nsg * 2u; + const uint64_t w_skip = (k_off / block_elems) * block_bytes; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)(inner + w_skip) atIndex:1]; + [enc setBuffer:xbuf + offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + x_elem_off * sizeof(float)) + atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:32 atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + rows_ptg - 1u) / rows_ptg, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4_K kslice matvec")) return 0; + return 1; + } + } + + fprintf(stderr, "ds4: Metal quant kslice has no kernel for type %u\n", + weight_type); + return 0; + } +} + +int ds4_gpu_attention_output_q8_tp_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t out_b_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups_total, + uint32_t group0, + uint32_t group_cnt, + uint64_t out_dim, + const ds4_gpu_tensor *heads) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !low || !heads || !model_map || + group_dim == 0 || rank == 0 || group_cnt == 0 || + group0 + group_cnt > n_groups_total || + (group_dim % 32u) != 0 || ((rank * group_cnt) % 32u) != 0 || + group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + const uint64_t low_dim_total = (uint64_t)n_groups_total * rank; + const uint64_t row_a_bytes = (group_dim / 32u) * 34u; + const uint64_t a_group_bytes = rank * row_a_bytes; + const uint64_t out_a_bytes = (uint64_t)n_groups_total * a_group_bytes; + if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset) { + fprintf(stderr, "ds4: Metal TP attention output weights are outside the mapped model\n"); + return 0; + } + /* The heads buffer holds only the owned groups, compact at its + * base (the head slice keeps q/attention output halves packed). */ + if (ds4_gpu_tensor_bytes(heads) < (uint64_t)group_cnt * group_dim * sizeof(float) || + ds4_gpu_tensor_bytes(low) < (uint64_t)group_cnt * rank * sizeof(float) || + ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) { + fprintf(stderr, "ds4: Metal TP attention output received undersized buffers\n"); + return 0; + } + + uint64_t out_a_inner = 0; + id out_a_buf = + ds4_gpu_wrap_model_range(model_map, model_size, + out_a_offset, out_a_bytes, &out_a_inner); + if (!out_a_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) return 0; + + /* Low projection for the owned groups only: identical dispatch to + * the single-node direct path with the weight base, heads input and + * group count shifted to the slice. The owned low half lands + * compactly at low[0 .. group_cnt*rank). */ + ds4_gpu_mul_mv_id_args args = { + .nei0 = (int32_t)group_cnt, + .nei1 = 1, + .nbi1 = 0, + .ne00 = (int32_t)group_dim, + .ne01 = (int32_t)rank, + .ne02 = (int32_t)group_cnt, + .nb00 = 34, + .nb01 = row_a_bytes, + .nb02 = a_group_bytes, + .ne10 = (int32_t)group_dim, + .ne11 = (int32_t)group_cnt, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)group_dim * sizeof(float), + .nb12 = (uint64_t)group_cnt * group_dim * sizeof(float), + .ne0 = (int32_t)rank, + .ne1 = (int32_t)group_cnt, + .nb1 = (uint64_t)rank * sizeof(float), + .nr0 = 2, + }; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4); + int ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, + pipeline, + &args, + out_a_buf, + (NSUInteger)(out_a_inner + (uint64_t)group0 * a_group_bytes), + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + 32u * 2u * sizeof(float), + 4, + true); + if (!ok) return 0; + + /* Expand projection over the owned k window only; the result is this + * rank's partial attention block output. */ + return ds4_gpu_matmul_q8_0_kslice_tensor(out, model_map, model_size, + out_b_offset, + low_dim_total, + (uint64_t)group0 * rank, + (uint64_t)group_cnt * rank, + out_dim, low, 0); + } +} + +int ds4_gpu_attention_output_low_q8_tensor( + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t n_groups, + const ds4_gpu_tensor *heads) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!low || !heads || !model_map || group_dim == 0 || rank == 0 || + n_groups == 0 || group_dim > UINT32_MAX || rank > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + const uint64_t low_dim = (uint64_t)n_groups * rank; + if ((group_dim % 32u) != 0 || low_dim > UINT32_MAX) { + fprintf(stderr, "ds4: Metal attention output low received invalid q8 dimensions\n"); + return 0; + } + + const uint64_t row_a_bytes = (group_dim / 32u) * 34u; + const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes; + if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset) { + fprintf(stderr, "ds4: Metal attention output low weights are outside the mapped model\n"); + return 0; + } + + const uint64_t heads_bytes = (uint64_t)n_groups * group_dim * sizeof(float); + const uint64_t low_bytes = low_dim * sizeof(float); + if (ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(low) < low_bytes) { + fprintf(stderr, "ds4: Metal attention output low received undersized buffers\n"); + return 0; + } + + uint64_t out_a_inner = 0; + id out_a_buf = + ds4_gpu_wrap_model_range(model_map, model_size, + out_a_offset, out_a_bytes, + &out_a_inner); + if (!out_a_buf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + bool ok = true; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) { + ok = false; + } + + if (ok) { + ds4_gpu_mul_mv_id_args args = { + .nei0 = (int32_t)n_groups, + .nei1 = 1, + .nbi1 = 0, + .ne00 = (int32_t)group_dim, + .ne01 = (int32_t)rank, + .ne02 = (int32_t)n_groups, + .nb00 = 34, + .nb01 = row_a_bytes, + .nb02 = (uint64_t)rank * row_a_bytes, + .ne10 = (int32_t)group_dim, + .ne11 = (int32_t)n_groups, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)group_dim * sizeof(float), + .nb12 = (uint64_t)n_groups * group_dim * sizeof(float), + .ne0 = (int32_t)rank, + .ne1 = (int32_t)n_groups, + .nb1 = (uint64_t)rank * sizeof(float), + .nr0 = 2, + }; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4); + ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, + pipeline, + &args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + 32u * 2u * sizeof(float), + 4, + true) != 0; + } + + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + return ok ? 1 : 0; + } +} + +int ds4_gpu_attention_output_low_q4_K_slice_tensor( + ds4_gpu_tensor *low, + const void *model_map, + uint64_t model_size, + uint64_t out_a_offset, + uint64_t group_dim, + uint64_t rank, + uint32_t group0, + uint32_t group_cnt, + const ds4_gpu_tensor *heads) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!low || !heads || !model_map || group_dim == 0 || rank == 0 || + group_cnt == 0 || group_dim > UINT32_MAX || rank > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + if ((group_dim % 256u) != 0) { + fprintf(stderr, "ds4: Metal attention output low received invalid Q4_K dimensions\n"); + return 0; + } + + uint64_t row_a_bytes = 0; + if (!ds4_gpu_quant_row_bytes(DS4_METAL_TENSOR_Q4_K, + (uint32_t)group_dim, + &row_a_bytes)) { + return 0; + } + if (rank > UINT64_MAX / row_a_bytes) return 0; + const uint64_t group_weight_bytes = rank * row_a_bytes; + if (group0 > UINT64_MAX / group_weight_bytes || + group_cnt > UINT64_MAX / group_weight_bytes) { + return 0; + } + const uint64_t group_skip = (uint64_t)group0 * group_weight_bytes; + const uint64_t out_a_bytes = (uint64_t)group_cnt * group_weight_bytes; + if (out_a_offset > UINT64_MAX - group_skip || + out_a_offset + group_skip > model_size || + out_a_bytes > model_size - (out_a_offset + group_skip)) { + fprintf(stderr, "ds4: Metal Q4 attention output low weights are outside the mapped model\n"); + return 0; + } + + const uint64_t heads_bytes = (uint64_t)group_cnt * group_dim * sizeof(float); + const uint64_t low_bytes = (uint64_t)group_cnt * rank * sizeof(float); + if (ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(low) < low_bytes) { + fprintf(stderr, "ds4: Metal Q4 attention output low received undersized buffers\n"); + return 0; + } + + uint64_t out_a_inner = 0; + id out_a_buf = + ds4_gpu_wrap_model_range(model_map, model_size, + out_a_offset + group_skip, + out_a_bytes, + &out_a_inner); + if (!out_a_buf) return 0; + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + + bool ok = true; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb || owned) { + ok = false; + } + + if (ok) { + ds4_gpu_mul_mv_id_args args = { + .nei0 = (int32_t)group_cnt, + .nei1 = 1, + .nbi1 = 0, + .ne00 = (int32_t)group_dim, + .ne01 = (int32_t)rank, + .ne02 = (int32_t)group_cnt, + .nb00 = 1, + .nb01 = row_a_bytes, + .nb02 = group_weight_bytes, + .ne10 = (int32_t)group_dim, + .ne11 = (int32_t)group_cnt, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = (uint64_t)group_dim * sizeof(float), + .nb12 = (uint64_t)group_cnt * group_dim * sizeof(float), + .ne0 = (int32_t)rank, + .ne1 = (int32_t)group_cnt, + .nb1 = (uint64_t)rank * sizeof(float), + .nr0 = 2, + }; + const NSUInteger nsg = 2; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q4_K_f32", (int16_t)nsg); + ok = ds4_gpu_encode_attn_out_low_q8_direct(cb, + pipeline, + &args, + out_a_buf, + (NSUInteger)out_a_inner, + ds4_gpu_tensor_buffer(heads), + ds4_gpu_tensor_offset(heads), + ds4_gpu_tensor_buffer(low), + ds4_gpu_tensor_offset(low), + 32u, + nsg, + false) != 0; + } + + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + return ok ? 1 : 0; + } +} + +static NSUInteger ds4_gpu_align_up_ns(NSUInteger value, NSUInteger align) { + return (value + align - 1u) & ~(align - 1u); +} + +static int ds4_gpu_encode_cpy_f32_f32_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n) { + if (!cb || !src || !dst || n == 0) return 0; + + ds4_gpu_cpy_args args = + ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(float)); + const NSUInteger nth = ds4_gpu_cpy_threads(n, g_cpy_f32_f32_pipeline); + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f32_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f32_f32_3d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride) { + if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; + + ds4_gpu_cpy_args args = { + .nk0 = (int64_t)cols, + .ne00 = (int64_t)cols, + .ne01 = (int64_t)rows, + .ne02 = (int64_t)planes, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src_row_stride, + .nb02 = src_plane_stride, + .nb03 = (uint64_t)planes * src_plane_stride, + .ne0 = (int64_t)cols, + .ne1 = (int64_t)rows, + .ne2 = (int64_t)planes, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = dst_row_stride, + .nb2 = dst_plane_stride, + .nb3 = (uint64_t)planes * dst_plane_stride, + }; + const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f32_pipeline); + const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f32_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f32_f32_3d_src_strided( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_col_stride, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride) { + if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; + + ds4_gpu_cpy_args args = { + .nk0 = (int64_t)cols, + .ne00 = (int64_t)cols, + .ne01 = (int64_t)rows, + .ne02 = (int64_t)planes, + .ne03 = 1, + .nb00 = src_col_stride, + .nb01 = src_row_stride, + .nb02 = src_plane_stride, + .nb03 = (uint64_t)planes * src_plane_stride, + .ne0 = (int64_t)cols, + .ne1 = (int64_t)rows, + .ne2 = (int64_t)planes, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = dst_row_stride, + .nb2 = dst_plane_stride, + .nb3 = (uint64_t)planes * dst_plane_stride, + }; + const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f32_pipeline); + const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f32_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f32_f16_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n) { + if (!cb || !src || !dst || n == 0) return 0; + + const int use_contiguous = + ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0; + id pipeline = use_contiguous + ? g_cpy_contig_f32_f16_pipeline + : g_cpy_f32_f16_pipeline; + const NSUInteger work_items = use_contiguous + ? ((NSUInteger)n + 3u) / 4u + : (NSUInteger)n; + const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline); + const NSUInteger groups = (work_items + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + if (use_contiguous) { + [enc setBytes:&n length:sizeof(n) atIndex:0]; + } else { + ds4_gpu_cpy_args args = + ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(uint16_t)); + [enc setBytes:&args length:sizeof(args) atIndex:0]; + } + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f32_f16_2d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint64_t src_row_stride, + uint64_t dst_row_stride) { + if (!cb || !src || !dst || cols == 0 || rows == 0) return 0; + + ds4_gpu_cpy_args args = { + .nk0 = (int64_t)cols, + .ne00 = (int64_t)cols, + .ne01 = (int64_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src_row_stride, + .nb02 = (uint64_t)rows * src_row_stride, + .nb03 = (uint64_t)rows * src_row_stride, + .ne0 = (int64_t)cols, + .ne1 = (int64_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(uint16_t), + .nb1 = dst_row_stride, + .nb2 = (uint64_t)rows * dst_row_stride, + .nb3 = (uint64_t)rows * dst_row_stride, + }; + const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f16_pipeline); + const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f32_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f32_f16_3d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride) { + if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; + + ds4_gpu_cpy_args args = { + .nk0 = (int64_t)cols, + .ne00 = (int64_t)cols, + .ne01 = (int64_t)rows, + .ne02 = (int64_t)planes, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src_row_stride, + .nb02 = src_plane_stride, + .nb03 = (uint64_t)planes * src_plane_stride, + .ne0 = (int64_t)cols, + .ne1 = (int64_t)rows, + .ne2 = (int64_t)planes, + .ne3 = 1, + .nb0 = sizeof(uint16_t), + .nb1 = dst_row_stride, + .nb2 = dst_plane_stride, + .nb3 = (uint64_t)planes * dst_plane_stride, + }; + const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f16_pipeline); + const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f32_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f16_f16_3d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t cols, + uint32_t rows, + uint32_t planes, + uint64_t src_row_stride, + uint64_t src_plane_stride, + uint64_t dst_row_stride, + uint64_t dst_plane_stride) { + if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0; + + ds4_gpu_cpy_args args = { + .nk0 = (int64_t)cols, + .ne00 = (int64_t)cols, + .ne01 = (int64_t)rows, + .ne02 = (int64_t)planes, + .ne03 = 1, + .nb00 = sizeof(uint16_t), + .nb01 = src_row_stride, + .nb02 = src_plane_stride, + .nb03 = (uint64_t)planes * src_plane_stride, + .ne0 = (int64_t)cols, + .ne1 = (int64_t)rows, + .ne2 = (int64_t)planes, + .ne3 = 1, + .nb0 = sizeof(uint16_t), + .nb1 = dst_row_stride, + .nb2 = dst_plane_stride, + .nb3 = (uint64_t)planes * dst_plane_stride, + }; + const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f16_f16_pipeline); + const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_f16_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_cpy_f16_f32_1d( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t n) { + if (!cb || !src || !dst || n == 0) return 0; + + const int use_contiguous = + ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0; + id pipeline = use_contiguous + ? g_cpy_contig_f16_f32_pipeline + : g_cpy_f16_f32_pipeline; + const NSUInteger work_items = use_contiguous + ? ((NSUInteger)n + 3u) / 4u + : (NSUInteger)n; + const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline); + const NSUInteger groups = (work_items + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + if (use_contiguous) { + [enc setBytes:&n length:sizeof(n) atIndex:0]; + } else { + ds4_gpu_cpy_args args = + ds4_gpu_make_cpy_1d_args(n, sizeof(uint16_t), sizeof(float)); + [enc setBytes:&args length:sizeof(args) atIndex:0]; + } + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_copy_to_f16_1d( + id cb, + id src, + NSUInteger src_off, + bool src_is_f16, + id dst, + NSUInteger dst_off, + uint32_t n) { + if (!cb || !src || !dst) return 0; + if (n == 0) return 1; + if (!src_is_f16) { + return ds4_gpu_encode_cpy_f32_f16_1d(cb, src, src_off, dst, dst_off, n); + } + + if (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") <= 0) { + const NSUInteger work_items = ((NSUInteger)n + 3u) / 4u; + const NSUInteger nth = ds4_gpu_cpy_threads( + (uint32_t)work_items, + g_cpy_contig_f16_f16_pipeline); + const NSUInteger groups = (work_items + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_cpy_contig_f16_f16_pipeline]; + [enc setBytes:&n length:sizeof(n) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + + if (g_batch_cb && cb == g_batch_cb) ds4_gpu_close_batch_encoder(); + id blit = [cb blitCommandEncoder]; + if (!blit) return 0; + [blit copyFromBuffer:src + sourceOffset:src_off + toBuffer:dst + destinationOffset:dst_off + size:(NSUInteger)n * sizeof(uint16_t)]; + [blit endEncoding]; + return 1; +} + +static int ds4_gpu_encode_copy_raw_ring_to_f16( + id cb, + id raw, + NSUInteger raw_offset, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_raw, + uint32_t head_dim, + id dst, + NSUInteger dst_offset) { + if (!cb || !raw || !dst || raw_cap == 0 || raw_start >= raw_cap || + n_raw == 0 || n_raw > raw_cap || head_dim == 0) { + return 0; + } + + const NSUInteger raw_row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger dst_row_bytes = (NSUInteger)head_dim * sizeof(uint16_t); + const uint32_t tail_rows = raw_cap - raw_start < n_raw + ? raw_cap - raw_start + : n_raw; + const uint32_t head_rows = n_raw - tail_rows; + const uint64_t tail_count = (uint64_t)tail_rows * head_dim; + const uint64_t head_count = (uint64_t)head_rows * head_dim; + const uint64_t raw_inner = (uint64_t)raw_start * raw_row_bytes; + const uint64_t dst_inner = (uint64_t)tail_rows * dst_row_bytes; + if (tail_count > UINT32_MAX || head_count > UINT32_MAX || + raw_inner > NSUIntegerMax - raw_offset || + dst_inner > NSUIntegerMax - dst_offset) { + return 0; + } + + if (tail_rows && + !ds4_gpu_encode_cpy_f32_f16_1d( + cb, + raw, + raw_offset + (NSUInteger)raw_inner, + dst, + dst_offset, + (uint32_t)tail_count)) { + return 0; + } + if (head_rows && + !ds4_gpu_encode_cpy_f32_f16_1d( + cb, + raw, + raw_offset, + dst, + dst_offset + (NSUInteger)dst_inner, + (uint32_t)head_count)) { + return 0; + } + return 1; +} + +static int ds4_gpu_encode_flash_kv_stage_f16( + id cb, + id raw, + NSUInteger raw_offset, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_raw, + id comp, + NSUInteger comp_offset, + bool comp_is_f16, + uint32_t n_comp, + uint32_t head_dim, + id dst, + NSUInteger dst_offset, + id mask, + NSUInteger mask_offset, + id pad, + NSUInteger pad_offset, + bool fuse_pad, + bool shared_pad, + bool *did_fuse_pad) { + if (did_fuse_pad) *did_fuse_pad = false; + if (!cb || !raw || !comp || !dst || raw_cap == 0 || + raw_start >= raw_cap || n_raw == 0 || n_raw > raw_cap || + n_comp == 0 || head_dim == 0) { + return 0; + } + + const bool force = + getenv("DS4_METAL_ENABLE_GATHERED_KV_STAGE") != NULL; + const bool disabled = + getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_STAGE") != NULL; + const bool require = + getenv("DS4_METAL_REQUIRE_GATHERED_KV_STAGE") != NULL; + const bool supported_shape = + comp_is_f16 && head_dim == 512u && raw_cap <= UINT32_MAX / 128u; + const uint64_t row_vecs64 = 128u; + const uint64_t comp_count64 = (uint64_t)n_comp * head_dim; + const uint64_t dst_comp_inner64 = + (uint64_t)n_raw * head_dim * sizeof(uint16_t); + if (comp_count64 > UINT32_MAX || + dst_comp_inner64 > NSUIntegerMax - dst_offset) { + return 0; + } + const uint64_t total_vecs64 = + ((uint64_t)n_raw + n_comp) * row_vecs64; + const bool valid_grid = + total_vecs64 != 0 && total_vecs64 <= UINT32_MAX; + const bool eligible = + supported_shape && valid_grid && !g_quality_mode && !disabled && + g_flash_kv_stage_f16_pipeline != nil && + (ds4_gpu_device_name_contains("M3") || + ds4_gpu_device_name_contains("M5") || force); + const bool component_disabled = eligible && + (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") > 0 || + ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") > 0); + const bool use_fusion = eligible && !component_disabled; + const bool use_pad_fusion = + use_fusion && fuse_pad && mask != nil && pad != nil && + getenv("DS4_METAL_DISABLE_M3_GATHERED_KV_PAD_FUSION") == NULL; + if (require && supported_shape && !use_fusion) { + fprintf(stderr, + "ds4: required Metal gathered KV staging kernel was not selected\n"); + return 0; + } + + if (use_fusion) { + ds4_gpu_flash_kv_stage_f16_args args = { + .raw_cap = raw_cap, + .raw_start = raw_start, + .n_raw = n_raw, + .n_comp = n_comp, + .pad_rows = use_pad_fusion ? 32u : 0u, + .shared_pad = use_pad_fusion && shared_pad ? 1u : 0u, + }; + const NSUInteger total_vecs = (NSUInteger)total_vecs64 + + (use_pad_fusion ? 32u * 128u + 32u : 0u); + NSUInteger nth = + g_flash_kv_stage_f16_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > total_vecs) nth = total_vecs; + if (nth == 0) return 0; + const NSUInteger groups = (total_vecs + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_flash_kv_stage_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:raw offset:raw_offset atIndex:1]; + [enc setBuffer:comp offset:comp_offset atIndex:2]; + [enc setBuffer:dst offset:dst_offset atIndex:3]; + [enc setBuffer:(use_pad_fusion ? mask : dst) + offset:(use_pad_fusion ? mask_offset : dst_offset) + atIndex:4]; + [enc setBuffer:(use_pad_fusion ? pad : dst) + offset:(use_pad_fusion ? pad_offset : dst_offset) + atIndex:5]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (did_fuse_pad) *did_fuse_pad = use_pad_fusion; + return 1; + } + + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, + raw, + raw_offset, + raw_cap, + raw_start, + n_raw, + head_dim, + dst, + dst_offset)) { + return 0; + } + return ds4_gpu_encode_copy_to_f16_1d( + cb, + comp, + comp_offset, + comp_is_f16, + dst, + dst_offset + (NSUInteger)dst_comp_inner64, + (uint32_t)comp_count64); +} + +int ds4_gpu_flash_kv_stage_f16_tensor( + ds4_gpu_tensor *dst, + const ds4_gpu_tensor *raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_raw, + const ds4_gpu_tensor *comp, + uint32_t comp_is_f16, + uint32_t n_comp, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!dst || !raw || !comp || raw_cap == 0 || raw_start >= raw_cap || + n_raw == 0 || n_raw > raw_cap || n_comp == 0 || + comp_is_f16 == 0 || head_dim != 512u) { + return 0; + } + + @autoreleasepool { + const uint64_t raw_bytes = + (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = + (uint64_t)n_comp * head_dim * + (comp_is_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t dst_bytes = + ((uint64_t)n_raw + n_comp) * head_dim * sizeof(uint16_t); + id rawbuf = ds4_gpu_tensor_buffer(raw); + id compbuf = ds4_gpu_tensor_buffer(comp); + id dstbuf = ds4_gpu_tensor_buffer(dst); + if (!rawbuf || !compbuf || !dstbuf || + ds4_gpu_tensor_bytes(raw) < raw_bytes || + ds4_gpu_tensor_bytes(comp) < comp_bytes || + ds4_gpu_tensor_bytes(dst) < dst_bytes) { + fprintf(stderr, + "ds4: Metal gathered KV staging received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (!ds4_gpu_encode_flash_kv_stage_f16( + cb, + rawbuf, + ds4_gpu_tensor_offset(raw), + raw_cap, + raw_start, + n_raw, + compbuf, + ds4_gpu_tensor_offset(comp), + comp_is_f16 != 0, + n_comp, + head_dim, + dstbuf, + ds4_gpu_tensor_offset(dst), + nil, + 0, + nil, + 0, + false, + false, + NULL)) { + return 0; + } + if (!ds4_gpu_finish_command_buffer( + cb, owned, "gathered KV staging")) { + return 0; + } + } + return 1; +} + +static int ds4_gpu_encode_fill_f16_1d( + id cb, + id buf, + NSUInteger offset, + uint32_t n, + float value) { + if (!cb || !buf || n == 0) return 0; + + ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(n, 1, 0, 0.0f, 0.0f); + args.val = value; + + NSUInteger nth = (NSUInteger)n; + const NSUInteger max_threads = g_unary_fill_f16_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth > 256u) nth = 256u; + if (nth == 0) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_unary_fill_f16_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:buf offset:offset atIndex:1]; + [enc setBuffer:buf offset:offset atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_flash_attention_raw_heads( + id cb, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_head, + uint32_t head_dim) { + if (head_dim != 512 || n_head == 0 || n_raw == 0 || raw_cap < n_raw) { + return 0; + } + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t heads_bytes = q_bytes; + if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(heads) < heads_bytes) { + fprintf(stderr, "ds4: Metal DS4 FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t ncpsg = 32; + const uint32_t nwg = 32; + const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_raw, nwg, ncpsg); + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_raw * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_raw * row_bytes_f16; + const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 + + (NSUInteger)ncpsg * sizeof(uint16_t); + const NSUInteger nrows = (NSUInteger)n_head; + const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + + nrows * (2u * (NSUInteger)nwg) * sizeof(float); + + id mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, + &g_flash_attn_tmp_bytes, + tmp_bytes, + "ds4_flash_attn_tmp")) { + return 0; + } + memset([mask_buffer contents], 0, mask_bytes); + + id pad_pipeline = nil; + if ((n_raw % ncpsg) != 0) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id vec_pipeline = + ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", + true, true, false, false, (n_raw % ncpsg) != 0, + false, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg, + (int32_t)nwg); + id reduce_pipeline = + ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); + if (!vec_pipeline || !reduce_pipeline) return 0; + + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + raw_cap, + raw_start, + n_raw, + head_dim, + g_flash_attn_kv_buffer, + 0)) { + return 0; + } + + if ((n_raw % ncpsg) != 0) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_raw, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_raw * row_bytes_f16, + .nb13 = (uint64_t)n_raw * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_raw * row_bytes_f16, + .nb23 = (uint64_t)n_raw * row_bytes_f16, + .ne31 = 1, + .ne32 = 1, + .ne33 = 1, + .nb31 = mask_bytes, + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + ds4_gpu_flash_attn_vec_args vec_args = { + .ne01 = 1, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_head * row_bytes, + .ne11 = (int32_t)n_raw, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_raw * row_bytes_f16, + .nb13 = (uint64_t)n_raw * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_raw * row_bytes_f16, + .nb23 = (uint64_t)n_raw * row_bytes_f16, + .ne31 = 1, + .ne32 = 1, + .ne33 = 1, + .nb31 = mask_bytes, + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = 1, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + + 4u * ncpsg + + 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:vec_pipeline]; + [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, n_head, nwg) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_flash_attn_reduce_args reduce_args = { + .nrows = (int32_t)nrows, + }; + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:reduce_pipeline]; + [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +/* Rectangular causal/window mask for raw prefill: q covers rows + * [q_row0, q_row0 + n_q) of an n_kv-token block whose keys all live in + * rows [0, n_kv). The square prefill case is q_row0 == 0, n_q == n_kv. */ +static void ds4_gpu_fill_raw_prefill_mask( + uint16_t *mask, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_kv, + uint32_t window) { + const uint16_t neg_inf_half = 0xfc00u; + for (uint32_t q = 0; q < n_q; q++) { + const uint32_t qpos = q_row0 + q; + uint16_t *row = mask + (uint64_t)q * n_kv; + for (uint32_t k = 0; k < n_kv; k++) { + const bool causal = k <= qpos; + const bool in_window = window == 0 || qpos - k < window; + row[k] = causal && in_window ? 0u : neg_inf_half; + } + } +} + +static void ds4_gpu_fill_glm_prefill_mask( + uint16_t *mask, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len) { + const uint16_t neg_inf_half = 0xfc00u; + for (uint32_t q = 0; q < n_tokens; q++) { + const uint32_t qpos = pos0 + q; + uint16_t *row = mask + (uint64_t)q * cache_len; + for (uint32_t k = 0; k < cache_len; k++) { + row[k] = k <= qpos ? 0u : neg_inf_half; + } + } +} + +static id ds4_gpu_glm_prefill_mask_buffer( + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + NSUInteger mask_bytes) { + const int same_shape = + g_glm_flash_attn_mask_valid && + g_glm_flash_attn_mask_buffer && + g_glm_flash_attn_mask_bytes >= mask_bytes && + g_glm_flash_attn_mask_pos0 == pos0 && + g_glm_flash_attn_mask_tokens == n_tokens && + g_glm_flash_attn_mask_cache_len == cache_len; + if (same_shape) return g_glm_flash_attn_mask_buffer; + + if (g_glm_flash_attn_mask_buffer) { + [g_transient_buffers addObject:g_glm_flash_attn_mask_buffer]; + g_glm_flash_attn_mask_buffer = nil; + } + g_glm_flash_attn_mask_bytes = 0; + g_glm_flash_attn_mask_valid = 0; + if (!ds4_gpu_ensure_scratch_buffer(&g_glm_flash_attn_mask_buffer, + &g_glm_flash_attn_mask_bytes, + mask_bytes, + "ds4_glm_flash_attn_mask")) { + return nil; + } + + ds4_gpu_fill_glm_prefill_mask((uint16_t *)[g_glm_flash_attn_mask_buffer contents], + pos0, + n_tokens, + cache_len); + g_glm_flash_attn_mask_pos0 = pos0; + g_glm_flash_attn_mask_tokens = n_tokens; + g_glm_flash_attn_mask_cache_len = cache_len; + g_glm_flash_attn_mask_valid = 1; + return g_glm_flash_attn_mask_buffer; +} + +static void ds4_gpu_fill_raw_decode_batch_mask( + uint16_t *mask, + uint32_t n_tokens, + uint32_t n_raw, + uint32_t pos0, + uint32_t window) { + const uint16_t neg_inf_half = 0xfc00u; + const uint32_t last_pos = pos0 + n_tokens - 1u; + /* The caller has already copied the SWA ring into logical order when it + * wraps, so key row k represents first_raw_pos + k. */ + const uint32_t first_raw_pos = last_pos + 1u - n_raw; + for (uint32_t q = 0; q < n_tokens; q++) { + const uint32_t qpos = pos0 + q; + uint16_t *row = mask + (uint64_t)q * n_raw; + for (uint32_t k = 0; k < n_raw; k++) { + const uint32_t kpos = first_raw_pos + k; + const bool causal = kpos <= qpos; + const bool in_window = causal && (window == 0 || qpos - kpos < window); + row[k] = causal && in_window ? 0u : neg_inf_half; + } + } +} + +static void ds4_gpu_fill_mixed_decode_batch_mask( + uint16_t *mask, + uint32_t n_tokens, + uint32_t n_raw, + uint32_t n_comp, + uint32_t pos0, + uint32_t window, + uint32_t ratio) { + const uint16_t neg_inf_half = 0xfc00u; + const uint32_t n_keys = n_raw + n_comp; + const uint32_t last_pos = pos0 + n_tokens - 1u; + /* Raw keys are laid out by logical position; compressed keys follow them. */ + const uint32_t first_raw_pos = last_pos + 1u - n_raw; + for (uint32_t q = 0; q < n_tokens; q++) { + const uint32_t qpos = pos0 + q; + uint16_t *row = mask + (uint64_t)q * n_keys; + for (uint32_t k = 0; k < n_raw; k++) { + const uint32_t kpos = first_raw_pos + k; + const bool causal = kpos <= qpos; + const bool in_window = causal && (window == 0 || qpos - kpos < window); + row[k] = causal && in_window ? 0u : neg_inf_half; + } + const uint32_t n_visible = (qpos + 1u) / ratio; + for (uint32_t c = 0; c < n_comp; c++) { + row[n_raw + c] = c < n_visible ? 0u : neg_inf_half; + } + } +} + +/* Rectangular causal/window + compressed-key visibility mask: q covers rows + * [q_row0, q_row0 + n_q) of an n_tokens-token chunk whose raw keys all stay + * resident, followed by n_comp compressed keys. The square prefill case is + * q_row0 == 0, n_q == n_tokens. */ +static void ds4_gpu_fill_static_mixed_prefill_mask( + uint16_t *mask, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio) { + const uint16_t neg_inf_half = 0xfc00u; + const uint32_t n_keys = n_tokens + n_comp; + for (uint32_t q = 0; q < n_q; q++) { + const uint32_t qpos = q_row0 + q; + uint16_t *row = mask + (uint64_t)q * n_keys; + for (uint32_t k = 0; k < n_tokens; k++) { + const bool causal = k <= qpos; + const bool in_window = window == 0 || qpos - k < window; + row[k] = causal && in_window ? 0u : neg_inf_half; + } + + const uint32_t n_visible = (qpos + 1u) / ratio; + for (uint32_t c = 0; c < n_comp; c++) { + row[n_tokens + c] = c < n_visible ? 0u : neg_inf_half; + } + } +} + +/* Static-mixed prefill FlashAttention over a rectangular problem: q holds + * n_q query rows for token positions [q_row0, q_row0 + n_q) of the chunk, + * while the keys stay full (all n_tokens raw rows plus n_comp compressed + * rows). The classic square prefill is q_row0 == 0, n_q == n_tokens. */ +static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long( + id __strong *cbp, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!cbp || !*cbp) return 0; + id cb = *cbp; + if (head_dim != 512 || n_head == 0 || n_q == 0 || n_tokens == 0 || ratio == 0) { + return 0; + } + + const uint32_t n_keys = n_tokens + n_comp; + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf; + id maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * + (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t comp_mask_bytes = use_comp_mask + ? (uint64_t)n_comp * (q_row0 + n_q) * sizeof(float) : 0u; + if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || + (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal prefill static mixed DS4 non-vector FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t nqptg = 8; + const uint32_t ncpsg = 64; + const uint32_t nsg = head_dim >= 512 ? 8u : 4u; + const bool has_kvpad = (n_keys % ncpsg) != 0; + const bool bc_mask = (n_q % nqptg) != 0; + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_q * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_q * sizeof(uint16_t)) + : 1u; + const NSUInteger nblk0 = ((NSUInteger)n_keys + ncpsg - 1u) / ncpsg; + const NSUInteger nblk1 = ((NSUInteger)n_q + nqptg - 1u) / nqptg; + const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); + + const uint32_t mask_cache_kind = ratio == 4u + ? DS4_GPU_PREFILL_MASK_CACHE_RATIO4 + : (ratio == 128u ? DS4_GPU_PREFILL_MASK_CACHE_RATIO128 : 0u); + bool mask_cache_created = false; + ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache = + use_comp_mask == 0u && mask_cache_kind != 0u && + q_row0 == 0u && n_q == n_tokens + ? ds4_gpu_get_zero_prefix_prefill_mask_cache(mask_cache_kind, + n_tokens, + n_comp, + n_keys, + window, + ratio, + nqptg, + ncpsg, + has_kvpad, + bc_mask, + mask_bytes, + blk_bytes, + &mask_cache_created) + : NULL; + id mask_buffer = mask_cache + ? mask_cache->mask + : ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (mask_cache && mask_cache_created) { + ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], + 0u, + n_tokens, + n_tokens, + n_comp, + window, + ratio); + mask_cache->valid = true; + } + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + (!mask_cache && + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_flash_attn_blk"))) { + return 0; + } + id blk_buffer = mask_cache + ? mask_cache->blk + : g_flash_attn_blk_buffer; + + const bool flash_stage_profile = + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; + double flash_stage_t0 = 0.0; + if (flash_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + int profile_owned = 0; + cb = ds4_gpu_command_buffer(&profile_owned); + if (!cb || profile_owned) return 0; + *cbp = cb; + flash_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ + if (flash_stage_profile) { \ + if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ + "static_mixed_nonvec", (name), n_q, n_comp, n_keys, \ + n_head, head_dim, window, ratio, &flash_stage_t0)) { \ + return 0; \ + } \ + cb = *cbp; \ + } \ + } while (0) + + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + g_flash_attn_kv_buffer, + 0, + n_tokens * head_dim)) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); + if (n_comp && + !ds4_gpu_encode_copy_to_f16_1d(cb, + compbuf, + ds4_gpu_tensor_offset(comp_kv), + comp_kv_f16 != 0, + g_flash_attn_kv_buffer, + (NSUInteger)n_tokens * row_bytes_f16, + n_comp * head_dim)) { + return 0; + } + if (n_comp) { + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp"); + } + + if (!mask_cache) { + ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], + q_row0, + n_q, + n_tokens, + n_comp, + window, + ratio); + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); + if (use_comp_mask && n_comp != 0) { + if (!ds4_gpu_encode_cpy_f32_f16_2d(cb, + maskbuf, + ds4_gpu_tensor_offset(comp_mask) + + (NSUInteger)((uint64_t)q_row0 * n_comp * sizeof(float)), + mask_buffer, + (NSUInteger)n_tokens * sizeof(uint16_t), + n_comp, + n_q, + (uint64_t)n_comp * sizeof(float), + (uint64_t)n_keys * sizeof(uint16_t))) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_comp_copy"); + } + + id pad_pipeline = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id blk_pipeline = + ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); + id attn_pipeline = + ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", + true, true, false, false, has_kvpad, bc_mask, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg); + if (!blk_pipeline || !attn_pipeline) return 0; + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); + } + + ds4_gpu_flash_attn_blk_args blk_args = { + .ne01 = (int32_t)n_q, + .ne30 = (int32_t)n_keys, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = nil; + if (!mask_cache || !mask_cache->blk_ready) { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (mask_cache) mask_cache->blk_ready = true; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map"); + + ds4_gpu_flash_attn_vec_args args = { + .ne01 = (int32_t)n_q, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_q * n_head * row_bytes, + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_q, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); + const NSUInteger shared_elems = (NSUInteger)nqptg * + ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:blk_buffer offset:0 atIndex:7]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention"); + +#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE + return 1; +} + +static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec( + id __strong *cbp, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!cbp || !*cbp) return 0; + id cb = *cbp; + if (head_dim != 512 || n_head == 0 || n_tokens == 0 || ratio == 0) { + return 0; + } + + const uint32_t n_keys = n_tokens + n_comp; + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf; + id maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * + (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t comp_mask_bytes = use_comp_mask ? (uint64_t)n_comp * n_tokens * sizeof(float) : 0u; + if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || + (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal prefill static mixed DS4 FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t ncpsg = 32; + const uint32_t nwg = 32; + const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_keys, nwg, ncpsg); + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_tokens * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; + const bool has_kvpad = (n_keys % ncpsg) != 0; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t)) + : 1u; + const NSUInteger nrows = (NSUInteger)n_tokens * n_head; + const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + + nrows * (2u * (NSUInteger)nwg) * sizeof(float); + + id mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, + &g_flash_attn_tmp_bytes, + tmp_bytes, + "ds4_flash_attn_tmp")) { + return 0; + } + + const bool flash_stage_profile = + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; + double flash_stage_t0 = 0.0; + if (flash_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + int profile_owned = 0; + cb = ds4_gpu_command_buffer(&profile_owned); + if (!cb || profile_owned) return 0; + *cbp = cb; + flash_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ + if (flash_stage_profile) { \ + if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ + "static_mixed_vec", (name), n_tokens, n_comp, n_keys, \ + n_head, head_dim, window, ratio, &flash_stage_t0)) { \ + return 0; \ + } \ + cb = *cbp; \ + } \ + } while (0) + + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + g_flash_attn_kv_buffer, + 0, + n_tokens * head_dim)) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); + if (n_comp) { + if (!ds4_gpu_encode_copy_to_f16_1d(cb, + compbuf, + ds4_gpu_tensor_offset(comp_kv), + comp_kv_f16 != 0, + g_flash_attn_kv_buffer, + (NSUInteger)n_tokens * row_bytes_f16, + n_comp * head_dim)) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp"); + } + + ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents], + 0, + n_tokens, + n_tokens, + n_comp, + window, + ratio); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); + if (use_comp_mask && n_comp != 0) { + if (!ds4_gpu_encode_cpy_f32_f16_2d(cb, + maskbuf, + ds4_gpu_tensor_offset(comp_mask), + mask_buffer, + (NSUInteger)n_tokens * sizeof(uint16_t), + n_comp, + n_tokens, + (uint64_t)n_comp * sizeof(float), + (uint64_t)n_keys * sizeof(uint16_t))) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_comp_copy"); + } + + id pad_pipeline = nil; + id enc = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id vec_pipeline = + ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", + true, true, false, false, has_kvpad, + false, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg, + (int32_t)nwg); + id reduce_pipeline = + ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); + if (!vec_pipeline || !reduce_pipeline) return 0; + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); + } + + ds4_gpu_flash_attn_vec_args vec_args = { + .ne01 = (int32_t)n_tokens, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_tokens * n_head * row_bytes, + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + + 4u * ncpsg + + 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:vec_pipeline]; + [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_head, nwg) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_vec"); + + ds4_gpu_flash_attn_reduce_args reduce_args = { + .nrows = (int32_t)nrows, + }; + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:reduce_pipeline]; + [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_reduce"); + +#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE + return 1; +} + +static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec( + id __strong *cbp, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + /* The vector sibling below only handles the small square case; every + * rectangular (TP row-split) problem goes through the long path. */ + if (n_tokens >= 20 || q_row0 != 0 || n_q != n_tokens) { + return ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long(cbp, + heads, + sinks_buf, + sinks_offset, + q, + raw_kv, + comp_kv, + comp_kv_f16, + comp_mask, + use_comp_mask, + q_row0, + n_q, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim); + } + return ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec(cbp, + heads, + sinks_buf, + sinks_offset, + q, + raw_kv, + comp_kv, + comp_kv_f16, + comp_mask, + use_comp_mask, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim); +} + +/* Raw prefill FlashAttention over a rectangular problem: q holds n_q query + * rows that correspond to token positions [q_row0, q_row0 + n_q) of the + * chunk, raw_kv holds all n_kv key rows, and heads receives one output row + * per query row. The classic square prefill is q_row0 == 0, n_q == n_kv. */ +static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec( + id __strong *cbp, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_kv, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + if (!cbp || !*cbp) return 0; + id cb = *cbp; + if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) { + return 0; + } + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)n_kv * head_dim * sizeof(float); + if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal prefill raw DS4 non-vector FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t nqptg = 8; + const uint32_t ncpsg = 64; + const uint32_t nsg = head_dim >= 512 ? 8u : 4u; + const bool has_kvpad = (n_kv % ncpsg) != 0; + const bool bc_mask = (n_q % nqptg) != 0; + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_q * (NSUInteger)n_kv * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_kv * row_bytes_f16; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_q * sizeof(uint16_t)) + : 1u; + const NSUInteger nblk0 = ((NSUInteger)n_kv + ncpsg - 1u) / ncpsg; + const NSUInteger nblk1 = ((NSUInteger)n_q + nqptg - 1u) / nqptg; + const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); + + bool mask_cache_created = false; + ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache = + q_row0 == 0u && n_q == n_kv + ? ds4_gpu_get_zero_prefix_prefill_mask_cache( + DS4_GPU_PREFILL_MASK_CACHE_RAW, + n_kv, + 0u, + n_kv, + window, + 0u, + nqptg, + ncpsg, + has_kvpad, + bc_mask, + mask_bytes, + blk_bytes, + &mask_cache_created) + : NULL; + id mask_buffer = mask_cache + ? mask_cache->mask + : ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (mask_cache && mask_cache_created) { + ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], + 0u, n_kv, n_kv, window); + mask_cache->valid = true; + } + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + (!mask_cache && + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_flash_attn_blk"))) { + return 0; + } + id blk_buffer = mask_cache + ? mask_cache->blk + : g_flash_attn_blk_buffer; + + const bool flash_stage_profile = + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; + double flash_stage_t0 = 0.0; + if (flash_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + int profile_owned = 0; + cb = ds4_gpu_command_buffer(&profile_owned); + if (!cb || profile_owned) return 0; + *cbp = cb; + flash_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ + if (flash_stage_profile) { \ + if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ + "raw_nonvec", (name), n_q, 0, n_kv, \ + n_head, head_dim, window, 0, &flash_stage_t0)) { \ + return 0; \ + } \ + cb = *cbp; \ + } \ + } while (0) + + if (!mask_cache) { + ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], + q_row0, n_q, n_kv, window); + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); + + id pad_pipeline = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id blk_pipeline = + ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); + id attn_pipeline = + ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", + true, true, false, false, has_kvpad, bc_mask, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg); + if (!blk_pipeline || !attn_pipeline) return 0; + + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + g_flash_attn_kv_buffer, + 0, + n_kv * head_dim)) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_kv, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_kv * row_bytes_f16, + .nb13 = (uint64_t)n_kv * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_kv * row_bytes_f16, + .nb23 = (uint64_t)n_kv * row_bytes_f16, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_kv * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); + } + + ds4_gpu_flash_attn_blk_args blk_args = { + .ne01 = (int32_t)n_q, + .ne30 = (int32_t)n_kv, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_kv * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = nil; + if (!mask_cache || !mask_cache->blk_ready) { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + if (mask_cache) mask_cache->blk_ready = true; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map"); + + ds4_gpu_flash_attn_vec_args args = { + .ne01 = (int32_t)n_q, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_q * n_head * row_bytes, + .ne11 = (int32_t)n_kv, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_kv * row_bytes_f16, + .nb13 = (uint64_t)n_kv * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_kv * row_bytes_f16, + .nb23 = (uint64_t)n_kv * row_bytes_f16, + .ne31 = (int32_t)n_q, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_kv * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_q, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); + const NSUInteger shared_elems = (NSUInteger)nqptg * + ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:blk_buffer offset:0 atIndex:7]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention"); + +#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE + return 1; +} + +static int ds4_gpu_encode_flash_attention_prefill_raw_heads( + id __strong *cbp, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_kv, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + if (!cbp || !*cbp) return 0; + id cb = *cbp; + if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) { + return 0; + } + /* The vector sibling below only handles the small square case; every + * rectangular (TP row-split) problem goes through the non-vector path. */ + if (n_kv >= 20 || q_row0 != 0 || n_q != n_kv) { + return ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec(cbp, + heads, + sinks_buf, + sinks_offset, + q, + raw_kv, + q_row0, + n_q, + n_kv, + window, + n_head, + head_dim); + } + const uint32_t n_tokens = n_q; + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float); + if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal prefill raw DS4 FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t ncpsg = 32; + const uint32_t nwg = 32; + const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_tokens, nwg, ncpsg); + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_tokens * (NSUInteger)n_tokens * sizeof(uint16_t); + const NSUInteger kv_f16_offset = 0; + const NSUInteger kv_f16_bytes = (NSUInteger)n_tokens * row_bytes_f16; + const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 + + (NSUInteger)ncpsg * (NSUInteger)n_tokens * sizeof(uint16_t); + const NSUInteger nrows = (NSUInteger)n_tokens * n_head; + const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + + nrows * (2u * (NSUInteger)nwg) * sizeof(float); + + id mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_f16_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, + &g_flash_attn_tmp_bytes, + tmp_bytes, + "ds4_flash_attn_tmp")) { + return 0; + } + + const bool flash_stage_profile = + getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil; + double flash_stage_t0 = 0.0; + if (flash_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + int profile_owned = 0; + cb = ds4_gpu_command_buffer(&profile_owned); + if (!cb || profile_owned) return 0; + *cbp = cb; + flash_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \ + if (flash_stage_profile) { \ + if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \ + "raw_vec", (name), n_tokens, 0, n_tokens, \ + n_head, head_dim, window, 0, &flash_stage_t0)) { \ + return 0; \ + } \ + cb = *cbp; \ + } \ + } while (0) + + ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], 0, n_tokens, n_tokens, window); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill"); + + id pad_pipeline = nil; + if ((n_tokens % ncpsg) != 0) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id vec_pipeline = + ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", + true, true, false, false, true, + false, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg, + (int32_t)nwg); + id reduce_pipeline = + ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); + if (!vec_pipeline || !reduce_pipeline) return 0; + + if (!ds4_gpu_encode_cpy_f32_f16_1d(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + g_flash_attn_kv_buffer, + kv_f16_offset, + n_tokens * head_dim)) { + return 0; + } + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw"); + + if ((n_tokens % ncpsg) != 0) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_tokens, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_tokens * row_bytes_f16, + .nb13 = (uint64_t)n_tokens * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_tokens * row_bytes_f16, + .nb23 = (uint64_t)n_tokens * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_tokens * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad"); + } + + ds4_gpu_flash_attn_vec_args vec_args = { + .ne01 = (int32_t)n_tokens, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_tokens * n_head * row_bytes, + .ne11 = (int32_t)n_tokens, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_tokens * row_bytes_f16, + .nb13 = (uint64_t)n_tokens * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_tokens * row_bytes_f16, + .nb23 = (uint64_t)n_tokens * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_tokens * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + + 4u * ncpsg + + 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:vec_pipeline]; + [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_head, nwg) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_vec"); + + ds4_gpu_flash_attn_reduce_args reduce_args = { + .nrows = (int32_t)nrows, + }; + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:reduce_pipeline]; + [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_reduce"); + +#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE + return 1; +} + +static int ds4_gpu_encode_flash_attention_gathered_heads( + id cb, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_comp, + const ds4_gpu_tensor *comp_mask, + uint32_t use_mask, + uint32_t n_head, + uint32_t head_dim) { + const uint32_t n_keys = n_raw + n_comp; + if (head_dim != 512 || n_head == 0 || n_raw == 0 || n_keys == 0 || + raw_cap < n_raw || n_keys < n_raw) { + return 0; + } + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : nil; + id headsbuf = ds4_gpu_tensor_buffer(heads); + id maskbuf = use_mask ? ds4_gpu_tensor_buffer(comp_mask) : nil; + const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * + (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t comp_mask_bytes = use_mask ? (uint64_t)n_comp * sizeof(float) : 0u; + if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + (n_comp && !compbuf) || + (use_mask && !maskbuf) || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || + ds4_gpu_tensor_bytes(heads) < q_bytes || + (use_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes)) { + fprintf(stderr, "ds4: Metal gathered DS4 FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t ncpsg = 32; + const uint32_t nwg = 32; + const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_keys, nwg, ncpsg); + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_keys * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; + const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 + + (NSUInteger)ncpsg * sizeof(uint16_t); + const NSUInteger nrows = (NSUInteger)n_head; + const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) + + nrows * (2u * (NSUInteger)nwg) * sizeof(float); + + const bool use_persistent_zero_mask = + use_mask == 0u && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_PERSISTENT_ZERO_ATTN_MASK") != NULL) && + getenv("DS4_METAL_DISABLE_M3_PERSISTENT_ZERO_ATTN_MASK") == NULL; + + if (!(use_persistent_zero_mask + ? ds4_gpu_ensure_zero_attention_mask(mask_bytes) + : ds4_gpu_ensure_scratch_buffer(&g_flash_attn_mask_buffer, + &g_flash_attn_mask_bytes, + mask_bytes, + "ds4_flash_attn_mask")) || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer, + &g_flash_attn_tmp_bytes, + tmp_bytes, + "ds4_flash_attn_tmp")) { + return 0; + } + id flash_mask_buffer = use_persistent_zero_mask + ? g_flash_attn_zero_mask_buffer + : g_flash_attn_mask_buffer; + + const bool has_kvpad = (n_keys % ncpsg) != 0; + const bool use_shared_kvpad = + has_kvpad && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_SHARED_KV_PAD") != NULL) && + getenv("DS4_METAL_DISABLE_M3_SHARED_KV_PAD") == NULL; + id vec_pipeline = + ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512", + true, true, false, false, has_kvpad, + use_shared_kvpad, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg, + (int32_t)nwg); + id reduce_pipeline = + ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg); + if (!vec_pipeline || !reduce_pipeline) return 0; + + if (!use_persistent_zero_mask && + !ds4_gpu_encode_fill_f16_1d(cb, flash_mask_buffer, 0, n_keys, 0.0f)) { + return 0; + } + if (use_mask && n_comp && + !ds4_gpu_encode_cpy_f32_f16_1d(cb, + maskbuf, + ds4_gpu_tensor_offset(comp_mask), + flash_mask_buffer, + (NSUInteger)n_raw * sizeof(uint16_t), + n_comp)) { + return 0; + } + + bool pad_fused = false; + if (!ds4_gpu_encode_flash_kv_stage_f16( + cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + raw_cap, + raw_start, + n_raw, + compbuf, + ds4_gpu_tensor_offset(comp_kv), + comp_kv_f16 != 0, + n_comp, + head_dim, + g_flash_attn_kv_buffer, + 0, + flash_mask_buffer, + 0, + g_flash_attn_pad_buffer, + 0, + has_kvpad, + use_shared_kvpad, + &pad_fused)) { + return 0; + } + + id pad_pipeline = nil; + if (has_kvpad && !pad_fused) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + + if (has_kvpad && !pad_fused) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = 1, + .ne32 = 1, + .ne33 = 1, + .nb31 = mask_bytes, + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:flash_mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + ds4_gpu_flash_attn_vec_args vec_args = { + .ne01 = 1, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_head * row_bytes, + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = 1, + .ne32 = 1, + .ne33 = 1, + .nb31 = mask_bytes, + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = 1, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) + + 4u * ncpsg + + 2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg; + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:vec_pipeline]; + [enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:flash_mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, n_head, nwg) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_flash_attn_reduce_args reduce_args = { + .nrows = (int32_t)nrows, + }; + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:reduce_pipeline]; + [enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0]; + [enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_flash_attention_decode_raw_batch_heads( + id cb, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t window, + uint32_t n_head, + uint32_t head_dim, + bool noncausal) { + if (head_dim != 512 || n_head == 0 || n_tokens == 0 || + n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) { + return 0; + } + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + if (!qbuf || !rawbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal decode raw batch FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t nqptg = 8; + const uint32_t ncpsg = 64; + const uint32_t nsg = head_dim >= 512 ? 8u : 4u; + const bool has_kvpad = (n_raw % ncpsg) != 0; + const bool bc_mask = (n_tokens % nqptg) != 0; + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_raw * (NSUInteger)n_tokens * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_raw * row_bytes_f16; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t)) + : 1u; + const NSUInteger nblk0 = ((NSUInteger)n_raw + ncpsg - 1u) / ncpsg; + const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; + const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); + + id mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_flash_attn_blk")) { + return 0; + } + + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + raw_cap, + raw_start, + n_raw, + head_dim, + g_flash_attn_kv_buffer, + 0)) { + return 0; + } + + if (noncausal) { + memset([mask_buffer contents], 0, mask_bytes); + } else { + ds4_gpu_fill_raw_decode_batch_mask((uint16_t *)[mask_buffer contents], + n_tokens, + n_raw, + pos0, + window); + } + + id pad_pipeline = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id blk_pipeline = + ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); + id attn_pipeline = + ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", + true, true, false, false, has_kvpad, bc_mask, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg); + if (!blk_pipeline || !attn_pipeline) return 0; + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_raw, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_raw * row_bytes_f16, + .nb13 = (uint64_t)n_raw * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_raw * row_bytes_f16, + .nb23 = (uint64_t)n_raw * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_raw * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + ds4_gpu_flash_attn_blk_args blk_args = { + .ne01 = (int32_t)n_tokens, + .ne30 = (int32_t)n_raw, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_raw * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_flash_attn_vec_args args = { + .ne01 = (int32_t)n_tokens, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_tokens * n_head * row_bytes, + .ne11 = (int32_t)n_raw, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_raw * row_bytes_f16, + .nb13 = (uint64_t)n_raw * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_raw * row_bytes_f16, + .nb23 = (uint64_t)n_raw * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_raw * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); + const NSUInteger shared_elems = (NSUInteger)nqptg * + ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +static int ds4_gpu_encode_flash_attention_decode_mixed_batch_heads( + id cb, + ds4_gpu_tensor *heads, + id sinks_buf, + NSUInteger sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (n_comp == 0) { + return ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb, + heads, + sinks_buf, + sinks_offset, + q, + raw_kv, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + window, + n_head, + head_dim, + false); + } + if (head_dim != 512 || n_head == 0 || n_tokens == 0 || + n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || + ratio == 0 || !comp_kv || (use_comp_mask && !comp_mask)) { + return 0; + } + + const uint32_t n_keys = n_raw + n_comp; + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = ds4_gpu_tensor_buffer(comp_kv); + id maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * + (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t comp_mask_bytes = use_comp_mask ? (uint64_t)n_comp * n_tokens * sizeof(float) : 0u; + if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(comp_kv) < comp_bytes || + (use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal decode mixed batch FlashAttention received undersized buffers\n"); + return 0; + } + + const uint32_t nqptg = 8; + const uint32_t ncpsg = 64; + const uint32_t nsg = head_dim >= 512 ? 8u : 4u; + const bool has_kvpad = (n_keys % ncpsg) != 0; + const bool bc_mask = (n_tokens % nqptg) != 0; + const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float); + const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_tokens * sizeof(uint16_t); + const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t)) + : 1u; + const NSUInteger nblk0 = ((NSUInteger)n_keys + ncpsg - 1u) / ncpsg; + const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; + const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); + + id mask_buffer = + ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask"); + if (!mask_buffer || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_bytes, + "ds4_flash_attn_kv_f16") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_flash_attn_blk")) { + return 0; + } + + if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb, + rawbuf, + ds4_gpu_tensor_offset(raw_kv), + raw_cap, + raw_start, + n_raw, + head_dim, + g_flash_attn_kv_buffer, + 0) || + !ds4_gpu_encode_copy_to_f16_1d(cb, + compbuf, + ds4_gpu_tensor_offset(comp_kv), + comp_kv_f16 != 0, + g_flash_attn_kv_buffer, + (NSUInteger)n_raw * row_bytes_f16, + n_comp * head_dim)) { + return 0; + } + + ds4_gpu_fill_mixed_decode_batch_mask((uint16_t *)[mask_buffer contents], + n_tokens, + n_raw, + n_comp, + pos0, + window, + ratio); + if (use_comp_mask) { + if (!ds4_gpu_encode_cpy_f32_f16_2d(cb, + maskbuf, + ds4_gpu_tensor_offset(comp_mask), + mask_buffer, + (NSUInteger)n_raw * sizeof(uint16_t), + n_comp, + n_tokens, + (uint64_t)n_comp * sizeof(float), + (uint64_t)n_keys * sizeof(uint16_t))) { + return 0; + } + } + + id pad_pipeline = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id blk_pipeline = + ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); + id attn_pipeline = + ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512", + true, true, false, false, has_kvpad, bc_mask, + (int32_t)head_dim, + (int32_t)head_dim, + (int32_t)nsg); + if (!blk_pipeline || !attn_pipeline) return 0; + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + ds4_gpu_flash_attn_blk_args blk_args = { + .ne01 = (int32_t)n_tokens, + .ne30 = (int32_t)n_keys, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_flash_attn_vec_args args = { + .ne01 = (int32_t)n_tokens, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * row_bytes, + .nb02 = row_bytes, + .nb03 = (uint64_t)n_tokens * n_head * row_bytes, + .ne11 = (int32_t)n_keys, + .ne_12_2 = 1, + .ne_12_3 = 1, + .ns10 = (int32_t)head_dim, + .nb11 = row_bytes_f16, + .nb12 = (uint64_t)n_keys * row_bytes_f16, + .nb13 = (uint64_t)n_keys * row_bytes_f16, + .ns20 = (int32_t)head_dim, + .nb21 = row_bytes_f16, + .nb22 = (uint64_t)n_keys * row_bytes_f16, + .nb23 = (uint64_t)n_keys * row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)n_keys * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)head_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u); + const NSUInteger shared_elems = (NSUInteger)nqptg * + ((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:sinks_buf offset:sinks_offset atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +int ds4_gpu_attention_prefill_raw_heads_range_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_kv, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !model_map || n_q == 0 || n_kv == 0) return 0; + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_prefill_raw_heads(&cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + q_row0, + n_q, + n_kv, + window, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill raw attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_prefill_raw_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + return ds4_gpu_attention_prefill_raw_heads_range_tensor(heads, + model_map, + model_size, + sinks_offset, + q, + raw_kv, + 0, + n_tokens, + n_tokens, + window, + n_head, + head_dim); +} + +int ds4_gpu_attention_decode_raw_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t window, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 || + n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) { + return 0; + } + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + window, + n_head, + head_dim, + false)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode raw batch attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_noncausal_raw_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_tokens, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !model_map || + n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || + raw_start >= raw_cap || n_head == 0 || head_dim == 0) { + return 0; + } + + @autoreleasepool { + const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float); + if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal noncausal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + sink_bytes, + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + n_tokens, + 0, + n_raw, + raw_cap, + raw_start, + 0, + n_head, + head_dim, + true)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph noncausal raw batch attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_decode_mixed_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t use_comp_mask, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 || + n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || + ratio == 0 || (n_comp != 0 && !comp_kv) || + (use_comp_mask != 0 && !comp_mask)) { + return 0; + } + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + comp_kv, + comp_kv_f16, + comp_mask, + use_comp_mask, + n_tokens, + pos0, + n_raw, + raw_cap, + raw_start, + n_comp, + window, + ratio, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode mixed batch attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *topk, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + uint32_t n_comp, + uint32_t top_k, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !model_map || !q || !raw_kv || !comp_kv || !topk || + n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap || + n_comp == 0 || top_k == 0 || top_k > n_comp || (top_k & (top_k - 1u)) != 0 || + ratio == 0 || n_head == 0 || head_dim != 512) { + return 0; + } + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal indexed attention sinks range is outside the mapped model\n"); + return 0; + } + + const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float); + const uint64_t row_bytes_f16 = (uint64_t)head_dim * sizeof(uint16_t); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * row_bytes; + const uint64_t raw_bytes = (uint64_t)raw_cap * row_bytes; + const uint64_t comp_bytes = (uint64_t)n_comp * (comp_kv_f16 ? row_bytes_f16 : row_bytes); + const uint64_t topk_bytes = (uint64_t)top_k * n_tokens * sizeof(int32_t); + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = ds4_gpu_tensor_buffer(comp_kv); + id topkbuf = ds4_gpu_tensor_buffer(topk); + id headsbuf = ds4_gpu_tensor_buffer(heads); + if (!qbuf || !rawbuf || !compbuf || !topkbuf || !headsbuf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + ds4_gpu_tensor_bytes(comp_kv) < comp_bytes || + ds4_gpu_tensor_bytes(topk) < topk_bytes || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal indexed mixed attention received undersized buffers\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + id sort_pipeline = + ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline, + "kernel_dsv4_sort_i32_rows_asc"); + const bool decode_one_token = n_tokens == 1u; + id attn_pipeline = + decode_one_token ? + ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_rb16_pipeline, + "kernel_dsv4_indexed_mixed_attention_heads8_rb16") : + ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_pipeline, + "kernel_dsv4_indexed_mixed_attention_heads8"); + if (!sort_pipeline || !attn_pipeline) return 0; + if ((NSUInteger)top_k > sort_pipeline.maxTotalThreadsPerThreadgroup) { + fprintf(stderr, "ds4: Metal indexed attention top-k exceeds sort threadgroup limit\n"); + return 0; + } + /* + * Fast decode attends to the same full top-k compressed rows but keeps + * them in score order, avoiding a chronological sort dispatch. + * --quality restores the sorted order for stricter reproducibility. + */ + const bool skip_decode_sort = !g_quality_mode && decode_one_token; + if (!skip_decode_sort && + !ds4_gpu_ensure_scratch_buffer(&g_indexed_topk_buffer, + &g_indexed_topk_bytes, + (NSUInteger)topk_bytes, + "ds4_indexed_topk_sorted")) { + return 0; + } + + ds4_gpu_dsv4_topk_mask_args sort_args = { + .ne00 = (int64_t)top_k, + .ne01 = (int64_t)n_tokens, + .nb00 = sizeof(int32_t), + .nb01 = (uint64_t)top_k * sizeof(int32_t), + .ne0 = (int64_t)top_k, + .ne1 = (int64_t)n_tokens, + .nb0 = sizeof(int32_t), + .nb1 = (uint64_t)top_k * sizeof(int32_t), + }; + ds4_gpu_dsv4_indexed_attention_args attn_args = { + .n_tokens = n_tokens, + .n_head = n_head, + .n_raw = n_raw, + .raw_cap = raw_cap, + .raw_start = raw_start, + .n_comp = n_comp, + .top_k = top_k, + .pos0 = pos0, + .window = window, + .ratio = ratio, + .comp_kv_f16 = comp_kv_f16 ? 1u : 0u, + .pad0 = 0, + .q_token_stride = (uint64_t)n_head * row_bytes, + .q_head_stride = row_bytes, + .raw_row_stride = row_bytes, + .comp_row_stride = comp_kv_f16 ? row_bytes_f16 : row_bytes, + .topk_token_stride = (uint64_t)top_k * sizeof(int32_t), + .dst_token_stride = (uint64_t)n_head * row_bytes, + .dst_head_stride = row_bytes, + .scale = 1.0f / sqrtf((float)head_dim), + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = nil; + if (!skip_decode_sort) { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:sort_pipeline]; + [enc setBytes:&sort_args length:sizeof(sort_args) atIndex:0]; + [enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1]; + [enc setBuffer:g_indexed_topk_buffer offset:0 atIndex:2]; + [enc setThreadgroupMemoryLength:(NSUInteger)top_k * sizeof(int32_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(top_k, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_kv) atIndex:2]; + [enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(comp_kv) atIndex:3]; + [enc setBuffer:skip_decode_sort ? topkbuf : g_indexed_topk_buffer + offset:skip_decode_sort ? ds4_gpu_tensor_offset(topk) : 0 + atIndex:4]; + [enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:5]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:6]; + [enc setThreadgroupMemoryLength:(decode_one_token ? 16u : 1u) * + 128u * 4u * sizeof(uint16_t) + atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, ((NSUInteger)n_head + 7u) / 8u, 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph indexed mixed attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_prefill_static_mixed_heads_range_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t q_row0, + uint32_t n_q, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !model_map || n_q == 0 || n_tokens == 0 || + ratio == 0 || (n_comp != 0 && !comp_kv)) { + return 0; + } + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(&cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + comp_kv, + comp_kv_f16, + NULL, + 0, + q_row0, + n_q, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill static mixed attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_prefill_static_mixed_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + return ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(heads, + model_map, + model_size, + sinks_offset, + q, + raw_kv, + comp_kv, + comp_kv_f16, + 0, + n_tokens, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim); +} + +int ds4_gpu_attention_prefill_masked_mixed_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + const ds4_gpu_tensor *comp_mask, + uint32_t n_tokens, + uint32_t n_comp, + uint32_t window, + uint32_t ratio, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !raw_kv || !comp_kv || !comp_mask || !model_map || + n_tokens == 0 || n_comp == 0 || ratio == 0) { + return 0; + } + + @autoreleasepool { + if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, + sinks_offset, + (uint64_t)n_head * sizeof(float), + &sinks_inner); + if (!sinks_buf) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(&cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + comp_kv, + comp_kv_f16, + comp_mask, + 1, + 0, + n_tokens, + n_tokens, + n_comp, + window, + ratio, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill masked mixed attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_attention_decode_heads_tensor( + ds4_gpu_tensor *heads, + const void *model_map, + uint64_t model_size, + uint64_t sinks_offset, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *raw_kv, + uint32_t n_raw, + uint32_t raw_cap, + uint32_t raw_start, + const ds4_gpu_tensor *comp_kv, + uint32_t comp_kv_f16, + uint32_t n_comp, + const ds4_gpu_tensor *comp_mask, + uint32_t use_mask, + uint32_t n_head, + uint32_t head_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !model_map || !q || !raw_kv || + n_raw == 0 || n_head == 0 || head_dim == 0 || + raw_cap < n_raw || raw_start >= raw_cap || + n_raw > UINT32_MAX - n_comp || n_raw + n_comp > 8192u || + (n_comp != 0 && !comp_kv) || + (use_mask != 0 && !comp_mask)) { + return 0; + } + + @autoreleasepool { + const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float); + const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * + (comp_kv_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float); + if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) { + fprintf(stderr, "ds4: Metal graph attention heads sink range is outside the mapped model\n"); + return 0; + } + + id qbuf = ds4_gpu_tensor_buffer(q); + id rawbuf = ds4_gpu_tensor_buffer(raw_kv); + id compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf; + id maskbuf = use_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf; + id headsbuf = ds4_gpu_tensor_buffer(heads); + const uint64_t comp_mask_bytes = use_mask ? (uint64_t)n_comp * sizeof(float) : 0u; + if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(raw_kv) < raw_bytes || + (n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) || + (use_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) || + ds4_gpu_tensor_bytes(heads) < q_bytes) { + fprintf(stderr, "ds4: Metal graph attention heads received undersized buffers\n"); + return 0; + } + + uint64_t sinks_inner = 0; + id sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, sinks_offset, sink_bytes, &sinks_inner); + if (!sinks_buf) return 0; + + if (n_comp == 0) { + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_raw_heads(cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + n_raw, + raw_cap, + raw_start, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph raw attention heads")) return 0; + return 1; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!ds4_gpu_encode_flash_attention_gathered_heads(cb, + heads, + sinks_buf, + (NSUInteger)sinks_inner, + q, + raw_kv, + n_raw, + raw_cap, + raw_start, + comp_kv, + comp_kv_f16, + n_comp, + comp_mask, + use_mask, + n_head, + head_dim)) { + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "graph attention heads")) return 0; + } + + return 1; +} + +int ds4_gpu_swiglu_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *gate, + const ds4_gpu_tensor *up, + uint32_t n, + float clamp, + float weight) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !gate || !up || n == 0) return 0; + if (!isfinite(clamp) || clamp < 0.0f || !isfinite(weight)) return 0; + + @autoreleasepool { + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t bytes = (uint64_t)n * sizeof(float); + if (!gatebuf || !upbuf || !outbuf || + ds4_gpu_tensor_bytes(gate) < bytes || + ds4_gpu_tensor_bytes(up) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal SwiGLU received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glu_args args = { + .ne00 = (int32_t)n, + .nb01 = (uint64_t)n * sizeof(float), + .ne10 = (int32_t)n, + .nb11 = (uint64_t)n * sizeof(float), + .ne0 = (int32_t)n, + .nb1 = (uint64_t)n * sizeof(float), + .i00 = 0, + .i10 = 0, + .alpha = weight, + .limit = clamp, + }; + NSUInteger nth = g_swiglu_flat_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + if (nth == 0u) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_swiglu_flat_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:1]; + [enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "SwiGLU")) return 0; + } + + return 1; +} + +int ds4_gpu_add_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + uint32_t n) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !a || !b || n == 0) return 0; + + @autoreleasepool { + id abuf = ds4_gpu_tensor_buffer(a); + id bbuf = ds4_gpu_tensor_buffer(b); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t bytes = (uint64_t)n * sizeof(float); + if (!abuf || !bbuf || !outbuf || + ds4_gpu_tensor_bytes(a) < bytes || + ds4_gpu_tensor_bytes(b) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal tensor add received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_add_flat_args args = { .n = n }; + NSUInteger nth = g_add2_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + if (nth == 0u) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_add2_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1]; + [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "tensor add")) return 0; + } + + return 1; +} + +int ds4_gpu_add3_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *a, + const ds4_gpu_tensor *b, + const ds4_gpu_tensor *c, + uint32_t n) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !a || !b || !c || n == 0) return 0; + + @autoreleasepool { + id abuf = ds4_gpu_tensor_buffer(a); + id bbuf = ds4_gpu_tensor_buffer(b); + id cbuf = ds4_gpu_tensor_buffer(c); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t bytes = (uint64_t)n * sizeof(float); + if (!abuf || !bbuf || !cbuf || !outbuf || + ds4_gpu_tensor_bytes(a) < bytes || + ds4_gpu_tensor_bytes(b) < bytes || + ds4_gpu_tensor_bytes(c) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal tensor add3 received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_add_flat_args args = { .n = n }; + NSUInteger nth = g_add3_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + if (nth == 0u) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_add3_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1]; + [enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2]; + [enc setBuffer:cbuf offset:ds4_gpu_tensor_offset(c) atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "tensor add3")) return 0; + } + + return 1; +} + +typedef struct { + uint32_t width; + uint32_t rows; + uint32_t layer; + uint32_t n_threads; + float scale; +} ds4_gpu_directional_steering_project_args; + +int ds4_gpu_directional_steering_project_tensor( + ds4_gpu_tensor *x, + const ds4_gpu_tensor *directions, + uint32_t layer, + uint32_t width, + uint32_t rows, + float scale) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || !directions || width == 0 || rows == 0 || scale == 0.0f) return 0; + + @autoreleasepool { + id pipeline = + ds4_gpu_get_pipeline("kernel_dsv4_directional_steering_project_f32"); + if (!pipeline) return 0; + + id xbuf = ds4_gpu_tensor_buffer(x); + id dbuf = ds4_gpu_tensor_buffer(directions); + const uint64_t x_bytes = (uint64_t)width * rows * sizeof(float); + const uint64_t dir_bytes = (uint64_t)(layer + 1u) * width * sizeof(float); + if (!xbuf || !dbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(directions) < dir_bytes) { + fprintf(stderr, "ds4: Metal directional steering received undersized buffers\n"); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + while (nth > width && nth > 1u) nth >>= 1; + if (nth == 0) nth = 1; + + ds4_gpu_directional_steering_project_args args = { + .width = width, + .rows = rows, + .layer = layer, + .n_threads = (uint32_t)nth, + .scale = scale, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc setBuffer:dbuf offset:ds4_gpu_tensor_offset(directions) atIndex:2]; + [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "directional steering")) return 0; + } + + return 1; +} + +static NSUInteger ds4_gpu_bin_threads(uint32_t width, id pipeline) { + NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup; + if (nth_max > 256u) nth_max = 256u; + NSUInteger nth = 1u; + while (2u * nth < (NSUInteger)width && nth < nth_max) nth *= 2u; + return nth ? nth : 1u; +} + +static int ds4_gpu_encode_unary_f32_rows( + id cb, + id pipeline, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t width, + uint32_t rows, + int c4, + float min, + float max) { + if (!cb || !pipeline || !src || !dst || width == 0 || rows == 0) return 0; + if (c4 && (width & 3u) != 0) return 0; + + ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(width, rows, c4, 0.0f, 0.0f); + args.min = min; + args.max = max; + + NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup; + if (nth_max > 256u) nth_max = 256u; + NSUInteger nth = (NSUInteger)args.ne00; + if (nth > nth_max) nth = nth_max; + if (nth == 0) nth = 1u; + const NSUInteger nk0 = ((NSUInteger)args.ne00 + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nk0 * (NSUInteger)args.ne01, + (NSUInteger)args.ne02, + (NSUInteger)args.ne03) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_bin_f32_rows( + id cb, + id pipeline, + const ds4_gpu_bin_args *args, + id a, + NSUInteger a_off, + id b, + NSUInteger b_off, + id out, + NSUInteger out_off) { + if (!cb || !pipeline || !args || !a || !b || !out || args->ne0 <= 0 || args->ne1 <= 0) { + return 0; + } + + const NSUInteger nth = ds4_gpu_bin_threads((uint32_t)args->ne0, pipeline); + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:a offset:a_off atIndex:1]; + [enc setBuffer:b offset:b_off atIndex:2]; + [enc setBuffer:out offset:out_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)args->ne1, + (NSUInteger)args->ne2, + (NSUInteger)args->ne3) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static ds4_gpu_bin_args ds4_gpu_make_bin_rowwise_scalar_args(uint32_t width, uint32_t rows) { + const uint64_t lhs_row_bytes = (uint64_t)width * sizeof(float); + const uint64_t rhs_row_bytes = sizeof(float); + return (ds4_gpu_bin_args) { + .ne00 = (int32_t)width, + .ne01 = (int32_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = lhs_row_bytes, + .nb02 = (uint64_t)rows * lhs_row_bytes, + .nb03 = (uint64_t)rows * lhs_row_bytes, + .ne10 = 1, + .ne11 = (int32_t)rows, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = rhs_row_bytes, + .nb12 = (uint64_t)rows * rhs_row_bytes, + .nb13 = (uint64_t)rows * rhs_row_bytes, + .ne0 = (int32_t)width, + .ne1 = (int32_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = lhs_row_bytes, + .nb2 = (uint64_t)rows * lhs_row_bytes, + .nb3 = (uint64_t)rows * lhs_row_bytes, + .offs = 0, + .o1 = { 0 }, + }; +} + +static ds4_gpu_mul_mv_id_args ds4_gpu_make_mul_mv_id_args( + uint32_t src0_cols, + uint32_t src0_rows, + uint32_t src0_experts, + uint64_t src0_row_bytes, + uint64_t src0_expert_bytes, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens, + uint32_t nr0) { + const uint64_t src1_row_bytes = (uint64_t)src0_cols * sizeof(float); + const uint64_t src0_blocks = src0_cols / 256u; + const uint64_t src0_block_bytes = src0_blocks ? src0_row_bytes / src0_blocks : 1u; + return (ds4_gpu_mul_mv_id_args) { + .nei0 = (int32_t)selected_experts, + .nei1 = (int32_t)n_tokens, + .nbi1 = (uint64_t)selected_experts * sizeof(int32_t), + .ne00 = (int32_t)src0_cols, + .ne01 = (int32_t)src0_rows, + .ne02 = (int32_t)src0_experts, + .nb00 = src0_block_bytes, + .nb01 = src0_row_bytes, + .nb02 = src0_expert_bytes, + .ne10 = (int32_t)src0_cols, + .ne11 = (int32_t)src1_expert_rows, + .ne12 = (int32_t)n_tokens, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = src1_row_bytes, + .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes, + .ne0 = (int32_t)src0_rows, + .ne1 = (int32_t)selected_experts, + .nb1 = (uint64_t)src0_rows * sizeof(float), + .nr0 = (int32_t)nr0, + }; +} + +static ds4_gpu_mul_mm_id_map_args ds4_gpu_make_mul_mm_id_map_args( + uint32_t src0_cols, + uint32_t src0_experts, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens) { + const uint64_t src1_row_bytes = (uint64_t)src0_cols * sizeof(float); + return (ds4_gpu_mul_mm_id_map_args) { + .ne02 = (int32_t)src0_experts, + .ne10 = (int32_t)src0_cols, + .ne11 = (int32_t)src1_expert_rows, + .nb11 = src1_row_bytes, + .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes, + .ne21 = (int32_t)n_tokens, + .ne20 = (int32_t)selected_experts, + .nb21 = (uint64_t)selected_experts * sizeof(int32_t), + }; +} + +static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args( + uint32_t src0_cols, + uint32_t src0_rows, + uint32_t src0_experts, + uint64_t src0_row_bytes, + uint64_t src0_expert_bytes, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens) { + return ds4_gpu_make_mul_mm_id_args_src1_size(src0_cols, + src0_rows, + src0_experts, + src0_row_bytes, + src0_expert_bytes, + src1_expert_rows, + selected_experts, + n_tokens, + sizeof(float)); +} + +static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args_src1_size( + uint32_t src0_cols, + uint32_t src0_rows, + uint32_t src0_experts, + uint64_t src0_row_bytes, + uint64_t src0_expert_bytes, + uint32_t src1_expert_rows, + uint32_t selected_experts, + uint32_t n_tokens, + uint32_t src1_elem_size) { + const uint64_t src1_row_bytes = (uint64_t)src0_cols * src1_elem_size; + return (ds4_gpu_mul_mm_id_args) { + .ne00 = (int32_t)src0_cols, + .ne02 = (int32_t)src0_experts, + .nb01 = src0_row_bytes, + .nb02 = src0_expert_bytes, + .nb03 = (uint64_t)src0_experts * src0_expert_bytes, + .ne11 = (int32_t)src1_expert_rows, + .nb10 = src1_elem_size, + .nb11 = src1_row_bytes, + .nb12 = (uint64_t)src1_expert_rows * src1_row_bytes, + .nb13 = (uint64_t)n_tokens * (uint64_t)src1_expert_rows * src1_row_bytes, + .ne20 = (int32_t)selected_experts, + .ne21 = (int32_t)n_tokens, + .ne0 = (int32_t)src0_rows, + .ne1 = (int32_t)selected_experts, + .r2 = 1, + .r3 = 1, + }; +} + +static uint32_t ds4_gpu_routed_mv_nr0(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q8_0: return 2; + case DS4_METAL_TENSOR_Q8_K: return 2; + case DS4_METAL_TENSOR_Q4_K: return 2; + case DS4_METAL_TENSOR_Q2_K: + case DS4_METAL_TENSOR_IQ2_XXS: return 4; + default: return 0; + } +} + +static const char *ds4_gpu_metal_tensor_type_name(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_IQ2_XXS: return "iq2_xxs"; + case DS4_METAL_TENSOR_Q2_K: return "q2_k"; + case DS4_METAL_TENSOR_Q4_K: return "q4_k"; + case DS4_METAL_TENSOR_Q5_K: return "q5_k"; + case DS4_METAL_TENSOR_Q6_K: return "q6_k"; + default: return "unknown"; + } +} + +static const char *ds4_gpu_trim_env_value(const char *env, size_t *len_out) { + if (len_out) *len_out = 0; + if (!env) return NULL; + + while (isspace((unsigned char)*env)) env++; + size_t n = strlen(env); + while (n > 0 && isspace((unsigned char)env[n - 1])) n--; + if (len_out) *len_out = n; + return env; +} + +static bool ds4_gpu_profile_layer_value_match(const char *env, uint32_t layer_index) { + size_t env_len = 0; + env = ds4_gpu_trim_env_value(env, &env_len); + if (!env || env_len == 0) return true; + if (ds4_gpu_env_value_eq(env, env_len, "all")) return true; + + const char *p = env; + const char *end_env = env + env_len; + while (p < end_env) { + while (p < end_env && (*p == ' ' || *p == '\t' || *p == ',')) p++; + if (p >= end_env) break; + + char *end = NULL; + const unsigned long first = strtoul(p, &end, 10); + if (end == p || end > end_env || first > UINT32_MAX) return false; + + unsigned long last = first; + p = end; + if (p < end_env && *p == '-') { + p++; + last = strtoul(p, &end, 10); + if (end == p || end > end_env || last > UINT32_MAX) return false; + p = end; + } + + if (first <= layer_index && layer_index <= last) return true; + while (p < end_env && (*p == ' ' || *p == '\t')) p++; + if (p < end_env && *p != ',') return false; + } + return false; +} + +static bool ds4_gpu_stage_profile_enabled_for_layer(const char *flag_env_name, + const char *layer_env_name, + uint32_t layer_index) { + size_t flag_len = 0; + const char *flag = ds4_gpu_trim_env_value(getenv(flag_env_name), &flag_len); + if (!flag) return false; + + size_t layer_len = 0; + const char *layer = ds4_gpu_trim_env_value(getenv(layer_env_name), &layer_len); + const bool has_layer_filter = layer && layer_len != 0; + + if (flag_len != 0) { + if (ds4_gpu_env_value_eq(flag, flag_len, "0") || + ds4_gpu_env_value_eq(flag, flag_len, "false") || + ds4_gpu_env_value_eq(flag, flag_len, "no") || + ds4_gpu_env_value_eq(flag, flag_len, "off")) { + return false; + } + if (!has_layer_filter && + !ds4_gpu_env_value_eq(flag, flag_len, "1") && + !ds4_gpu_env_value_eq(flag, flag_len, "true") && + !ds4_gpu_env_value_eq(flag, flag_len, "yes") && + !ds4_gpu_env_value_eq(flag, flag_len, "on") && + !ds4_gpu_env_value_eq(flag, flag_len, "all")) { + return ds4_gpu_profile_layer_value_match(flag, layer_index); + } + } + + return ds4_gpu_profile_layer_value_match(layer, layer_index); +} + +static NSUInteger ds4_gpu_routed_mv_smem(uint32_t type) { + if (type == DS4_METAL_TENSOR_Q8_0) { + return 32u * 2u * sizeof(float); + } + if (type == DS4_METAL_TENSOR_IQ2_XXS) { + return 256u * sizeof(uint64_t) + 128u * sizeof(uint8_t); + } + return 0; +} + +static NSUInteger ds4_gpu_routed_mv_nsg(uint32_t type) { + return type == DS4_METAL_TENSOR_Q8_0 ? 4u : 2u; +} + +static bool ds4_gpu_routed_mv_rows_per_group_is_nr0(uint32_t type) { + return type == DS4_METAL_TENSOR_Q8_0; +} + +static id ds4_gpu_routed_mv_pipeline(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q8_0: + return ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_0_f32", 4); + case DS4_METAL_TENSOR_Q8_K: + return ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_K_f32", 2); + case DS4_METAL_TENSOR_IQ2_XXS: return g_moe_mul_mv_id_iq2_xxs_pipeline; + case DS4_METAL_TENSOR_Q2_K: return g_moe_mul_mv_id_q2_k_pipeline; + case DS4_METAL_TENSOR_Q4_K: return g_moe_mul_mv_id_q4_k_pipeline; + default: return nil; + } +} + +static id ds4_gpu_routed_mm_pipeline(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q8_0: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false); + case DS4_METAL_TENSOR_Q8_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_K_f32", false); + case DS4_METAL_TENSOR_IQ2_XXS: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f32", false); + case DS4_METAL_TENSOR_Q2_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q2_K_f32", false); + case DS4_METAL_TENSOR_Q4_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false); + case DS4_METAL_TENSOR_Q5_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f32", false); + case DS4_METAL_TENSOR_Q6_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f32", false); + default: + return nil; + } +} + +static id ds4_gpu_routed_mm_addr_pipeline(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q2_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f32", false); + case DS4_METAL_TENSOR_Q4_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f32", false); + default: + return nil; + } +} + +static id ds4_gpu_routed_mm_f16_rhs_pipeline(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q8_0: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f16", false); + case DS4_METAL_TENSOR_Q8_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_K_f16", false); + case DS4_METAL_TENSOR_IQ2_XXS: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f16", false); + case DS4_METAL_TENSOR_Q2_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q2_K_f16", false); + case DS4_METAL_TENSOR_Q4_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f16", false); + case DS4_METAL_TENSOR_Q5_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f16", false); + case DS4_METAL_TENSOR_Q6_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f16", false); + default: + return nil; + } +} + +static id ds4_gpu_routed_mm_addr_f16_rhs_pipeline(uint32_t type) { + switch (type) { + case DS4_METAL_TENSOR_Q2_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f16", false); + case DS4_METAL_TENSOR_Q4_K: + return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f16", false); + default: + return nil; + } +} + +static int ds4_gpu_encode_mul_mv_id( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:src0 offset:src0_off atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_attn_out_low_q8_direct( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !src0 || !src1 || !dst || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { + return 0; + } + + /* Two row conventions in the classic matvec family: Q8_0 k-splits one + * nr0-row group across all nsg simdgroups (cross-simdgroup reduce), so a + * threadgroup covers nr0 rows; Q4_K gives each simdgroup its own nr0 + * rows, covering nr0*nsg. Dispatching Q8 with the Q4 stride leaves + * (nsg-1)/nsg of the output rows unwritten. */ + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? + (NSUInteger)args->nr0 : (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:src0 offset:src0_off atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_id_pair( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0_a, + NSUInteger src0_a_off, + id src0_b, + NSUInteger src0_b_off, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:src0_a offset:src0_a_off atIndex:1]; + [enc setBuffer:src0_b offset:src0_b_off atIndex:2]; + [enc setBuffer:src1 offset:src1_off atIndex:3]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:4]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:5]; + [enc setBuffer:ids offset:ids_off atIndex:6]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_id_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + id src0_a, + NSUInteger src0_a_off, + id src0_b, + NSUInteger src0_b_off, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || + !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBuffer:src0_a offset:src0_a_off atIndex:2]; + [enc setBuffer:src0_b offset:src0_b_off atIndex:3]; + [enc setBuffer:src1 offset:src1_off atIndex:4]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; + [enc setBuffer:ids offset:ids_off atIndex:8]; + [enc setBuffer:weights offset:weights_off atIndex:9]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_table_q4_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + DS4MetalQ4ExpertTable *gate_table, + DS4MetalQ4ExpertTable *up_table, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0, + bool queue_residency) { + if (!cb || !pipeline || !args || !act || !gate_table || !up_table || + !gate_table.argumentBuffer || !up_table.argumentBuffer || + !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBuffer:gate_table.argumentBuffer offset:0 atIndex:2]; + [enc setBuffer:up_table.argumentBuffer offset:0 atIndex:3]; + [enc setBuffer:src1 offset:src1_off atIndex:4]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; + [enc setBuffer:ids offset:ids_off atIndex:8]; + [enc setBuffer:weights offset:weights_off atIndex:9]; + if (!ds4_gpu_bind_q4_expert_table_anchors(enc, gate_table, 10, 6) || + !ds4_gpu_bind_q4_expert_table_anchors(enc, up_table, 16, 6)) { + ds4_gpu_end_compute_encoder(cb, enc); + return 0; + } + ds4_gpu_use_q4_expert_table_resources(cb, enc, gate_table, queue_residency); + ds4_gpu_use_q4_expert_table_resources(cb, enc, up_table, queue_residency); + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_q4_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + DS4MetalQ4ExpertTable *gate_table, + DS4MetalQ4ExpertTable *up_table, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !gate_table || !up_table || + !gate_table.addressBuffer || !up_table.addressBuffer || + !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBuffer:gate_table.addressBuffer offset:0 atIndex:2]; + [enc setBuffer:up_table.addressBuffer offset:0 atIndex:3]; + [enc setBuffer:src1 offset:src1_off atIndex:4]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; + [enc setBuffer:ids offset:ids_off atIndex:8]; + [enc setBuffer:weights offset:weights_off atIndex:9]; + if (!ds4_gpu_bind_q4_expert_table_anchors(enc, gate_table, 10, 6) || + !ds4_gpu_bind_q4_expert_table_anchors(enc, up_table, 16, 6)) { + ds4_gpu_end_compute_encoder(cb, enc); + return 0; + } + if (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL) { + ds4_gpu_use_q4_expert_table_resources(cb, enc, gate_table, false); + ds4_gpu_use_q4_expert_table_resources(cb, enc, up_table, false); + } + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_table_q4_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + DS4MetalQ4ExpertTable *table, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool queue_residency) { + if (!cb || !pipeline || !args || !table || !table.argumentBuffer || + !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:table.argumentBuffer offset:0 atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + if (!ds4_gpu_bind_q4_expert_table_anchors(enc, table, 5, 6)) { + ds4_gpu_end_compute_encoder(cb, enc); + return 0; + } + ds4_gpu_use_q4_expert_table_resources(cb, enc, table, queue_residency); + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_q4_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + DS4MetalQ4ExpertTable *table, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !table || !table.addressBuffer || + !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:table.addressBuffer offset:0 atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + if (!ds4_gpu_bind_q4_expert_table_anchors(enc, table, 5, 6)) { + ds4_gpu_end_compute_encoder(cb, enc); + return 0; + } + if (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL) { + ds4_gpu_use_q4_expert_table_resources(cb, enc, table, false); + } + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static uint32_t ds4_gpu_q4_expert_group_size(uint32_t n_total_expert) { + uint32_t group_size = 32; + const char *env = getenv("DS4_METAL_Q4_EXPERT_GROUP_SIZE"); + if (env && env[0]) { + char *end = NULL; + unsigned long v = strtoul(env, &end, 10); + if (end != env && *end == '\0' && v > 0 && v <= UINT32_MAX) { + group_size = (uint32_t)v; + } + } + if (group_size == 0) group_size = 1; + if (group_size > n_total_expert) group_size = n_total_expert; + return group_size; +} + +static int ds4_gpu_encode_mul_mv_group_q4_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + const ds4_gpu_moe_expert_group_args *group, + id src0_a, + NSUInteger src0_a_off, + id src0_b, + NSUInteger src0_b_off, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !group || + !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0 || + group->expert_count == 0) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBytes:group length:sizeof(*group) atIndex:2]; + [enc setBuffer:src0_a offset:src0_a_off atIndex:3]; + [enc setBuffer:src0_b offset:src0_b_off atIndex:4]; + [enc setBuffer:src1 offset:src1_off atIndex:5]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:6]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:7]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:8]; + [enc setBuffer:ids offset:ids_off atIndex:9]; + [enc setBuffer:weights offset:weights_off atIndex:10]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group_q4_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_moe_expert_group_args *group, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !group || !src0 || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + group->expert_count == 0) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:group length:sizeof(*group) atIndex:1]; + [enc setBuffer:src0 offset:src0_off atIndex:2]; + [enc setBuffer:src1 offset:src1_off atIndex:3]; + [enc setBuffer:dst offset:dst_off atIndex:4]; + [enc setBuffer:ids offset:ids_off atIndex:5]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_id_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + id add_in, + NSUInteger add_in_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || + args->nei0 <= 0 || args->nei0 > 8 || args->nei1 <= 0) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:src0 offset:src0_off atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + [enc setBuffer:(add_in ? add_in : dst) offset:(add_in ? add_in_off : dst_off) atIndex:5]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_q4_gather_slots6( + id cb, + id pipeline, + const ds4_gpu_q4_gather_slots6_args *args, + __unsafe_unretained id src_groups[6], + const NSUInteger src_group_offsets[6], + id ids, + NSUInteger ids_off, + id dst, + NSUInteger dst_off) { + if (!cb || !pipeline || !args || !src_groups || !src_group_offsets || !ids || !dst || + args->expert_bytes == 0 || (args->expert_bytes & 15u) != 0 || + args->group_size == 0 || args->n_slots == 0 || args->n_slots > 6) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if (!src_groups[i]) return 0; + } + + const uint64_t chunks = args->expert_bytes >> 4; + if (chunks == 0 || chunks > NSUIntegerMax) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src_groups[i] offset:src_group_offsets[i] atIndex:1 + i]; + } + [enc setBuffer:ids offset:ids_off atIndex:7]; + [enc setBuffer:dst offset:dst_off atIndex:8]; + const NSUInteger threads = 256u; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)chunks + threads - 1u) / threads, + (NSUInteger)args->n_slots, + 1) + threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_slots6_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + __unsafe_unretained id src0_a[6], + const NSUInteger src0_a_off[6], + __unsafe_unretained id src0_b[6], + const NSUInteger src0_b_off[6], + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off || + !src1 || !dst_a || !dst_b || !dst_mid || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if (!src0_a[i] || !src0_b[i]) return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i]; + } + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:8 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:14]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:15]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:16]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:17]; + [enc setBuffer:weights offset:weights_off atIndex:18]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_slots6_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + __unsafe_unretained id src0[6], + const NSUInteger src0_off[6], + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if (!src0[i]) return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:7]; + [enc setBuffer:dst offset:dst_off atIndex:8]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group6_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + __unsafe_unretained id src0_a[6], + const NSUInteger src0_a_off[6], + __unsafe_unretained id src0_b[6], + const NSUInteger src0_b_off[6], + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off || + !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if (!src0_a[i] || !src0_b[i]) return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i]; + } + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:8 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:14]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:15]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:16]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:17]; + [enc setBuffer:ids offset:ids_off atIndex:18]; + [enc setBuffer:weights offset:weights_off atIndex:19]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group6_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + __unsafe_unretained id src0[6], + const NSUInteger src0_off[6], + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if (!src0[i]) return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 6; i++) { + [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:7]; + [enc setBuffer:dst offset:dst_off atIndex:8]; + [enc setBuffer:ids offset:ids_off atIndex:9]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + ds4_gpu_stream_expert_cache_entry * const *entries, + uint32_t n_entries, + id gate_addrs, + id up_addrs, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0, + id overflow_gate, + id overflow_up) { + if (!cb || !pipeline || !args || !act || !entries || + (n_entries == 0 && !overflow_gate) || + !gate_addrs || !up_addrs || + !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || + args->nei0 <= 0 || args->nei0 > DS4_METAL_MAX_ROUTED_EXPERT_USED || + args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + for (uint32_t i = 0; i < n_entries; i++) { + if (!entries[i] || !entries[i]->gate_buffer || !entries[i]->up_buffer) return 0; + } + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, + n_entries, + 0)) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBuffer:gate_addrs offset:0 atIndex:2]; + [enc setBuffer:up_addrs offset:0 atIndex:3]; + [enc setBuffer:src1 offset:src1_off atIndex:4]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:5]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:6]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:7]; + [enc setBuffer:ids offset:ids_off atIndex:8]; + [enc setBuffer:weights offset:weights_off atIndex:9]; + for (uint32_t i = 0; i < n_entries; i++) { + [enc useResource:entries[i]->gate_buffer usage:MTLResourceUsageRead]; + [enc useResource:entries[i]->up_buffer usage:MTLResourceUsageRead]; + } + /* Overflow experts are addressed straight into the mapped model views + * when a layer's unique selected set exceeds the cache budget. */ + if (overflow_gate) [enc useResource:overflow_gate usage:MTLResourceUsageRead]; + if (overflow_up) [enc useResource:overflow_up usage:MTLResourceUsageRead]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_iq2( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + ds4_gpu_stream_expert_cache_entry * const *entries, + uint32_t n_entries, + id addrs, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !entries || n_entries == 0 || + n_entries > DS4_METAL_MAX_ROUTED_EXPERT_USED || + !addrs || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || + args->nei0 <= 0 || args->nei0 > DS4_METAL_MAX_ROUTED_EXPERT_USED || + args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + for (uint32_t i = 0; i < n_entries; i++) { + if (!entries[i] || !entries[i]->down_buffer) return 0; + } + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, + n_entries, + 0)) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:addrs offset:0 atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + for (uint32_t i = 0; i < n_entries; i++) { + [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead]; + } + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_q2_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + ds4_gpu_stream_expert_cache_entry * const *entries, + uint32_t n_entries, + id addrs, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + id overflow_down) { + if (!cb || !pipeline || !args || !entries || + (n_entries == 0 && !overflow_down) || + !addrs || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + for (uint32_t i = 0; i < n_entries; i++) { + if (!entries[i] || !entries[i]->down_buffer) return 0; + } + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, + n_entries, + 0)) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBuffer:addrs offset:0 atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setBuffer:ids offset:ids_off atIndex:4]; + for (uint32_t i = 0; i < n_entries; i++) { + [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead]; + } + if (overflow_down) [enc useResource:overflow_down usage:MTLResourceUsageRead]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + const ds4_gpu_stream_expert_split_args *split, + ds4_gpu_stream_expert_cache_entry * const entries[6], + id gate_addrs, + id up_addrs, + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !split || !entries || + !gate_addrs || !up_addrs || !src1 || !dst_a || !dst_b || !dst_mid || + !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if ((split->active_mask & (1u << i)) == 0) continue; + if (!entries[i] || !entries[i]->gate_buffer || !entries[i]->up_buffer) return 0; + } + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, + 6, + split->active_mask)) { + return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + [enc setBytes:split length:sizeof(*split) atIndex:2]; + [enc setBuffer:gate_addrs offset:0 atIndex:3]; + [enc setBuffer:up_addrs offset:0 atIndex:4]; + [enc setBuffer:src1 offset:src1_off atIndex:5]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:6]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:7]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:8]; + [enc setBuffer:ids offset:ids_off atIndex:9]; + [enc setBuffer:weights offset:weights_off atIndex:10]; + for (uint32_t i = 0; i < 6; i++) { + if ((split->active_mask & (1u << i)) == 0) continue; + [enc useResource:entries[i]->gate_buffer usage:MTLResourceUsageRead]; + [enc useResource:entries[i]->up_buffer usage:MTLResourceUsageRead]; + } + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_addr_q2_sum6_masked( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_stream_expert_split_args *split, + ds4_gpu_stream_expert_cache_entry * const entries[6], + id addrs, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !split || !entries || !addrs || + !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 <= 0 || args->ne02 > 384) { + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + if ((split->active_mask & (1u << i)) == 0) continue; + if (!entries[i] || !entries[i]->down_buffer) return 0; + } + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries, + 6, + split->active_mask)) { + return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:split length:sizeof(*split) atIndex:1]; + [enc setBuffer:addrs offset:0 atIndex:2]; + [enc setBuffer:src1 offset:src1_off atIndex:3]; + [enc setBuffer:dst offset:dst_off atIndex:4]; + [enc setBuffer:ids offset:ids_off atIndex:5]; + for (uint32_t i = 0; i < 6; i++) { + if ((split->active_mask & (1u << i)) == 0) continue; + [enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead]; + } + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group8_pair_swiglu( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act, + __unsafe_unretained id src0_a[8], + const NSUInteger src0_a_off[8], + __unsafe_unretained id src0_b[8], + const NSUInteger src0_b_off[8], + id src1, + NSUInteger src1_off, + id dst_a, + NSUInteger dst_a_off, + id dst_b, + NSUInteger dst_b_off, + id dst_mid, + NSUInteger dst_mid_off, + id ids, + NSUInteger ids_off, + id weights, + NSUInteger weights_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off || + !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 8; i++) { + if (!src0_a[i] || !src0_b[i]) return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + [enc setBytes:act length:sizeof(*act) atIndex:1]; + for (uint32_t i = 0; i < 8; i++) { + [enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i]; + } + for (uint32_t i = 0; i < 8; i++) { + [enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:10 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:18]; + [enc setBuffer:dst_a offset:dst_a_off atIndex:19]; + [enc setBuffer:dst_b offset:dst_b_off atIndex:20]; + [enc setBuffer:dst_mid offset:dst_mid_off atIndex:21]; + [enc setBuffer:ids offset:ids_off atIndex:22]; + [enc setBuffer:weights offset:weights_off atIndex:23]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group8_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + __unsafe_unretained id src0[8], + const NSUInteger src0_off[8], + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 8; i++) { + if (!src0[i]) return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 8; i++) { + [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:9]; + [enc setBuffer:dst offset:dst_off atIndex:10]; + [enc setBuffer:ids offset:ids_off atIndex:11]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group24_id( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + __unsafe_unretained id src0[24], + const NSUInteger src0_off[24], + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg, + bool rows_per_group_is_nr0) { + if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 24; i++) { + if (!src0[i]) return 0; + } + + const NSUInteger nr0 = (NSUInteger)args->nr0; + const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 24; i++) { + [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:25]; + [enc setBuffer:dst offset:dst_off atIndex:26]; + [enc setBuffer:ids offset:ids_off atIndex:27]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mv_group24_sum6( + id cb, + id pipeline, + const ds4_gpu_mul_mv_id_args *args, + __unsafe_unretained id src0[24], + const NSUInteger src0_off[24], + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off, + NSUInteger threadgroup_bytes, + NSUInteger nsg) { + if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids || + args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 || + args->ne02 != 384) { + return 0; + } + for (uint32_t i = 0; i < 24; i++) { + if (!src0[i]) return 0; + } + + const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg; + const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:args length:sizeof(*args) atIndex:0]; + for (uint32_t i = 0; i < 24; i++) { + [enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i]; + } + [enc setBuffer:src1 offset:src1_off atIndex:25]; + [enc setBuffer:dst offset:dst_off atIndex:26]; + [enc setBuffer:ids offset:ids_off atIndex:27]; + if (threadgroup_bytes != 0) { + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mm_id( + id cb, + id map_pipeline, + id mm_pipeline, + const ds4_gpu_mul_mm_id_map_args *map_args, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + id ids, + NSUInteger ids_off) { + if (!cb || !map_pipeline || !mm_pipeline || !map_args || !mm_args || + !src0 || !src1 || !dst || !ids || + mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || + mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { + return 0; + } + + return ds4_gpu_encode_mul_mm_id_map(cb, + map_pipeline, + map_args, + mm_args, + ids, + ids_off) && + ds4_gpu_encode_mul_mm_id_mapped(cb, + mm_pipeline, + mm_args, + src0, + src0_off, + src1, + src1_off, + dst, + dst_off); +} + +static int ds4_gpu_encode_mul_mm_id_map( + id cb, + id map_pipeline, + const ds4_gpu_mul_mm_id_map_args *map_args, + const ds4_gpu_mul_mm_id_args *mm_args, + id ids, + NSUInteger ids_off) { + if (!cb || !map_pipeline || !map_args || !mm_args || !ids || + mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { + return 0; + } + + const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); + const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); + if (tpe_bytes > NSUIntegerMax - hids_bytes) return 0; + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_id_map_buffer, + &g_moe_id_map_bytes, + tpe_bytes + hids_bytes, + "ds4_moe_id_map")) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:map_pipeline]; + [enc setBytes:map_args length:sizeof(*map_args) atIndex:0]; + [enc setBuffer:ids offset:ids_off atIndex:1]; + [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:2]; + [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:3]; + [enc setThreadgroupMemoryLength:(NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne20 * sizeof(uint16_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake((NSUInteger)mm_args->ne02, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mm_id_mapped_tile( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes) { + if (!cb || !mm_pipeline || !mm_args || !src0 || !src1 || !dst || + !g_moe_id_map_buffer || + mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || + mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { + return 0; + } + /* + * The routed MoE grouped matmul uses the legacy 32-token expert-major tile. + * The removed TensorOps variant was not semantically stable on evals, so keep + * this encoder tied to the tested simdgroup kernel shape. + */ + const NSUInteger tile_n = 32u; + const bool use_resource_hints = + getenv("DS4_METAL_MOE_MM_ID_USE_RESOURCES") != NULL && + getenv("DS4_METAL_DISABLE_MOE_MM_ID_USE_RESOURCES") == NULL; + + const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); + const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); + if (tpe_bytes > NSUIntegerMax - hids_bytes || + g_moe_id_map_bytes < tpe_bytes + hids_bytes) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:mm_pipeline]; + [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; + [enc setBuffer:src0 offset:src0_off atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3]; + [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4]; + [enc setBuffer:dst offset:dst_off atIndex:5]; + if (use_resource_hints) { + [enc useResource:src0 usage:MTLResourceUsageRead]; + } + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + tile_n - 1u) / tile_n, + ((NSUInteger)mm_args->ne0 + 63u) / 64u, + (NSUInteger)mm_args->ne02) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0_addrs, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off, + NSUInteger threadgroup_bytes, + ds4_gpu_stream_expert_cache_entry * const *resources, + uint32_t resource_count, + uint32_t resource_kind, + id overflow_resource) { + if (!cb || !mm_pipeline || !mm_args || !src0_addrs || !src1 || !dst || + !g_moe_id_map_buffer || + mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || + mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { + return 0; + } + + const NSUInteger tile_n = 32u; + const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); + const NSUInteger hids_bytes = + (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); + if (tpe_bytes > NSUIntegerMax - hids_bytes || + g_moe_id_map_bytes < tpe_bytes + hids_bytes) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:mm_pipeline]; + [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; + [enc setBuffer:src0_addrs offset:0 atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3]; + [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4]; + [enc setBuffer:dst offset:dst_off atIndex:5]; + [enc useResource:src0_addrs usage:MTLResourceUsageRead]; + for (uint32_t i = 0; resources && i < resource_count; i++) { + ds4_gpu_stream_expert_cache_entry *entry = resources[i]; + if (!entry) continue; + id b = + resource_kind == 0 ? entry->gate_buffer : + resource_kind == 1 ? entry->up_buffer : + entry->down_buffer; + if (b) [enc useResource:b usage:MTLResourceUsageRead]; + } + if (overflow_resource) { + [enc useResource:overflow_resource usage:MTLResourceUsageRead]; + } + [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + tile_n - 1u) / tile_n, + ((NSUInteger)mm_args->ne0 + 63u) / 64u, + (NSUInteger)mm_args->ne02) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16( + id cb, + id pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + const ds4_gpu_dsv4_moe_swiglu_weight_args *act_args, + id gate_src0, + NSUInteger gate_src0_off, + id up_src0, + NSUInteger up_src0_off, + id src1, + NSUInteger src1_off, + id mid, + NSUInteger mid_off, + id weights, + NSUInteger weights_off) { + if (!cb || !pipeline || !mm_args || !act_args || + !gate_src0 || !up_src0 || !src1 || !mid || !weights || + !g_moe_id_map_buffer || + mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || + mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) { + return 0; + } + + const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t); + const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t); + if (tpe_bytes > NSUIntegerMax - hids_bytes || + g_moe_id_map_bytes < tpe_bytes + hids_bytes) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; + [enc setBytes:act_args length:sizeof(*act_args) atIndex:1]; + [enc setBuffer:gate_src0 offset:gate_src0_off atIndex:2]; + [enc setBuffer:up_src0 offset:up_src0_off atIndex:3]; + [enc setBuffer:src1 offset:src1_off atIndex:4]; + [enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:5]; + [enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:6]; + [enc setBuffer:mid offset:mid_off atIndex:7]; + [enc setBuffer:weights offset:weights_off atIndex:8]; + [enc setThreadgroupMemoryLength:16384u atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + 31u) / 32u, + ((NSUInteger)mm_args->ne0 + 63u) / 64u, + (NSUInteger)mm_args->ne02) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_mul_mm_id_mapped( + id cb, + id mm_pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off) { + return ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + mm_pipeline, + mm_args, + src0, + src0_off, + src1, + src1_off, + dst, + dst_off, + 8192u); +} + +static int ds4_gpu_encode_attn_out_low_q8_mpp( + id cb, + id pipeline, + const ds4_gpu_mul_mm_id_args *mm_args, + id src0, + NSUInteger src0_off, + id src1, + NSUInteger src1_off, + id dst, + NSUInteger dst_off) { + if (!cb || !pipeline || !mm_args || !src0 || !src1 || !dst || + mm_args->ne00 <= 0 || mm_args->ne0 <= 0 || + mm_args->ne02 <= 0 || mm_args->ne1 <= 0 || mm_args->ne21 <= 0) { + return 0; + } + + const uint32_t tile_n = DS4_METAL_ATTN_OUT_MPP_TILE_N; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0]; + [enc setBuffer:src0 offset:src0_off atIndex:1]; + [enc setBuffer:src1 offset:src1_off atIndex:2]; + [enc setBuffer:dst offset:dst_off atIndex:3]; + [enc setThreadgroupMemoryLength:8192u atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + (NSUInteger)tile_n - 1u) / (NSUInteger)tile_n, + ((NSUInteger)mm_args->ne0 + 63u) / 64u, + (NSUInteger)mm_args->ne02) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_swiglu_flat( + id cb, + id gate, + NSUInteger gate_off, + id up, + NSUInteger up_off, + id out, + NSUInteger out_off, + uint32_t n) { + if (!cb || !gate || !up || !out || n == 0) return 0; + + ds4_gpu_glu_args args = { + .ne00 = (int32_t)n, + .nb01 = (uint64_t)n * sizeof(float), + .ne10 = (int32_t)n, + .nb11 = (uint64_t)n * sizeof(float), + .ne0 = (int32_t)n, + .nb1 = (uint64_t)n * sizeof(float), + .i00 = 0, + .i10 = 0, + .alpha = 1.0f, + .limit = 0.0f, + }; + NSUInteger nth = g_swiglu_flat_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n) nth = (NSUInteger)n; + if (nth == 0u) nth = 1u; + const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_swiglu_flat_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate offset:gate_off atIndex:1]; + [enc setBuffer:up offset:up_off atIndex:2]; + [enc setBuffer:out offset:out_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_moe_swiglu_weight( + id cb, + id gate, + NSUInteger gate_off, + id up, + NSUInteger up_off, + id mid, + NSUInteger mid_off, + id weights, + NSUInteger weights_off, + uint32_t width, + uint32_t rows, + float clamp_value, + bool mid_f16) { + if (!cb || !gate || !up || !mid || !weights || width == 0 || rows == 0) return 0; + + id pipeline = + ds4_gpu_get_pipeline(mid_f16 ? "kernel_dsv4_moe_swiglu_weight_f16" : + "kernel_dsv4_moe_swiglu_weight"); + if (!pipeline) return 0; + + ds4_gpu_dsv4_moe_swiglu_weight_args args = { + .width = width, + .rows = rows, + .gate_row_stride = (uint64_t)width * sizeof(float), + .up_row_stride = (uint64_t)width * sizeof(float), + .mid_row_stride = (uint64_t)width * (mid_f16 ? sizeof(uint16_t) : sizeof(float)), + .weight_stride = sizeof(float), + .write_clamped = getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL ? 1u : 0u, + .clamp_value = clamp_value, + }; + + NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > width) nth = width; + if (nth == 0) nth = 1u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate offset:gate_off atIndex:1]; + [enc setBuffer:up offset:up_off atIndex:2]; + [enc setBuffer:mid offset:mid_off atIndex:3]; + [enc setBuffer:weights offset:weights_off atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_moe_sum6( + id cb, + id experts, + NSUInteger experts_off, + id out, + NSUInteger out_off, + uint32_t out_dim, + uint32_t n_tokens) { + if (!cb || !experts || !out || out_dim == 0 || n_tokens == 0) return 0; + + if (!g_moe_sum6_pipeline) return 0; + + const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float); + ds4_gpu_dsv4_moe_sum6_args args = { + .width = out_dim, + .tokens = n_tokens, + .src_token_stride = 6u * out_row_bytes, + .dst_token_stride = out_row_bytes, + }; + + NSUInteger nth = g_moe_sum6_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > out_dim) nth = out_dim; + if (nth == 0) nth = 1u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_moe_sum6_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:experts offset:experts_off atIndex:1]; + [enc setBuffer:out offset:out_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_moe_sum8( + id cb, + id experts, + NSUInteger experts_off, + id out, + NSUInteger out_off, + uint32_t out_dim, + uint32_t n_tokens) { + if (!cb || !experts || !out || out_dim == 0 || n_tokens == 0) return 0; + + if (!g_moe_sum8_pipeline) return 0; + + const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float); + ds4_gpu_dsv4_moe_sum6_args args = { + .width = out_dim, + .tokens = n_tokens, + .src_token_stride = 8u * out_row_bytes, + .dst_token_stride = out_row_bytes, + }; + + NSUInteger nth = g_moe_sum8_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > out_dim) nth = out_dim; + if (nth == 0) nth = 1u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_moe_sum8_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:experts offset:experts_off atIndex:1]; + [enc setBuffer:out offset:out_off atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static ds4_gpu_bin_args ds4_gpu_make_moe_add_args( + uint32_t out_dim, + uint32_t n_tokens, + uint64_t src0_token_stride, + uint64_t src1_token_stride, + uint64_t dst_token_stride) { + return (ds4_gpu_bin_args) { + .ne00 = (int32_t)out_dim, + .ne01 = (int32_t)n_tokens, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src0_token_stride, + .nb02 = (uint64_t)n_tokens * src0_token_stride, + .nb03 = (uint64_t)n_tokens * src0_token_stride, + .ne10 = (int32_t)out_dim, + .ne11 = (int32_t)n_tokens, + .ne12 = 1, + .ne13 = 1, + .nb10 = sizeof(float), + .nb11 = src1_token_stride, + .nb12 = (uint64_t)n_tokens * src1_token_stride, + .nb13 = (uint64_t)n_tokens * src1_token_stride, + .ne0 = (int32_t)out_dim, + .ne1 = (int32_t)n_tokens, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = dst_token_stride, + .nb2 = (uint64_t)n_tokens * dst_token_stride, + .nb3 = (uint64_t)n_tokens * dst_token_stride, + .offs = 0, + .o1 = { 0 }, + }; +} + +static int ds4_gpu_encode_moe_sum_experts( + id cb, + id experts, + NSUInteger experts_off, + id out, + NSUInteger out_off, + uint32_t out_dim, + uint32_t n_expert, + uint32_t n_tokens) { + if (!cb || !experts || !out || out_dim == 0 || n_expert < 2 || n_tokens == 0) return 0; + + const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float); + const uint64_t expert_token_stride = (uint64_t)n_expert * out_row_bytes; + + if (n_expert == 6 && + ds4_gpu_encode_moe_sum6(cb, + experts, + experts_off, + out, + out_off, + out_dim, + n_tokens)) { + return 1; + } + + if (n_expert == 8 && + ds4_gpu_encode_moe_sum8(cb, + experts, + experts_off, + out, + out_off, + out_dim, + n_tokens)) { + return 1; + } + + ds4_gpu_bin_args first = + ds4_gpu_make_moe_add_args(out_dim, n_tokens, expert_token_stride, expert_token_stride, out_row_bytes); + if (!ds4_gpu_encode_bin_f32_rows(cb, + g_add_pipeline, + &first, + experts, + experts_off, + experts, + experts_off + (NSUInteger)out_row_bytes, + out, + out_off)) { + return 0; + } + + ds4_gpu_bin_args accum = + ds4_gpu_make_moe_add_args(out_dim, n_tokens, out_row_bytes, expert_token_stride, out_row_bytes); + for (uint32_t slot = 2; slot < n_expert; slot++) { + if (!ds4_gpu_encode_bin_f32_rows(cb, + g_add_pipeline, + &accum, + out, + out_off, + experts, + experts_off + (NSUInteger)((uint64_t)slot * out_row_bytes), + out, + out_off)) { + return 0; + } + } + return 1; +} + +static int ds4_gpu_encode_get_rows_i32_token_rows( + id cb, + id table, + NSUInteger table_off, + id tokens, + NSUInteger tokens_off, + const int32_t *token_inline, + id selected, + NSUInteger selected_off, + uint32_t hash_rows, + uint32_t n_cols, + uint32_t n_tokens) { + if (!cb || !table || !selected || hash_rows == 0 || n_cols == 0 || n_tokens == 0) return 0; + if (!tokens && !token_inline) return 0; + + const uint64_t table_row_bytes = (uint64_t)n_cols * sizeof(int32_t); + const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t); + ds4_gpu_get_rows_args args = { + .ne00t = (int64_t)n_cols, + .ne00 = (int64_t)n_cols, + .nb01 = table_row_bytes, + .nb02 = (uint64_t)hash_rows * table_row_bytes, + .nb03 = (uint64_t)hash_rows * table_row_bytes, + .ne10 = (int32_t)n_tokens, + .nb10 = sizeof(int32_t), + .nb11 = token_bytes, + .nb12 = token_bytes, + .nb1 = table_row_bytes, + .nb2 = (uint64_t)n_tokens * table_row_bytes, + .nb3 = (uint64_t)n_tokens * table_row_bytes, + }; + + NSUInteger nth = (NSUInteger)n_cols; + const NSUInteger max_threads = g_get_rows_i32_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > max_threads) nth = max_threads; + if (nth == 0) nth = 1u; + const NSUInteger nw0 = ((NSUInteger)n_cols + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_get_rows_i32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:table offset:table_off atIndex:1]; + if (tokens) { + [enc setBuffer:tokens offset:tokens_off atIndex:2]; + } else { + [enc setBytes:token_inline length:sizeof(*token_inline) atIndex:2]; + } + [enc setBuffer:selected offset:selected_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(nw0 * n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_get_rows_f32_router_weights( + id cb, + id probs, + NSUInteger probs_off, + id selected, + NSUInteger selected_off, + id weights, + NSUInteger weights_off, + uint32_t n_expert, + uint32_t n_expert_used, + uint32_t n_tokens) { + if (!cb || !probs || !selected || !weights || n_expert == 0 || n_expert_used == 0 || n_tokens == 0) return 0; + + const uint64_t probs_token_bytes = (uint64_t)n_expert * sizeof(float); + const uint64_t selected_row_bytes = (uint64_t)n_expert_used * sizeof(int32_t); + const uint64_t weights_row_bytes = (uint64_t)n_expert_used * sizeof(float); + ds4_gpu_get_rows_args args = { + .ne00t = 1, + .ne00 = 1, + .nb01 = sizeof(float), + .nb02 = probs_token_bytes, + .nb03 = (uint64_t)n_tokens * probs_token_bytes, + .ne10 = (int64_t)n_expert_used, + .nb10 = sizeof(int32_t), + .nb11 = selected_row_bytes, + .nb12 = (uint64_t)n_tokens * selected_row_bytes, + .nb1 = sizeof(float), + .nb2 = weights_row_bytes, + .nb3 = (uint64_t)n_tokens * weights_row_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_get_rows_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:probs offset:probs_off atIndex:1]; + [enc setBuffer:selected offset:selected_off atIndex:2]; + [enc setBuffer:weights offset:weights_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_expert_used, n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_sum_rows_f32( + id cb, + id src, + NSUInteger src_off, + id dst, + NSUInteger dst_off, + uint32_t width, + uint32_t rows) { + if (!cb || !src || !dst || width == 0 || rows == 0) return 0; + + const uint64_t src_row_bytes = (uint64_t)width * sizeof(float); + ds4_gpu_kargs_sum_rows args = { + .ne00 = (int64_t)width, + .ne01 = (int64_t)rows, + .ne02 = 1, + .ne03 = 1, + .nb00 = sizeof(float), + .nb01 = src_row_bytes, + .nb02 = (uint64_t)rows * src_row_bytes, + .nb03 = (uint64_t)rows * src_row_bytes, + .ne0 = 1, + .ne1 = (int64_t)rows, + .ne2 = 1, + .ne3 = 1, + .nb0 = sizeof(float), + .nb1 = sizeof(float), + .nb2 = (uint64_t)rows * sizeof(float), + .nb3 = (uint64_t)rows * sizeof(float), + }; + + NSUInteger nth = 32u; + const NSUInteger max_threads = g_sum_rows_f32_f32_pipeline.maxTotalThreadsPerThreadgroup; + while (nth < (NSUInteger)args.ne00 && nth < max_threads) nth *= 2u; + if (nth > max_threads) nth = max_threads; + if (nth > (NSUInteger)args.ne00) nth = (NSUInteger)args.ne00; + if (nth == 0) nth = 1u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_sum_rows_f32_f32_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:src offset:src_off atIndex:1]; + [enc setBuffer:dst offset:dst_off atIndex:2]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; +} + +static int ds4_gpu_encode_router_select( + id cb, + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + id logitsbuf, + NSUInteger logits_off, + id biasbuf, + NSUInteger bias_off, + id hashbuf, + NSUInteger hash_off, + id tokensbuf, + NSUInteger tokens_off, + const int32_t *single_token, + uint32_t hash_rows, + uint32_t n_tokens, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + bool has_bias, + bool hash_mode) { + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + const NSUInteger selected_off = ds4_gpu_tensor_offset(selected); + const NSUInteger weights_off = ds4_gpu_tensor_offset(weights); + const NSUInteger probs_off = ds4_gpu_tensor_offset(probs); + + if (!cb || !selectedbuf || !weightsbuf || !probsbuf || !logitsbuf || + n_tokens == 0 || n_expert == 0 || n_expert_used == 0) return 0; + + const NSUInteger probs_bytes = (NSUInteger)n_tokens * (NSUInteger)n_expert * sizeof(float); + const bool flash_router_fast_path = + n_expert == 256u && + n_expert_used == 6u && + fabsf(expert_weight_scale - 1.5f) <= 1.0e-6f; + + int ok = 0; + if (flash_router_fast_path && + !g_quality_mode && n_tokens == 1 && + getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL) { + const bool force_simd_weights_fusion = + getenv("DS4_METAL_ENABLE_ROUTER_SIMD_WEIGHTS_FUSION") != NULL; + const bool use_simd_finalize = + !hash_mode && + g_dsv4_router_finalize_one_simd_pipeline != nil && + g_dsv4_router_finalize_one_simd_pipeline.threadExecutionWidth == 32u && + g_dsv4_router_finalize_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u && + (ds4_gpu_device_name_contains("M3") || + ds4_gpu_device_name_contains("M5") || + getenv("DS4_METAL_ENABLE_ROUTER_SIMD_FINALIZE") != NULL || + force_simd_weights_fusion) && + getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_FINALIZE") == NULL; + const bool use_simd_weights_fusion = + use_simd_finalize && + g_dsv4_router_finalize_weights_one_simd_pipeline != nil && + g_dsv4_router_finalize_weights_one_simd_pipeline.threadExecutionWidth == 32u && + g_dsv4_router_finalize_weights_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u && + (ds4_gpu_device_name_contains("M3") || + ds4_gpu_device_name_contains("M5") || + force_simd_weights_fusion) && + getenv("DS4_METAL_DISABLE_M3_ROUTER_SIMD_WEIGHTS_FUSION") == NULL; + id softplus_sqrt_pipeline = + ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline, + "kernel_dsv4_softplus_sqrt_f32_4"); + id router_finalize_pipeline = + ds4_gpu_hot_pipeline( + use_simd_weights_fusion + ? g_dsv4_router_finalize_weights_one_simd_pipeline + : use_simd_finalize + ? g_dsv4_router_finalize_one_simd_pipeline + : g_dsv4_router_finalize_one_pipeline, + use_simd_weights_fusion + ? "kernel_dsv4_router_finalize_weights_one_simd" + : use_simd_finalize + ? "kernel_dsv4_router_finalize_one_simd" + : "kernel_dsv4_router_finalize_one"); + id router_weights_pipeline = use_simd_weights_fusion + ? nil + : ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline, + "kernel_dsv4_router_weights_one"); + if (!softplus_sqrt_pipeline || !router_finalize_pipeline || + (!use_simd_weights_fusion && !router_weights_pipeline)) return 0; + + ok = ds4_gpu_encode_unary_f32_rows(cb, + softplus_sqrt_pipeline, + logitsbuf, + logits_off, + probsbuf, + probs_off, + n_expert, + 1, + 1, + 0.0f, + 0.0f); + if (!ok) return 0; + + const bool use_token_buffer = single_token == NULL; + ds4_gpu_dsv4_router_select_one_args args = { + .has_bias = has_bias ? 1u : 0u, + .hash_mode = hash_mode ? 1u : 0u, + .use_token_buffer = use_token_buffer ? 1u : 0u, + .token = single_token ? (uint32_t)*single_token : 0u, + .hash_rows = hash_rows, + }; + + const float zero_f32 = 0.0f; + const int32_t zero_i32 = 0; + if ((has_bias && !biasbuf) || + (hash_mode && !hashbuf) || + (use_token_buffer && !tokensbuf)) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:router_finalize_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:probsbuf offset:probs_off atIndex:1]; + if (has_bias) { + [enc setBuffer:biasbuf offset:bias_off atIndex:2]; + } else { + [enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:2]; + } + if (hash_mode) { + [enc setBuffer:hashbuf offset:hash_off atIndex:3]; + } else { + [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:3]; + } + if (use_token_buffer) { + [enc setBuffer:tokensbuf offset:tokens_off atIndex:4]; + } else { + [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:4]; + } + [enc setBuffer:selectedbuf offset:selected_off atIndex:5]; + if (use_simd_weights_fusion) { + [enc setBuffer:weightsbuf offset:weights_off atIndex:6]; + } + const NSUInteger router_finalize_scratch_bytes = use_simd_finalize + ? 2u * (256u * sizeof(float) + 256u * sizeof(int32_t)) + : 256u * sizeof(float) + 256u * sizeof(int32_t); + [enc setThreadgroupMemoryLength:router_finalize_scratch_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (use_simd_weights_fusion) return 1; + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:router_weights_pipeline]; + [enc setBuffer:probsbuf offset:probs_off atIndex:0]; + [enc setBuffer:selectedbuf offset:selected_off atIndex:1]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:2]; + [enc dispatchThreads:MTLSizeMake(6, 1, 1) + threadsPerThreadgroup:MTLSizeMake(6, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + + if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) { + id softplus_sqrt_pipeline = + ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline, + "kernel_dsv4_softplus_sqrt_f32_4"); + ok = softplus_sqrt_pipeline && + ds4_gpu_encode_unary_f32_rows(cb, + softplus_sqrt_pipeline, + logitsbuf, + logits_off, + probsbuf, + probs_off, + n_expert, + 1, + 1, + 0.0f, + 0.0f); + } else { + ok = ds4_gpu_encode_unary_f32_rows(cb, + g_unary_softplus_pipeline, + logitsbuf, + logits_off, + probsbuf, + probs_off, + n_expert, + n_tokens, + 1, + 0.0f, + 0.0f) && + ds4_gpu_encode_unary_f32_rows(cb, + g_unary_sqrt_pipeline, + probsbuf, + probs_off, + probsbuf, + probs_off, + n_expert, + n_tokens, + 1, + 0.0f, + 0.0f); + } + if (!ok) return 0; + + if (hash_mode) { + ok = ds4_gpu_encode_get_rows_i32_token_rows(cb, + hashbuf, + hash_off, + tokensbuf, + tokens_off, + single_token, + selectedbuf, + selected_off, + hash_rows, + n_expert_used, + n_tokens); + } else { + ds4_gpu_tensor *score_tensor = probs; + DS4MetalTensor *selection_view = nil; + + if (has_bias) { + if (!biasbuf || + !ds4_gpu_ensure_scratch_buffer(&g_router_selection_buffer, + &g_router_selection_bytes, + probs_bytes, + "ds4_router_selection")) { + return 0; + } + + ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_expert, n_tokens, n_expert); + ok = ds4_gpu_encode_bin_f32_rows(cb, + g_add_pipeline, + &add_args, + probsbuf, + probs_off, + biasbuf, + bias_off, + g_router_selection_buffer, + 0); + if (!ok) return 0; + + selection_view = [DS4MetalTensor new]; + selection_view.buffer = g_router_selection_buffer; + selection_view.offset = 0; + selection_view.bytes = probs_bytes; + selection_view.owner = 0; + score_tensor = (__bridge ds4_gpu_tensor *)selection_view; + } + + ok = ds4_gpu_indexer_topk_tensor(selected, score_tensor, n_expert, n_tokens, n_expert_used) != 0; + } + if (!ok) return 0; + + const bool use_batch_weights_fusion = + flash_router_fast_path && !g_quality_mode && n_tokens > 1u && + g_dsv4_router_weights_batch_pipeline != nil && + (ds4_gpu_device_name_contains("M3") || + getenv("DS4_METAL_ENABLE_ROUTER_WEIGHTS_BATCH_FUSION") != NULL) && + getenv("DS4_METAL_DISABLE_M3_ROUTER_WEIGHTS_BATCH_FUSION") == NULL && + getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL; + if (use_batch_weights_fusion) { + const float scale = expert_weight_scale; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_dsv4_router_weights_batch_pipeline]; + [enc setBytes:&scale length:sizeof(scale) atIndex:0]; + [enc setBuffer:probsbuf offset:probs_off atIndex:1]; + [enc setBuffer:selectedbuf offset:selected_off atIndex:2]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:3]; + [enc setThreadgroupMemoryLength:40u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(6, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + + if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) { + id router_weights_pipeline = + ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline, + "kernel_dsv4_router_weights_one"); + if (!router_weights_pipeline) return 0; + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:router_weights_pipeline]; + [enc setBuffer:probsbuf offset:probs_off atIndex:0]; + [enc setBuffer:selectedbuf offset:selected_off atIndex:1]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:2]; + [enc dispatchThreads:MTLSizeMake(6, 1, 1) + threadsPerThreadgroup:MTLSizeMake(6, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return 1; + } + + const NSUInteger sum_bytes = (NSUInteger)n_tokens * sizeof(float); + if (!ds4_gpu_ensure_scratch_buffer(&g_router_weight_sum_buffer, + &g_router_weight_sum_bytes, + sum_bytes, + "ds4_router_weight_sum")) { + return 0; + } + + ok = ds4_gpu_encode_get_rows_f32_router_weights(cb, + probsbuf, + probs_off, + selectedbuf, + selected_off, + weightsbuf, + weights_off, + n_expert, + n_expert_used, + n_tokens) && + ds4_gpu_encode_sum_rows_f32(cb, + weightsbuf, + weights_off, + g_router_weight_sum_buffer, + 0, + n_expert_used, + n_tokens) && + ds4_gpu_encode_unary_f32_rows(cb, + g_unary_clamp_pipeline, + g_router_weight_sum_buffer, + 0, + g_router_weight_sum_buffer, + 0, + 1, + n_tokens, + 0, + 6.103515625e-5f, + ds4_gpu_positive_infinity()); + if (!ok) return 0; + + ds4_gpu_bin_args div_args = ds4_gpu_make_bin_rowwise_scalar_args(n_expert_used, n_tokens); + const float scale = expert_weight_scale; + ds4_gpu_bin_args scale_args = ds4_gpu_make_bin_rows_args(n_expert_used, n_tokens, 1); + + ok = ds4_gpu_encode_bin_f32_rows(cb, + g_bin_div_row_pipeline, + &div_args, + weightsbuf, + weights_off, + g_router_weight_sum_buffer, + 0, + weightsbuf, + weights_off); + if (!ok) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_bin_mul_scalar_pipeline]; + [enc setBytes:&scale_args length:sizeof(scale_args) atIndex:0]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:1]; + [enc setBytes:&scale length:sizeof(scale) atIndex:2]; + [enc setBuffer:weightsbuf offset:weights_off atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)scale_args.ne1, + (NSUInteger)scale_args.ne2, + (NSUInteger)scale_args.ne3) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_bin_threads(n_expert_used, g_bin_mul_scalar_pipeline), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + return 1; +} + +int ds4_gpu_glm_kv_lora_rms_norm_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *kv_raw, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !kv_raw || !model_map || + n_tokens == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || + kv_lora_dim > kv_raw_dim || (kv_lora_dim & 3u) != 0 || + !isfinite(eps) || eps < 0.0f) { + return 0; + } + + @autoreleasepool { + id rawbuf = ds4_gpu_tensor_buffer(kv_raw); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float); + const uint64_t weight_bytes = (uint64_t)kv_lora_dim * sizeof(float); + if (!rawbuf || !outbuf || + ds4_gpu_tensor_bytes(kv_raw) < raw_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal GLM KV RMS norm received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal GLM KV RMS norm weight range is outside the mapped model\n"); + return 0; + } + + const bool exact_decode_weight_view = + n_tokens == 1u && + weight_bytes <= (1ull << 20) && + getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL; + uint64_t weight_inner = 0; + id weightbuf = exact_decode_weight_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner); + if (!weightbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_kv_lora_rms_norm_pipeline, + "kernel_glm_kv_lora_rms_norm"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_kv_lora_rms_norm_args args = { + .n_tokens = n_tokens, + .kv_raw_dim = kv_raw_dim, + .kv_lora_dim = kv_lora_dim, + .eps = eps, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(kv_lora_dim); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM KV RMS norm")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_k_b_project_typed_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *kv_norm, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_tokens, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t n_head) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !kv_norm || !model_map || + n_tokens == 0 || kv_lora_dim == 0 || + qk_nope == 0 || n_head == 0) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(out); + id kvbuf = ds4_gpu_tensor_buffer(kv_norm); + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) { + fprintf(stderr, "ds4: Metal GLM k_b projection received unsupported weight type\n"); + return 0; + } + const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim; + const uint64_t weight_bytes = weight_rows * row_bytes; + const uint64_t kv_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float); + if (!outbuf || !kvbuf || + ds4_gpu_tensor_bytes(kv_norm) < kv_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes) { + fprintf(stderr, "ds4: Metal GLM k_b projection received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal GLM k_b projection range is outside the mapped model\n"); + return 0; + } + const NSUInteger q_blocks = ((NSUInteger)qk_nope + 31u) / 32u; + if (q_blocks > 8u) { + fprintf(stderr, "ds4: Metal GLM k_b projection q width is too large for the tiled kernel\n"); + return 0; + } + + uint64_t weight_inner = 0; + id weightbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner); + if (!weightbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_k_b_project_pipeline, + "kernel_glm_k_b_project_q8_0"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_k_b_project_args args = { + .n_tokens = n_tokens, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .n_head = n_head, + .row_bytes = (uint32_t)row_bytes, + .weight_type = weight_type, + .pad1 = 0, + .pad2 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv_norm) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, q_blocks, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM k_b projection")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_k_b_project_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *kv_norm, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t n_head) { + return ds4_gpu_glm_k_b_project_typed_tensor(out, + kv_norm, + model_map, + model_size, + weight_offset, + DS4_METAL_TENSOR_Q8_0, + n_tokens, + kv_lora_dim, + qk_nope, + n_head); +} + +int ds4_gpu_glm_store_compact_kv_tensor( + ds4_gpu_tensor *kv_lora_cache, + ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *kv_norm, + const ds4_gpu_tensor *kv_raw, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!kv_lora_cache || !k_rope_cache || !kv_norm || !kv_raw || + n_tokens == 0 || cache_cap == 0 || + kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 || + kv_lora_dim > kv_raw_dim || + qk_rope > kv_raw_dim - kv_lora_dim || + pos0 > cache_cap || n_tokens > cache_cap - pos0) { + return 0; + } + + @autoreleasepool { + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id kvnormbuf = ds4_gpu_tensor_buffer(kv_norm); + id kvrawbuf = ds4_gpu_tensor_buffer(kv_raw); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t kv_cache_bytes = + (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = + (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t kv_norm_bytes = + (uint64_t)n_tokens * kv_lora_dim * sizeof(float); + const uint64_t kv_raw_bytes = + (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + if (!kvcachebuf || !ropecachebuf || !kvnormbuf || !kvrawbuf || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(kv_norm) < kv_norm_bytes || + ds4_gpu_tensor_bytes(kv_raw) < kv_raw_bytes) { + fprintf(stderr, "ds4: Metal GLM compact KV store received undersized buffers\n"); + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_store_compact_kv_pipeline, + "kernel_glm_store_compact_kv"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_store_compact_kv_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_cap = cache_cap, + .kv_raw_dim = kv_raw_dim, + .kv_lora_dim = kv_lora_dim, + .qk_rope = qk_rope, + .cache_f16 = cache_f16 ? 1u : 0u, + .pad1 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:kvnormbuf offset:ds4_gpu_tensor_offset(kv_norm) atIndex:1]; + [enc setBuffer:kvrawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 2, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM compact KV store")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_qkv_norm_store_compact_kv_tensor( + ds4_gpu_tensor *q_out, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t q_weight_offset, + uint32_t q_n, + ds4_gpu_tensor *kv_lora_cache, + ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *kv_raw, + uint64_t kv_weight_offset, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + bool cache_f16, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!q_out || !q || !kv_lora_cache || !k_rope_cache || !kv_raw || + !model_map || n_tokens == 0 || cache_cap == 0 || + q_n == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 || + (q_n & 3u) != 0 || (kv_lora_dim & 3u) != 0 || + kv_lora_dim > kv_raw_dim || + qk_rope > kv_raw_dim - kv_lora_dim || + pos0 > cache_cap || n_tokens > cache_cap - pos0) { + return 0; + } + + @autoreleasepool { + id qbuf = ds4_gpu_tensor_buffer(q); + id qoutbuf = ds4_gpu_tensor_buffer(q_out); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id kvrawbuf = ds4_gpu_tensor_buffer(kv_raw); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float); + const uint64_t kv_weight_bytes = (uint64_t)kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = + (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = + (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t kv_raw_bytes = + (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + if (!qbuf || !qoutbuf || !kvcachebuf || !ropecachebuf || !kvrawbuf || + ds4_gpu_tensor_bytes(q) < q_row_bytes * n_tokens || + ds4_gpu_tensor_bytes(q_out) < q_row_bytes * n_tokens || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(kv_raw) < kv_raw_bytes) { + fprintf(stderr, "ds4: Metal GLM fused q/kv norm compact store received undersized buffers\n"); + return 0; + } + if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset || + kv_weight_offset > model_size || kv_weight_bytes > model_size - kv_weight_offset) { + fprintf(stderr, "ds4: Metal GLM fused q/kv norm compact store weight range is outside the mapped model\n"); + return 0; + } + + uint64_t q_weight_inner = 0; + id q_weightbuf = ds4_gpu_wrap_model_range(model_map, model_size, + q_weight_offset, q_row_bytes, + &q_weight_inner); + if (!q_weightbuf) return 0; + uint64_t kv_weight_inner = 0; + id kv_weightbuf = ds4_gpu_wrap_model_range(model_map, model_size, + kv_weight_offset, kv_weight_bytes, + &kv_weight_inner); + if (!kv_weightbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_qkv_norm_store_compact_kv_pipeline, + "kernel_glm_qkv_norm_store_compact_kv"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_qkv_norm_store_compact_kv_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_cap = cache_cap, + .q_n = q_n, + .q_n4 = q_n / 4u, + .kv_raw_dim = kv_raw_dim, + .kv_lora_dim = kv_lora_dim, + .kv_lora_n4 = kv_lora_dim / 4u, + .qk_rope = qk_rope, + .cache_f16 = cache_f16 ? 1u : 0u, + .eps = eps, + .pad0 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:q_weightbuf offset:(NSUInteger)q_weight_inner atIndex:2]; + [enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:3]; + [enc setBuffer:kvrawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:4]; + [enc setBuffer:kv_weightbuf offset:(NSUInteger)kv_weight_inner atIndex:5]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:6]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:7]; + [enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 3, 1) + threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM fused q/kv norm compact store")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_store_indexer_k_tensor( + ds4_gpu_tensor *indexer_key_cache, + const ds4_gpu_tensor *raw_k, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t bias_offset, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t n_ctx_orig, + float eps, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!indexer_key_cache || !raw_k || !model_map || + n_tokens == 0 || cache_cap == 0 || + head_dim == 0 || rot_dim == 0 || + rot_dim > head_dim || (rot_dim & 1u) != 0 || + pos0 > cache_cap || n_tokens > cache_cap - pos0) { + return 0; + } + + @autoreleasepool { + id cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache); + id rawbuf = ds4_gpu_tensor_buffer(raw_k); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t cache_bytes = + (uint64_t)cache_cap * head_dim * cache_elem_bytes; + const uint64_t raw_bytes = + (uint64_t)n_tokens * head_dim * sizeof(float); + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + if (!cachebuf || !rawbuf || + ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes || + ds4_gpu_tensor_bytes(raw_k) < raw_bytes) { + fprintf(stderr, "ds4: Metal GLM indexer K store received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || norm_bytes > model_size - weight_offset || + bias_offset > model_size || norm_bytes > model_size - bias_offset) { + fprintf(stderr, "ds4: Metal GLM indexer K norm range is outside the mapped model\n"); + return 0; + } + + uint64_t weight_inner = 0; + uint64_t bias_inner = 0; + id weightbuf = + ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, norm_bytes, + &weight_inner); + id biasbuf = + ds4_gpu_wrap_model_range(model_map, model_size, + bias_offset, norm_bytes, + &bias_inner); + if (!weightbuf || !biasbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_store_indexer_k_pipeline, + "kernel_glm_store_indexer_k"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_store_indexer_k_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_cap = cache_cap, + .head_dim = head_dim, + .rot_dim = rot_dim, + .n_ctx_orig = n_ctx_orig, + .cache_f16 = cache_f16 ? 1u : 0u, + .pad0 = 0, + .eps = eps, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .pad1 = 0.0f, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_k) atIndex:1]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:2]; + [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:3]; + [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:4]; + [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer K store")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_build_kv_cache_tensor( + ds4_gpu_tensor *key_cache, + ds4_gpu_tensor *value_cache, + const ds4_gpu_tensor *kv_raw, + const ds4_gpu_tensor *k_nope, + const ds4_gpu_tensor *value, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!key_cache || !value_cache || !kv_raw || !k_nope || !value || + n_tokens == 0 || cache_cap == 0 || n_head == 0 || + kv_raw_dim == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || value_dim == 0 || + kv_lora_dim + qk_rope > kv_raw_dim || + qk_dim < qk_nope || (qk_rope & 1u) != 0 || + pos0 > cache_cap || n_tokens > cache_cap - pos0 || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id keybuf = ds4_gpu_tensor_buffer(key_cache); + id valbuf = ds4_gpu_tensor_buffer(value_cache); + id rawbuf = ds4_gpu_tensor_buffer(kv_raw); + id knbuf = ds4_gpu_tensor_buffer(k_nope); + id vbuf = ds4_gpu_tensor_buffer(value); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; + const uint64_t cache_value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; + const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + const uint64_t kn_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float); + const uint64_t value_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + if (!keybuf || !valbuf || !rawbuf || !knbuf || !vbuf || + ds4_gpu_tensor_bytes(key_cache) < key_bytes || + ds4_gpu_tensor_bytes(value_cache) < cache_value_bytes || + ds4_gpu_tensor_bytes(kv_raw) < raw_bytes || + ds4_gpu_tensor_bytes(k_nope) < kn_bytes || + ds4_gpu_tensor_bytes(value) < value_bytes) { + fprintf(stderr, "ds4: Metal GLM KV cache builder received undersized buffers\n"); + return 0; + } + + const bool decode_group4 = + n_tokens == 1u && + n_head >= 4u && + ds4_gpu_env_bool("DS4_METAL_DISABLE_GLM_DECODE_KV_GROUP4") <= 0; + id pipeline = + decode_group4 ? + ds4_gpu_hot_pipeline(g_glm_build_kv_cache_decode_group4_pipeline, + "kernel_glm_build_kv_cache_decode_group4") : + ds4_gpu_hot_pipeline(g_glm_build_kv_cache_pipeline, + "kernel_glm_build_kv_cache"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_build_kv_cache_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_cap = cache_cap, + .n_head = n_head, + .kv_raw_dim = kv_raw_dim, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = value_dim, + .n_ctx_orig = n_ctx_orig, + .cache_f16 = cache_f16 ? 1u : 0u, + .pad0 = 0u, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1]; + [enc setBuffer:knbuf offset:ds4_gpu_tensor_offset(k_nope) atIndex:2]; + [enc setBuffer:vbuf offset:ds4_gpu_tensor_offset(value) atIndex:3]; + [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:4]; + [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:5]; + const NSUInteger group_y = decode_group4 ? + ((NSUInteger)n_head + 3u) / 4u : + (NSUInteger)n_head; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, group_y, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM KV cache build")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_build_kv_cache_flash_tensor( + ds4_gpu_tensor *key_cache, + ds4_gpu_tensor *value_cache, + const ds4_gpu_tensor *kv_raw, + const ds4_gpu_tensor *k_nope, + const ds4_gpu_tensor *value, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t n_head, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!key_cache || !value_cache || !kv_raw || !k_nope || !value || + pos0 != 0 || n_tokens == 0 || cache_cap == 0 || n_head == 0 || + kv_raw_dim == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || value_dim == 0 || + kv_lora_dim + qk_rope > kv_raw_dim || + qk_dim < qk_nope || (qk_rope & 1u) != 0 || + n_tokens > cache_cap || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id keybuf = ds4_gpu_tensor_buffer(key_cache); + id valbuf = ds4_gpu_tensor_buffer(value_cache); + id rawbuf = ds4_gpu_tensor_buffer(kv_raw); + id knbuf = ds4_gpu_tensor_buffer(k_nope); + id vbuf = ds4_gpu_tensor_buffer(value); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; + const uint64_t cache_value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; + const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + const uint64_t kn_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float); + const uint64_t value_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + if (!keybuf || !valbuf || !rawbuf || !knbuf || !vbuf || + ds4_gpu_tensor_bytes(key_cache) < key_bytes || + ds4_gpu_tensor_bytes(value_cache) < cache_value_bytes || + ds4_gpu_tensor_bytes(kv_raw) < raw_bytes || + ds4_gpu_tensor_bytes(k_nope) < kn_bytes || + ds4_gpu_tensor_bytes(value) < value_bytes) { + fprintf(stderr, "ds4: Metal GLM staged KV cache builder received undersized buffers\n"); + return 0; + } + + const NSUInteger q_row_bytes_f16 = (NSUInteger)qk_dim * sizeof(uint16_t); + const NSUInteger v_row_bytes_f16 = (NSUInteger)value_dim * sizeof(uint16_t); + const NSUInteger key_f16_offset = 0; + const NSUInteger key_f16_bytes = + (NSUInteger)n_tokens * (NSUInteger)n_head * q_row_bytes_f16; + const NSUInteger value_f16_offset = key_f16_bytes; + const NSUInteger value_f16_bytes = + (NSUInteger)n_tokens * (NSUInteger)n_head * v_row_bytes_f16; + const NSUInteger kv_f16_bytes = key_f16_bytes + value_f16_bytes; + if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_f16_bytes, + "ds4_glm_flash_attn_kv_f16")) { + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_build_kv_cache_flash_pipeline, + "kernel_glm_build_kv_cache_flash"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_build_kv_cache_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_cap = cache_cap, + .n_head = n_head, + .kv_raw_dim = kv_raw_dim, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = value_dim, + .n_ctx_orig = n_ctx_orig, + .cache_f16 = cache_f16 ? 1u : 0u, + .pad0 = 0u, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1]; + [enc setBuffer:knbuf offset:ds4_gpu_tensor_offset(k_nope) atIndex:2]; + [enc setBuffer:vbuf offset:ds4_gpu_tensor_offset(value) atIndex:3]; + [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:4]; + [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:5]; + [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:6]; + [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:7]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM staged KV cache build")) return 0; + } + + return 1; +} + +static int ds4_gpu_glm_attention_flash_tensor_impl( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16, + int kv_pre_staged) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !key_cache || !value_cache || + n_tokens == 0 || cache_len == 0 || cache_cap == 0 || + n_head == 0 || qk_dim != 256u || value_dim != 256u || + cache_len > cache_cap || + pos0 > cache_len || n_tokens > cache_len - pos0 || + cache_len > ds4_gpu_glm_flash_attention_max_cache_len()) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id qbuf = ds4_gpu_tensor_buffer(q); + id keybuf = ds4_gpu_tensor_buffer(key_cache); + id valbuf = ds4_gpu_tensor_buffer(value_cache); + const uint64_t heads_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; + const uint64_t value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; + if (!headsbuf || !qbuf || !keybuf || !valbuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(key_cache) < key_bytes || + ds4_gpu_tensor_bytes(value_cache) < value_bytes) { + fprintf(stderr, "ds4: Metal GLM FlashAttention received undersized buffers\n"); + return 0; + } + const uint64_t key_elems = (uint64_t)cache_len * n_head * qk_dim; + const uint64_t value_elems = (uint64_t)cache_len * n_head * value_dim; + if (key_elems > UINT32_MAX || value_elems > UINT32_MAX) { + return 0; + } + + const uint32_t nqptg = 8; + const uint32_t ncpsg = 64; + const uint32_t nsg = 4; + const bool has_kvpad = (cache_len % ncpsg) != 0; + const bool bc_mask = (n_tokens % nqptg) != 0; + const NSUInteger q_row_bytes = (NSUInteger)qk_dim * sizeof(float); + const NSUInteger q_row_bytes_f16 = (NSUInteger)qk_dim * sizeof(uint16_t); + const NSUInteger v_row_bytes = (NSUInteger)value_dim * sizeof(float); + const NSUInteger v_row_bytes_f16 = (NSUInteger)value_dim * sizeof(uint16_t); + const NSUInteger mask_bytes = (NSUInteger)n_tokens * (NSUInteger)cache_len * sizeof(uint16_t); + const NSUInteger key_f16_offset = 0; + const NSUInteger key_f16_bytes = + (NSUInteger)cache_len * (NSUInteger)n_head * q_row_bytes_f16; + const NSUInteger value_f16_offset = key_f16_bytes; + const NSUInteger value_f16_bytes = + (NSUInteger)cache_len * (NSUInteger)n_head * v_row_bytes_f16; + const NSUInteger kv_f16_bytes = key_f16_bytes + value_f16_bytes; + const NSUInteger pad_bytes = has_kvpad + ? (NSUInteger)ncpsg * ((NSUInteger)n_head * (q_row_bytes_f16 + v_row_bytes_f16) + + (NSUInteger)n_tokens * sizeof(uint16_t)) + : 1u; + const NSUInteger nblk0 = ((NSUInteger)cache_len + ncpsg - 1u) / ncpsg; + const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg; + const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u); + + id mask_buffer = + ds4_gpu_glm_prefill_mask_buffer(pos0, n_tokens, cache_len, mask_bytes); + if (!mask_buffer) return 0; + if (kv_pre_staged) { + if (!g_flash_attn_kv_buffer || g_flash_attn_kv_bytes < kv_f16_bytes) { + fprintf(stderr, "ds4: GLM staged FlashAttention KV scratch is missing\n"); + return 0; + } + } else if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer, + &g_flash_attn_kv_bytes, + kv_f16_bytes, + "ds4_glm_flash_attn_kv_f16")) { + return 0; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer, + &g_flash_attn_pad_bytes, + pad_bytes, + "ds4_glm_flash_attn_pad") || + !ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer, + &g_flash_attn_blk_bytes, + blk_bytes, + "ds4_glm_flash_attn_blk")) { + return 0; + } + + id pad_pipeline = nil; + if (has_kvpad) { + pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg); + if (!pad_pipeline) return 0; + } + id blk_pipeline = + ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg); + id attn_pipeline = + ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk256_dv256", + true, false, false, false, has_kvpad, bc_mask, + (int32_t)qk_dim, + (int32_t)value_dim, + (int32_t)nsg); + if (!blk_pipeline || !attn_pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + if (!kv_pre_staged) { + const bool copied = cache_f16 ? + (ds4_gpu_encode_cpy_f16_f16_3d(cb, + keybuf, + ds4_gpu_tensor_offset(key_cache), + g_flash_attn_kv_buffer, + key_f16_offset, + qk_dim, + cache_len, + n_head, + (uint64_t)n_head * q_row_bytes_f16, + q_row_bytes_f16, + q_row_bytes_f16, + (uint64_t)cache_len * q_row_bytes_f16) && + ds4_gpu_encode_cpy_f16_f16_3d(cb, + valbuf, + ds4_gpu_tensor_offset(value_cache), + g_flash_attn_kv_buffer, + value_f16_offset, + value_dim, + cache_len, + n_head, + (uint64_t)n_head * v_row_bytes_f16, + v_row_bytes_f16, + v_row_bytes_f16, + (uint64_t)cache_len * v_row_bytes_f16)) : + (ds4_gpu_encode_cpy_f32_f16_3d(cb, + keybuf, + ds4_gpu_tensor_offset(key_cache), + g_flash_attn_kv_buffer, + key_f16_offset, + qk_dim, + cache_len, + n_head, + (uint64_t)n_head * q_row_bytes, + q_row_bytes, + q_row_bytes_f16, + (uint64_t)cache_len * q_row_bytes_f16) && + ds4_gpu_encode_cpy_f32_f16_3d(cb, + valbuf, + ds4_gpu_tensor_offset(value_cache), + g_flash_attn_kv_buffer, + value_f16_offset, + value_dim, + cache_len, + n_head, + (uint64_t)n_head * v_row_bytes, + v_row_bytes, + v_row_bytes_f16, + (uint64_t)cache_len * v_row_bytes_f16)); + if (!copied) { + return 0; + } + } + + if (has_kvpad) { + ds4_gpu_flash_attn_pad_args pad_args = { + .ne11 = (int32_t)cache_len, + .ne_12_2 = (int32_t)n_head, + .ne_12_3 = 1, + .nb11 = q_row_bytes_f16, + .nb12 = (uint64_t)cache_len * q_row_bytes_f16, + .nb13 = (uint64_t)cache_len * (uint64_t)n_head * q_row_bytes_f16, + .nb21 = v_row_bytes_f16, + .nb22 = (uint64_t)cache_len * v_row_bytes_f16, + .nb23 = (uint64_t)cache_len * (uint64_t)n_head * v_row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)cache_len * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pad_pipeline]; + [enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0]; + [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:2]; + [enc setBuffer:mask_buffer offset:0 atIndex:3]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(ncpsg, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + + ds4_gpu_flash_attn_blk_args blk_args = { + .ne01 = (int32_t)n_tokens, + .ne30 = (int32_t)cache_len, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)cache_len * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:blk_pipeline]; + [enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0]; + [enc setBuffer:mask_buffer offset:0 atIndex:1]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + ds4_gpu_flash_attn_vec_args args = { + .ne01 = (int32_t)n_tokens, + .ne02 = (int32_t)n_head, + .ne03 = 1, + .nb01 = (uint64_t)n_head * q_row_bytes, + .nb02 = q_row_bytes, + .nb03 = (uint64_t)n_tokens * n_head * q_row_bytes, + .ne11 = (int32_t)cache_len, + .ne_12_2 = (int32_t)n_head, + .ne_12_3 = 1, + .ns10 = (int32_t)qk_dim, + .nb11 = q_row_bytes_f16, + .nb12 = (uint64_t)cache_len * q_row_bytes_f16, + .nb13 = (uint64_t)cache_len * (uint64_t)n_head * q_row_bytes_f16, + .ns20 = (int32_t)value_dim, + .nb21 = v_row_bytes_f16, + .nb22 = (uint64_t)cache_len * v_row_bytes_f16, + .nb23 = (uint64_t)cache_len * (uint64_t)n_head * v_row_bytes_f16, + .ne31 = (int32_t)n_tokens, + .ne32 = 1, + .ne33 = 1, + .nb31 = (uint64_t)cache_len * sizeof(uint16_t), + .nb32 = mask_bytes, + .nb33 = mask_bytes, + .ne1 = (int32_t)n_head, + .ne2 = (int32_t)n_tokens, + .ne3 = 1, + .scale = 1.0f / sqrtf((float)qk_dim), + .max_bias = 0.0f, + .m0 = 0.0f, + .m1 = 0.0f, + .n_head_log2 = 0, + .logit_softcap = 0.0f, + }; + + const NSUInteger padded_v = ds4_gpu_align_up_ns(value_dim, 64u); + const NSUInteger shared_elems = (NSUInteger)nqptg * + ((NSUInteger)qk_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg)); + const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:attn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:2]; + [enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:3]; + [enc setBuffer:mask_buffer offset:0 atIndex:4]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:5]; + [enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6]; + [enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1) + threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM FlashAttention")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_flash_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16) { + return ds4_gpu_glm_attention_flash_tensor_impl(heads, + q, + key_cache, + value_cache, + pos0, + n_tokens, + cache_len, + cache_cap, + n_head, + qk_dim, + value_dim, + cache_f16, + 0); +} + +int ds4_gpu_glm_attention_flash_staged_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16) { + if (pos0 != 0 || n_tokens != cache_len) return 0; + return ds4_gpu_glm_attention_flash_tensor_impl(heads, + q, + key_cache, + value_cache, + pos0, + n_tokens, + cache_len, + cache_cap, + n_head, + qk_dim, + value_dim, + cache_f16, + 1); +} + +int ds4_gpu_glm_attention_full_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *key_cache, + const ds4_gpu_tensor *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !q || !key_cache || !value_cache || + n_tokens == 0 || cache_len == 0 || cache_cap == 0 || + n_head == 0 || qk_dim == 0 || value_dim == 0 || + (qk_dim & 3u) != 0 || + cache_len > cache_cap || + pos0 > cache_len || n_tokens > cache_len - pos0 || + cache_len > ds4_gpu_glm_full_attention_max_cache_len()) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id qbuf = ds4_gpu_tensor_buffer(q); + id keybuf = ds4_gpu_tensor_buffer(key_cache); + id valbuf = ds4_gpu_tensor_buffer(value_cache); + const uint64_t heads_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes; + const uint64_t value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes; + if (!headsbuf || !qbuf || !keybuf || !valbuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(key_cache) < key_bytes || + ds4_gpu_tensor_bytes(value_cache) < value_bytes) { + fprintf(stderr, "ds4: Metal GLM attention received undersized buffers\n"); + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_attention_full_pipeline, + "kernel_glm_attention_full"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const uint32_t full_attn_mode = 2u; + ds4_gpu_glm_attention_full_args args = { + .pos0 = pos0, + .n_tokens = n_tokens, + .cache_len = cache_len, + .cache_cap = cache_cap, + .n_head = n_head, + .qk_dim = qk_dim, + .value_dim = value_dim, + .pad0 = full_attn_mode, + .cache_f16 = cache_f16 ? 1u : 0u, + .pad1 = 0u, + .pad2 = 0u, + .scale = 1.0f / sqrtf((float)qk_dim), + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:2]; + [enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:3]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:4]; + [enc setThreadgroupMemoryLength:(256u + (NSUInteger)cache_len) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM full attention")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_fill_selected_range_tensor( + ds4_gpu_tensor *selected, + uint32_t n_selected) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || n_selected == 0) return 0; + + @autoreleasepool { + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t); + if (!selectedbuf || ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM selected range received undersized buffer\n"); + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_fill_selected_range_pipeline, + "kernel_glm_fill_selected_range"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_fill_selected_range_args args = { + .n_selected = n_selected, + }; + const NSUInteger nth = 256u; + const NSUInteger n_groups = ((NSUInteger)n_selected + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake(n_groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM selected range")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_fill_selected_range_batch_tensor( + ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + uint32_t pad_row) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || n_tokens == 0 || n_selected == 0) return 0; + + @autoreleasepool { + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t total = (uint64_t)n_tokens * n_selected; + if (n_tokens != 0 && total / n_tokens != n_selected) return 0; + if (total > UINT64_MAX / sizeof(uint32_t)) return 0; + const uint64_t selected_bytes = total * sizeof(uint32_t); + if (!selectedbuf || ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM selected range batch received undersized buffer\n"); + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_fill_selected_range_batch_pipeline, + "kernel_glm_fill_selected_range_batch"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_fill_selected_range_batch_args args = { + .n_tokens = n_tokens, + .pos0 = pos0, + .n_selected = n_selected, + .pad_row = pad_row, + }; + const NSUInteger nth = 256u; + const NSUInteger n_groups = ((NSUInteger)total + nth - 1u) / nth; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake(n_groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM selected range batch")) return 0; + } + + return 1; +} + +static int ds4_gpu_glm_rope_tail_offset_tensor( + ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t rot_offset, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + const char *label) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!x || n_tokens == 0 || n_head == 0 || head_dim == 0 || + rot_dim == 0 || rot_offset > head_dim || rot_dim > head_dim - rot_offset || + (rot_dim & 1u) != 0 || + pos0 > UINT32_MAX - n_tokens || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + const uint64_t bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) { + fprintf(stderr, "ds4: Metal %s received undersized buffer\n", + label ? label : "GLM RoPE"); + return 0; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_indexer_rope_tail_pipeline, + "kernel_glm_indexer_rope_tail_f32"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_rope_tail_args args = { + .n_tokens = n_tokens, + .n_head = n_head, + .head_dim = head_dim, + .rot_dim = rot_dim, + .rot_offset = rot_offset, + .pos0 = pos0, + .n_ctx_orig = n_ctx_orig, + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(rot_dim); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, (NSUInteger)n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, label ? label : "GLM RoPE")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (rot_dim > head_dim) return 0; + return ds4_gpu_glm_rope_tail_offset_tensor(x, + n_tokens, + n_head, + head_dim, + rot_dim, + head_dim - rot_dim, + pos0, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow, + "GLM RoPE"); +} + +int ds4_gpu_glm_indexer_rope_tail_tensor( + ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t n_head, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t pos0, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ds4_gpu_glm_rope_tail_offset_tensor(x, + n_tokens, + n_head, + head_dim, + rot_dim, + 0, + pos0, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow, + "GLM indexer RoPE"); +} + +int ds4_gpu_glm_indexer_score_one_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!scores || !q || !weights || !indexer_key_cache || + n_rows == 0 || n_head == 0 || head_dim == 0 || + !isfinite(scale) || scale <= 0.0f) { + return 0; + } + + @autoreleasepool { + id scoresbuf = ds4_gpu_tensor_buffer(scores); + id qbuf = ds4_gpu_tensor_buffer(q); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t score_bytes = (uint64_t)n_rows * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float); + const uint64_t weights_bytes = (uint64_t)n_head * sizeof(float); + const uint64_t cache_bytes = (uint64_t)n_rows * head_dim * cache_elem_bytes; + if (!scoresbuf || !qbuf || !weightsbuf || !cachebuf || + ds4_gpu_tensor_bytes(scores) < score_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes || + ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes) { + fprintf(stderr, "ds4: Metal GLM indexer score received undersized buffers\n"); + return 0; + } + + ds4_gpu_glm_indexer_score_one_args args = { + .n_rows = n_rows, + .n_head = n_head, + .head_dim = head_dim, + .cache_f16 = cache_f16 ? 1u : 0u, + .scale = scale, + }; + + if (n_head == 32u && head_dim == 128u) { + id direct_pipeline = + ds4_gpu_hot_pipeline(g_glm_indexer_score_one_direct_pipeline, + "kernel_glm_indexer_score_one_direct"); + if (!direct_pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:direct_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3]; + [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + [enc setThreadgroupMemoryLength:(128u + 4u) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer direct score")) return 0; + return 1; + } + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_indexer_score_one_pipeline, + "kernel_glm_indexer_score_one"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3]; + [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer score")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_indexer_scores_batch_tensor( + ds4_gpu_tensor *scores, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *weights, + const ds4_gpu_tensor *indexer_key_cache, + uint32_t n_rows, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + float scale, + bool cache_f16) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!scores || !q || !weights || !indexer_key_cache || + n_rows == 0 || n_tokens == 0 || n_head == 0 || head_dim != 128 || + pos0 >= n_rows || n_tokens > n_rows - pos0 || + !isfinite(scale) || scale <= 0.0f) { + return 0; + } + + @autoreleasepool { + id scoresbuf = ds4_gpu_tensor_buffer(scores); + id qbuf = ds4_gpu_tensor_buffer(q); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t score_bytes = (uint64_t)n_rows * n_tokens * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float); + const uint64_t weights_bytes = (uint64_t)n_tokens * n_head * sizeof(float); + const uint64_t cache_bytes = (uint64_t)n_rows * head_dim * cache_elem_bytes; + if (!scoresbuf || !qbuf || !weightsbuf || !cachebuf || + ds4_gpu_tensor_bytes(scores) < score_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes || + ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes) { + fprintf(stderr, "ds4: Metal GLM indexer batch scores received undersized buffers\n"); + return 0; + } + + const bool force_scalar = g_quality_mode; + const bool use_tiled_f32 = false; + const bool use_tiled = !force_scalar && n_tokens >= 8u && + n_head == 32u && head_dim == 128u; + id pipeline = + use_tiled + ? ds4_gpu_hot_pipeline(use_tiled_f32 ? g_glm_indexer_scores_tiled_f32_pipeline + : g_glm_indexer_scores_tiled_pipeline, + use_tiled_f32 ? "kernel_glm_indexer_scores_tiled_f32" + : "kernel_glm_indexer_scores_tiled") + : ds4_gpu_hot_pipeline(g_glm_indexer_scores_batch_pipeline, + "kernel_glm_indexer_scores_batch"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_indexer_scores_batch_args args = { + .n_rows = n_rows, + .n_tokens = n_tokens, + .n_head = n_head, + .head_dim = head_dim, + .pos0 = pos0, + .cache_f16 = cache_f16 ? 1u : 0u, + .q_token_stride = (uint64_t)n_head * head_dim * sizeof(float), + .q_head_stride = (uint64_t)head_dim * sizeof(float), + .weights_token_stride = (uint64_t)n_head * sizeof(float), + .score_token_stride = (uint64_t)n_rows * sizeof(float), + .scale = scale, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3]; + [enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4]; + if (use_tiled) { + const NSUInteger q_shared = 8u * 128u; + const NSUInteger k_shared = 32u * 128u; + const NSUInteger dot_shared = 8u * 32u; + if (use_tiled_f32) { + [enc setThreadgroupMemoryLength:(q_shared + k_shared + dot_shared) * + sizeof(float) atIndex:0]; + } else { + [enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) + + dot_shared * sizeof(float) atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_rows + 31u) / 32u, + ((NSUInteger)n_tokens + 7u) / 8u, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + } else { + [enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer batch scores")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_qk_lowrank_typed_tensor( + ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!qk_low || !q || !model_map || + n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_nope > qk_dim) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(qk_low); + id qbuf = ds4_gpu_tensor_buffer(q); + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) { + fprintf(stderr, "ds4: Metal GLM qk lowrank received unsupported weight type\n"); + return 0; + } + const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim; + const uint64_t weight_bytes = weight_rows * row_bytes; + const uint64_t out_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float); + if (!outbuf || !qbuf || + ds4_gpu_tensor_bytes(qk_low) < out_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes) { + fprintf(stderr, "ds4: Metal GLM qk lowrank received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal GLM qk lowrank range is outside the mapped model\n"); + return 0; + } + + uint64_t weight_inner = 0; + id weightbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner); + if (!weightbuf) return 0; + + const int use_glm52 = + n_head == 64u && + kv_lora_dim == 512u && + qk_nope == 192u && + qk_dim == 256u && + row_bytes == 204u && + weight_type == DS4_METAL_TENSOR_Q8_0; + /* Coalesced simdgroup variant: lanes split the 192-dot so weight + * reads coalesce, and 2048 threadgroups replace 64. The thread- + * per-row kernels measured ~7.5x off the weight-bandwidth floor + * (7.2ms of the decode token by skip-ablation). Covers the GLM 5.2 + * shape for both Q8_0 k_b and the DenseQ4 GGUF's Q4_0 k_b (the Q8 + * fast path above never engaged there). */ + const int use_glm52_sg = + n_head == 64u && + kv_lora_dim == 512u && + qk_nope == 192u && + qk_dim == 256u && + ((weight_type == DS4_METAL_TENSOR_Q8_0 && row_bytes == 204u) || + (weight_type == DS4_METAL_TENSOR_Q4_0 && row_bytes == 108u)) && + getenv("DS4_METAL_DISABLE_GLM_QKLOW_SG") == NULL && + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_sg_pipeline, + "kernel_glm_qk_lowrank_q8_0_glm52_sg") != nil; + if (getenv("DS4_METAL_GLM_QKLOW_DEBUG")) { + static int printed = 0; + if (!printed) { + printed = 1; + fprintf(stderr, "ds4: qk_lowrank decode path: use_glm52=%d sg=%d n_head=%u kv=%u nope=%u dim=%u rb=%llu type=%u\n", + use_glm52, use_glm52_sg, n_head, kv_lora_dim, qk_nope, qk_dim, + (unsigned long long)row_bytes, weight_type); + } + } + id pipeline = + use_glm52_sg ? + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_sg_pipeline, + "kernel_glm_qk_lowrank_q8_0_glm52_sg") : + use_glm52 ? + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_pipeline, + "kernel_glm_qk_lowrank_q8_0_glm52") : + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_pipeline, + "kernel_glm_qk_lowrank_q8_0"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_qk_lowrank_args args = { + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_dim = qk_dim, + .row_bytes = (uint32_t)row_bytes, + .weight_type = weight_type, + .pad1 = 0, + .pad2 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3]; + if (use_glm52_sg) { + /* 8 simdgroups x 2 rows per threadgroup: (64, 512/16) grid. */ + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, + (NSUInteger)(kv_lora_dim / 16u), + 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + } else { + if (use_glm52) { + [enc setThreadgroupMemoryLength:192u * sizeof(float) atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM qk lowrank")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_qk_lowrank_q8_0_tensor( + ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim) { + return ds4_gpu_glm_qk_lowrank_typed_tensor(qk_low, + q, + model_map, + model_size, + weight_offset, + DS4_METAL_TENSOR_Q8_0, + n_head, + kv_lora_dim, + qk_nope, + qk_dim); +} + +int ds4_gpu_glm_qk_lowrank_typed_batch_tensor( + ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!qk_low || !q || !model_map || + n_tokens == 0 || n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_nope > qk_dim) { + return 0; + } + + @autoreleasepool { + id outbuf = ds4_gpu_tensor_buffer(qk_low); + id qbuf = ds4_gpu_tensor_buffer(q); + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) { + fprintf(stderr, "ds4: Metal GLM batch qk lowrank received unsupported weight type\n"); + return 0; + } + const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim; + const uint64_t weight_bytes = weight_rows * row_bytes; + const uint64_t out_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t q_bytes = + (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + if (!outbuf || !qbuf || + ds4_gpu_tensor_bytes(qk_low) < out_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes) { + fprintf(stderr, "ds4: Metal GLM batch qk lowrank received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal GLM batch qk lowrank range is outside the mapped model\n"); + return 0; + } + + uint64_t weight_inner = 0; + id weightbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner); + if (!weightbuf) return 0; + + const int use_glm52_t4 = + n_tokens >= 4u && + n_head == 64u && + kv_lora_dim == 512u && + qk_nope == 192u && + qk_dim == 256u && + row_bytes == 204u && + weight_type == DS4_METAL_TENSOR_Q8_0; + id pipeline = + use_glm52_t4 ? + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_glm52_t4_pipeline, + "kernel_glm_qk_lowrank_q8_0_batch_glm52_t4") : + ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_pipeline, + "kernel_glm_qk_lowrank_q8_0_batch"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + uint32_t head_base = 0; + uint32_t head_count = n_head; + ds4_gpu_tp_attn_head_range(n_head, 8u, &head_base, &head_count); + ds4_gpu_glm_qk_lowrank_batch_args args = { + .n_tokens = n_tokens, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_dim = qk_dim, + .row_bytes = (uint32_t)row_bytes, + .weight_type = weight_type, + .head_base = head_base, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3]; + if (use_glm52_t4) { + [enc setThreadgroupMemoryLength:4u * 192u * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, + ((NSUInteger)n_tokens + 3u) / 4u, + 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } else { + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch qk lowrank")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_qk_lowrank_q8_0_batch_tensor( + ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim) { + return ds4_gpu_glm_qk_lowrank_typed_batch_tensor(qk_low, + q, + model_map, + model_size, + weight_offset, + DS4_METAL_TENSOR_Q8_0, + n_tokens, + n_head, + kv_lora_dim, + qk_nope, + qk_dim); +} + +int ds4_gpu_glm_value_project_typed_batch_heads_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *lora, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t weight_type, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t value_dim) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!heads || !lora || !model_map || + n_tokens == 0 || n_head == 0 || + kv_lora_dim == 0 || value_dim == 0) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id lorabuf = ds4_gpu_tensor_buffer(lora); + uint64_t row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(weight_type, kv_lora_dim, &row_bytes)) { + fprintf(stderr, "ds4: Metal GLM batch value project received unsupported weight type\n"); + return 0; + } + const uint64_t weight_rows = (uint64_t)n_head * value_dim; + const uint64_t weight_bytes = weight_rows * row_bytes; + const uint64_t heads_bytes = + (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + const uint64_t lora_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + if (!headsbuf || !lorabuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(lora) < lora_bytes) { + fprintf(stderr, "ds4: Metal GLM batch value project received undersized buffers\n"); + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal GLM batch value project range is outside the mapped model\n"); + return 0; + } + + uint64_t weight_inner = 0; + id weightbuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + weight_offset, + weight_bytes, + &weight_inner); + if (!weightbuf) return 0; + + const int use_mma = + n_head == 64u && + kv_lora_dim == 512u && + value_dim == 256u && + row_bytes == 544u && + weight_type == DS4_METAL_TENSOR_Q8_0 && + n_tokens >= 32u; + id pipeline = + use_mma ? + ds4_gpu_hot_pipeline(g_glm_value_project_q8_0_batch_heads_mma_pipeline, + "kernel_glm_value_project_q8_0_batch_heads_mma") : + ds4_gpu_hot_pipeline(g_glm_value_project_q8_0_batch_heads_pipeline, + "kernel_glm_value_project_q8_0_batch_heads"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + uint32_t head_base = 0; + uint32_t head_count = n_head; + ds4_gpu_tp_attn_head_range(n_head, 8u, &head_base, &head_count); + ds4_gpu_glm_qk_lowrank_batch_args args = { + .n_tokens = n_tokens, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = 0, + .qk_dim = value_dim, + .row_bytes = (uint32_t)row_bytes, + .weight_type = weight_type, + .head_base = head_base, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora) atIndex:2]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:3]; + if (use_mma) { + [enc setThreadgroupMemoryLength:16u * 1024u atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tokens + 31u) / 32u, + ((NSUInteger)value_dim + 63u) / 64u, + (NSUInteger)head_count) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + } else { + [enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float) + atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch value project")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_value_project_q8_0_batch_heads_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *lora, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t value_dim) { + return ds4_gpu_glm_value_project_typed_batch_heads_tensor(heads, + lora, + model_map, + model_size, + weight_offset, + DS4_METAL_TENSOR_Q8_0, + n_tokens, + n_head, + kv_lora_dim, + value_dim); +} + +int ds4_gpu_glm_attention_indexed_decode_typed_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + uint32_t value_weight_type, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + !model_map || !selected || + n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + value_dim == 0 || qk_dim < qk_nope || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id qbuf = ds4_gpu_tensor_buffer(q); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + uint64_t value_row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) { + fprintf(stderr, "ds4: Metal GLM indexed attention received unsupported value type\n"); + return 0; + } + const uint64_t value_weight_rows = (uint64_t)n_head * value_dim; + const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes; + const uint64_t heads_bytes = (uint64_t)n_head * value_dim * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float); + const uint64_t low_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t); + if (!headsbuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(qk_low) < low_bytes || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM indexed attention received undersized buffers\n"); + return 0; + } + if (value_weight_offset > model_size || + value_weight_bytes > model_size - value_weight_offset) { + fprintf(stderr, "ds4: Metal GLM indexed attention value range is outside the mapped model\n"); + return 0; + } + + uint64_t value_inner = 0; + id valuebuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + value_weight_offset, + value_weight_bytes, + &value_inner); + if (!valuebuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_attention_indexed_decode_pipeline, + "kernel_glm_attention_indexed_decode"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + ds4_gpu_glm_attention_indexed_decode_args args = { + .n_selected = n_selected, + .cache_cap = cache_cap, + .cache_f16 = cache_f16 ? 1u : 0u, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = value_dim, + .n_ctx_orig = n_ctx_orig, + .value_row_bytes = (uint32_t)value_row_bytes, + .scale = 1.0f / sqrtf((float)qk_dim), + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .value_type = value_weight_type, + }; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:5]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:7]; + const NSUInteger scratch_floats = + 256u + (NSUInteger)n_selected + (NSUInteger)kv_lora_dim; + [enc setThreadgroupMemoryLength:scratch_floats * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed attention decode")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_indexed_decode_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ds4_gpu_glm_attention_indexed_decode_typed_tensor(heads, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + model_map, + model_size, + value_weight_offset, + DS4_METAL_TENSOR_Q8_0, + selected, + n_selected, + cache_cap, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + value_dim, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow); +} + +int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *partial_lora, + ds4_gpu_tensor *partial_ms, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + uint32_t value_weight_type, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + bool selected_rows_valid, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + uint32_t block_rows, + uint32_t n_blocks, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + const uint32_t needed_blocks = + block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; + if (!heads || !partial_lora || !partial_ms || !q || !qk_low || + !kv_lora_cache || !k_rope_cache || !model_map || !selected || + n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + n_head == 0 || (n_head % 8u) != 0 || + kv_lora_dim != 512u || + qk_nope == 0 || qk_rope != 64u || + value_dim == 0 || qk_dim < qk_nope || + block_rows == 0u || needed_blocks == 0u || + n_blocks < needed_blocks || n_blocks > 64u || + !cache_f16 || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id partial_lorabuf = ds4_gpu_tensor_buffer(partial_lora); + id partial_msbuf = ds4_gpu_tensor_buffer(partial_ms); + id qbuf = ds4_gpu_tensor_buffer(q); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t cache_elem_bytes = sizeof(uint16_t); + uint64_t value_row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) { + fprintf(stderr, "ds4: Metal GLM split grouped indexed attention received unsupported value type\n"); + return 0; + } + const uint64_t value_weight_rows = (uint64_t)n_head * value_dim; + const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes; + const uint64_t heads_bytes = (uint64_t)n_head * value_dim * sizeof(float); + const uint64_t partial_lora_bytes = + (uint64_t)n_blocks * n_head * kv_lora_dim * sizeof(float); + const uint64_t partial_ms_bytes = + (uint64_t)n_blocks * n_head * 2u * sizeof(float); + const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float); + const uint64_t low_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t); + if (!headsbuf || !partial_lorabuf || !partial_msbuf || + !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(partial_lora) < partial_lora_bytes || + ds4_gpu_tensor_bytes(partial_ms) < partial_ms_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(qk_low) < low_bytes || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM split grouped indexed attention received undersized buffers\n"); + return 0; + } + if (value_weight_offset > model_size || + value_weight_bytes > model_size - value_weight_offset) { + fprintf(stderr, "ds4: Metal GLM split grouped indexed attention value range is outside the mapped model\n"); + return 0; + } + + uint64_t value_inner = 0; + id valuebuf = + ds4_gpu_wrap_model_range(model_map, + model_size, + value_weight_offset, + value_weight_bytes, + &value_inner); + if (!valuebuf) return 0; + + const bool use_valid_fullheads = + selected_rows_valid && (n_head % 8u) == 0u; + id partial_pipeline = + ds4_gpu_hot_pipeline(use_valid_fullheads ? + g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline : + g_glm_attention_indexed_decode_split_group8_partial_pipeline, + use_valid_fullheads ? + "kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads" : + "kernel_glm_attention_indexed_decode_split_group8_partial"); + const bool use_reduce16 = + n_blocks == 16u && block_rows == 128u && + n_selected == 2048u && value_dim == 256u; + id reduce_pipeline = + ds4_gpu_hot_pipeline(use_reduce16 ? + g_glm_attention_indexed_decode_split_group8_reduce16_pipeline : + g_glm_attention_indexed_decode_split_group8_reduce_pipeline, + use_reduce16 ? + "kernel_glm_attention_indexed_decode_split_group8_reduce16" : + "kernel_glm_attention_indexed_decode_split_group8_reduce"); + if (!partial_pipeline || !reduce_pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_attention_indexed_decode_split_args args = { + .n_selected = n_selected, + .cache_cap = cache_cap, + .cache_f16 = 1u, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = value_dim, + .n_ctx_orig = n_ctx_orig, + .value_row_bytes = (uint32_t)value_row_bytes, + .block_rows = block_rows, + .n_blocks = n_blocks, + .scale = 1.0f / sqrtf((float)qk_dim), + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .value_type = value_weight_type, + }; + const NSUInteger stage_rows = 16u; + const NSUInteger kv_vecs = (NSUInteger)kv_lora_dim / 4u; + const NSUInteger rope_vecs = (NSUInteger)qk_rope / 4u; + const NSUInteger partial_scratch_bytes = + stage_rows * kv_vecs * sizeof(uint16_t) * 4u + + stage_rows * rope_vecs * sizeof(float) * 4u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:partial_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; + [enc setBuffer:partial_lorabuf offset:ds4_gpu_tensor_offset(partial_lora) atIndex:6]; + [enc setBuffer:partial_msbuf offset:ds4_gpu_tensor_offset(partial_ms) atIndex:7]; + [enc setThreadgroupMemoryLength:partial_scratch_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head / 8u, (NSUInteger)n_blocks, 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + const NSUInteger reduce_scratch_floats = 256u + 64u + (NSUInteger)kv_lora_dim; + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:reduce_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:partial_lorabuf offset:ds4_gpu_tensor_offset(partial_lora) atIndex:1]; + [enc setBuffer:partial_msbuf offset:ds4_gpu_tensor_offset(partial_ms) atIndex:2]; + [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:3]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:4]; + [enc setThreadgroupMemoryLength:reduce_scratch_floats * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM split grouped indexed attention decode")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor( + ds4_gpu_tensor *heads, + ds4_gpu_tensor *partial_lora, + ds4_gpu_tensor *partial_ms, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + const ds4_gpu_tensor *selected, + uint32_t n_selected, + bool selected_rows_valid, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + uint32_t block_rows, + uint32_t n_blocks, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(heads, + partial_lora, + partial_ms, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + model_map, + model_size, + value_weight_offset, + DS4_METAL_TENSOR_Q8_0, + selected, + n_selected, + selected_rows_valid, + cache_cap, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + value_dim, + n_ctx_orig, + block_rows, + n_blocks, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow); +} + +int ds4_gpu_glm_attention_indexed_batch_typed_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + uint32_t value_weight_type, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + !model_map || !selected || + n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + value_dim == 0 || qk_dim < qk_nope || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id headsbuf = ds4_gpu_tensor_buffer(heads); + id qbuf = ds4_gpu_tensor_buffer(q); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + uint64_t value_row_bytes = 0; + if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) { + fprintf(stderr, "ds4: Metal GLM indexed batch attention received unsupported value type\n"); + return 0; + } + const uint64_t value_weight_rows = (uint64_t)n_head * value_dim; + const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes; + const uint64_t heads_bytes = + (uint64_t)n_tokens * n_head * value_dim * sizeof(float); + const uint64_t q_bytes = + (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + const uint64_t low_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t selected_bytes = + (uint64_t)n_tokens * n_selected * sizeof(uint32_t); + if (!headsbuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || + ds4_gpu_tensor_bytes(heads) < heads_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(qk_low) < low_bytes || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM indexed batch attention received undersized buffers\n"); + return 0; + } + if (value_weight_offset > model_size || + value_weight_bytes > model_size - value_weight_offset) { + fprintf(stderr, "ds4: Metal GLM indexed batch attention value range is outside the mapped model\n"); + return 0; + } + + uint64_t value_inner = 0; + id valuebuf = + ds4_gpu_wrap_model_range(model_map, model_size, + value_weight_offset, + value_weight_bytes, + &value_inner); + if (!valuebuf) return 0; + + const uint64_t grouped_attn = (n_tokens >= 128u) ? 8u : 2u; + const NSUInteger q2_bit_words = ((NSUInteger)cache_cap + 31u) / 32u; + const NSUInteger q2_scratch_bytes = + q2_bit_words * sizeof(uint32_t) + + 4u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) + + 8u * (NSUInteger)kv_lora_dim * sizeof(float); + const NSUInteger max_tg_mem = [g_device maxThreadgroupMemoryLength]; + const bool q2_fits = max_tg_mem == 0 || q2_scratch_bytes <= max_tg_mem; + const bool use_q2_group4 = grouped_attn == 4u && n_tokens >= 2u && q2_fits; + const bool use_group8 = + grouped_attn == 8u || + (grouped_attn == 4u && !use_q2_group4); + const bool use_group2 = grouped_attn == 1u || grouped_attn == 2u; + id pipeline = use_q2_group4 ? + ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_q2_group4_pipeline, + "kernel_glm_attention_indexed_batch_q2_group4") : + (use_group8 ? + ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group8_pipeline, + "kernel_glm_attention_indexed_batch_group8") : + (use_group2 ? + ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group2_pipeline, + "kernel_glm_attention_indexed_batch_group2") : + ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_pipeline, + "kernel_glm_attention_indexed_batch"))); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_attention_indexed_batch_args args = { + .n_tokens = n_tokens, + .n_selected = n_selected, + .cache_cap = cache_cap, + .cache_f16 = cache_f16 ? 1u : 0u, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = value_dim, + .n_ctx_orig = n_ctx_orig, + .value_row_bytes = (uint32_t)value_row_bytes, + .value_type = value_weight_type, + .scale = 1.0f / sqrtf((float)qk_dim), + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .head_base = 0, + }; + const NSUInteger scratch_bytes = use_q2_group4 ? + q2_scratch_bytes : + (use_group8 ? + (8u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) + + 8u * (NSUInteger)kv_lora_dim * sizeof(float)) : + ((use_group2 ? + (512u + 2u * (NSUInteger)n_selected + 2u * (NSUInteger)kv_lora_dim) : + (256u + (NSUInteger)n_selected + (NSUInteger)kv_lora_dim)) * sizeof(float))); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + [enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:5]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6]; + [enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:7]; + [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; + if (use_q2_group4) { + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 3u) / 4u, + ((NSUInteger)n_tokens + 1u) / 2u, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 4, 1)]; + } else if (use_group8) { + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 7u) / 8u, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + } else if (use_group2) { + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 1u) / 2u, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } else { + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, (NSUInteger)n_tokens, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + } + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed batch attention")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_indexed_batch_tensor( + ds4_gpu_tensor *heads, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const void *model_map, + uint64_t model_size, + uint64_t value_weight_offset, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t value_dim, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ds4_gpu_glm_attention_indexed_batch_typed_tensor(heads, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + model_map, + model_size, + value_weight_offset, + DS4_METAL_TENSOR_Q8_0, + selected, + n_tokens, + n_selected, + cache_cap, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + value_dim, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow); +} + +int ds4_gpu_sort_i32_rows_asc_tensor( + ds4_gpu_tensor *dst, + const ds4_gpu_tensor *src, + uint32_t row_width, + uint32_t n_rows) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!dst || !src || row_width == 0 || n_rows == 0 || + (row_width & (row_width - 1u)) != 0) { + return 0; + } + + @autoreleasepool { + id pipeline = + ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline, + "kernel_dsv4_sort_i32_rows_asc"); + if (!pipeline) return 0; + + const uint64_t bytes = (uint64_t)row_width * n_rows * sizeof(int32_t); + id srcbuf = ds4_gpu_tensor_buffer(src); + id dstbuf = ds4_gpu_tensor_buffer(dst); + if (!srcbuf || !dstbuf || + ds4_gpu_tensor_bytes(src) < bytes || + ds4_gpu_tensor_bytes(dst) < bytes) { + fprintf(stderr, "ds4: Metal row sort received undersized buffers\n"); + return 0; + } + + NSUInteger threads = (NSUInteger)row_width; + const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup; + if (max_threads != 0 && threads > max_threads) threads = max_threads; + if (threads == 0) return 0; + + const NSUInteger scratch_bytes = (NSUInteger)row_width * sizeof(int32_t); + const NSUInteger max_tg_mem = [g_device maxThreadgroupMemoryLength]; + if (max_tg_mem != 0 && scratch_bytes > max_tg_mem) { + fprintf(stderr, "ds4: Metal row sort scratch exceeds threadgroup memory limit\n"); + return 0; + } + + ds4_gpu_dsv4_topk_mask_args args = { + .ne00 = (int64_t)row_width, + .ne01 = (int64_t)n_rows, + .nb00 = sizeof(int32_t), + .nb01 = (uint64_t)row_width * sizeof(int32_t), + .ne0 = (int64_t)row_width, + .ne1 = (int64_t)n_rows, + .nb0 = sizeof(int32_t), + .nb1 = (uint64_t)row_width * sizeof(int32_t), + }; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:srcbuf offset:ds4_gpu_tensor_offset(src) atIndex:1]; + [enc setBuffer:dstbuf offset:ds4_gpu_tensor_offset(dst) atIndex:2]; + [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1) + threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "sort i32 rows asc")) return 0; + } + + return 1; +} + +static int ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor( + ds4_gpu_tensor *lora_out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool selected_rows_valid) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || !selected || + n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + qk_dim < qk_nope || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id lorabuf = ds4_gpu_tensor_buffer(lora_out); + id qbuf = ds4_gpu_tensor_buffer(q); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t lora_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t q_bytes = + (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + const uint64_t low_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = + (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = + (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + const uint64_t selected_bytes = + (uint64_t)n_tokens * n_selected * sizeof(uint32_t); + if (!lorabuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf || + ds4_gpu_tensor_bytes(lora_out) < lora_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(qk_low) < low_bytes || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes) { + fprintf(stderr, "ds4: Metal GLM indexed batch attention-lora received undersized buffers\n"); + return 0; + } + + const bool use_vec_lora = + cache_f16 && kv_lora_dim == 512u && qk_rope == 64u; + const bool full_head_groups = (n_head % 8u) == 0u; + id pipeline = nil; + if (use_vec_lora && selected_rows_valid && full_head_groups) { + pipeline = ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline, + "kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads"); + } else if (use_vec_lora && selected_rows_valid) { + pipeline = ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline, + "kernel_glm_attention_indexed_batch_lora_group8_vec_valid"); + } else if (use_vec_lora) { + pipeline = ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group8_vec_pipeline, + "kernel_glm_attention_indexed_batch_lora_group8_vec"); + } else { + pipeline = ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group8_pipeline, + "kernel_glm_attention_indexed_batch_group8"); + } + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_attention_indexed_batch_args args = { + .n_tokens = n_tokens, + .n_selected = n_selected, + .cache_cap = cache_cap, + .cache_f16 = cache_f16 ? 1u : 0u, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = kv_lora_dim, + .n_ctx_orig = n_ctx_orig, + .value_row_bytes = 0, + .value_type = 1u, + .scale = 1.0f / sqrtf((float)qk_dim), + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .head_base = 0, + }; + uint32_t head_count = n_head; + ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count); + const NSUInteger scratch_bytes = use_vec_lora ? + (16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u + + 16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u) : + (8u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) + + 8u * (NSUInteger)kv_lora_dim * sizeof(float)); + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + if (use_vec_lora) { + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:6]; + } else { + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:5]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6]; + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:7]; + } + [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed batch attention-lora")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor( + ds4_gpu_tensor *lora_out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + if (!g_initialized && !ds4_gpu_init()) return 0; + const uint32_t qk_dim = qk_nope + qk_rope; + if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap || + pos0 > n_selected || n_tokens > n_selected - pos0 || + n_head == 0 || kv_lora_dim != 512u || + qk_nope == 0 || qk_rope != 64u || + qk_dim < qk_nope || !cache_f16 || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + + @autoreleasepool { + id lorabuf = ds4_gpu_tensor_buffer(lora_out); + id qbuf = ds4_gpu_tensor_buffer(q); + id lowbuf = ds4_gpu_tensor_buffer(qk_low); + id kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache); + id ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache); + const uint64_t cache_elem_bytes = sizeof(uint16_t); + const uint64_t lora_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t q_bytes = + (uint64_t)n_tokens * n_head * qk_dim * sizeof(float); + const uint64_t low_bytes = + (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float); + const uint64_t kv_cache_bytes = + (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes; + const uint64_t rope_cache_bytes = + (uint64_t)cache_cap * qk_rope * cache_elem_bytes; + if (!lorabuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || + ds4_gpu_tensor_bytes(lora_out) < lora_bytes || + ds4_gpu_tensor_bytes(q) < q_bytes || + ds4_gpu_tensor_bytes(qk_low) < low_bytes || + ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes || + ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes) { + fprintf(stderr, "ds4: Metal GLM causal batch attention-lora received undersized buffers\n"); + return 0; + } + + const bool full_head_groups = (n_head % 8u) == 0u; + id pipeline = full_head_groups ? + ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline, + "kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads") : + ds4_gpu_hot_pipeline( + g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline, + "kernel_glm_attention_indexed_batch_lora_group8_vec_causal"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_attention_indexed_batch_args args = { + .n_tokens = n_tokens, + .n_selected = n_selected, + .cache_cap = cache_cap, + .cache_f16 = 1u, + .n_head = n_head, + .kv_lora_dim = kv_lora_dim, + .qk_nope = qk_nope, + .qk_rope = qk_rope, + .value_dim = kv_lora_dim, + .n_ctx_orig = n_ctx_orig, + .value_row_bytes = 0, + .value_type = 1u, + .pos0 = pos0, + .scale = 1.0f / sqrtf((float)qk_dim), + .freq_base = freq_base, + .freq_scale = freq_scale, + .ext_factor = ext_factor, + .attn_factor = attn_factor, + .beta_fast = beta_fast, + .beta_slow = beta_slow, + .head_base = 0, + }; + uint32_t head_count = n_head; + ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count); + const NSUInteger scratch_bytes = + 16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u + + 16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2]; + [enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3]; + [enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4]; + [enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:5]; + [enc setThreadgroupMemoryLength:scratch_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(32, 8, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM causal indexed batch attention-lora")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_attention_indexed_batch_lora_tensor( + ds4_gpu_tensor *lora_out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(lora_out, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + selected, + n_tokens, + n_selected, + cache_cap, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow, + false); +} + +int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor( + ds4_gpu_tensor *lora_out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *kv_lora_cache, + const ds4_gpu_tensor *k_rope_cache, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + bool cache_f16, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_rope, + uint32_t n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(lora_out, + q, + qk_low, + kv_lora_cache, + k_rope_cache, + selected, + n_tokens, + n_selected, + cache_cap, + cache_f16, + n_head, + kv_lora_dim, + qk_nope, + qk_rope, + n_ctx_orig, + freq_base, + freq_scale, + ext_factor, + attn_factor, + beta_fast, + beta_slow, + true); +} + +int ds4_gpu_glm_router_select_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + const ds4_gpu_tensor *logits, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !weights || !probs || !logits || !model_map || + n_expert == 0 || n_expert > 256u || + n_expert_used == 0 || n_expert_used > n_expert) { + return 0; + } + + @autoreleasepool { + id logitsbuf = ds4_gpu_tensor_buffer(logits); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || + ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int32_t) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) || + ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float)) { + fprintf(stderr, "ds4: Metal GLM router received undersized buffers\n"); + return 0; + } + + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + if (bias_offset > model_size || bias_bytes > model_size - bias_offset) { + fprintf(stderr, "ds4: Metal GLM router bias range is outside the mapped model\n"); + return 0; + } + const bool exact_bias_view = + bias_bytes <= (1ull << 20) && + getenv("DS4_METAL_DISABLE_DECODE_ROUTER_BIAS_EXACT_VIEWS") == NULL; + uint64_t bias_inner = 0; + id biasbuf = exact_bias_view ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + bias_offset, + bias_bytes, + &bias_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + bias_offset, + bias_bytes, + &bias_inner); + if (!biasbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_router_select_one_pipeline, + "kernel_glm_router_select_one"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_router_select_one_args args = { + .n_expert = n_expert, + .n_expert_used = n_expert_used, + .expert_weight_scale = expert_weight_scale, + .pad0 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:1]; + [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4]; + [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5]; + [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM router select")) return 0; + } + + return 1; +} + +int ds4_gpu_glm_router_select_batch_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + const ds4_gpu_tensor *logits, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !weights || !probs || !logits || !model_map || + n_tokens == 0 || + n_expert == 0 || n_expert > 256u || + n_expert_used == 0 || n_expert_used > n_expert) { + return 0; + } + + @autoreleasepool { + id logitsbuf = ds4_gpu_tensor_buffer(logits); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + const uint64_t logits_bytes = (uint64_t)n_tokens * n_expert * sizeof(float); + const uint64_t selected_bytes = (uint64_t)n_tokens * n_expert_used * sizeof(int32_t); + const uint64_t weights_bytes = (uint64_t)n_tokens * n_expert_used * sizeof(float); + const uint64_t probs_bytes = (uint64_t)n_tokens * n_expert * sizeof(float); + if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || + ds4_gpu_tensor_bytes(logits) < logits_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes || + ds4_gpu_tensor_bytes(probs) < probs_bytes) { + fprintf(stderr, "ds4: Metal GLM batch router received undersized buffers\n"); + return 0; + } + + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + if (bias_offset > model_size || bias_bytes > model_size - bias_offset) { + fprintf(stderr, "ds4: Metal GLM batch router bias range is outside the mapped model\n"); + return 0; + } + uint64_t bias_inner = 0; + id biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, + bias_offset, bias_bytes, + &bias_inner); + if (!biasbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_glm_router_select_one_pipeline, + "kernel_glm_router_select_one"); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + ds4_gpu_glm_router_select_one_args args = { + .n_expert = n_expert, + .n_expert_used = n_expert_used, + .expert_weight_scale = expert_weight_scale, + .pad0 = 0, + }; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:1]; + [enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4]; + [enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5]; + [enc setThreadgroupMemoryLength:256u * sizeof(float) + 256u * sizeof(int32_t) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch router select")) return 0; + } + + return 1; +} + +static bool ds4_gpu_glm_gate_pair_type_supported( + uint32_t gate_type, + uint32_t up_type) { + return gate_type == up_type && + (gate_type == DS4_METAL_TENSOR_Q2_K || + gate_type == DS4_METAL_TENSOR_Q4_K || + gate_type == DS4_METAL_TENSOR_Q5_K); +} + +static bool ds4_gpu_glm_down_type_supported(uint32_t down_type) { + return down_type == DS4_METAL_TENSOR_Q2_K || + down_type == DS4_METAL_TENSOR_Q4_K || + down_type == DS4_METAL_TENSOR_Q5_K || + down_type == DS4_METAL_TENSOR_Q6_K; +} + +int ds4_gpu_glm_routed_moe_one_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + bool force_resident) { + if (!g_initialized && !ds4_gpu_init()) return 0; + /* TP sharding: only the owned contiguous expert range is mapped, + * so bind from the owned base, validate only its bytes, and tell the + * kernels the first expert id present at that base. */ + uint32_t first_expert = 0; + uint32_t n_bind_expert = 0; + ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + const int32_t tp_expert_base_host = (int32_t)first_expert; + gate_offset += (uint64_t)first_expert * gate_expert_bytes; + up_offset += (uint64_t)first_expert * up_expert_bytes; + down_offset += (uint64_t)first_expert * down_expert_bytes; + + if (!out || !mid || !model_map || !selected || !weights || !x || + n_total_expert == 0 || n_expert == 0 || n_expert > 256u || + n_expert > n_total_expert || + expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 || + gate_expert_bytes == 0 || gate_row_bytes == 0 || + up_expert_bytes == 0 || up_row_bytes == 0 || + down_expert_bytes == 0 || down_row_bytes == 0 || + (expert_in_dim % 256u) != 0 || + (expert_mid_dim % 256u) != 0 || + !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) || + !ds4_gpu_glm_down_type_supported(down_type)) { + return 0; + } + + if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes || + (uint64_t)expert_mid_dim > UINT64_MAX / gate_row_bytes || + (uint64_t)expert_mid_dim > UINT64_MAX / up_row_bytes || + (uint64_t)out_dim > UINT64_MAX / down_row_bytes) { + fprintf(stderr, "ds4: Metal GLM routed MoE tensor byte size overflow\n"); + return 0; + } + + const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; + const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + if (gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes || + up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes || + down_expert_bytes != (uint64_t)out_dim * down_row_bytes) { + fprintf(stderr, "ds4: Metal GLM routed MoE received inconsistent expert strides\n"); + return 0; + } + if (layer_index == 3u) { + ds4_gpu_stream_expert_cache_note_decode_token(); + } + if (gate_offset > model_size || gate_tensor_bytes > model_size - gate_offset || + up_offset > model_size || up_tensor_bytes > model_size - up_offset || + down_offset > model_size || down_tensor_bytes > model_size - down_offset) { + fprintf(stderr, "ds4: Metal GLM routed MoE tensor range is outside the mapped model\n"); + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + const uint64_t x_bytes = (uint64_t)expert_in_dim * sizeof(float); + const uint64_t mid_bytes = (uint64_t)n_expert * expert_mid_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float); + if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert * sizeof(int32_t) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert * sizeof(float)) { + fprintf(stderr, "ds4: Metal GLM routed MoE received undersized activation buffers\n"); + return 0; + } + + const BOOL gate_pair_q2 = gate_type == DS4_METAL_TENSOR_Q2_K; + const BOOL gate_pair_q5 = gate_type == DS4_METAL_TENSOR_Q5_K; + const BOOL down_scalar_q2 = down_type == DS4_METAL_TENSOR_Q2_K; + const BOOL down_scalar_q4 = down_type == DS4_METAL_TENSOR_Q4_K; + const BOOL down_simd_q4 = down_scalar_q4; + const BOOL down_simd_q5 = down_type == DS4_METAL_TENSOR_Q5_K; + const BOOL down_simd_q6 = down_type == DS4_METAL_TENSOR_Q6_K; + const BOOL down_simd = down_simd_q4 || down_simd_q5 || down_simd_q6; + const BOOL stream_addr_q2 = + gate_pair_q2 && down_scalar_q2 && + g_glm_q2_k_addr_pair_swiglu2_f32_pipeline != nil && + g_glm_q2_k_addr_down_f32_pipeline != nil; + const BOOL stream_addr_q4 = + !gate_pair_q2 && !gate_pair_q5 && down_scalar_q4 && + g_glm_q4_k_addr_pair_swiglu_f32_pipeline != nil && + g_glm_q4_k_addr_down_f32_pipeline != nil; + BOOL use_stream_expert_addr_table = + g_ssd_streaming_mode && + !force_resident && + getenv("DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE") == NULL && + (stream_addr_q2 || stream_addr_q4) && + layer_index < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && + n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && + n_expert <= 8u && + ds4_gpu_stream_expert_cache_configured_budget() >= n_expert && + ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) && + ds4_gpu_stream_expert_cache_effective_cap(layer_index, + n_total_expert, + n_expert) != 0; + int32_t stream_selected_ids[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + ds4_gpu_stream_expert_cache_entry *stream_entries[8] = { + NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL + }; + uint64_t stream_gate_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint64_t stream_up_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint64_t stream_down_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + uint32_t stream_missing_mask = 0; + uint32_t stream_entry_count = 0; + uint32_t stream_resident_mask = 0; + BOOL use_stream_split_deferred = false; + id stream_gate_addr_buf = nil; + id stream_up_addr_buf = nil; + id stream_down_addr_buf = nil; + const int glm_stream_timing = + ds4_gpu_stream_expert_timing_summary_enabled(); + + if (use_stream_expert_addr_table) { + const int had_batch = g_batch_cb != nil; + int stream_ok = 1; + const int have_prefetched_selected = + ds4_gpu_glm_stream_selected_prefetch_take(model_map, + model_size, + layer_index, + n_total_expert, + n_expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes, + stream_selected_ids); + if (have_prefetched_selected) { + if (had_batch && g_batch_has_work && + g_stream_expert_pending_load.active && + ds4_gpu_flush_commands() == 0) { + stream_ok = 0; + } + } else { + if (had_batch && ds4_gpu_end_commands() == 0) { + stream_ok = 0; + } + if (stream_ok && + ds4_gpu_tensor_read(selected, + 0, + stream_selected_ids, + (uint64_t)n_expert * sizeof(stream_selected_ids[0])) == 0) { + stream_ok = 0; + } + } + for (uint32_t i = 0; stream_ok && i < n_expert; i++) { + if (stream_selected_ids[i] < 0 || + (uint32_t)stream_selected_ids[i] >= n_total_expert) { + fprintf(stderr, + "ds4: Metal GLM routed MoE selected expert id %d is outside 0..%u\n", + stream_selected_ids[i], + n_total_expert); + stream_ok = 0; + } + } + if (stream_ok) { + ds4_gpu_stream_expert_cache_note_selected_hotness(layer_index, + stream_selected_ids, + n_expert); + if (!ds4_gpu_moe_selected_hotlist_record(layer_index, + stream_selected_ids, + n_expert, + n_total_expert)) { + stream_ok = 0; + } + } + if (stream_ok) { + g_glm_stream_expert_addr_table_building++; + for (uint32_t i = 0; stream_ok && i < n_expert; i++) { + const uint64_t expert_id = (uint64_t)(uint32_t)stream_selected_ids[i]; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal GLM routed MoE selected expert offset overflow\n"); + stream_ok = 0; + break; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal GLM routed MoE selected expert offset overflow\n"); + stream_ok = 0; + break; + } + stream_gate_abs_offsets[i] = gate_offset + gate_rel; + stream_up_abs_offsets[i] = up_offset + gate_rel; + stream_down_abs_offsets[i] = down_offset + down_rel; + stream_entries[i] = + ds4_gpu_stream_expert_cache_peek(model_map, + model_size, + layer_index, + (uint32_t)stream_selected_ids[i], + n_total_expert, + n_expert, + stream_gate_abs_offsets[i], + stream_up_abs_offsets[i], + stream_down_abs_offsets[i], + gate_expert_bytes, + down_expert_bytes); + if (!stream_entries[i]) { + stream_missing_mask |= 1u << i; + } else { + stream_resident_mask |= 1u << i; + } + } + if (stream_ok && glm_stream_timing) { + ds4_gpu_stream_expert_timing_note_cache_class( + stream_resident_mask, + stream_missing_mask); + } + use_stream_split_deferred = + stream_ok && + getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_SPLIT") == NULL && + stream_missing_mask != 0 && + stream_resident_mask != 0 && + ds4_gpu_stream_expert_split_worthwhile(stream_resident_mask, + stream_missing_mask) && + ds4_gpu_stream_expert_split_ready() && + ((gate_pair_q2 && + g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline != nil) || + (!gate_pair_q2 && !gate_pair_q5 && + g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline != nil)); + if (use_stream_split_deferred) { + const ds4_gpu_stream_expert_table table = { + .model_map = model_map, + .model_size = model_size, + .layer = layer_index, + .n_total_expert = n_total_expert, + .gate_offset = gate_offset, + .up_offset = up_offset, + .down_offset = down_offset, + .gate_expert_bytes = gate_expert_bytes, + .down_expert_bytes = down_expert_bytes, + }; + if (!ds4_gpu_stream_expert_cache_begin_selected_load( + &table, + stream_selected_ids, + n_expert)) { + stream_ok = 0; + } + } + if (stream_ok && stream_missing_mask != 0 && + !use_stream_split_deferred && + !ds4_gpu_stream_expert_cache_load_selected_missing( + model_map, + model_size, + layer_index, + stream_selected_ids, + n_total_expert, + n_expert, + stream_gate_abs_offsets, + stream_up_abs_offsets, + stream_down_abs_offsets, + gate_expert_bytes, + down_expert_bytes, + stream_missing_mask, + stream_entries)) { + fprintf(stderr, + "ds4: Metal GLM streaming expert cache failed to load " + "layer=%u missing=0x%x budget=%u\n", + layer_index, + stream_missing_mask, + ds4_gpu_stream_expert_cache_configured_budget()); + stream_ok = 0; + } + for (uint32_t i = 0; stream_ok && i < n_expert; i++) { + ds4_gpu_stream_expert_cache_entry *entry = stream_entries[i]; + if (!entry) { + if (use_stream_split_deferred && + (stream_missing_mask & (1u << i)) != 0) { + continue; + } + stream_ok = 0; + break; + } + if (!ds4_gpu_stream_expert_cache_set_addr_slot( + layer_index, + (uint32_t)stream_selected_ids[i], + entry->gate_buffer, + entry->gate_inner, + entry->up_buffer, + entry->up_inner, + entry->down_buffer, + entry->down_inner)) { + stream_ok = 0; + break; + } + stream_entry_count++; + } + if (stream_ok && + !ds4_gpu_stream_expert_cache_addr_buffers(layer_index, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf)) { + stream_ok = 0; + } + g_glm_stream_expert_addr_table_building--; + } + if (!have_prefetched_selected && had_batch && + ds4_gpu_begin_commands() == 0) { + stream_ok = 0; + } + if (!stream_ok) return 0; + ds4_gpu_stream_expert_cache_prune_layer(layer_index, + n_total_expert, + n_expert, + stream_selected_ids, + n_expert); + ds4_gpu_stream_expert_cache_prune_global(layer_index, + stream_selected_ids, + n_expert); + } + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gatebuf = nil; + id upbuf = nil; + id downbuf = nil; + if (!use_stream_expert_addr_table) { + gatebuf = ds4_gpu_wrap_model_range(model_map, model_size, + gate_offset, gate_tensor_bytes, + &gate_inner); + upbuf = ds4_gpu_wrap_model_range(model_map, model_size, + up_offset, up_tensor_bytes, + &up_inner); + downbuf = ds4_gpu_wrap_model_range(model_map, model_size, + down_offset, down_tensor_bytes, + &down_inner); + if (!gatebuf || !upbuf || !downbuf) return 0; + } + + id pair_pipeline = + (use_stream_split_deferred ? + (gate_pair_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline, + "kernel_glm_q2_K_addr_pair_swiglu2_f32_masked") : + ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline, + "kernel_glm_q4_K_addr_pair_swiglu_f32_masked")) : + use_stream_expert_addr_table ? + (gate_pair_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_f32_pipeline, + "kernel_glm_q2_K_addr_pair_swiglu2_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_f32_pipeline, + "kernel_glm_q4_K_addr_pair_swiglu_f32")) : + gate_pair_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_pair_swiglu_f32_pipeline, + "kernel_glm_q2_K_pair_swiglu_f32") : + gate_pair_q5 ? + ds4_gpu_hot_pipeline(g_glm_q5_k_pair_swiglu_f32_pipeline, + "kernel_glm_q5_K_pair_swiglu_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu2_f32_pipeline, + "kernel_glm_q4_K_pair_swiglu2_f32")); + id down_pipeline = + (use_stream_expert_addr_table ? + (down_scalar_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_addr_down_f32_pipeline, + "kernel_glm_q2_K_addr_down_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_addr_down_f32_pipeline, + "kernel_glm_q4_K_addr_down_f32")) : + down_scalar_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_down_f32_pipeline, + "kernel_glm_q2_K_down_f32") : + down_scalar_q4 ? + ds4_gpu_hot_pipeline(g_glm_q4_k_down_f32_pipeline, + "kernel_glm_q4_K_down_f32") : + down_simd_q5 ? + ds4_gpu_hot_pipeline(g_glm_q5_k_down_f32_pipeline, + "kernel_glm_q5_K_down_f32") : + ds4_gpu_hot_pipeline(g_glm_q6_k_down_f32_pipeline, + "kernel_glm_q6_K_down_f32")); + if (!pair_pipeline || !down_pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const bool glm_moe_stage_profile = + g_batch_cb != nil && + ds4_gpu_stage_profile_enabled_for_layer( + "DS4_METAL_GLM_MOE_ONE_STAGE_PROFILE", + "DS4_METAL_GLM_MOE_ONE_STAGE_PROFILE_LAYER", + layer_index); + const char *glm_moe_stage_filter = + getenv("DS4_METAL_GLM_MOE_STAGE_PROFILE_FILTER"); + const char *glm_pair_path = + use_stream_expert_addr_table ? + (gate_pair_q2 ? "q2_stream_addr_swiglu" : + "q4_stream_addr_swiglu") : + gate_pair_q2 ? "q2_scalar_swiglu" : + gate_pair_q5 ? "q5_pair_simd_swiglu" : "q4_pair2_simd_swiglu"; + const char *glm_down_path = + use_stream_expert_addr_table ? + (down_scalar_q2 ? "q2_stream_addr_down" : + "q4_stream_addr_down_simd") : + down_scalar_q2 ? "q2_down_simd" : + down_scalar_q4 ? "q4_down_simd" : + down_simd_q5 ? "q5_down_simd" : "q6_down_simd"; + double glm_moe_stage_t0 = 0.0; + if (glm_moe_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + glm_moe_stage_t0 = ds4_gpu_now_ms(); + } + int ok = 1; +#define DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE(name) do { \ + if (ok && glm_moe_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ + strstr(stage_name, glm_moe_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal GLM routed MoE one stage layer=%u tokens=1 experts=%u " \ + "gate=%s down=%s pair=%s down_path=%s %s=%.3f ms\n", \ + layer_index, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + glm_pair_path, glm_down_path, \ + stage_name, now_ms - glm_moe_stage_t0); \ + } \ + glm_moe_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + + ds4_gpu_glm_routed_moe_args args = { + .tp_rank = g_tp_split_rank, + .tp_world = g_tp_split_world, + .tp_expert_base = tp_expert_base_host, + .in_dim = expert_in_dim, + .mid_dim = expert_mid_dim, + .out_dim = out_dim, + .n_total_expert = n_total_expert, + .n_expert_used = n_expert, + .n_tokens = 1, + .mid_token_stride = n_expert * expert_mid_dim, + .down_type = down_type, + .gate_expert_bytes = gate_expert_bytes, + .gate_row_bytes = gate_row_bytes, + .up_expert_bytes = up_expert_bytes, + .up_row_bytes = up_row_bytes, + .down_expert_bytes = down_expert_bytes, + .down_row_bytes = down_row_bytes, + }; + const NSUInteger pair_x_groups = + gate_pair_q2 ? (use_stream_expert_addr_table ? + (NSUInteger)((expert_mid_dim + 1u) / 2u) : + (NSUInteger)((expert_mid_dim + 7u) / 8u)) : + gate_pair_q5 ? (NSUInteger)((expert_mid_dim + 7u) / 8u) : + use_stream_expert_addr_table ? (NSUInteger)((expert_mid_dim + 3u) / 4u) : + (NSUInteger)((expert_mid_dim + 1u) / 2u); + const NSUInteger pair_threadgroup_bytes = 0u; + const NSUInteger pair_threads = 64u; + const NSUInteger down_x_groups = + down_scalar_q2 ? (NSUInteger)((out_dim + 7u) / 8u) : + down_simd_q4 ? (NSUInteger)((out_dim + 3u) / 4u) : + down_simd_q5 ? (NSUInteger)((out_dim + 3u) / 4u) : + down_simd_q6 ? (NSUInteger)((out_dim + 3u) / 4u) : + (NSUInteger)out_dim; + const NSUInteger down_threadgroup_bytes = + (down_scalar_q2 || down_simd) ? 0u : 256u * sizeof(float); + const NSUInteger down_threads = + (down_scalar_q2 || down_simd) ? 64u : 256u; + if (use_stream_expert_addr_table && + !ds4_gpu_stream_expert_cache_mark_entries_inflight( + stream_entries, + use_stream_split_deferred ? n_expert : stream_entry_count, + use_stream_split_deferred ? stream_resident_mask : 0)) { + return 0; + } + + const int glm_stream_split_timing = + use_stream_split_deferred && glm_stream_timing; + double glm_stream_split_t0 = + glm_stream_split_timing ? ds4_gpu_now_ms() : 0.0; + double glm_stream_split_resident_ms = 0.0; + double glm_stream_split_missing_load_ms = 0.0; + double glm_stream_split_missing_wait_ms = 0.0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pair_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + if (use_stream_split_deferred) { + [enc setBytes:&stream_resident_mask length:sizeof(stream_resident_mask) atIndex:1]; + [enc setBuffer:stream_gate_addr_buf offset:0u atIndex:2]; + [enc setBuffer:stream_up_addr_buf offset:0u atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:6]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:7]; + } else { + [enc setBuffer:use_stream_expert_addr_table ? stream_gate_addr_buf : gatebuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)gate_inner + atIndex:1]; + [enc setBuffer:use_stream_expert_addr_table ? stream_up_addr_buf : upbuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)up_inner + atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:4]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:5]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; + } + if (use_stream_expert_addr_table) { + const uint32_t use_count = + use_stream_split_deferred ? n_expert : stream_entry_count; + const uint32_t use_mask = + use_stream_split_deferred ? stream_resident_mask : 0; + for (uint32_t i = 0; i < use_count; i++) { + if (use_mask != 0 && (use_mask & (1u << i)) == 0) continue; + [enc useResource:stream_entries[i]->gate_buffer usage:MTLResourceUsageRead]; + [enc useResource:stream_entries[i]->up_buffer usage:MTLResourceUsageRead]; + } + } + if (pair_threadgroup_bytes != 0u) { + [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups, + (NSUInteger)n_expert, + 1) + threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE("pair"); + + if (ok && use_stream_split_deferred) { + id resident_cb = nil; + const int resident_owned = owned; + if (owned) { + resident_cb = cb; + [resident_cb commit]; + cb = nil; + } else { + ok = ds4_gpu_flush_commands(); + } + if (glm_stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + glm_stream_split_resident_ms = now_ms - glm_stream_split_t0; + glm_stream_split_t0 = now_ms; + } + if (ok) { + ok = ds4_gpu_stream_expert_pending_load_finish(stream_entries); + } + if (glm_stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + glm_stream_split_missing_load_ms = now_ms - glm_stream_split_t0; + glm_stream_split_t0 = now_ms; + } + if (ok) { + ds4_gpu_stream_expert_cache_prune_layer(layer_index, + n_total_expert, + n_expert, + stream_selected_ids, + n_expert); + ds4_gpu_stream_expert_cache_prune_global(layer_index, + stream_selected_ids, + n_expert); + ok = ds4_gpu_stream_expert_cache_addr_buffers(layer_index, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf); + } + if (ok) { + if (resident_owned) { + ok = ds4_gpu_wait_command_buffer( + resident_cb, + "GLM streaming split resident pair"); + ds4_gpu_stream_expert_cache_note_owned_completed(); + } else { + ok = ds4_gpu_wait_pending_command_buffers( + "GLM streaming split resident pair"); + } + } + if (glm_stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + glm_stream_split_missing_wait_ms = now_ms - glm_stream_split_t0; + glm_stream_split_t0 = now_ms; + ds4_gpu_stream_expert_timing_note_split( + stream_resident_mask, + stream_missing_mask, + glm_stream_split_resident_ms, + glm_stream_split_missing_load_ms + + glm_stream_split_missing_wait_ms); + ds4_gpu_stream_expert_timing_note_split_missing_detail( + glm_stream_split_missing_load_ms, + 0.0, + 0.0, + 0.0, + glm_stream_split_missing_wait_ms); + } + if (ok && + !ds4_gpu_stream_expert_cache_mark_entries_inflight(stream_entries, + n_expert, + stream_missing_mask)) { + ok = 0; + } + if (ok) { + cb = ds4_gpu_command_buffer(&owned); + if (!cb) ok = 0; + } + if (ok) { + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pair_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBytes:&stream_missing_mask length:sizeof(stream_missing_mask) atIndex:1]; + [enc setBuffer:stream_gate_addr_buf offset:0u atIndex:2]; + [enc setBuffer:stream_up_addr_buf offset:0u atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:6]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:7]; + for (uint32_t i = 0; i < n_expert; i++) { + if ((stream_missing_mask & (1u << i)) == 0) continue; + [enc useResource:stream_entries[i]->gate_buffer usage:MTLResourceUsageRead]; + [enc useResource:stream_entries[i]->up_buffer usage:MTLResourceUsageRead]; + } + if (pair_threadgroup_bytes != 0u) { + [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups, + (NSUInteger)n_expert, + 1) + threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + } + } + + if (!ok) return 0; + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:down_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:use_stream_expert_addr_table ? stream_down_addr_buf : downbuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)down_inner + atIndex:1]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + if (use_stream_expert_addr_table) { + const uint32_t use_count = + use_stream_split_deferred ? n_expert : stream_entry_count; + for (uint32_t i = 0; i < use_count; i++) { + [enc useResource:stream_entries[i]->down_buffer usage:MTLResourceUsageRead]; + } + } + if (down_threadgroup_bytes != 0u) { + [enc setThreadgroupMemoryLength:down_threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(down_x_groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(down_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE("down"); + + if (!ok) return 0; + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM routed MoE")) return 0; +#undef DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE + } + + return 1; +} + +static bool ds4_gpu_glm_grouped_moe_fast_default(void) { + return !g_quality_mode; +} + +static bool ds4_gpu_glm_routed_moe_batch_grouped_available( + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint32_t n_expert, + uint32_t n_tokens) { + if (n_tokens < 32u || + !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) || + !ds4_gpu_glm_down_type_supported(down_type) || + ds4_gpu_mul_mm_id_map0_name(n_expert) == NULL) { + return false; + } + const bool fast_default = ds4_gpu_glm_grouped_moe_fast_default(); + if (!fast_default) { + return false; + } + if (n_tokens < 96u) return false; + + return ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)) != nil && + ds4_gpu_routed_mm_pipeline(gate_type) != nil && + ds4_gpu_routed_mm_pipeline(up_type) != nil && + ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) != nil; +} + +static bool ds4_gpu_glm_grouped_moe_layer_enabled(uint32_t layer_index) { + (void)layer_index; + return true; +} + +static int ds4_gpu_glm_routed_moe_batch_grouped_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens) { + if (!ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, + up_type, + down_type, + n_expert, + n_tokens)) { + return 0; + } + if (n_expert > UINT32_MAX / n_tokens || + (uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { + return 0; + } + + const uint32_t pair_rows = n_tokens * n_expert; + if ((uint64_t)pair_rows > UINT64_MAX / expert_mid_dim || + (uint64_t)pair_rows > UINT64_MAX / out_dim || + (uint64_t)n_tokens > UINT64_MAX / expert_in_dim || + (uint64_t)n_tokens > UINT64_MAX / out_dim) { + return 0; + } + + const bool mid_f16 = true; + const NSUInteger mm_id_threadgroup_bytes = 8192u; + const uint64_t compact_mid_values = (uint64_t)pair_rows * expert_mid_dim; + const uint64_t down_values = (uint64_t)pair_rows * out_dim; + const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim; + const uint64_t out_values = (uint64_t)n_tokens * out_dim; + if (compact_mid_values > UINT64_MAX / sizeof(float) || + compact_mid_values > UINT64_MAX / (mid_f16 ? sizeof(uint16_t) : sizeof(float)) || + down_values > UINT64_MAX / sizeof(float) || + x_values > UINT64_MAX / sizeof(float) || + out_values > UINT64_MAX / sizeof(float)) { + return 0; + } + + const uint64_t gate_scratch_bytes = compact_mid_values * sizeof(float); + const uint64_t mid_bytes = compact_mid_values * (mid_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t down_scratch_bytes = down_values * sizeof(float); + const uint64_t x_bytes = x_values * sizeof(float); + const uint64_t out_bytes = out_values * sizeof(float); + const uint64_t selected_values = (uint64_t)n_tokens * n_expert; + const uint64_t selected_bytes = selected_values * sizeof(int32_t); + const uint64_t weights_bytes = selected_values * sizeof(float); + if (gate_scratch_bytes > UINT64_MAX - gate_scratch_bytes || + gate_scratch_bytes > NSUIntegerMax || + gate_scratch_bytes * 2ull > NSUIntegerMax || + down_scratch_bytes > NSUIntegerMax) { + return 0; + } + + uint32_t first_expert = 0; + uint32_t n_bind_expert = 0; + ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + gate_offset += (uint64_t)first_expert * gate_expert_bytes; + up_offset += (uint64_t)first_expert * up_expert_bytes; + down_offset += (uint64_t)first_expert * down_expert_bytes; + const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; + const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes) { + fprintf(stderr, "ds4: Metal GLM grouped routed MoE received undersized activation buffers\n"); + return 0; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_gate_scratch_buffer, + &g_moe_gate_scratch_bytes, + (NSUInteger)(gate_scratch_bytes * 2ull), + "ds4_glm_moe_gate_up_scratch")) { + return 0; + } + if (n_expert > 1 && + !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, + &g_moe_down_scratch_bytes, + (NSUInteger)down_scratch_bytes, + "ds4_glm_moe_down_scratch")) { + return 0; + } + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gatebuf = ds4_gpu_wrap_model_range(model_map, model_size, + gate_offset, gate_tensor_bytes, + &gate_inner); + id upbuf = ds4_gpu_wrap_model_range(model_map, model_size, + up_offset, up_tensor_bytes, + &up_inner); + id downbuf = ds4_gpu_wrap_model_range(model_map, model_size, + down_offset, down_tensor_bytes, + &down_inner); + if (!gatebuf || !upbuf || !downbuf) return 0; + + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); + id gate_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); + id up_pipeline = ds4_gpu_routed_mm_pipeline(up_type); + id down_pipeline = + ds4_gpu_routed_mm_f16_rhs_pipeline(down_type); + if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) { + return 0; + } + + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args gate_args = + ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + gate_row_bytes, gate_expert_bytes, + 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args up_args = + ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + up_row_bytes, up_expert_bytes, + 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args down_args = + ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert, + down_row_bytes, down_expert_bytes, + n_expert, n_expert, n_tokens, + mid_f16 ? sizeof(uint16_t) : sizeof(float)); + gate_args.tp_rank = g_tp_split_rank; + gate_args.tp_world = g_tp_split_world; + gate_args.tp_expert_base = (int32_t)first_expert; + up_args.tp_rank = g_tp_split_rank; + up_args.tp_world = g_tp_split_world; + up_args.tp_expert_base = (int32_t)first_expert; + down_args.tp_rank = g_tp_split_rank; + down_args.tp_world = g_tp_split_world; + down_args.tp_expert_base = (int32_t)first_expert; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const bool glm_moe_stage_profile = false; + const char *glm_moe_stage_filter = NULL; + double glm_moe_stage_t0 = 0.0; + if (glm_moe_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + glm_moe_stage_t0 = ds4_gpu_now_ms(); + } + + int ok = 1; +#define DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE(name) do { \ + if (ok && glm_moe_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ + strstr(stage_name, glm_moe_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal GLM grouped routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \ + "gate=%s down=%s mid=%s %s=%.3f ms\n", \ + layer_index, n_tokens, pair_rows, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + mid_f16 ? "f16" : "f32", \ + stage_name, now_ms - glm_moe_stage_t0); \ + } \ + glm_moe_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + + ok = ds4_gpu_encode_mul_mm_id_map(cb, + map_pipeline, + &map_args, + &gate_args, + selectedbuf, + ds4_gpu_tensor_offset(selected)); + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("map"); + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + gate_pipeline, + &gate_args, + gatebuf, + (NSUInteger)gate_inner, + xbuf, + ds4_gpu_tensor_offset(x), + g_moe_gate_scratch_buffer, + 0, + mm_id_threadgroup_bytes); + } + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("gate"); + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + up_pipeline, + &up_args, + upbuf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + g_moe_gate_scratch_buffer, + (NSUInteger)gate_scratch_bytes, + mm_id_threadgroup_bytes); + } + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("up"); + if (ok) { + ok = ds4_gpu_encode_moe_swiglu_weight(cb, + g_moe_gate_scratch_buffer, + 0, + g_moe_gate_scratch_buffer, + (NSUInteger)gate_scratch_bytes, + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + expert_mid_dim, + pair_rows, + 0.0f, + mid_f16); + } + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("activation_weight"); + + id down_dst = n_expert == 1 ? outbuf : g_moe_down_scratch_buffer; + NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : 0; + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + down_pipeline, + &down_args, + downbuf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + mm_id_threadgroup_bytes); + } + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("down"); + if (ok && n_expert > 1) { + ok = ds4_gpu_encode_moe_sum_experts(cb, + down_dst, + down_dst_off, + outbuf, + ds4_gpu_tensor_offset(out), + out_dim, + n_expert, + n_tokens); + } + DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("sum"); + if (!ok) return 0; + + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM grouped routed batch MoE")) { + return 0; + } +#undef DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE + } + + return 1; +} + +static int ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + id gate_addrs, + id up_addrs, + id down_addrs, + id overflow_gate, + id overflow_up, + id overflow_down, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + ds4_gpu_stream_expert_cache_entry * const *resources, + uint32_t resource_count) { + if (n_expert > UINT32_MAX / n_tokens || + !gate_addrs || !up_addrs || !down_addrs || + !ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, + up_type, + down_type, + n_expert, + n_tokens)) { + return 0; + } + + const uint32_t pair_rows = n_tokens * n_expert; + if ((uint64_t)pair_rows > UINT64_MAX / expert_mid_dim || + (uint64_t)pair_rows > UINT64_MAX / out_dim || + (uint64_t)n_tokens > UINT64_MAX / expert_in_dim || + (uint64_t)n_tokens > UINT64_MAX / out_dim) { + return 0; + } + + const bool mid_f16 = true; + const NSUInteger mm_id_threadgroup_bytes = 8192u; + const uint64_t compact_mid_values = (uint64_t)pair_rows * expert_mid_dim; + const uint64_t down_values = (uint64_t)pair_rows * out_dim; + const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim; + const uint64_t out_values = (uint64_t)n_tokens * out_dim; + if (compact_mid_values > UINT64_MAX / sizeof(float) || + compact_mid_values > UINT64_MAX / sizeof(uint16_t) || + down_values > UINT64_MAX / sizeof(float) || + x_values > UINT64_MAX / sizeof(float) || + out_values > UINT64_MAX / sizeof(float)) { + return 0; + } + + const uint64_t gate_scratch_bytes = compact_mid_values * sizeof(float); + const uint64_t mid_bytes = compact_mid_values * sizeof(uint16_t); + const uint64_t down_scratch_bytes = down_values * sizeof(float); + const uint64_t x_bytes = x_values * sizeof(float); + const uint64_t out_bytes = out_values * sizeof(float); + const uint64_t selected_values = (uint64_t)n_tokens * n_expert; + const uint64_t selected_bytes = selected_values * sizeof(int32_t); + const uint64_t weights_bytes = selected_values * sizeof(float); + if (gate_scratch_bytes > UINT64_MAX - gate_scratch_bytes || + gate_scratch_bytes > NSUIntegerMax || + gate_scratch_bytes * 2ull > NSUIntegerMax || + down_scratch_bytes > NSUIntegerMax) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes) { + fprintf(stderr, "ds4: Metal GLM grouped-address routed MoE received undersized activation buffers\n"); + return 0; + } + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_gate_scratch_buffer, + &g_moe_gate_scratch_bytes, + (NSUInteger)(gate_scratch_bytes * 2ull), + "ds4_glm_moe_gate_up_scratch")) { + return 0; + } + if (n_expert > 1 && + !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, + &g_moe_down_scratch_bytes, + (NSUInteger)down_scratch_bytes, + "ds4_glm_moe_down_scratch")) { + return 0; + } + + id map_pipeline = + ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); + id gate_pipeline = + ds4_gpu_routed_mm_addr_pipeline(gate_type); + id up_pipeline = + ds4_gpu_routed_mm_addr_pipeline(up_type); + id down_pipeline = + ds4_gpu_routed_mm_addr_f16_rhs_pipeline(down_type); + if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) { + return 0; + } + + ds4_gpu_mul_mm_id_map_args map_args = + ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args gate_args = + ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + gate_row_bytes, gate_expert_bytes, + 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args up_args = + ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + up_row_bytes, up_expert_bytes, + 1, n_expert, n_tokens); + ds4_gpu_mul_mm_id_args down_args = + ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert, + down_row_bytes, down_expert_bytes, + n_expert, n_expert, n_tokens, + sizeof(uint16_t)); + if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(resources, + resource_count, + 0)) { + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const bool glm_moe_stage_profile = false; + const char *glm_moe_stage_filter = NULL; + double glm_moe_stage_t0 = 0.0; + if (glm_moe_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + glm_moe_stage_t0 = ds4_gpu_now_ms(); + } + + int ok = 1; +#define DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE(name) do { \ + if (ok && glm_moe_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ + strstr(stage_name, glm_moe_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal GLM grouped-address routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \ + "gate=%s down=%s mid=f16 %s=%.3f ms\n", \ + layer_index, n_tokens, pair_rows, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + stage_name, now_ms - glm_moe_stage_t0); \ + } \ + glm_moe_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + + ok = ds4_gpu_encode_mul_mm_id_map(cb, + map_pipeline, + &map_args, + &gate_args, + selectedbuf, + ds4_gpu_tensor_offset(selected)); + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("map"); + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb, + gate_pipeline, + &gate_args, + gate_addrs, + xbuf, + ds4_gpu_tensor_offset(x), + g_moe_gate_scratch_buffer, + 0, + mm_id_threadgroup_bytes, + resources, + resource_count, + 0, + overflow_gate); + } + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("gate"); + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb, + up_pipeline, + &up_args, + up_addrs, + xbuf, + ds4_gpu_tensor_offset(x), + g_moe_gate_scratch_buffer, + (NSUInteger)gate_scratch_bytes, + mm_id_threadgroup_bytes, + resources, + resource_count, + 1, + overflow_up); + } + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("up"); + if (ok) { + ok = ds4_gpu_encode_moe_swiglu_weight(cb, + g_moe_gate_scratch_buffer, + 0, + g_moe_gate_scratch_buffer, + (NSUInteger)gate_scratch_bytes, + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + expert_mid_dim, + pair_rows, + 0.0f, + mid_f16); + } + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("activation_weight"); + + id down_dst = n_expert == 1 ? outbuf : g_moe_down_scratch_buffer; + NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : 0; + if (ok) { + ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb, + down_pipeline, + &down_args, + down_addrs, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + mm_id_threadgroup_bytes, + resources, + resource_count, + 2, + overflow_down); + } + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("down"); + if (ok && n_expert > 1) { + ok = ds4_gpu_encode_moe_sum_experts(cb, + down_dst, + down_dst_off, + outbuf, + ds4_gpu_tensor_offset(out), + out_dim, + n_expert, + n_tokens); + } + DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("sum"); + if (!ok) return 0; + + if (!ds4_gpu_finish_command_buffer(cb, owned, + "GLM grouped-address routed batch MoE")) { + return 0; + } +#undef DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE + } + + return 1; +} + +static int ds4_gpu_glm_routed_moe_batch_tensor_impl( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride, + bool allow_grouped, + bool force_scalar_q4_pair) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !mid || !model_map || !selected || !weights || !x || + n_tokens == 0 || + n_total_expert == 0 || n_expert == 0 || n_expert > 256u || + n_expert > n_total_expert || + expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 || + gate_expert_bytes == 0 || gate_row_bytes == 0 || + up_expert_bytes == 0 || up_row_bytes == 0 || + down_expert_bytes == 0 || down_row_bytes == 0 || + (expert_in_dim % 256u) != 0 || + (expert_mid_dim % 256u) != 0 || + !ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) || + !ds4_gpu_glm_down_type_supported(down_type)) { + return 0; + } + + const uint64_t per_token_mid = (uint64_t)n_expert * expert_mid_dim; + if ((uint64_t)mid_token_stride < per_token_mid || + (uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes || + (uint64_t)expert_mid_dim > UINT64_MAX / gate_row_bytes || + (uint64_t)expert_mid_dim > UINT64_MAX / up_row_bytes || + (uint64_t)out_dim > UINT64_MAX / down_row_bytes || + (uint64_t)n_tokens > UINT64_MAX / expert_in_dim || + (uint64_t)n_tokens > UINT64_MAX / out_dim) { + fprintf(stderr, "ds4: Metal GLM routed batch MoE tensor byte size overflow\n"); + return 0; + } + + const uint64_t full_gate_tensor_bytes = + (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t full_up_tensor_bytes = + (uint64_t)n_total_expert * up_expert_bytes; + const uint64_t full_down_tensor_bytes = + (uint64_t)n_total_expert * down_expert_bytes; + if (gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes || + up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes || + down_expert_bytes != (uint64_t)out_dim * down_row_bytes) { + fprintf(stderr, "ds4: Metal GLM routed batch MoE received inconsistent expert strides\n"); + return 0; + } + if (gate_offset > model_size || full_gate_tensor_bytes > model_size - gate_offset || + up_offset > model_size || full_up_tensor_bytes > model_size - up_offset || + down_offset > model_size || full_down_tensor_bytes > model_size - down_offset) { + fprintf(stderr, "ds4: Metal GLM routed batch MoE tensor range is outside the mapped model\n"); + return 0; + } + + if (allow_grouped && + (!g_ssd_streaming_mode || + ds4_gpu_glm_streaming_prefill_full_layer_active()) && + ds4_gpu_glm_grouped_moe_layer_enabled(layer_index) && + ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, + up_type, + down_type, + n_expert, + n_tokens)) { + return ds4_gpu_glm_routed_moe_batch_grouped_tensor(out, + mid, + model_map, + model_size, + gate_offset, + up_offset, + down_offset, + gate_type, + up_type, + down_type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + selected, + weights, + n_total_expert, + n_expert, + layer_index, + x, + n_tokens); + } + + uint32_t first_expert = 0; + uint32_t n_bind_expert = 0; + ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + gate_offset += (uint64_t)first_expert * gate_expert_bytes; + up_offset += (uint64_t)first_expert * up_expert_bytes; + down_offset += (uint64_t)first_expert * down_expert_bytes; + const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; + const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim; + const uint64_t out_values = (uint64_t)n_tokens * out_dim; + const uint64_t mid_values = + (uint64_t)(n_tokens - 1u) * mid_token_stride + per_token_mid; + const uint64_t selected_values = (uint64_t)n_tokens * n_expert; + if (x_values > UINT64_MAX / sizeof(float) || + out_values > UINT64_MAX / sizeof(float) || + mid_values > UINT64_MAX / sizeof(float) || + selected_values > UINT64_MAX / sizeof(int32_t)) { + fprintf(stderr, "ds4: Metal GLM routed batch MoE activation byte size overflow\n"); + return 0; + } + const uint64_t x_bytes = x_values * sizeof(float); + const uint64_t mid_bytes = mid_values * sizeof(float); + const uint64_t out_bytes = out_values * sizeof(float); + const uint64_t selected_bytes = selected_values * sizeof(int32_t); + const uint64_t weights_bytes = selected_values * sizeof(float); + if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes) { + fprintf(stderr, "ds4: Metal GLM routed batch MoE received undersized activation buffers\n"); + return 0; + } + + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gatebuf = nil; + id upbuf = nil; + id downbuf = nil; + id stream_gate_addr_buf = nil; + id stream_up_addr_buf = nil; + id stream_down_addr_buf = nil; + id stream_overflow_gate = nil; + id stream_overflow_up = nil; + id stream_overflow_down = nil; + ds4_gpu_stream_expert_cache_entry + *stream_resources[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { NULL }; + uint32_t stream_resource_count = 0; + uint32_t stream_unique = 0; + + const BOOL gate_pair_q2 = gate_type == DS4_METAL_TENSOR_Q2_K; + const BOOL gate_pair_q5 = gate_type == DS4_METAL_TENSOR_Q5_K; + const BOOL down_scalar_q2 = down_type == DS4_METAL_TENSOR_Q2_K; + const BOOL down_scalar_q4 = down_type == DS4_METAL_TENSOR_Q4_K; + const BOOL down_simd_q4 = down_scalar_q4; + const BOOL down_simd_q5 = down_type == DS4_METAL_TENSOR_Q5_K; + const BOOL down_simd_q6 = down_type == DS4_METAL_TENSOR_Q6_K; + const BOOL down_simd = down_simd_q4 || down_simd_q5 || down_simd_q6; + const BOOL stream_addr_q2 = + gate_pair_q2 && down_scalar_q2 && + g_glm_q2_k_addr_pair_swiglu2_f32_pipeline != nil && + g_glm_q2_k_addr_down_f32_pipeline != nil; + const BOOL stream_addr_q4 = + !gate_pair_q2 && !gate_pair_q5 && down_scalar_q4 && + g_glm_q4_k_addr_pair_swiglu_f32_pipeline != nil && + g_glm_q4_k_addr_down_f32_pipeline != nil; + BOOL use_stream_expert_addr_table = + g_ssd_streaming_mode && + !ds4_gpu_glm_streaming_prefill_full_layer_active() && + n_tokens > 1 && + (stream_addr_q2 || stream_addr_q4) && + layer_index < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER && + n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && + n_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED && + ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes) && + ds4_gpu_stream_expert_cache_effective_cap(layer_index, + n_total_expert, + n_expert) != 0; + const BOOL use_stream_grouped_addr_table = + use_stream_expert_addr_table && + allow_grouped && + getenv("DS4_METAL_GLM_DISABLE_STREAMING_GROUPED_ADDR_PREFILL") == NULL && + ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type, + up_type, + down_type, + n_expert, + n_tokens); + const BOOL enable_q4_pair4 = true; + const BOOL q4_scalar_pair = false; + const BOOL q4_pair2 = + !use_stream_expert_addr_table && + !gate_pair_q5 && !q4_scalar_pair && + (force_scalar_q4_pair || !enable_q4_pair4); + id pair_pipeline = + use_stream_expert_addr_table ? + (gate_pair_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_f32_pipeline, + "kernel_glm_q2_K_addr_pair_swiglu2_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_f32_pipeline, + "kernel_glm_q4_K_addr_pair_swiglu_f32")) : + gate_pair_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_pair_swiglu_f32_pipeline, + "kernel_glm_q2_K_pair_swiglu_f32") : + gate_pair_q5 ? + ds4_gpu_hot_pipeline(g_glm_q5_k_pair_swiglu_f32_pipeline, + "kernel_glm_q5_K_pair_swiglu_f32") : + (q4_scalar_pair ? + ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu_f32_pipeline, + "kernel_glm_q4_K_pair_swiglu_f32") : + q4_pair2 ? + ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu2_f32_pipeline, + "kernel_glm_q4_K_pair_swiglu2_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu4_f32_pipeline, + "kernel_glm_q4_K_pair_swiglu4_f32")); + id down_pipeline = + use_stream_expert_addr_table ? + (down_scalar_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_addr_down_f32_pipeline, + "kernel_glm_q2_K_addr_down_f32") : + ds4_gpu_hot_pipeline(g_glm_q4_k_addr_down_f32_pipeline, + "kernel_glm_q4_K_addr_down_f32")) : + down_scalar_q2 ? + ds4_gpu_hot_pipeline(g_glm_q2_k_down_f32_pipeline, + "kernel_glm_q2_K_down_f32") : + down_scalar_q4 ? + ds4_gpu_hot_pipeline(g_glm_q4_k_down_f32_pipeline, + "kernel_glm_q4_K_down_f32") : + down_simd_q5 ? + ds4_gpu_hot_pipeline(g_glm_q5_k_down_f32_pipeline, + "kernel_glm_q5_K_down_f32") : + ds4_gpu_hot_pipeline(g_glm_q6_k_down_f32_pipeline, + "kernel_glm_q6_K_down_f32"); + if (!pair_pipeline || !down_pipeline) return 0; + + if (use_stream_expert_addr_table) { + const int had_batch = g_batch_cb != nil; + if (had_batch && ds4_gpu_end_commands() == 0) { + return 0; + } + if (!ds4_gpu_stream_expert_cache_prepare_selected_batch( + model_map, + model_size, + layer_index, + selected, + n_tokens, + n_total_expert, + n_expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf, + stream_resources, + &stream_resource_count, + &stream_unique, + &stream_overflow_gate, + &stream_overflow_up, + &stream_overflow_down)) { + return 0; + } + if (stream_unique == 0) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + if (had_batch && ds4_gpu_begin_commands() == 0) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + if (use_stream_grouped_addr_table) { + return ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor( + out, + mid, + stream_gate_addr_buf, + stream_up_addr_buf, + stream_down_addr_buf, + stream_overflow_gate, + stream_overflow_up, + stream_overflow_down, + gate_type, + up_type, + down_type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + selected, + weights, + n_total_expert, + n_expert, + layer_index, + x, + n_tokens, + stream_resources, + stream_resource_count); + } + } else { + gatebuf = ds4_gpu_wrap_model_range(model_map, model_size, + gate_offset, gate_tensor_bytes, + &gate_inner); + upbuf = ds4_gpu_wrap_model_range(model_map, model_size, + up_offset, up_tensor_bytes, + &up_inner); + downbuf = ds4_gpu_wrap_model_range(model_map, model_size, + down_offset, down_tensor_bytes, + &down_inner); + if (!gatebuf || !upbuf || !downbuf) return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + const bool glm_moe_stage_profile = false; + const char *glm_moe_stage_filter = NULL; + const char *glm_pair_path = use_stream_expert_addr_table ? + (gate_pair_q2 ? "q2_stream_addr_swiglu" : + "q4_stream_addr_swiglu") : + gate_pair_q2 ? "q2_scalar_swiglu" : + gate_pair_q5 ? "q5_pair_simd_swiglu" : + (q4_scalar_pair ? "q4_scalar_swiglu" : + (q4_pair2 ? "q4_pair2_simd_swiglu" : + "q4_pair4_simd_swiglu")); + const char *glm_down_path = + use_stream_expert_addr_table ? + (down_scalar_q2 ? "q2_stream_addr_down" : "q4_stream_addr_down_simd") : + down_scalar_q2 ? "q2_down_scalar" : + down_scalar_q4 ? "q4_down_simd" : + down_simd_q5 ? "q5_down_simd" : "q6_down_simd"; + double glm_moe_stage_t0 = 0.0; + if (glm_moe_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + glm_moe_stage_t0 = ds4_gpu_now_ms(); + } + int ok = 1; +#define DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE(name) do { \ + if (ok && glm_moe_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !glm_moe_stage_filter || !glm_moe_stage_filter[0] || \ + strstr(stage_name, glm_moe_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal GLM routed MoE batch stage layer=%u tokens=%u experts=%u " \ + "gate=%s down=%s pair=%s down_path=%s %s=%.3f ms\n", \ + layer_index, n_tokens, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + glm_pair_path, glm_down_path, \ + stage_name, now_ms - glm_moe_stage_t0); \ + } \ + glm_moe_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + + ds4_gpu_glm_routed_moe_args args = { + .tp_rank = g_tp_split_rank, + .tp_world = g_tp_split_world, + .tp_expert_base = (int32_t)first_expert, + .in_dim = expert_in_dim, + .mid_dim = expert_mid_dim, + .out_dim = out_dim, + .n_total_expert = n_total_expert, + .n_expert_used = n_expert, + .n_tokens = n_tokens, + .mid_token_stride = mid_token_stride, + .down_type = down_type, + .gate_expert_bytes = gate_expert_bytes, + .gate_row_bytes = gate_row_bytes, + .up_expert_bytes = up_expert_bytes, + .up_row_bytes = up_row_bytes, + .down_expert_bytes = down_expert_bytes, + .down_row_bytes = down_row_bytes, + }; + const NSUInteger pair_x_groups = + gate_pair_q2 ? (use_stream_expert_addr_table ? + (NSUInteger)((expert_mid_dim + 1u) / 2u) : + (NSUInteger)((expert_mid_dim + 7u) / 8u)) : + gate_pair_q5 ? (NSUInteger)((expert_mid_dim + 7u) / 8u) : + use_stream_expert_addr_table ? (NSUInteger)((expert_mid_dim + 3u) / 4u) : + q4_scalar_pair ? (NSUInteger)expert_mid_dim : + q4_pair2 ? (NSUInteger)((expert_mid_dim + 1u) / 2u) : + (NSUInteger)((expert_mid_dim + 7u) / 8u); + const NSUInteger pair_threadgroup_bytes = + q4_scalar_pair ? 512u * sizeof(float) : 0u; + const NSUInteger pair_threads = + q4_scalar_pair ? 256u : 64u; + const NSUInteger down_x_groups = + down_scalar_q2 ? (NSUInteger)((out_dim + 7u) / 8u) : + down_simd_q4 ? (NSUInteger)((out_dim + 3u) / 4u) : + down_simd_q5 ? (NSUInteger)((out_dim + 3u) / 4u) : + down_simd_q6 ? (NSUInteger)((out_dim + 3u) / 4u) : + (NSUInteger)out_dim; + const NSUInteger down_threadgroup_bytes = + down_simd ? 0u : 256u * sizeof(float); + const NSUInteger down_threads = + down_simd ? 64u : 256u; + if (use_stream_expert_addr_table && + !ds4_gpu_stream_expert_cache_mark_entries_inflight(stream_resources, + stream_resource_count, + 0)) { + return 0; + } + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pair_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:use_stream_expert_addr_table ? stream_gate_addr_buf : gatebuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)gate_inner + atIndex:1]; + [enc setBuffer:use_stream_expert_addr_table ? stream_up_addr_buf : upbuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)up_inner + atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:4]; + [enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:5]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6]; + if (use_stream_expert_addr_table) { + for (uint32_t i = 0; i < stream_resource_count; i++) { + [enc useResource:stream_resources[i]->gate_buffer usage:MTLResourceUsageRead]; + [enc useResource:stream_resources[i]->up_buffer usage:MTLResourceUsageRead]; + } + if (stream_overflow_gate) [enc useResource:stream_overflow_gate usage:MTLResourceUsageRead]; + if (stream_overflow_up) [enc useResource:stream_overflow_up usage:MTLResourceUsageRead]; + } + if (pair_threadgroup_bytes != 0u) { + [enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(pair_x_groups, + (NSUInteger)n_expert, + (NSUInteger)n_tokens) + threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE("pair"); + + enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:down_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:use_stream_expert_addr_table ? stream_down_addr_buf : downbuf + offset:use_stream_expert_addr_table ? 0u : (NSUInteger)down_inner + atIndex:1]; + [enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + if (use_stream_expert_addr_table) { + for (uint32_t i = 0; i < stream_resource_count; i++) { + [enc useResource:stream_resources[i]->down_buffer usage:MTLResourceUsageRead]; + } + if (stream_overflow_down) [enc useResource:stream_overflow_down usage:MTLResourceUsageRead]; + } + if (down_threadgroup_bytes != 0u) { + [enc setThreadgroupMemoryLength:down_threadgroup_bytes atIndex:0]; + } + [enc dispatchThreadgroups:MTLSizeMake(down_x_groups, + (NSUInteger)n_tokens, + 1) + threadsPerThreadgroup:MTLSizeMake(down_threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE("down"); + + if (!ok) return 0; + if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM routed batch MoE")) return 0; +#undef DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE + } + + return 1; +} + +int ds4_gpu_glm_routed_moe_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride, + bool force_resident) { + (void)force_resident; + return ds4_gpu_glm_routed_moe_batch_tensor_impl(out, + mid, + model_map, + model_size, + gate_offset, + up_offset, + down_offset, + gate_type, + up_type, + down_type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + selected, + weights, + n_total_expert, + n_expert, + layer_index, + x, + n_tokens, + mid_token_stride, + true, + false); +} + +int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *mid, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t up_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + uint32_t layer_index, + const ds4_gpu_tensor *x, + uint32_t n_tokens, + uint32_t mid_token_stride) { + return ds4_gpu_glm_routed_moe_batch_tensor_impl(out, + mid, + model_map, + model_size, + gate_offset, + up_offset, + down_offset, + gate_type, + up_type, + down_type, + gate_expert_bytes, + gate_row_bytes, + up_expert_bytes, + up_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + selected, + weights, + n_total_expert, + n_expert, + layer_index, + x, + n_tokens, + mid_token_stride, + false, + false); +} + +int ds4_gpu_router_select_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + uint64_t hash_offset, + uint32_t hash_rows, + uint32_t token, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_expert_groups, + uint32_t n_group_used, + bool has_bias, + bool hash_mode, + const ds4_gpu_tensor *logits) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !weights || !probs || !logits || !model_map || + n_expert == 0 || n_expert_used == 0) return 0; + if (hash_mode && token >= hash_rows) return 0; + if (n_expert_groups > 1u || n_group_used > 0u) { + fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n"); + return 0; + } + + @autoreleasepool { + id logitsbuf = ds4_gpu_tensor_buffer(logits); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || + ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) || + ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float)) { + fprintf(stderr, "ds4: Metal router select received undersized buffers\n"); + return 0; + } + + uint64_t bias_inner = 0; + uint64_t hash_inner = 0; + id biasbuf = nil; + id hashbuf = nil; + NSUInteger bias_set_offset = 0; + NSUInteger hash_set_offset = 0; + if (has_bias && !hash_mode) { + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner); + if (!biasbuf) return 0; + bias_set_offset = (NSUInteger)bias_inner; + } + if (hash_mode) { + const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t); + hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner); + if (!hashbuf) return 0; + hash_set_offset = (NSUInteger)hash_inner; + } + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + const int32_t token_i32 = (int32_t)token; + int ok = cb && + ds4_gpu_encode_router_select(cb, + selected, + weights, + probs, + logitsbuf, + ds4_gpu_tensor_offset(logits), + biasbuf, + bias_set_offset, + hashbuf, + hash_set_offset, + nil, + 0, + &token_i32, + hash_rows, + 1, + n_expert, + n_expert_used, + expert_weight_scale, + has_bias && !hash_mode, + hash_mode); + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + if (!ok) return 0; + } + + return 1; +} + +int ds4_gpu_router_select_batch_tensor( + ds4_gpu_tensor *selected, + ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, + const void *model_map, + uint64_t model_size, + uint64_t bias_offset, + uint64_t hash_offset, + uint32_t hash_rows, + uint32_t n_expert_groups, + uint32_t n_group_used, + bool has_bias, + bool hash_mode, + const ds4_gpu_tensor *logits, + const ds4_gpu_tensor *tokens, + uint32_t n_expert, + uint32_t n_expert_used, + float expert_weight_scale, + uint32_t n_tokens) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!selected || !weights || !probs || !logits || !tokens || !model_map || + n_expert == 0 || n_expert_used == 0 || n_tokens == 0) return 0; + if (n_expert_groups > 1u || n_group_used > 0u) { + fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n"); + return 0; + } + + @autoreleasepool { + id logitsbuf = ds4_gpu_tensor_buffer(logits); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + id probsbuf = ds4_gpu_tensor_buffer(probs); + id tokensbuf = ds4_gpu_tensor_buffer(tokens); + if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || !tokensbuf || + ds4_gpu_tensor_bytes(logits) < (uint64_t)n_tokens * n_expert * sizeof(float) || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_tokens * n_expert_used * sizeof(int) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_tokens * n_expert_used * sizeof(float) || + ds4_gpu_tensor_bytes(probs) < (uint64_t)n_tokens * n_expert * sizeof(float) || + ds4_gpu_tensor_bytes(tokens) < (uint64_t)n_tokens * sizeof(int32_t)) { + fprintf(stderr, "ds4: Metal router batch select received undersized buffers\n"); + return 0; + } + + uint64_t bias_inner = 0; + uint64_t hash_inner = 0; + id biasbuf = nil; + id hashbuf = nil; + NSUInteger bias_set_offset = 0; + NSUInteger hash_set_offset = 0; + if (has_bias && !hash_mode) { + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner); + if (!biasbuf) return 0; + bias_set_offset = (NSUInteger)bias_inner; + } + if (hash_mode) { + const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t); + hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner); + if (!hashbuf) return 0; + hash_set_offset = (NSUInteger)hash_inner; + } + + const bool had_batch = g_batch_cb != nil; + if (!had_batch && ds4_gpu_begin_commands() == 0) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + int ok = cb && + ds4_gpu_encode_router_select(cb, + selected, + weights, + probs, + logitsbuf, + ds4_gpu_tensor_offset(logits), + biasbuf, + bias_set_offset, + hashbuf, + hash_set_offset, + tokensbuf, + ds4_gpu_tensor_offset(tokens), + NULL, + hash_rows, + n_tokens, + n_expert, + n_expert_used, + expert_weight_scale, + has_bias && !hash_mode, + hash_mode); + if (!had_batch) { + ok = ds4_gpu_end_commands() != 0 && ok; + } + if (!ok) return 0; + } + + return 1; +} + +int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected) { + if (n_selected > DS4_METAL_MAX_ROUTED_EXPERT_USED || + (!selected && n_selected != 0)) return 0; + for (uint32_t i = 0; i < n_selected; i++) { + g_routed_moe_selected_override[i] = selected[i]; + } + g_routed_moe_selected_override_n = n_selected; + return 1; +} + +int ds4_gpu_routed_moe_one_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *experts, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + float clamp, + const ds4_gpu_tensor *x, + const ds4_gpu_tensor *add_in, + uint32_t layer_index, + bool force_resident) { + if (!g_initialized && !ds4_gpu_init()) return 0; + /* TP sharding: only the owned contiguous expert range is mapped, + * so bind from the owned base, validate only its bytes, and tell the + * kernels the first expert id present at that base. */ + uint32_t first_expert = 0; + uint32_t n_bind_expert = 0; + ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + const int32_t tp_expert_base_host = (int32_t)first_expert; + gate_offset += (uint64_t)first_expert * gate_expert_bytes; + up_offset += (uint64_t)first_expert * gate_expert_bytes; + down_offset += (uint64_t)first_expert * down_expert_bytes; + + if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights || + n_total_expert == 0 || n_expert == 0 || + n_expert > DS4_METAL_MAX_ROUTED_EXPERT_USED || + gate_expert_bytes == 0 || down_expert_bytes == 0 || + gate_row_bytes == 0 || down_row_bytes == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32044); + return 0; + } + if ((expert_in_dim % 256u) != 0 || (expert_mid_dim % 256u) != 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32047); return 0; } + ds4_gpu_stream_expert_cache_note_token(layer_index); + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id expertsbuf = ds4_gpu_tensor_buffer(experts); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id selected_exec_buf = selectedbuf; + NSUInteger selected_exec_off = ds4_gpu_tensor_offset(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + const uint64_t x_bytes = (uint64_t)expert_in_dim * sizeof(float); + const uint64_t mid_bytes = (uint64_t)n_expert * expert_mid_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float); + if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(gate) < mid_bytes || + ds4_gpu_tensor_bytes(up) < mid_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert * sizeof(int) || + ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert * sizeof(float)) { + fprintf(stderr, "ds4: Metal routed tensor MoE received undersized activation buffers\n"); + return 0; + } + if (n_expert > 1 && + (!expertsbuf || + ds4_gpu_tensor_bytes(experts) < (uint64_t)n_expert * out_dim * sizeof(float))) { + fprintf(stderr, "ds4: Metal routed tensor MoE received undersized expert output buffer\n"); + return 0; + } + + if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed MoE tensor byte size overflow\n"); + return 0; + } + const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gate_buf = nil; + id up_buf = nil; + id down_buf = nil; + __unsafe_unretained id gate_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil }; + __unsafe_unretained id up_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil }; + __unsafe_unretained id down_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil }; + ds4_gpu_stream_expert_cache_entry *stream_slot_entries[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { NULL }; + NSUInteger gate_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + NSUInteger up_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + NSUInteger down_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + uint64_t stream_gate_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + uint64_t stream_up_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + uint64_t stream_down_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + id stream_gate_addr_buf = nil; + id stream_up_addr_buf = nil; + id stream_down_addr_buf = nil; + bool use_stream_expert_addr_table = false; + bool use_stream_expert_masked_addr_table = false; + bool use_stream_compact_addr_table = false; + bool use_stream_expert_cache = false; + bool use_stream_expert_split_candidate = false; + bool use_stream_expert_split_deferred = false; + bool stream_expert_split_completed = false; + uint32_t stream_expert_resident_mask = 0; + uint32_t stream_expert_missing_mask = 0; + __unsafe_unretained id gate_group6_bufs[6] = { nil, nil, nil, nil, nil, nil }; + __unsafe_unretained id up_group6_bufs[6] = { nil, nil, nil, nil, nil, nil }; + __unsafe_unretained id down_group6_bufs[6] = { nil, nil, nil, nil, nil, nil }; + NSUInteger gate_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + NSUInteger up_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + NSUInteger down_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 }; + __unsafe_unretained id gate_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil }; + __unsafe_unretained id up_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil }; + __unsafe_unretained id down_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil }; + NSUInteger gate_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 }; + NSUInteger up_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 }; + NSUInteger down_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 }; + __unsafe_unretained id gate_group24_bufs[24] = { nil }; + __unsafe_unretained id up_group24_bufs[24] = { nil }; + __unsafe_unretained id down_group24_bufs[24] = { nil }; + NSUInteger gate_group24_offsets[24] = { 0 }; + NSUInteger up_group24_offsets[24] = { 0 }; + NSUInteger down_group24_offsets[24] = { 0 }; + DS4MetalQ4ExpertTable *gate_table = nil; + DS4MetalQ4ExpertTable *up_table = nil; + DS4MetalQ4ExpertTable *down_table = nil; + id q4_table_layer_residency = nil; + int32_t selected_ids[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 }; + + const uint32_t n_tokens = 1; + const uint32_t pair_rows = n_tokens * n_expert; + const uint64_t down_scratch_bytes = (uint64_t)pair_rows * out_dim * sizeof(float); + if ((n_expert > 1 && !expertsbuf && + !ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, + &g_moe_down_scratch_bytes, + (NSUInteger)down_scratch_bytes, + "ds4_moe_down_scratch"))) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32148); + return 0; + } + + const uint32_t gate_nr0 = ds4_gpu_routed_mv_nr0(gate_type); + const uint32_t down_nr0 = ds4_gpu_routed_mv_nr0(down_type); + id gate_mv_pipeline = ds4_gpu_routed_mv_pipeline(gate_type); + id down_mv_pipeline = ds4_gpu_routed_mv_pipeline(down_type); + if (gate_nr0 == 0 || down_nr0 == 0 || !gate_mv_pipeline || !down_mv_pipeline) { + fprintf(stderr, "ds4: unsupported Metal routed MoE quant types gate=%u down=%u\n", + gate_type, down_type); + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32158); + return 0; + } + + ds4_gpu_mul_mv_id_args gate_args = + ds4_gpu_make_mul_mv_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + gate_row_bytes, gate_expert_bytes, + 1, n_expert, n_tokens, gate_nr0); + ds4_gpu_mul_mv_id_args down_args = + ds4_gpu_make_mul_mv_id_args(expert_mid_dim, out_dim, n_total_expert, + down_row_bytes, down_expert_bytes, + n_expert, n_expert, n_tokens, down_nr0); + /* Tensor-parallel expert ownership; non-TP calls keep tp_world at 1. */ + gate_args.tp_rank = g_tp_split_rank; + gate_args.tp_world = g_tp_split_world; + gate_args.tp_expert_base = tp_expert_base_host; + down_args.tp_rank = g_tp_split_rank; + down_args.tp_world = g_tp_split_world; + down_args.tp_addend = add_in != NULL; + down_args.tp_expert_base = tp_expert_base_host; + + const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type); + const NSUInteger down_smem = ds4_gpu_routed_mv_smem(down_type); + const NSUInteger gate_nsg = ds4_gpu_routed_mv_nsg(gate_type); + const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type); + const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type); + const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type); + int ok = 1; + const bool write_clamped_moe = + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL; + id pair_swiglu_pipeline = nil; + if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) { + pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; + } else if (gate_type == DS4_METAL_TENSOR_Q4_K) { + pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; + } + const bool fuse_pair_swiglu = + !g_quality_mode && + !write_clamped_moe && + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && + pair_swiglu_pipeline != nil; + id down_sum6_pipeline = nil; + if (down_type == DS4_METAL_TENSOR_Q2_K) { + down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline; + } else if (down_type == DS4_METAL_TENSOR_Q4_K) { + down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline; + } else if (down_type == DS4_METAL_TENSOR_IQ2_XXS && + g_tp_split_world == 2) { + /* IQ2 down-sum exists for the GLM TP resident split only; the + * streaming paths must keep their addr/masked chain (expert + * bytes are not at their model-map offsets when streamed). */ + down_sum6_pipeline = g_moe_mul_mv_id_iq2_xxs_sum6_pipeline; + } + const bool direct_down_sum = + !g_quality_mode && + (n_expert == 6 || (n_expert == 8 && g_tp_split_world == 2)) && + n_tokens == 1 && + down_sum6_pipeline != nil; + /* The expert-ownership split lives only in the fused id pair+sum6 + * kernels; every other routed variant would silently compute full + * sums on both ranks and double the combine. Fail fast instead. */ + if ((g_tp_split_world > 1 || add_in) && !(fuse_pair_swiglu && direct_down_sum)) { + fprintf(stderr, "ds4: tensor-parallel routed MoE requires the fused pair+sum6 decode path\n"); + return 0; + } + const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; + /* + * The grouped Q4 experiment keeps selected IDs on GPU, but it also walks + * every expert window in the layer. On PRO Q4 this measured far slower + * than the active selected-slot path, so keep it opt-in for profiling. + */ + const bool enable_q4_grouped_experts = + getenv("DS4_METAL_ENABLE_Q4_GROUPED_EXPERTS") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUPED_EXPERTS") == NULL; + const bool use_q4_grouped_experts = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert >= 128 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group_q4_k_sum6_pipeline != nil && + enable_q4_grouped_experts; + const uint32_t q4_expert_group_size = + use_q4_grouped_experts ? ds4_gpu_q4_expert_group_size(n_total_expert) : 0; + const bool q4_grouped_boundary = + use_q4_grouped_experts && + g_batch_cb != nil && + getenv("DS4_METAL_DISABLE_Q4_GROUPED_BOUNDARY") == NULL; + const bool q4_grouped_cache_views = + getenv("DS4_METAL_Q4_GROUPED_CACHE_VIEWS") != NULL; + const uint32_t q4_group6_expert_group_size = 64; + const bool use_q4_group6_experts = + !use_q4_grouped_experts && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == q4_group6_expert_group_size * 6u && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group6_q4_k_sum6_pipeline != nil && + getenv("DS4_METAL_ENABLE_Q4_GROUP6_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP6_EXPERT_TABLE") == NULL; + const uint32_t q4_group8_expert_group_size = 48; + const bool use_q4_group8_experts = + !use_q4_grouped_experts && + !use_q4_group6_experts && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == q4_group8_expert_group_size * 8u && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group8_q4_k_sum6_pipeline != nil && + getenv("DS4_METAL_ENABLE_Q4_GROUP8_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP8_EXPERT_TABLE") == NULL; + const uint32_t q4_group24_expert_group_size = 16; + const bool use_q4_group24_experts = + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == q4_group24_expert_group_size * 24u && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + direct_down_sum && + g_moe_mul_mv_group24_q4_k_id_pipeline != nil && + g_moe_mul_mv_group24_q4_k_sum6_pipeline != nil && + getenv("DS4_METAL_ENABLE_Q4_GROUP24_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP24_EXPERT_TABLE") == NULL; + const bool q4_group24_exact_views = + use_q4_group24_experts && + getenv("DS4_METAL_Q4_GROUP24_EXACT_VIEWS") != NULL && + getenv("DS4_METAL_Q4_GROUP24_BASE_VIEWS") == NULL; + const uint64_t max_buffer_len = g_device ? (uint64_t)[g_device maxBufferLength] : 0; + const bool can_wrap_q4_exact_tensors = + max_buffer_len != 0 && + gate_tensor_bytes <= max_buffer_len && + down_tensor_bytes <= max_buffer_len; + /* + * The full-tensor Q4 ID path is the closest arithmetic analogue to IQ2, + * but PRO Q4 routed tensors are multi-GiB. Keep it opt-in: even with + * per-layer command boundaries it is much slower than binding only the + * six active experts on current M3 Ultra Metal. + */ + const bool enable_q4_exact_tensor_id = + getenv("DS4_METAL_ENABLE_Q4_EXACT_TENSOR_ID") != NULL && + getenv("DS4_METAL_DISABLE_Q4_EXACT_TENSOR_ID") == NULL; + const bool use_q4_exact_tensor_id = + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + !use_q4_group24_experts && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == 384 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + can_wrap_q4_exact_tensors && + fuse_pair_swiglu && + direct_down_sum && + enable_q4_exact_tensor_id; + const bool q4_exact_boundary = + use_q4_exact_tensor_id && + g_batch_cb != nil && + getenv("DS4_METAL_DISABLE_Q4_EXACT_BOUNDARY") == NULL; + const bool q4_expert_table_auto = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes); + const bool q4_expert_address_auto = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + ds4_gpu_pro_q4_expert_address_auto_enabled(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes); + const bool q4_table_queue_residency = + ds4_gpu_q4_table_queue_residency_enabled(q4_expert_table_auto || + q4_expert_address_auto); + const bool use_q4_expert_address_table = + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + !use_q4_group24_experts && + !use_q4_exact_tensor_id && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == 384 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil && + (getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL || + q4_expert_address_auto) && + (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL || + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL || + q4_table_queue_residency || + ds4_gpu_q4_table_model_residency_enabled() || + getenv("DS4_METAL_USE_QUEUE_RESIDENCY_SET") != NULL) && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL; + const bool enable_q4_expert_table = + getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || + q4_expert_table_auto; + const bool use_q4_expert_table = + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + !use_q4_group24_experts && + !use_q4_exact_tensor_id && + !use_q4_expert_address_table && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == 384 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_table_q4_k_sum6_pipeline != nil && + g_moe_table_q4_pair_gate_encoder != nil && + g_moe_table_q4_pair_up_encoder != nil && + g_moe_table_q4_sum_down_encoder != nil && + enable_q4_expert_table && + (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") != NULL || + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL || + q4_table_queue_residency || + ds4_gpu_q4_table_model_residency_enabled() || + getenv("DS4_METAL_USE_QUEUE_RESIDENCY_SET") != NULL) && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL; + const bool q4_table_boundary = + use_q4_expert_table && + g_batch_cb != nil && + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL && + !q4_table_queue_residency && + getenv("DS4_METAL_DISABLE_Q4_TABLE_BOUNDARY") == NULL; + const bool enable_q4_gather_slots = + getenv("DS4_METAL_ENABLE_Q4_GATHER_SLOTS") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GATHER_SLOTS") == NULL; + const bool use_q4_gather_slots = + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + !use_q4_group24_experts && + !use_q4_exact_tensor_id && + !use_q4_expert_address_table && + !use_q4_expert_table && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert == q4_group6_expert_group_size * 6u && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + fuse_pair_swiglu && + direct_down_sum && + g_moe_q4_gather_slots6_pipeline != nil && + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil && + enable_q4_gather_slots; + const bool use_q4_selected_slots = + !force_resident && + !use_q4_grouped_experts && + !use_q4_group6_experts && + !use_q4_group8_experts && + !use_q4_group24_experts && + !use_q4_exact_tensor_id && + !use_q4_expert_address_table && + !use_q4_expert_table && + !use_q4_gather_slots && + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens == 1 && + n_total_expert >= 128 && + (g_ssd_streaming_mode || + (gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes)) && + fuse_pair_swiglu && + direct_down_sum && + ds4_gpu_q4_selected_paths_allowed() && + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil && + getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL; + const bool use_iq2_selected_slots = + !force_resident && + g_ssd_streaming_mode && + gate_type == DS4_METAL_TENSOR_IQ2_XXS && + down_type == DS4_METAL_TENSOR_Q2_K && + n_expert == 6 && + n_tokens == 1 && + fuse_pair_swiglu && + direct_down_sum && + g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_slots6_q2_k_sum6_pipeline != nil && + getenv("DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS") == NULL; + const bool use_iq2_stream_addr_table = + !force_resident && + g_ssd_streaming_mode && + gate_type == DS4_METAL_TENSOR_IQ2_XXS && + down_type == DS4_METAL_TENSOR_IQ2_XXS && + n_expert <= DS4_METAL_MAX_ROUTED_EXPERT_USED && + n_tokens == 1 && + fuse_pair_swiglu && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_iq2_xxs_pipeline != nil && + getenv("DS4_METAL_DISABLE_IQ2_STREAM_ADDR_TABLE") == NULL; + const bool use_selected_slots = + use_q4_selected_slots || use_iq2_selected_slots || use_iq2_stream_addr_table; + id slots_pair_swiglu_pipeline = + use_iq2_selected_slots ? g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline : + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline; + id slots_sum6_pipeline = + use_iq2_selected_slots ? g_moe_mul_mv_slots6_q2_k_sum6_pipeline : + g_moe_mul_mv_slots6_q4_k_sum6_pipeline; + const char *selected_profile_env = getenv("DS4_METAL_SELECTED_PROFILE"); + if (!selected_profile_env) { + selected_profile_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE"); + } + const char *selected_profile_layer_env = getenv("DS4_METAL_SELECTED_PROFILE_LAYER"); + if (!selected_profile_layer_env) { + selected_profile_layer_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE_LAYER"); + } + bool selected_profile_layer_match = true; + if (selected_profile_layer_env && selected_profile_layer_env[0]) { + char *end = NULL; + long layer = strtol(selected_profile_layer_env, &end, 10); + selected_profile_layer_match = + end && *end == '\0' && layer >= 0 && (uint32_t)layer == layer_index; + } + const bool selected_profile = + use_selected_slots && + selected_profile_env != NULL && + selected_profile_layer_match; + const bool q4_selected_shared_event = + use_q4_selected_slots && + getenv("DS4_METAL_Q4_SELECTED_SHARED_EVENT") != NULL; + const bool q4_selected_base_views = + use_q4_selected_slots && + getenv("DS4_METAL_Q4_SELECTED_USE_BASE_VIEWS") != NULL && + getenv("DS4_METAL_Q4_SELECTED_EXACT_VIEWS") == NULL; + const bool q4_selected_transient_views = + use_q4_selected_slots && + !q4_selected_base_views && + getenv("DS4_METAL_Q4_SELECTED_TRANSIENT_VIEWS") != NULL; + const char *q4_selected_view_mode = + q4_selected_base_views ? "base" : + (q4_selected_transient_views ? "transient" : "cached"); + if (!use_selected_slots) { + g_routed_moe_selected_override_n = 0; + } + if (use_q4_expert_address_table) { + gate_table = ds4_gpu_q4_expert_address_table(model_map, + model_size, + gate_offset, + gate_expert_bytes, + n_total_expert); + up_table = ds4_gpu_q4_expert_address_table(model_map, + model_size, + up_offset, + gate_expert_bytes, + n_total_expert); + down_table = ds4_gpu_q4_expert_address_table(model_map, + model_size, + down_offset, + down_expert_bytes, + n_total_expert); + if (!gate_table || !up_table || !down_table) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32546); + return 0; + } + q4_table_layer_residency = + ds4_gpu_q4_expert_layer_residency_set(gate_table, + up_table, + down_table, + q4_expert_address_auto); + } else if (use_q4_expert_table) { + gate_table = ds4_gpu_q4_expert_table(model_map, + model_size, + gate_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_gate_encoder); + up_table = ds4_gpu_q4_expert_table(model_map, + model_size, + up_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_up_encoder); + down_table = ds4_gpu_q4_expert_table(model_map, + model_size, + down_offset, + down_expert_bytes, + n_total_expert, + g_moe_table_q4_sum_down_encoder); + if (!gate_table || !up_table || !down_table) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32573); + return 0; + } + q4_table_layer_residency = + ds4_gpu_q4_expert_layer_residency_set(gate_table, + up_table, + down_table, + q4_expert_table_auto); + } else if (use_q4_exact_tensor_id) { + gate_buf = ds4_gpu_wrap_model_exact_range(model_map, + model_size, + gate_offset, + gate_tensor_bytes, + &gate_inner); + up_buf = ds4_gpu_wrap_model_exact_range(model_map, + model_size, + up_offset, + gate_tensor_bytes, + &up_inner); + down_buf = ds4_gpu_wrap_model_exact_range(model_map, + model_size, + down_offset, + down_tensor_bytes, + &down_inner); + if (!gate_buf || !up_buf || !down_buf) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32601); return 0; } + } else if (use_q4_group6_experts || use_q4_gather_slots) { + if ((uint64_t)q4_group6_expert_group_size > UINT64_MAX / gate_expert_bytes || + (uint64_t)q4_group6_expert_group_size > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group6 byte size overflow\n"); + return 0; + } + const uint64_t gate_group_bytes = + (uint64_t)q4_group6_expert_group_size * gate_expert_bytes; + const uint64_t down_group_bytes = + (uint64_t)q4_group6_expert_group_size * down_expert_bytes; + for (uint32_t i = 0; i < 6; i++) { + const uint64_t gate_rel = (uint64_t)i * gate_group_bytes; + const uint64_t down_rel = (uint64_t)i * down_group_bytes; + if (gate_rel > gate_tensor_bytes || + gate_group_bytes > gate_tensor_bytes - gate_rel || + down_rel > down_tensor_bytes || + down_group_bytes > down_tensor_bytes - down_rel || + gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group6 offset overflow\n"); + return 0; + } + + uint64_t group_inner = 0; + gate_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset + gate_rel, + gate_group_bytes, + &group_inner); + gate_group6_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + up_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset + gate_rel, + gate_group_bytes, + &group_inner); + up_group6_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + down_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset + down_rel, + down_group_bytes, + &group_inner); + down_group6_offsets[i] = (NSUInteger)group_inner; + if (!gate_group6_bufs[i] || !up_group6_bufs[i] || !down_group6_bufs[i]) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32643); + return 0; + } + } + if (use_q4_gather_slots) { + if (gate_expert_bytes > NSUIntegerMax || + down_expert_bytes > NSUIntegerMax || + gate_expert_bytes > UINT64_MAX / 6u || + down_expert_bytes > UINT64_MAX / 6u || + 6ull * gate_expert_bytes > NSUIntegerMax || + 6ull * down_expert_bytes > NSUIntegerMax) { + fprintf(stderr, "ds4: Metal routed MoE Q4 gather scratch byte size overflow\n"); + return 0; + } + const NSUInteger gate_slots_bytes = (NSUInteger)(6ull * gate_expert_bytes); + const NSUInteger down_slots_bytes = (NSUInteger)(6ull * down_expert_bytes); + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_q4_gate_slots_buffer, + &g_moe_q4_gate_slots_bytes, + gate_slots_bytes, + "ds4_moe_q4_gate_slots") || + !ds4_gpu_ensure_scratch_buffer(&g_moe_q4_up_slots_buffer, + &g_moe_q4_up_slots_bytes, + gate_slots_bytes, + "ds4_moe_q4_up_slots") || + !ds4_gpu_ensure_scratch_buffer(&g_moe_q4_down_slots_buffer, + &g_moe_q4_down_slots_bytes, + down_slots_bytes, + "ds4_moe_q4_down_slots")) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32670); + return 0; + } + for (uint32_t i = 0; i < 6; i++) { + gate_slot_bufs[i] = g_moe_q4_gate_slots_buffer; + up_slot_bufs[i] = g_moe_q4_up_slots_buffer; + down_slot_bufs[i] = g_moe_q4_down_slots_buffer; + gate_slot_offsets[i] = (NSUInteger)((uint64_t)i * gate_expert_bytes); + up_slot_offsets[i] = (NSUInteger)((uint64_t)i * gate_expert_bytes); + down_slot_offsets[i] = (NSUInteger)((uint64_t)i * down_expert_bytes); + } + } + } else if (use_q4_group8_experts) { + if ((uint64_t)q4_group8_expert_group_size > UINT64_MAX / gate_expert_bytes || + (uint64_t)q4_group8_expert_group_size > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group8 byte size overflow\n"); + return 0; + } + const uint64_t gate_group_bytes = + (uint64_t)q4_group8_expert_group_size * gate_expert_bytes; + const uint64_t down_group_bytes = + (uint64_t)q4_group8_expert_group_size * down_expert_bytes; + for (uint32_t i = 0; i < 8; i++) { + const uint64_t gate_rel = (uint64_t)i * gate_group_bytes; + const uint64_t down_rel = (uint64_t)i * down_group_bytes; + if (gate_rel > gate_tensor_bytes || + gate_group_bytes > gate_tensor_bytes - gate_rel || + down_rel > down_tensor_bytes || + down_group_bytes > down_tensor_bytes - down_rel || + gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group8 offset overflow\n"); + return 0; + } + + uint64_t group_inner = 0; + gate_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset + gate_rel, + gate_group_bytes, + &group_inner); + gate_group8_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + up_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset + gate_rel, + gate_group_bytes, + &group_inner); + up_group8_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + down_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset + down_rel, + down_group_bytes, + &group_inner); + down_group8_offsets[i] = (NSUInteger)group_inner; + if (!gate_group8_bufs[i] || !up_group8_bufs[i] || !down_group8_bufs[i]) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32727); + return 0; + } + } + } else if (use_q4_group24_experts) { + if ((uint64_t)q4_group24_expert_group_size > UINT64_MAX / gate_expert_bytes || + (uint64_t)q4_group24_expert_group_size > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group24 byte size overflow\n"); + return 0; + } + const uint64_t gate_group_bytes = + (uint64_t)q4_group24_expert_group_size * gate_expert_bytes; + const uint64_t down_group_bytes = + (uint64_t)q4_group24_expert_group_size * down_expert_bytes; + for (uint32_t i = 0; i < 24; i++) { + const uint64_t gate_rel = (uint64_t)i * gate_group_bytes; + const uint64_t down_rel = (uint64_t)i * down_group_bytes; + if (gate_rel > gate_tensor_bytes || + gate_group_bytes > gate_tensor_bytes - gate_rel || + down_rel > down_tensor_bytes || + down_group_bytes > down_tensor_bytes - down_rel || + gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal routed MoE Q4 group24 offset overflow\n"); + return 0; + } + + uint64_t group_inner = 0; + gate_group24_bufs[i] = q4_group24_exact_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + gate_offset + gate_rel, + gate_group_bytes, + &group_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset + gate_rel, + gate_group_bytes, + &group_inner); + gate_group24_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + up_group24_bufs[i] = q4_group24_exact_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + up_offset + gate_rel, + gate_group_bytes, + &group_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset + gate_rel, + gate_group_bytes, + &group_inner); + up_group24_offsets[i] = (NSUInteger)group_inner; + group_inner = 0; + down_group24_bufs[i] = q4_group24_exact_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + down_offset + down_rel, + down_group_bytes, + &group_inner) : + ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset + down_rel, + down_group_bytes, + &group_inner); + down_group24_offsets[i] = (NSUInteger)group_inner; + if (!gate_group24_bufs[i] || !up_group24_bufs[i] || !down_group24_bufs[i]) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32794); + return 0; + } + } + } else if (use_selected_slots) { + const bool selected_timing = + selected_profile || + ds4_gpu_stream_expert_timing_summary_enabled(); + double selected_t0 = selected_timing ? ds4_gpu_now_ms() : 0.0; + double selected_read_ms = 0.0; + double selected_sync_ms = 0.0; + double selected_copy_ms = 0.0; + double selected_wrap_ms = 0.0; + uint64_t selected_cache_hits0 = g_stream_expert_cache_hits; + uint64_t selected_cache_misses0 = g_stream_expert_cache_misses; + uint64_t selected_cache_wraps0 = g_stream_expert_cache_wraps; + uint64_t selected_cache_evictions0 = g_stream_expert_cache_evictions; + const char *selected_id_source = "readback"; + bool selected_ids_available = true; + bool selected_exec_ids_from_host = false; + const int stream_expert_cache_size_known = + ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes, + down_expert_bytes); + const bool use_iq2_full_expert_addr_table = + use_iq2_selected_slots && + ds4_gpu_stream_full_expert_addr_table_requested() && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil; + use_stream_expert_cache = + !use_iq2_full_expert_addr_table && + (use_iq2_selected_slots || use_iq2_stream_addr_table || use_q4_selected_slots) && + stream_expert_cache_size_known && + ds4_gpu_stream_expert_cache_effective_cap(layer_index, + n_total_expert, + n_expert) != 0; + if (use_iq2_stream_addr_table && !use_stream_expert_cache) { + fprintf(stderr, + "ds4: Metal IQ2/IQ2 streaming decode requires a non-empty expert cache\n"); + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32831); + return 0; + } + const bool stream_split_ready = + use_stream_expert_cache && + ds4_gpu_stream_expert_split_ready(); + const bool use_stream_compact_addr = + use_stream_expert_cache && + use_iq2_selected_slots && + ds4_gpu_stream_compact_addr_requested() && + !stream_split_ready && + !ds4_gpu_stream_expert_masked_addr_requested() && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil; + use_stream_expert_split_candidate = + use_stream_expert_cache && + use_iq2_selected_slots && + !use_stream_compact_addr && + stream_split_ready && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline != nil; + const bool use_stream_hit_validator = + use_stream_expert_cache && + use_iq2_selected_slots && + ds4_gpu_stream_expert_hit_validator_requested() && + g_moe_stream_expert_cache_validate_pipeline != nil && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil && + ds4_gpu_stream_expert_cache_addr_buffers(layer_index, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf); + if (use_iq2_full_expert_addr_table) { + selected_id_source = "gpu-full-addr"; + selected_ids_available = false; + g_routed_moe_selected_override_n = 0; + ds4_gpu_stream_expert_cache_entry *full_entry = NULL; + if (!ds4_gpu_stream_full_expert_addr_table_prepare(model_map, + model_size, + layer_index, + n_total_expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf, + &full_entry)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32880); + return 0; + } + for (uint32_t i = 0; i < n_expert; i++) { + stream_slot_entries[i] = full_entry; + } + use_stream_expert_addr_table = true; + } else { + const int replayed_selected_ids = + ds4_gpu_moe_selected_trace_replay(selected_ids, n_expert); + if (replayed_selected_ids < 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32890); + return 0; + } else if (replayed_selected_ids > 0) { + selected_id_source = "replay"; + selected_exec_ids_from_host = true; + g_routed_moe_selected_override_n = 0; + } else if (g_routed_moe_selected_override_n == n_expert) { + memcpy(selected_ids, + g_routed_moe_selected_override, + (size_t)n_expert * sizeof(selected_ids[0])); + selected_id_source = "override"; + selected_exec_ids_from_host = true; + g_routed_moe_selected_override_n = 0; + } else if (use_stream_hit_validator) { + g_routed_moe_selected_override_n = 0; + uint32_t validator_all_cached = 0; + uint32_t validator_miss_mask = 0; + uint32_t validator_invalid_mask = 0; + if (!ds4_gpu_stream_expert_cache_validate_selected(selected, + stream_gate_addr_buf, + stream_up_addr_buf, + stream_down_addr_buf, + n_total_expert, + n_expert, + selected_ids, + &validator_all_cached, + &validator_miss_mask, + &validator_invalid_mask)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32917); + return 0; + } + selected_id_source = + validator_invalid_mask != 0 ? "validator-invalid" : + (validator_miss_mask != 0 ? "validator-miss" : + (validator_all_cached != 0 ? "validator-hit" : "validator-miss")); + } else { + g_routed_moe_selected_override_n = 0; + if (g_batch_cb != nil) { + double selected_boundary_t0 = + selected_timing ? ds4_gpu_now_ms() : 0.0; + if (q4_selected_shared_event) { + if (ds4_gpu_signal_batch_and_wait_event("selected-id readback") == 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32942); return 0; } + } else if (ds4_gpu_end_commands() == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32931); + return 0; + } + if (selected_timing) { + selected_sync_ms += + ds4_gpu_now_ms() - selected_boundary_t0; + } + double selected_copy_t0 = + selected_timing ? ds4_gpu_now_ms() : 0.0; + if (ds4_gpu_tensor_read(selected, + 0, + selected_ids, + (uint64_t)n_expert * sizeof(selected_ids[0])) == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32943); + return 0; + } + if (selected_timing) { + selected_copy_ms += + ds4_gpu_now_ms() - selected_copy_t0; + } + if (!q4_selected_shared_event) { + selected_boundary_t0 = + selected_timing ? ds4_gpu_now_ms() : 0.0; + if (ds4_gpu_begin_commands() == 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32967); return 0; } + if (selected_timing) { + selected_sync_ms += + ds4_gpu_now_ms() - selected_boundary_t0; + } + } + } else { + double selected_copy_t0 = + selected_timing ? ds4_gpu_now_ms() : 0.0; + if (ds4_gpu_tensor_read(selected, + 0, + selected_ids, + (uint64_t)n_expert * sizeof(selected_ids[0])) == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32965); + return 0; + } + if (selected_timing) { + selected_copy_ms += + ds4_gpu_now_ms() - selected_copy_t0; + } + } + } + } + if (selected_timing) { + selected_read_ms = selected_sync_ms + selected_copy_ms; + selected_t0 = ds4_gpu_now_ms(); + } + + if (selected_ids_available) { + for (uint32_t i = 0; i < n_expert; i++) { + if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) { + fprintf(stderr, + "ds4: Metal routed MoE selected expert id %d is outside 0..%u\n", + selected_ids[i], + n_total_expert); + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32986); + return 0; + } + } + ds4_gpu_stream_expert_cache_note_selected_hotness(layer_index, + selected_ids, + n_expert); + if (!ds4_gpu_moe_selected_trace_record(selected_ids, n_expert) || + !ds4_gpu_moe_selected_hotlist_record(layer_index, + selected_ids, + n_expert, + n_total_expert)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32997); + return 0; + } + + for (uint32_t i = 0; i < n_expert; i++) { + const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i]; + if (expert_id > UINT64_MAX / gate_expert_bytes || + expert_id > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed MoE selected expert offset overflow\n"); + return 0; + } + const uint64_t gate_rel = expert_id * gate_expert_bytes; + const uint64_t down_rel = expert_id * down_expert_bytes; + if (gate_rel > UINT64_MAX - gate_offset || + gate_rel > UINT64_MAX - up_offset || + down_rel > UINT64_MAX - down_offset) { + fprintf(stderr, "ds4: Metal routed MoE selected expert offset overflow\n"); + return 0; + } + stream_gate_abs_offsets[i] = gate_offset + gate_rel; + stream_up_abs_offsets[i] = up_offset + gate_rel; + stream_down_abs_offsets[i] = down_offset + down_rel; + + if (use_stream_expert_cache) { + ds4_gpu_stream_expert_cache_entry *entry = NULL; + entry = ds4_gpu_stream_expert_cache_peek(model_map, + model_size, + layer_index, + (uint32_t)selected_ids[i], + n_total_expert, + n_expert, + stream_gate_abs_offsets[i], + stream_up_abs_offsets[i], + stream_down_abs_offsets[i], + gate_expert_bytes, + down_expert_bytes); + if (!entry) { + stream_expert_missing_mask |= 1u << i; + continue; + } + stream_expert_resident_mask |= 1u << i; + stream_slot_entries[i] = entry; + gate_slot_bufs[i] = entry->gate_buffer; + gate_slot_offsets[i] = entry->gate_inner; + up_slot_bufs[i] = entry->up_buffer; + up_slot_offsets[i] = entry->up_inner; + down_slot_bufs[i] = entry->down_buffer; + down_slot_offsets[i] = entry->down_inner; + continue; + } + + uint64_t slot_inner = 0; + gate_slot_bufs[i] = q4_selected_base_views ? + ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset + gate_rel, + gate_expert_bytes, + &slot_inner) : + (q4_selected_transient_views ? + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + gate_offset + gate_rel, + gate_expert_bytes, + &slot_inner) : + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + gate_offset + gate_rel, + gate_expert_bytes, + &slot_inner)); + gate_slot_offsets[i] = (NSUInteger)slot_inner; + slot_inner = 0; + up_slot_bufs[i] = q4_selected_base_views ? + ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset + gate_rel, + gate_expert_bytes, + &slot_inner) : + (q4_selected_transient_views ? + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + up_offset + gate_rel, + gate_expert_bytes, + &slot_inner) : + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + up_offset + gate_rel, + gate_expert_bytes, + &slot_inner)); + up_slot_offsets[i] = (NSUInteger)slot_inner; + slot_inner = 0; + down_slot_bufs[i] = q4_selected_base_views ? + ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset + down_rel, + down_expert_bytes, + &slot_inner) : + (q4_selected_transient_views ? + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + down_offset + down_rel, + down_expert_bytes, + &slot_inner) : + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + down_offset + down_rel, + down_expert_bytes, + &slot_inner)); + down_slot_offsets[i] = (NSUInteger)slot_inner; + if (!gate_slot_bufs[i] || !up_slot_bufs[i] || !down_slot_bufs[i]) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33105); + return 0; + } + } + } + if (use_stream_expert_cache) { + use_stream_expert_addr_table = + ((use_iq2_selected_slots && + ds4_gpu_stream_expert_addr_table_kernel_requested() && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil) || + use_iq2_stream_addr_table) && + ds4_gpu_stream_expert_cache_addr_buffers(layer_index, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf); + if (use_iq2_stream_addr_table && !use_stream_expert_addr_table) { + fprintf(stderr, + "ds4: Metal IQ2/IQ2 streaming decode could not prepare expert address buffers\n"); + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33123); + return 0; + } + use_stream_expert_masked_addr_table = + use_stream_expert_addr_table && + use_iq2_selected_slots && + ds4_gpu_stream_expert_masked_addr_requested() && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline != nil; + use_stream_expert_split_deferred = + use_stream_expert_split_candidate && + use_stream_expert_masked_addr_table && + stream_expert_resident_mask != 0 && + stream_expert_missing_mask != 0 && + ds4_gpu_stream_expert_split_worthwhile(stream_expert_resident_mask, + stream_expert_missing_mask) && + g_batch_cb != nil && + getenv("DS4_METAL_MOE_ONE_STAGE_PROFILE") == NULL; + if (use_stream_expert_split_deferred) { + const ds4_gpu_stream_expert_table table = { + .model_map = model_map, + .model_size = model_size, + .layer = layer_index, + .n_total_expert = n_total_expert, + .gate_offset = gate_offset, + .up_offset = up_offset, + .down_offset = down_offset, + .gate_expert_bytes = gate_expert_bytes, + .down_expert_bytes = down_expert_bytes, + }; + if (!ds4_gpu_stream_expert_cache_begin_selected_load( + &table, + selected_ids, + n_expert)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33156); + return 0; + } + } + if (stream_expert_missing_mask != 0 && + !use_stream_expert_split_deferred) { + if (!ds4_gpu_stream_expert_cache_load_selected_missing( + model_map, + model_size, + layer_index, + selected_ids, + n_total_expert, + n_expert, + stream_gate_abs_offsets, + stream_up_abs_offsets, + stream_down_abs_offsets, + gate_expert_bytes, + down_expert_bytes, + stream_expert_missing_mask, + stream_slot_entries)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33175); + return 0; + } + for (uint32_t i = 0; i < n_expert; i++) { + if ((stream_expert_missing_mask & (1u << i)) == 0) continue; + ds4_gpu_stream_expert_cache_entry *entry = stream_slot_entries[i]; + if (!entry) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33202); return 0; } + gate_slot_bufs[i] = entry->gate_buffer; + gate_slot_offsets[i] = entry->gate_inner; + up_slot_bufs[i] = entry->up_buffer; + up_slot_offsets[i] = entry->up_inner; + down_slot_bufs[i] = entry->down_buffer; + down_slot_offsets[i] = entry->down_inner; + } + } + if (use_iq2_stream_addr_table && use_stream_expert_addr_table) { + for (uint32_t i = 0; i < n_expert; i++) { + ds4_gpu_stream_expert_cache_entry *entry = stream_slot_entries[i]; + if (!entry) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33214); return 0; } + if (!ds4_gpu_stream_expert_cache_set_addr_slot_raw( + layer_index, + (uint32_t)selected_ids[i], + entry->gate_buffer, + entry->gate_inner, + entry->up_buffer, + entry->up_inner, + entry->down_buffer, + entry->down_inner)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33202); + return 0; + } + } + } + ds4_gpu_stream_expert_cache_prune_layer(layer_index, + n_total_expert, + n_expert, + selected_ids, + n_expert); + ds4_gpu_stream_expert_cache_prune_global(layer_index, + selected_ids, + n_expert); + if (use_stream_compact_addr) { + id compact_selected = nil; + if (!ds4_gpu_stream_compact_addr_prepare(layer_index, + stream_slot_entries, + n_expert, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf, + &compact_selected)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33223); + return 0; + } + selected_exec_buf = compact_selected; + selected_exec_off = 0; + use_stream_expert_addr_table = true; + use_stream_expert_masked_addr_table = false; + use_stream_compact_addr_table = true; + } + if (use_stream_expert_addr_table && + !use_stream_compact_addr_table && + selected_exec_ids_from_host) { + if (!ds4_gpu_stream_selected_ids_prepare(layer_index, + selected_ids, + n_expert, + &selected_exec_buf, + &selected_exec_off)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33239); + return 0; + } + } + } + if (selected_timing) { + selected_wrap_ms = ds4_gpu_now_ms() - selected_t0; + if (use_stream_expert_cache) { + ds4_gpu_stream_expert_timing_note_cache_class( + stream_expert_resident_mask, + stream_expert_missing_mask); + } + ds4_gpu_stream_expert_timing_note_selected(selected_sync_ms, + selected_copy_ms, + selected_wrap_ms); + } + if (selected_profile) { + const uint64_t selected_cache_hits = + g_stream_expert_cache_hits - selected_cache_hits0; + const uint64_t selected_cache_misses = + g_stream_expert_cache_misses - selected_cache_misses0; + const uint64_t selected_cache_wraps = + g_stream_expert_cache_wraps - selected_cache_wraps0; + const uint64_t selected_cache_evictions = + g_stream_expert_cache_evictions - selected_cache_evictions0; + const char *selected_path = + use_iq2_stream_addr_table ? "iq2/iq2" : + (use_iq2_selected_slots ? "iq2/q2" : "q4/q4"); + const char *selected_view_mode = + use_stream_expert_split_deferred ? "stream-split" : + use_stream_expert_masked_addr_table ? "stream-addr-mask" : + use_stream_compact_addr_table ? "stream-compact-addr" : + use_stream_expert_addr_table ? "stream-addr" : + (use_stream_expert_cache ? "stream-cache" : + (use_iq2_selected_slots ? "exact-cache" : q4_selected_view_mode)); + fprintf(stderr, + "ds4: Metal selected views layer=%u path=%s mode=%s ids=%s " + "experts=%d,%d,%d,%d,%d,%d expert_gate=%.2f MiB " + "expert_down=%.2f MiB read=%.3f ms bind=%.3f ms " + "cache_hits=%llu cache_misses=%llu cache_wraps=%llu cache_evictions=%llu\n", + layer_index, + selected_path, + selected_view_mode, + selected_id_source, + selected_ids_available ? selected_ids[0] : -1, + selected_ids_available ? selected_ids[1] : -1, + selected_ids_available ? selected_ids[2] : -1, + selected_ids_available ? selected_ids[3] : -1, + selected_ids_available ? selected_ids[4] : -1, + selected_ids_available ? selected_ids[5] : -1, + ds4_gpu_mib(gate_expert_bytes), + ds4_gpu_mib(down_expert_bytes), + selected_read_ms, + selected_wrap_ms, + (unsigned long long)selected_cache_hits, + (unsigned long long)selected_cache_misses, + (unsigned long long)selected_cache_wraps, + (unsigned long long)selected_cache_evictions); + } + } else if (!use_q4_grouped_experts) { + gate_buf = ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, gate_tensor_bytes, &gate_inner); + up_buf = ds4_gpu_wrap_model_range(model_map, model_size, up_offset, gate_tensor_bytes, &up_inner); + down_buf = ds4_gpu_wrap_model_range(model_map, model_size, down_offset, down_tensor_bytes, &down_inner); + if (!gate_buf || !up_buf || !down_buf) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33326); return 0; } + } + if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33305); + return 0; + } + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33337); return 0; } + if ((use_q4_expert_address_table || use_q4_expert_table) && + !ds4_gpu_use_model_residency_set(cb)) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33314); + return 0; + } + if (!q4_table_queue_residency && + q4_table_layer_residency && + [cb respondsToSelector:@selector(useResidencySet:)]) { + [cb useResidencySet:q4_table_layer_residency]; + } + + const bool moe_one_stage_profile = + g_batch_cb != nil && + ds4_gpu_stage_profile_enabled_for_layer("DS4_METAL_MOE_ONE_STAGE_PROFILE", + "DS4_METAL_MOE_ONE_STAGE_PROFILE_LAYER", + layer_index); + const char *moe_one_stage_filter = getenv("DS4_METAL_MOE_STAGE_PROFILE_FILTER"); + const char *moe_one_path = + use_q4_grouped_experts ? "q4_grouped_pair_swiglu" : + use_q4_group6_experts ? "q4_group6_pair_swiglu" : + use_q4_group8_experts ? "q4_group8_pair_swiglu" : + use_q4_group24_experts ? "q4_group24_split_gate_up" : + use_q4_exact_tensor_id ? "q4_exact_pair_swiglu" : + use_q4_expert_address_table ? "q4_addr_pair_swiglu" : + use_q4_expert_table ? "q4_table_pair_swiglu" : + use_q4_gather_slots ? "q4_gather_slots6_pair_swiglu" : + use_stream_expert_split_deferred ? "iq2_stream_split_pair_swiglu" : + use_stream_expert_masked_addr_table ? "iq2_stream_addr_mask_pair_swiglu" : + use_stream_expert_addr_table ? "iq2_stream_addr_pair_swiglu" : + use_iq2_selected_slots ? "iq2_slots6_pair_swiglu" : + use_q4_selected_slots ? "q4_slots6_pair_swiglu" : + (fuse_pair_swiglu ? "pair_swiglu" : + ((!g_quality_mode && + ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) || + (gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline))) ? "pair" : "single")); + double moe_one_stage_t0 = moe_one_stage_profile ? ds4_gpu_now_ms() : 0.0; + if (moe_one_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33349); + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33380); return 0; } + moe_one_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_MOE_ONE_STAGE(name) do { \ + if (ok && moe_one_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !moe_one_stage_filter || !moe_one_stage_filter[0] || \ + strstr(stage_name, moe_one_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal routed MoE one stage layer=%u pairs=%u experts=%u " \ + "gate=%s down=%s path=%s %s=%.3f ms\n", \ + layer_index, pair_rows, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + moe_one_path, \ + stage_name, now_ms - moe_one_stage_t0); \ + } \ + moe_one_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + if (use_q4_gather_slots) { + ds4_gpu_q4_gather_slots6_args gate_gather_args = { + .expert_bytes = gate_expert_bytes, + .group_size = q4_group6_expert_group_size, + .n_slots = n_expert, + }; + ds4_gpu_q4_gather_slots6_args down_gather_args = { + .expert_bytes = down_expert_bytes, + .group_size = q4_group6_expert_group_size, + .n_slots = n_expert, + }; + ok = ds4_gpu_encode_q4_gather_slots6(cb, + g_moe_q4_gather_slots6_pipeline, + &gate_gather_args, + gate_group6_bufs, + gate_group6_offsets, + selectedbuf, + ds4_gpu_tensor_offset(selected), + g_moe_q4_gate_slots_buffer, + 0) && + ds4_gpu_encode_q4_gather_slots6(cb, + g_moe_q4_gather_slots6_pipeline, + &gate_gather_args, + up_group6_bufs, + up_group6_offsets, + selectedbuf, + ds4_gpu_tensor_offset(selected), + g_moe_q4_up_slots_buffer, + 0) && + ds4_gpu_encode_q4_gather_slots6(cb, + g_moe_q4_gather_slots6_pipeline, + &down_gather_args, + down_group6_bufs, + down_group6_offsets, + selectedbuf, + ds4_gpu_tensor_offset(selected), + g_moe_q4_down_slots_buffer, + 0); + } + if (use_q4_gather_slots) { + DS4_METAL_PROFILE_MOE_ONE_STAGE("q4_gather"); + } + if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33455); return 0; } + if (use_q4_grouped_experts) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + for (uint32_t expert_base = 0; ok && expert_base < n_total_expert; expert_base += q4_expert_group_size) { + const uint32_t expert_count = + q4_expert_group_size < n_total_expert - expert_base ? + q4_expert_group_size : n_total_expert - expert_base; + if ((uint64_t)expert_base > UINT64_MAX / gate_expert_bytes || + (uint64_t)expert_count > UINT64_MAX / gate_expert_bytes) { + ok = 0; + break; + } + const uint64_t group_rel = (uint64_t)expert_base * gate_expert_bytes; + const uint64_t group_bytes = (uint64_t)expert_count * gate_expert_bytes; + if (group_rel > UINT64_MAX - gate_offset || + group_rel > UINT64_MAX - up_offset) { + ok = 0; + break; + } + uint64_t gate_group_inner = 0; + uint64_t up_group_inner = 0; + id gate_group_buf = + q4_grouped_cache_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + gate_offset + group_rel, + group_bytes, + &gate_group_inner) : + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + gate_offset + group_rel, + group_bytes, + &gate_group_inner); + id up_group_buf = + q4_grouped_cache_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + up_offset + group_rel, + group_bytes, + &up_group_inner) : + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + up_offset + group_rel, + group_bytes, + &up_group_inner); + if (!gate_group_buf || !up_group_buf) { + ok = 0; + break; + } + ds4_gpu_moe_expert_group_args group_args = { + .expert_base = expert_base, + .expert_count = expert_count, + .accumulate = 0, + .pad0 = 0, + }; + ok = ds4_gpu_encode_mul_mv_group_q4_pair_swiglu(cb, + g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + &group_args, + gate_group_buf, + (NSUInteger)gate_group_inner, + up_group_buf, + (NSUInteger)up_group_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } + } else if (use_q4_expert_address_table) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_addr_q4_pair_swiglu(cb, + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_table, + up_table, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } else if (use_q4_expert_table) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(cb, + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_table, + up_table, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false, + q4_table_queue_residency); + } else if (use_q4_group6_experts) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_group6_pair_swiglu(cb, + g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_group6_bufs, + gate_group6_offsets, + up_group6_bufs, + up_group6_offsets, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } else if (use_q4_group8_experts) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_group8_pair_swiglu(cb, + g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_group8_bufs, + gate_group8_offsets, + up_group8_bufs, + up_group8_offsets, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } else if (use_q4_group24_experts) { + ok = ds4_gpu_encode_mul_mv_group24_id(cb, + g_moe_mul_mv_group24_q4_k_id_pipeline, + &gate_args, + gate_group24_bufs, + gate_group24_offsets, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + 2, + false) && + ds4_gpu_encode_mul_mv_group24_id(cb, + g_moe_mul_mv_group24_q4_k_id_pipeline, + &gate_args, + up_group24_bufs, + up_group24_offsets, + xbuf, + ds4_gpu_tensor_offset(x), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + 2, + false); + } else if (use_q4_gather_slots || use_selected_slots) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + if (use_stream_expert_addr_table) { + if (use_stream_expert_masked_addr_table) { + if (use_stream_expert_split_deferred) { + const bool stream_split_profile = + getenv("DS4_METAL_STREAMING_EXPERT_SPLIT_PROFILE") != NULL; + const bool stream_split_timing = + stream_split_profile || + ds4_gpu_stream_expert_timing_summary_enabled(); + double stream_split_t0 = + stream_split_timing ? ds4_gpu_now_ms() : 0.0; + ds4_gpu_stream_expert_split_args resident_pair_args = { + .active_mask = stream_expert_resident_mask, + .accumulate = 0u, + }; + ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb, + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline, + &gate_args, + &act_args, + &resident_pair_args, + stream_slot_entries, + stream_gate_addr_buf, + stream_up_addr_buf, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selected_exec_buf, + selected_exec_off, + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + if (ok) { + ok = ds4_gpu_flush_commands(); + if (ok) { + cb = ds4_gpu_command_buffer(&owned); + if (!cb) ok = 0; + } + } + const double stream_split_resident_ms = + stream_split_timing ? ds4_gpu_now_ms() - stream_split_t0 : 0.0; + if (stream_split_timing) stream_split_t0 = ds4_gpu_now_ms(); + const double stream_split_missing_start_ms = stream_split_t0; + double stream_split_missing_load_ms = 0.0; + double stream_split_missing_slot_ms = 0.0; + double stream_split_missing_prune_ms = 0.0; + double stream_split_missing_addr_ms = 0.0; + double stream_split_missing_wait_ms = 0.0; + if (ok) { + ok = ds4_gpu_stream_expert_cache_load_selected_missing( + model_map, + model_size, + layer_index, + selected_ids, + n_total_expert, + n_expert, + stream_gate_abs_offsets, + stream_up_abs_offsets, + stream_down_abs_offsets, + gate_expert_bytes, + down_expert_bytes, + stream_expert_missing_mask, + stream_slot_entries); + if (stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + stream_split_missing_load_ms = + now_ms - stream_split_t0; + stream_split_t0 = now_ms; + } + if (ok) { + for (uint32_t i = 0; i < n_expert; i++) { + if ((stream_expert_missing_mask & (1u << i)) == 0) continue; + ds4_gpu_stream_expert_cache_entry *entry = + stream_slot_entries[i]; + if (!entry) { + ok = 0; + break; + } + gate_slot_bufs[i] = entry->gate_buffer; + gate_slot_offsets[i] = entry->gate_inner; + up_slot_bufs[i] = entry->up_buffer; + up_slot_offsets[i] = entry->up_inner; + down_slot_bufs[i] = entry->down_buffer; + down_slot_offsets[i] = entry->down_inner; + } + } + } + if (stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + stream_split_missing_slot_ms = + now_ms - stream_split_t0; + stream_split_t0 = now_ms; + } + if (ok) { + ds4_gpu_stream_expert_cache_prune_layer(layer_index, + n_total_expert, + n_expert, + selected_ids, + n_expert); + ds4_gpu_stream_expert_cache_prune_global(layer_index, + selected_ids, + n_expert); + if (stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + stream_split_missing_prune_ms = + now_ms - stream_split_t0; + stream_split_t0 = now_ms; + } + ok = ds4_gpu_stream_expert_cache_addr_buffers(layer_index, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf); + if (stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + stream_split_missing_addr_ms = + now_ms - stream_split_t0; + stream_split_t0 = now_ms; + } + } + if (ok) { + /* + * The resident stage was submitted before the + * CPU read of missing experts so I/O can overlap + * with GPU work. The missing stage reuses the same + * gate/up/mid scratch buffers, so it must not + * execute until the resident command buffer has + * finished. The down/sum pass is issued once after + * all six mid slots exist; this keeps the final + * accumulation order stable regardless of the + * resident/missing split. + */ + ok = ds4_gpu_wait_pending_command_buffers( + "streaming expert split resident"); + if (stream_split_timing) { + const double now_ms = ds4_gpu_now_ms(); + stream_split_missing_wait_ms = + now_ms - stream_split_t0; + stream_split_t0 = now_ms; + } + } + const double stream_split_missing_ms = + stream_split_timing ? + ds4_gpu_now_ms() - stream_split_missing_start_ms : + 0.0; + ds4_gpu_stream_expert_split_args missing_pair_args = { + .active_mask = stream_expert_missing_mask, + .accumulate = 0u, + }; + ds4_gpu_stream_expert_split_args all_down_args = { + .active_mask = 0x3fu, + .accumulate = 0u, + }; + if (ok) { + ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb, + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline, + &gate_args, + &act_args, + &missing_pair_args, + stream_slot_entries, + stream_gate_addr_buf, + stream_up_addr_buf, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selected_exec_buf, + selected_exec_off, + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } + if (ok) { + ok = ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(cb, + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline, + &down_args, + &all_down_args, + stream_slot_entries, + stream_down_addr_buf, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selected_exec_buf, + selected_exec_off, + down_smem, + 2); + } + stream_expert_split_completed = ok; + if (stream_split_timing) { + ds4_gpu_stream_expert_timing_note_split( + stream_expert_resident_mask, + stream_expert_missing_mask, + stream_split_resident_ms, + stream_split_missing_ms); + ds4_gpu_stream_expert_timing_note_split_missing_detail( + stream_split_missing_load_ms, + stream_split_missing_slot_ms, + stream_split_missing_prune_ms, + stream_split_missing_addr_ms, + stream_split_missing_wait_ms); + } + if (stream_split_profile) { + fprintf(stderr, + "ds4: Metal streaming expert split layer=%u " + "resident=0x%02x missing=0x%02x resident_submit=%.3f ms " + "missing_bind=%.3f ms\n", + layer_index, + stream_expert_resident_mask, + stream_expert_missing_mask, + stream_split_resident_ms, + stream_split_missing_ms); + } + } else { + ds4_gpu_stream_expert_split_args split_args = { + .active_mask = 0x3fu, + .accumulate = 0u, + }; + ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb, + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline, + &gate_args, + &act_args, + &split_args, + stream_slot_entries, + stream_gate_addr_buf, + stream_up_addr_buf, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selected_exec_buf, + selected_exec_off, + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } + } else { + ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu(cb, + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline, + &gate_args, + &act_args, + stream_slot_entries, + n_expert, + stream_gate_addr_buf, + stream_up_addr_buf, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selected_exec_buf, + selected_exec_off, + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false, + nil, + nil); + } + } else { + ok = (!use_stream_expert_cache || + ds4_gpu_stream_expert_cache_mark_entries_inflight( + stream_slot_entries, + n_expert, + 0)) && + ds4_gpu_encode_mul_mv_slots6_pair_swiglu(cb, + slots_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_slot_bufs, + gate_slot_offsets, + up_slot_bufs, + up_slot_offsets, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } + } else if (fuse_pair_swiglu) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_id_pair_swiglu(cb, + pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } else if (!g_quality_mode && + gate_type == DS4_METAL_TENSOR_IQ2_XXS && + g_moe_mul_mv_id_iq2_xxs_pair_pipeline) { + ok = ds4_gpu_encode_mul_mv_id_pair(cb, + g_moe_mul_mv_id_iq2_xxs_pair_pipeline, + &gate_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + 2, + false); + } else if (!g_quality_mode && + gate_type == DS4_METAL_TENSOR_Q4_K && + g_moe_mul_mv_id_q4_k_pair_pipeline) { + ok = ds4_gpu_encode_mul_mv_id_pair(cb, + g_moe_mul_mv_id_q4_k_pair_pipeline, + &gate_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + 2, + false); + } else { + ok = ds4_gpu_encode_mul_mv_id(cb, + gate_mv_pipeline, + &gate_args, + gate_buf, + (NSUInteger)gate_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + gate_nsg, + gate_rows_per_group_is_nr0) && + ds4_gpu_encode_mul_mv_id(cb, + gate_mv_pipeline, + &gate_args, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + gate_nsg, + gate_rows_per_group_is_nr0); + } + DS4_METAL_PROFILE_MOE_ONE_STAGE("gate_up"); + if (ok && (!fuse_pair_swiglu || use_q4_group24_experts)) { + ok = ds4_gpu_encode_moe_swiglu_weight(cb, + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + expert_mid_dim, + pair_rows, + clamp, + false); + } + DS4_METAL_PROFILE_MOE_ONE_STAGE("activation_weight"); + + id down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer); + NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : + (expertsbuf ? ds4_gpu_tensor_offset(experts) : 0); + if (ok && stream_expert_split_completed) { + /* The split path already wrote the resident partial output and + * accumulated the missing experts into out. */ + } else if (ok && use_q4_grouped_experts) { + bool first_group = true; + for (uint32_t expert_base = 0; ok && expert_base < n_total_expert; expert_base += q4_expert_group_size) { + const uint32_t expert_count = + q4_expert_group_size < n_total_expert - expert_base ? + q4_expert_group_size : n_total_expert - expert_base; + if ((uint64_t)expert_base > UINT64_MAX / down_expert_bytes || + (uint64_t)expert_count > UINT64_MAX / down_expert_bytes) { + ok = 0; + break; + } + const uint64_t group_rel = (uint64_t)expert_base * down_expert_bytes; + const uint64_t group_bytes = (uint64_t)expert_count * down_expert_bytes; + if (group_rel > UINT64_MAX - down_offset) { + ok = 0; + break; + } + uint64_t down_group_inner = 0; + id down_group_buf = + q4_grouped_cache_views ? + ds4_gpu_wrap_model_exact_range(model_map, + model_size, + down_offset + group_rel, + group_bytes, + &down_group_inner) : + ds4_gpu_wrap_model_exact_range_transient(model_map, + model_size, + down_offset + group_rel, + group_bytes, + &down_group_inner); + if (!down_group_buf) { + ok = 0; + break; + } + ds4_gpu_moe_expert_group_args group_args = { + .expert_base = expert_base, + .expert_count = expert_count, + .accumulate = first_group ? 0u : 1u, + .pad0 = 0, + }; + ok = ds4_gpu_encode_mul_mv_group_q4_sum6(cb, + g_moe_mul_mv_group_q4_k_sum6_pipeline, + &down_args, + &group_args, + down_group_buf, + (NSUInteger)down_group_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2); + first_group = false; + } + } else if (ok && use_q4_expert_address_table) { + ok = ds4_gpu_encode_mul_mv_addr_q4_sum6(cb, + g_moe_mul_mv_addr_q4_k_sum6_pipeline, + &down_args, + down_table, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2); + } else if (ok && use_q4_expert_table) { + ok = ds4_gpu_encode_mul_mv_table_q4_sum6(cb, + g_moe_mul_mv_table_q4_k_sum6_pipeline, + &down_args, + down_table, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2, + q4_table_queue_residency); + } else if (ok && use_q4_group6_experts) { + ok = ds4_gpu_encode_mul_mv_group6_sum6(cb, + g_moe_mul_mv_group6_q4_k_sum6_pipeline, + &down_args, + down_group6_bufs, + down_group6_offsets, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2); + } else if (ok && use_q4_group8_experts) { + ok = ds4_gpu_encode_mul_mv_group8_sum6(cb, + g_moe_mul_mv_group8_q4_k_sum6_pipeline, + &down_args, + down_group8_bufs, + down_group8_offsets, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2); + } else if (ok && use_q4_group24_experts) { + ok = ds4_gpu_encode_mul_mv_group24_sum6(cb, + g_moe_mul_mv_group24_q4_k_sum6_pipeline, + &down_args, + down_group24_bufs, + down_group24_offsets, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2); + } else if (ok && (use_q4_gather_slots || use_selected_slots)) { + if (use_stream_expert_addr_table) { + if (down_type == DS4_METAL_TENSOR_IQ2_XXS) { + ok = ds4_gpu_encode_mul_mv_addr_iq2(cb, + g_moe_mul_mv_addr_iq2_xxs_pipeline, + &down_args, + stream_slot_entries, + n_expert, + stream_down_addr_buf, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + selected_exec_buf, + selected_exec_off, + down_smem, + 2, + false); + } else if (use_stream_expert_masked_addr_table) { + ds4_gpu_stream_expert_split_args split_args = { + .active_mask = 0x3fu, + .accumulate = 0u, + }; + ok = ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(cb, + g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline, + &down_args, + &split_args, + stream_slot_entries, + stream_down_addr_buf, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selected_exec_buf, + selected_exec_off, + down_smem, + 2); + } else { + ok = ds4_gpu_encode_mul_mv_addr_q2_sum6(cb, + g_moe_mul_mv_addr_q2_k_sum6_pipeline, + &down_args, + stream_slot_entries, + n_expert, + stream_down_addr_buf, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selected_exec_buf, + selected_exec_off, + down_smem, + 2, + nil); + } + } else { + ok = (!use_stream_expert_cache || + ds4_gpu_stream_expert_cache_mark_entries_inflight( + stream_slot_entries, + n_expert, + 0)) && + ds4_gpu_encode_mul_mv_slots6_sum6(cb, + slots_sum6_pipeline, + &down_args, + down_slot_bufs, + down_slot_offsets, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + down_smem, + 2); + } + } else if (ok && direct_down_sum) { + ok = ds4_gpu_encode_mul_mv_id_sum6(cb, + down_sum6_pipeline, + &down_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + add_in ? ds4_gpu_tensor_buffer(add_in) : nil, + add_in ? ds4_gpu_tensor_offset(add_in) : 0, + down_smem, + 2); + } else if (ok) { + ok = ds4_gpu_encode_mul_mv_id(cb, + down_mv_pipeline, + &down_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + down_nsg, + down_rows_per_group_is_nr0); + } + DS4_METAL_PROFILE_MOE_ONE_STAGE("down"); + if (ok && n_expert > 1 && !direct_down_sum && !stream_expert_split_completed) { + ok = ds4_gpu_encode_moe_sum_experts(cb, + down_dst, + down_dst_off, + outbuf, + ds4_gpu_tensor_offset(out), + out_dim, + n_expert, + n_tokens); + } + DS4_METAL_PROFILE_MOE_ONE_STAGE("sum"); + if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34395); return 0; } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "routed tensor MoE")) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34397); return 0; } + if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34372); + return 0; + } + } +#undef DS4_METAL_PROFILE_MOE_ONE_STAGE + } + + return 1; +} + +int ds4_gpu_routed_moe_batch_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + ds4_gpu_tensor *experts, + const void *model_map, + uint64_t model_size, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint32_t gate_type, + uint32_t down_type, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t down_expert_bytes, + uint64_t down_row_bytes, + uint32_t expert_in_dim, + uint32_t expert_mid_dim, + uint32_t out_dim, + const ds4_gpu_tensor *selected, + const ds4_gpu_tensor *weights, + uint32_t n_total_expert, + uint32_t n_expert, + float clamp, + const ds4_gpu_tensor *x, + uint32_t layer_index, + uint32_t n_tokens, + bool *mid_is_f16, + bool force_resident) { + (void)force_resident; + if (!g_initialized && !ds4_gpu_init()) return 0; + /* TP sharding (see ds4_gpu_routed_moe_one_tensor): bind from the owned + * expert range and rebase ids in the kernels. */ + uint32_t first_expert = 0; + uint32_t n_bind_expert = 0; + ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert); + const int32_t tp_expert_base_host = (int32_t)first_expert; + gate_offset += (uint64_t)first_expert * gate_expert_bytes; + up_offset += (uint64_t)first_expert * gate_expert_bytes; + down_offset += (uint64_t)first_expert * down_expert_bytes; + if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights || + n_tokens == 0 || n_total_expert == 0 || n_expert == 0 || + n_expert > DS4_METAL_MAX_ROUTED_EXPERT_USED) { + return 0; + } + if (gate_expert_bytes == 0 || down_expert_bytes == 0 || + gate_row_bytes == 0 || down_row_bytes == 0) { + return 0; + } + if ((expert_in_dim % 256u) != 0 || (expert_mid_dim % 256u) != 0) return 0; + if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: Metal routed batch MoE tensor byte size overflow\n"); + return 0; + } + const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes; + + /* + * PRO Q4 routed expert tensors are multi-GiB per layer. A one-token + * layer-slice prefill should use the one-token path so it can either bind + * exact tensor views with GPU-selected IDs or fall back to selected-expert + * views. Keep this guarded by tensor size so Flash and mixed-Flash Q4 keep + * their existing fast path. + */ + const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull; + const bool can_single_token_q4_grouped = + getenv("DS4_METAL_ENABLE_Q4_GROUPED_EXPERTS") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUPED_EXPERTS") == NULL && + g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group_q4_k_sum6_pipeline != nil; + const bool can_single_token_q4_selected_slots = + ds4_gpu_q4_selected_paths_allowed() && + getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL && + g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil; + const bool can_single_token_q4_group6 = + n_total_expert == 384 && + getenv("DS4_METAL_ENABLE_Q4_GROUP6_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP6_EXPERT_TABLE") == NULL && + g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group6_q4_k_sum6_pipeline != nil; + const bool can_single_token_q4_group8 = + n_total_expert == 384 && + getenv("DS4_METAL_ENABLE_Q4_GROUP8_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP8_EXPERT_TABLE") == NULL && + g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_group8_q4_k_sum6_pipeline != nil; + const bool can_single_token_q4_group24 = + n_total_expert == 384 && + getenv("DS4_METAL_ENABLE_Q4_GROUP24_EXPERT_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_GROUP24_EXPERT_TABLE") == NULL && + g_moe_mul_mv_group24_q4_k_id_pipeline != nil && + g_moe_mul_mv_group24_q4_k_sum6_pipeline != nil; + const uint64_t max_buffer_len = g_device ? (uint64_t)[g_device maxBufferLength] : 0; + const bool can_single_token_q4_exact_tensor_id = + n_total_expert == 384 && + max_buffer_len != 0 && + gate_tensor_bytes <= max_buffer_len && + down_tensor_bytes <= max_buffer_len && + getenv("DS4_METAL_ENABLE_Q4_EXACT_TENSOR_ID") != NULL && + getenv("DS4_METAL_DISABLE_Q4_EXACT_TENSOR_ID") == NULL; + const bool enable_single_token_q4_expert_table = + getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || + ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes); + const bool can_single_token_q4_expert_table = + n_total_expert == 384 && + enable_single_token_q4_expert_table && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL && + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_table_q4_k_sum6_pipeline != nil && + g_moe_table_q4_pair_gate_encoder != nil && + g_moe_table_q4_pair_up_encoder != nil && + g_moe_table_q4_sum_down_encoder != nil; + const bool can_single_token_q4_expert_address_table = + n_total_expert == 384 && + getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL && + g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil; + const bool use_single_token_q4_one_tensor = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_tokens == 1 && + n_expert == 6 && + n_total_expert >= 128 && + (g_ssd_streaming_mode || + (gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes)) && + !g_quality_mode && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && + (can_single_token_q4_grouped || + can_single_token_q4_group6 || + can_single_token_q4_group8 || + can_single_token_q4_group24 || + can_single_token_q4_exact_tensor_id || + can_single_token_q4_expert_address_table || + can_single_token_q4_expert_table || + can_single_token_q4_selected_slots); + if (use_single_token_q4_one_tensor) { + if (mid_is_f16) *mid_is_f16 = false; + return ds4_gpu_routed_moe_one_tensor(out, + gate, + up, + mid, + experts, + model_map, + model_size, + gate_offset, + up_offset, + down_offset, + gate_type, + down_type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + expert_mid_dim, + out_dim, + selected, + weights, + n_total_expert, + n_expert, + clamp, + x, + NULL, + layer_index, + false); + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id gatebuf = ds4_gpu_tensor_buffer(gate); + id upbuf = ds4_gpu_tensor_buffer(up); + id midbuf = ds4_gpu_tensor_buffer(mid); + id outbuf = ds4_gpu_tensor_buffer(out); + id expertsbuf = ds4_gpu_tensor_buffer(experts); + id selectedbuf = ds4_gpu_tensor_buffer(selected); + id weightsbuf = ds4_gpu_tensor_buffer(weights); + const uint64_t x_bytes = (uint64_t)n_tokens * expert_in_dim * sizeof(float); + const uint64_t mid_bytes = (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float); + const uint64_t selected_bytes = (uint64_t)n_tokens * n_expert * sizeof(int); + const uint64_t weights_bytes = (uint64_t)n_tokens * n_expert * sizeof(float); + if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(gate) < mid_bytes || + ds4_gpu_tensor_bytes(up) < mid_bytes || + ds4_gpu_tensor_bytes(mid) < mid_bytes || + ds4_gpu_tensor_bytes(out) < out_bytes || + ds4_gpu_tensor_bytes(selected) < selected_bytes || + ds4_gpu_tensor_bytes(weights) < weights_bytes) { + fprintf(stderr, "ds4: Metal routed batch MoE received undersized activation buffers\n"); + return 0; + } + uint64_t gate_inner = 0; + uint64_t up_inner = 0; + uint64_t down_inner = 0; + id gate_buf = nil; + id up_buf = nil; + id down_buf = nil; + DS4MetalQ4ExpertTable *gate_table = nil; + DS4MetalQ4ExpertTable *up_table = nil; + DS4MetalQ4ExpertTable *down_table = nil; + id q4_table_layer_residency = nil; + id stream_gate_addr_buf = nil; + id stream_up_addr_buf = nil; + id stream_down_addr_buf = nil; + id stream_overflow_gate = nil; + id stream_overflow_up = nil; + id stream_overflow_down = nil; + ds4_gpu_stream_expert_cache_entry + *stream_resources[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { NULL }; + uint32_t stream_resource_count = 0; + uint32_t stream_unique = 0; + + const uint32_t pair_rows = n_tokens * n_expert; + const uint64_t down_scratch_bytes = (uint64_t)pair_rows * out_dim * sizeof(float); + + const uint32_t gate_nr0 = ds4_gpu_routed_mv_nr0(gate_type); + const uint32_t down_nr0 = ds4_gpu_routed_mv_nr0(down_type); + id gate_mv_pipeline = ds4_gpu_routed_mv_pipeline(gate_type); + id down_mv_pipeline = ds4_gpu_routed_mv_pipeline(down_type); + id gate_mm_pipeline = nil; + id up_mm_pipeline = nil; + id down_mm_pipeline = nil; + id pair_swiglu_mm_pipeline = nil; + if (gate_nr0 == 0 || down_nr0 == 0 || !gate_mv_pipeline || !down_mv_pipeline) { + fprintf(stderr, "ds4: unsupported Metal routed batch MoE quant types gate=%u down=%u\n", + gate_type, down_type); + return 0; + } + const bool use_iq2_batch_selected_addr = + ds4_gpu_stream_prefill_batch_selected_addr_enabled(n_tokens, + n_total_expert, + n_expert, + gate_type, + down_type) && + n_tokens > 1 && + n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT && + !g_quality_mode && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil && + g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil; + + ds4_gpu_mul_mv_id_args gate_args = + ds4_gpu_make_mul_mv_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + gate_row_bytes, gate_expert_bytes, + 1, n_expert, n_tokens, gate_nr0); + gate_args.tp_rank = g_tp_split_rank; + gate_args.tp_world = g_tp_split_world; + gate_args.tp_expert_base = tp_expert_base_host; + ds4_gpu_mul_mv_id_args down_args = + ds4_gpu_make_mul_mv_id_args(expert_mid_dim, out_dim, n_total_expert, + down_row_bytes, down_expert_bytes, + n_expert, n_expert, n_tokens, down_nr0); + down_args.tp_rank = g_tp_split_rank; + down_args.tp_world = g_tp_split_world; + down_args.tp_expert_base = tp_expert_base_host; + const bool q4_batch_expert_table_auto = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert, + n_expert, + gate_tensor_bytes, + down_tensor_bytes); + const bool q4_batch_table_queue_residency = + ds4_gpu_q4_table_queue_residency_enabled(q4_batch_expert_table_auto); + const bool enable_q4_batch_expert_table = + getenv("DS4_METAL_ENABLE_Q4_BATCH_EXPERT_TABLE") != NULL || + getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL || + q4_batch_expert_table_auto; + const bool use_q4_batch_expert_table = + gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K && + n_expert == 6 && + n_tokens > 1 && + n_total_expert == 384 && + gate_tensor_bytes >= q4_selected_min_tensor_bytes && + down_tensor_bytes >= q4_selected_min_tensor_bytes && + !g_quality_mode && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && + getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL && + getenv("DS4_METAL_DISABLE_Q4_BATCH_EXPERT_TABLE") == NULL && + getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL && + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil && + g_moe_mul_mv_table_q4_k_sum6_pipeline != nil && + g_moe_table_q4_pair_gate_encoder != nil && + g_moe_table_q4_pair_up_encoder != nil && + g_moe_table_q4_sum_down_encoder != nil && + enable_q4_batch_expert_table && + (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") != NULL || + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL || + q4_batch_table_queue_residency || + ds4_gpu_q4_table_model_residency_enabled()); + const bool use_mm_id = + !use_q4_batch_expert_table && + !use_iq2_batch_selected_addr && + n_tokens >= 32u && + ds4_gpu_mul_mm_id_map0_name(n_expert) != NULL; + /* + * MTP verification is neither normal decode nor large prefill: the + * target model must verify a tiny suffix (up to DSpark's 5-token + * block) in one layer-major pass. For that shape the prefill + * expert-major GEMM path + * is too large, but the decode pair kernels are exactly the right + * primitive: they read the same activation once and compute routed + * gate/up together for every selected expert row. Keep this limited to + * tiny batches so ordinary prefill keeps using the higher-throughput + * grouped matmul path. + */ + const bool use_tiny_pair_mv = + !g_quality_mode && + n_tokens <= 5u && + !use_q4_batch_expert_table && + !use_mm_id && + ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) || + (gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline)); + id tiny_pair_swiglu_pipeline = nil; + if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) { + tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline; + } else if (gate_type == DS4_METAL_TENSOR_Q4_K) { + tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline; + } + const bool use_tiny_pair_swiglu = + use_tiny_pair_mv && + tiny_pair_swiglu_pipeline != nil && + getenv("DS4_METAL_DISABLE_TINY_PAIR_SWIGLU_FUSION") == NULL && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL; + ds4_gpu_mul_mm_id_map_args gate_map_args = { 0 }; + ds4_gpu_mul_mm_id_args gate_mm_args = { 0 }; + ds4_gpu_mul_mm_id_args down_mm_args = { 0 }; + id map_pipeline = nil; + /* + * The grouped routed-MoE matmul loads activation tiles as half before + * using SIMD-group MMA. Store the SwiGLU/route-weight intermediate in + * that same precision so the down projection avoids a large F32 mid + * write/read. --quality keeps the older F32 intermediate. + */ + const bool request_mid_f16 = + !g_quality_mode && + !use_q4_batch_expert_table && + !use_iq2_batch_selected_addr; + /* + * Fused gate+up grouped matmul with the SwiGLU epilogue. The IQ2 + * variant stays opt-in behind its env flag; the Q4_K variant is the + * default path — same MMA accumulation order and epilogue math as the + * separate GEMMs + swiglu pass, so the mid tensor is bit-identical. + */ + const bool use_mm_id_pair_swiglu = + use_mm_id && + g_tp_split_world != 2 && /* pair-swiglu mm kernel lacks expert ownership */ + request_mid_f16 && + n_expert == 6 && + ((gate_type == DS4_METAL_TENSOR_IQ2_XXS && + down_type == DS4_METAL_TENSOR_Q2_K && + getenv("DS4_METAL_ENABLE_MOE_MM_ID_PAIR_SWIGLU") != NULL) || + (gate_type == DS4_METAL_TENSOR_Q4_K && + down_type == DS4_METAL_TENSOR_Q4_K)) && + getenv("DS4_METAL_DISABLE_MOE_MM_ID_PAIR_SWIGLU") == NULL && + getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL && + getenv("DS4_METAL_GRAPH_DUMP_PREFIX") == NULL; + if (use_mm_id) { + gate_map_args = + ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens); + gate_mm_args = + ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert, + gate_row_bytes, gate_expert_bytes, + 1, n_expert, n_tokens); + down_mm_args = + ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert, + down_row_bytes, down_expert_bytes, + n_expert, n_expert, n_tokens, + request_mid_f16 ? sizeof(uint16_t) : sizeof(float)); + gate_mm_args.tp_rank = g_tp_split_rank; + gate_mm_args.tp_world = g_tp_split_world; + gate_mm_args.tp_expert_base = tp_expert_base_host; + down_mm_args.tp_rank = g_tp_split_rank; + down_mm_args.tp_world = g_tp_split_world; + down_mm_args.tp_expert_base = tp_expert_base_host; + + map_pipeline = ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)); + gate_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); + up_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type); + down_mm_pipeline = request_mid_f16 ? + ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) : + ds4_gpu_routed_mm_pipeline(down_type); + if (use_mm_id_pair_swiglu) { + pair_swiglu_mm_pipeline = + ds4_gpu_get_pipeline(gate_type == DS4_METAL_TENSOR_Q4_K ? + "kernel_mul_mm_id_q4_K_pair_swiglu_f16" : + "kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16"); + } + if (!map_pipeline || !gate_mm_pipeline || !up_mm_pipeline || !down_mm_pipeline || + (use_mm_id_pair_swiglu && !pair_swiglu_mm_pipeline)) { + return 0; + } + } + + if (use_iq2_batch_selected_addr) { + const int had_batch = g_batch_cb != nil; + if (had_batch && ds4_gpu_end_commands() == 0) { + return 0; + } + g_stream_prefill_batch_selected_addr_building++; + if (!ds4_gpu_stream_expert_cache_prepare_selected_batch( + model_map, + model_size, + layer_index, + selected, + n_tokens, + n_total_expert, + n_expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes, + &stream_gate_addr_buf, + &stream_up_addr_buf, + &stream_down_addr_buf, + stream_resources, + &stream_resource_count, + &stream_unique, + &stream_overflow_gate, + &stream_overflow_up, + &stream_overflow_down)) { + g_stream_prefill_batch_selected_addr_building--; + return 0; + } + g_stream_prefill_batch_selected_addr_building--; + if (stream_unique == 0) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "produced no resident experts\n", + layer_index); + return 0; + } + for (uint32_t i = 0; i < stream_resource_count; i++) { + ds4_gpu_stream_expert_cache_entry *entry = stream_resources[i]; + if (!entry || + !entry->valid || + !entry->gate_buffer || + !entry->up_buffer || + !entry->down_buffer) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "lost resident expert resource %u/%u during preparation\n", + layer_index, + i, + stream_resource_count); + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + } + if (had_batch && ds4_gpu_begin_commands() == 0) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "failed to reopen command batch after preparation\n", + layer_index); + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + } + + if (use_q4_batch_expert_table) { + gate_table = ds4_gpu_q4_expert_table(model_map, + model_size, + gate_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_gate_encoder); + up_table = ds4_gpu_q4_expert_table(model_map, + model_size, + up_offset, + gate_expert_bytes, + n_total_expert, + g_moe_table_q4_pair_up_encoder); + down_table = ds4_gpu_q4_expert_table(model_map, + model_size, + down_offset, + down_expert_bytes, + n_total_expert, + g_moe_table_q4_sum_down_encoder); + if (!gate_table || !up_table || !down_table) { + return 0; + } + q4_table_layer_residency = + ds4_gpu_q4_expert_layer_residency_set(gate_table, + up_table, + down_table, + q4_batch_expert_table_auto); + if ((q4_batch_table_queue_residency || + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL) && + !q4_table_layer_residency && + !ds4_gpu_q4_table_model_residency_enabled()) { + fprintf(stderr, "ds4: Metal Q4 batch expert table residency set is not available\n"); + return 0; + } + } else if (use_iq2_batch_selected_addr) { + if (n_expert > 1 && (!expertsbuf || + ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) { + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, + &g_moe_down_scratch_bytes, + (NSUInteger)down_scratch_bytes, + "ds4_moe_down_scratch")) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + } + } else { + if (n_expert > 1 && (!expertsbuf || + ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) { + if (!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer, + &g_moe_down_scratch_bytes, + (NSUInteger)down_scratch_bytes, + "ds4_moe_down_scratch")) { + return 0; + } + } + gate_buf = ds4_gpu_wrap_model_range(model_map, + model_size, + gate_offset, + gate_tensor_bytes, + &gate_inner); + up_buf = ds4_gpu_wrap_model_range(model_map, + model_size, + up_offset, + gate_tensor_bytes, + &up_inner); + down_buf = ds4_gpu_wrap_model_range(model_map, + model_size, + down_offset, + down_tensor_bytes, + &down_inner); + if (!gate_buf || !up_buf || !down_buf) return 0; + if (getenv("DS4_GLM_TP_DEBUG") && layer_index == 3) { + fprintf(stderr, + "ds4: batch mv binds l=%u rank=%d base=%d gate=%llu+%llu " + "up=%llu down=%llu inner=%llu/%llu/%llu ne02=%d nei0=%d nr0=%d\n", + layer_index, g_tp_split_rank, tp_expert_base_host, + (unsigned long long)gate_offset, + (unsigned long long)gate_tensor_bytes, + (unsigned long long)up_offset, + (unsigned long long)down_offset, + (unsigned long long)gate_inner, + (unsigned long long)up_inner, + (unsigned long long)down_inner, + gate_args.ne02, gate_args.nei0, gate_args.nr0); + } + } + + const bool q4_batch_table_boundary = + use_q4_batch_expert_table && + g_batch_cb != nil && + getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL && + !q4_batch_table_queue_residency && + getenv("DS4_METAL_DISABLE_Q4_TABLE_BOUNDARY") == NULL; + if (q4_batch_table_boundary) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + if (use_q4_batch_expert_table && !ds4_gpu_use_model_residency_set(cb)) { + return 0; + } + if (!q4_batch_table_queue_residency && + q4_table_layer_residency && + [cb respondsToSelector:@selector(useResidencySet:)]) { + [cb useResidencySet:q4_table_layer_residency]; + } + const bool moe_stage_profile = + g_batch_cb != nil && + ds4_gpu_stage_profile_enabled_for_layer("DS4_METAL_MOE_STAGE_PROFILE", + "DS4_METAL_MOE_STAGE_PROFILE_LAYER", + layer_index); + const char *moe_stage_filter = getenv("DS4_METAL_MOE_STAGE_PROFILE_FILTER"); + const char *moe_path = + use_q4_batch_expert_table ? "q4_table_pair_swiglu" : + use_iq2_batch_selected_addr ? "iq2_batch_stream_addr" : + use_mm_id_pair_swiglu ? "mm_id_pair_swiglu" : + use_mm_id ? "mm_id" : + use_tiny_pair_swiglu ? "tiny_pair_swiglu" : + (use_tiny_pair_mv ? "tiny_pair_mv" : "mv"); + double moe_stage_t0 = moe_stage_profile ? ds4_gpu_now_ms() : 0.0; + if (moe_stage_profile) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + moe_stage_t0 = ds4_gpu_now_ms(); + } +#define DS4_METAL_PROFILE_MOE_STAGE(name) do { \ + if (ok && moe_stage_profile) { \ + if (ds4_gpu_end_commands() == 0) { \ + ok = 0; \ + } else { \ + const char *stage_name = (name); \ + const double now_ms = ds4_gpu_now_ms(); \ + const int print_stage = \ + !moe_stage_filter || !moe_stage_filter[0] || \ + strstr(stage_name, moe_stage_filter) != NULL; \ + if (print_stage) { \ + fprintf(stderr, \ + "ds4: Metal routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \ + "gate=%s down=%s path=%s mid=%s %s=%.3f ms\n", \ + layer_index, n_tokens, pair_rows, n_expert, \ + ds4_gpu_metal_tensor_type_name(gate_type), \ + ds4_gpu_metal_tensor_type_name(down_type), \ + moe_path, \ + request_mid_f16 ? "f16" : "f32", \ + stage_name, now_ms - moe_stage_t0); \ + } \ + moe_stage_t0 = now_ms; \ + if (ds4_gpu_begin_commands() == 0) { \ + ok = 0; \ + } else { \ + cb = ds4_gpu_command_buffer(&owned); \ + if (!cb) ok = 0; \ + } \ + } \ + } \ + } while (0) + + const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type); + const NSUInteger down_smem = ds4_gpu_routed_mv_smem(down_type); + const NSUInteger gate_nsg = ds4_gpu_routed_mv_nsg(gate_type); + const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type); + const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type); + const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type); + id down_sum6_pipeline = nil; + if (down_type == DS4_METAL_TENSOR_Q2_K) { + down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline; + } else if (down_type == DS4_METAL_TENSOR_Q4_K) { + down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline; + } + const bool direct_down_sum = + !g_quality_mode && + !use_q4_batch_expert_table && + !use_mm_id && + n_expert == 6 && + n_tokens <= 4u && + down_sum6_pipeline != nil; + int ok = 0; + if (use_iq2_batch_selected_addr) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu( + cb, + g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline, + &gate_args, + &act_args, + stream_resources, + stream_resource_count, + stream_gate_addr_buf, + stream_up_addr_buf, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false, + stream_overflow_gate, + stream_overflow_up); + if (!ok) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "failed to encode gate/up path tokens=%u unique=%u\n", + layer_index, + n_tokens, + stream_unique); + } + } else if (use_q4_batch_expert_table) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(cb, + g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_table, + up_table, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false, + q4_batch_table_queue_residency); + } else if (use_mm_id) { + /* + * The routed pair ids are the same for gate, up, and down. Build + * the expert-major work map once, then reuse it for all three + * batched expert matmuls. + */ + ok = ds4_gpu_encode_mul_mm_id_map(cb, + map_pipeline, + &gate_map_args, + &gate_mm_args, + selectedbuf, + ds4_gpu_tensor_offset(selected)); + DS4_METAL_PROFILE_MOE_STAGE("map"); + if (ok && use_mm_id_pair_swiglu) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(uint16_t), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16(cb, + pair_swiglu_mm_pipeline, + &gate_mm_args, + &act_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights)); + DS4_METAL_PROFILE_MOE_STAGE("gate_up_fused"); + } else if (ok) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + gate_mm_pipeline, + &gate_mm_args, + gate_buf, + (NSUInteger)gate_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + 8192u); + DS4_METAL_PROFILE_MOE_STAGE("gate"); + } + if (ok && !use_mm_id_pair_swiglu) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + up_mm_pipeline, + &gate_mm_args, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + upbuf, + ds4_gpu_tensor_offset(up), + 8192u); + DS4_METAL_PROFILE_MOE_STAGE("up"); + } + } else if (use_tiny_pair_swiglu) { + ds4_gpu_dsv4_moe_swiglu_weight_args act_args = { + .width = expert_mid_dim, + .rows = pair_rows, + .gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .up_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float), + .weight_stride = sizeof(float), + .write_clamped = 0, + .clamp_value = clamp, + }; + ok = ds4_gpu_encode_mul_mv_id_pair_swiglu(cb, + tiny_pair_swiglu_pipeline, + &gate_args, + &act_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + selectedbuf, + ds4_gpu_tensor_offset(selected), + weightsbuf, + ds4_gpu_tensor_offset(weights), + gate_smem, + 2, + false); + } else if (use_tiny_pair_mv) { + id pair_pipeline = + gate_type == DS4_METAL_TENSOR_IQ2_XXS ? + g_moe_mul_mv_id_iq2_xxs_pair_pipeline : + g_moe_mul_mv_id_q4_k_pair_pipeline; + ok = ds4_gpu_encode_mul_mv_id_pair(cb, + pair_pipeline, + &gate_args, + gate_buf, + (NSUInteger)gate_inner, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + 2, + false); + } else { + ok = ds4_gpu_encode_mul_mv_id(cb, + gate_mv_pipeline, + &gate_args, + gate_buf, + (NSUInteger)gate_inner, + xbuf, + ds4_gpu_tensor_offset(x), + gatebuf, + ds4_gpu_tensor_offset(gate), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + gate_nsg, + gate_rows_per_group_is_nr0) && + ds4_gpu_encode_mul_mv_id(cb, + gate_mv_pipeline, + &gate_args, + up_buf, + (NSUInteger)up_inner, + xbuf, + ds4_gpu_tensor_offset(x), + upbuf, + ds4_gpu_tensor_offset(up), + selectedbuf, + ds4_gpu_tensor_offset(selected), + gate_smem, + gate_nsg, + gate_rows_per_group_is_nr0); + } + DS4_METAL_PROFILE_MOE_STAGE("gate_up"); + const bool use_fused_activation = !g_quality_mode && !use_q4_batch_expert_table; + const bool use_mid_f16 = + use_mm_id && + use_fused_activation && + request_mid_f16; + if (mid_is_f16) *mid_is_f16 = use_mid_f16; + if (ok && use_iq2_batch_selected_addr) { + /* The address-table pair kernel already wrote weighted SwiGLU rows into mid. */ + } else if (ok && use_q4_batch_expert_table) { + /* The table pair kernel already wrote weighted SwiGLU rows into mid. */ + } else if (ok && use_mm_id_pair_swiglu) { + /* The fused batch mm_id pair kernel already wrote weighted f16 SwiGLU rows into mid. */ + } else if (ok && use_tiny_pair_swiglu) { + /* The fused tiny pair kernel already wrote weighted F32 SwiGLU rows into mid. */ + } else if (ok && use_fused_activation) { + ok = ds4_gpu_encode_moe_swiglu_weight(cb, + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + expert_mid_dim, + pair_rows, + clamp, + use_mid_f16); + } else if (ok && clamp > 1.0e-6f) { + ok = ds4_gpu_encode_unary_f32_rows(cb, + g_unary_clamp_pipeline, + gatebuf, + ds4_gpu_tensor_offset(gate), + gatebuf, + ds4_gpu_tensor_offset(gate), + expert_mid_dim, + pair_rows, + 0, + -FLT_MAX, + clamp); + if (ok) { + ok = ds4_gpu_encode_unary_f32_rows(cb, + g_unary_silu_pipeline, + gatebuf, + ds4_gpu_tensor_offset(gate), + midbuf, + ds4_gpu_tensor_offset(mid), + expert_mid_dim, + pair_rows, + 1, + 0.0f, + 0.0f); + } + if (ok) { + ok = ds4_gpu_encode_unary_f32_rows(cb, + g_unary_clamp_pipeline, + upbuf, + ds4_gpu_tensor_offset(up), + upbuf, + ds4_gpu_tensor_offset(up), + expert_mid_dim, + pair_rows, + 0, + -clamp, + clamp); + } + if (ok) { + ds4_gpu_bin_args mul_args = + ds4_gpu_make_bin_same_rows_args(expert_mid_dim, pair_rows); + ok = ds4_gpu_encode_bin_f32_rows(cb, + g_mul_pipeline, + &mul_args, + midbuf, + ds4_gpu_tensor_offset(mid), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid)); + } + } else if (ok) { + ok = ds4_gpu_encode_swiglu_flat(cb, + gatebuf, + ds4_gpu_tensor_offset(gate), + upbuf, + ds4_gpu_tensor_offset(up), + midbuf, + ds4_gpu_tensor_offset(mid), + (uint32_t)((uint64_t)pair_rows * expert_mid_dim)); + } + if (ok && !use_fused_activation && !use_q4_batch_expert_table) { + ds4_gpu_bin_args weight_args = + ds4_gpu_make_bin_rowwise_scalar_args(expert_mid_dim, pair_rows); + ok = ds4_gpu_encode_bin_f32_rows(cb, + g_bin_mul_scalar_pipeline, + &weight_args, + midbuf, + ds4_gpu_tensor_offset(mid), + weightsbuf, + ds4_gpu_tensor_offset(weights), + midbuf, + ds4_gpu_tensor_offset(mid)); + } + DS4_METAL_PROFILE_MOE_STAGE("activation_weight"); + + id down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer); + NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : + (expertsbuf ? ds4_gpu_tensor_offset(experts) : 0); + if (ok) { + if (use_iq2_batch_selected_addr) { + ok = ds4_gpu_encode_mul_mv_addr_q2_sum6( + cb, + g_moe_mul_mv_addr_q2_k_sum6_pipeline, + &down_args, + stream_resources, + stream_resource_count, + stream_down_addr_buf, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2, + stream_overflow_down); + if (!ok) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "failed to encode down path tokens=%u unique=%u\n", + layer_index, + n_tokens, + stream_unique); + } + } else if (use_q4_batch_expert_table) { + ok = ds4_gpu_encode_mul_mv_table_q4_sum6(cb, + g_moe_mul_mv_table_q4_k_sum6_pipeline, + &down_args, + down_table, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + 2, + q4_batch_table_queue_residency); + } else if (direct_down_sum) { + ok = ds4_gpu_encode_mul_mv_id_sum6(cb, + down_sum6_pipeline, + &down_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + outbuf, + ds4_gpu_tensor_offset(out), + selectedbuf, + ds4_gpu_tensor_offset(selected), + nil, + 0, + down_smem, + 2); + } else if (use_mm_id) { + ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb, + down_mm_pipeline, + &down_mm_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + 8192u); + } else { + ok = ds4_gpu_encode_mul_mv_id(cb, + down_mv_pipeline, + &down_args, + down_buf, + (NSUInteger)down_inner, + midbuf, + ds4_gpu_tensor_offset(mid), + down_dst, + down_dst_off, + selectedbuf, + ds4_gpu_tensor_offset(selected), + down_smem, + down_nsg, + down_rows_per_group_is_nr0); + } + } + DS4_METAL_PROFILE_MOE_STAGE("down"); + if (ok && + n_expert > 1 && + !direct_down_sum && + !use_q4_batch_expert_table && + !use_iq2_batch_selected_addr) { + ok = ds4_gpu_encode_moe_sum_experts(cb, + down_dst, + down_dst_off, + outbuf, + ds4_gpu_tensor_offset(out), + out_dim, + n_expert, + n_tokens); + } + DS4_METAL_PROFILE_MOE_STAGE("sum"); + if (!ok) { + fprintf(stderr, + "ds4: Metal routed batch MoE failed before submit layer=%u tokens=%u " + "gate=%s down=%s path=%s\n", + layer_index, + n_tokens, + ds4_gpu_metal_tensor_type_name(gate_type), + ds4_gpu_metal_tensor_type_name(down_type), + moe_path); + return 0; + } + + if (!ds4_gpu_finish_command_buffer(cb, owned, "routed batch MoE")) { + if (use_iq2_batch_selected_addr) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + } + return 0; + } + if (use_iq2_batch_selected_addr) { + if (!owned) { + if (ds4_gpu_end_commands() == 0) { + ds4_gpu_stream_expert_cache_clear_layer(layer_index); + return 0; + } + if (ds4_gpu_begin_commands() == 0) { + fprintf(stderr, + "ds4: Metal streaming prefill batch selected addr layer=%u " + "failed to reopen command batch after execution\n", + layer_index); + return 0; + } + } + } + if (q4_batch_table_boundary) { + if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) { + return 0; + } + } +#undef DS4_METAL_PROFILE_MOE_STAGE + } + + return 1; +} + +int ds4_gpu_hc_split_sinkhorn_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *mix, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (n_hc == 0 || n_hc > 16) return 0; + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t mix_bytes = mix_hc * sizeof(float); + const uint64_t scale_bytes = 3ull * sizeof(float); + + @autoreleasepool { + id mixbuf = ds4_gpu_tensor_buffer(mix); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t mix_tensor_bytes = ds4_gpu_tensor_bytes(mix); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); + if (!mixbuf || !outbuf || + mix_tensor_bytes < mix_bytes || + out_tensor_bytes < mix_bytes) { + fprintf(stderr, "ds4: Metal HC split received undersized activation buffers\n"); + return 0; + } + if (scale_offset > model_size || scale_bytes > model_size - scale_offset || + base_offset > model_size || mix_bytes > model_size - base_offset) { + fprintf(stderr, "ds4: Metal HC split parameter range is outside the mapped model\n"); + return 0; + } + + uint64_t scale_inner = 0; + uint64_t base_inner = 0; + id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner); + id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner); + if (!scalebuf || !basebuf) return 0; + + uint64_t n_rows64 = mix_tensor_bytes / mix_bytes; + const uint64_t out_rows64 = out_tensor_bytes / mix_bytes; + if (out_rows64 < n_rows64) n_rows64 = out_rows64; + if (n_rows64 == 0 || n_rows64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC split row count is outside supported range\n"); + return 0; + } + + ds4_gpu_hc_split_args args = { + .n_hc = (int32_t)n_hc, + .sinkhorn_iters = (int32_t)sinkhorn_iters, + .n_rows = (int64_t)n_rows64, + .mix_hc = (int64_t)mix_hc, + .nb01 = mix_bytes, + .nb1 = mix_bytes, + .eps = eps, + }; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_rows64)); + const NSUInteger n_tg = ((NSUInteger)n_rows64 + nth - 1u) / nth; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_hc_split_sinkhorn_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sinkhorn")) return 0; + } + + return 1; +} + +static int ds4_gpu_hc_weighted_sum_strided( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *weights, + uint64_t weight_offset, + uint64_t weight_row_stride, + uint32_t n_embd, + uint32_t n_hc, + const char *label) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !residual_hc || !weights || n_embd == 0 || n_hc == 0 || + weight_row_stride < (uint64_t)n_hc * sizeof(float)) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(residual_hc); + id wbuf = ds4_gpu_tensor_buffer(weights); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); + if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || out_tensor_bytes % out_row_bytes != 0) { + fprintf(stderr, "ds4: Metal HC weighted sum output size is not a whole token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / out_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC weighted sum token count is outside supported range\n"); + return 0; + } + + const uint64_t x_row_values = (uint64_t)n_hc * n_embd; + if (x_row_values == 0 || + x_row_values > UINT64_MAX / sizeof(float) || + n_tokens64 > UINT64_MAX / (x_row_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / ((uint64_t)n_hc * sizeof(float))) { + fprintf(stderr, "ds4: Metal HC weighted sum activation size overflow\n"); + return 0; + } + + const uint64_t x_bytes = n_tokens64 * x_row_values * sizeof(float); + const uint64_t w_last = weight_offset + + (n_tokens64 - 1u) * weight_row_stride + + (uint64_t)n_hc * sizeof(float); + if (!xbuf || !wbuf || !outbuf || + ds4_gpu_tensor_bytes(residual_hc) < x_bytes || + ds4_gpu_tensor_bytes(weights) < w_last) { + fprintf(stderr, "ds4: Metal HC weighted sum received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_weighted_sum_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_x0 = sizeof(float), + .nb_x1 = (uint64_t)n_embd * sizeof(float), + .nb_x2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_w0 = sizeof(float), + .nb_w1 = weight_row_stride, + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + }; + const uint64_t n_elem = (uint64_t)n_embd * n_tokens64; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_hc_weighted_sum_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) + (NSUInteger)weight_offset atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, label)) return 0; + } + + return 1; +} + +int ds4_gpu_hc_weighted_sum_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *weights, + uint32_t n_embd, + uint32_t n_hc) { + return ds4_gpu_hc_weighted_sum_strided(out, + residual_hc, + weights, + 0, + (uint64_t)n_hc * sizeof(float), + n_embd, + n_hc, + "HC weighted sum"); +} + +int ds4_gpu_hc_weighted_sum_norm_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *weights, + const void *model_map, + uint64_t model_size, + uint64_t norm_weight_offset, + uint32_t n_embd, + uint32_t n_hc, + float norm_eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !norm_out || !residual_hc || !weights || !model_map) return 0; + + const bool force = + getenv("DS4_METAL_ENABLE_OUTPUT_HC_SUM_NORM_FUSION") != NULL; + const bool disabled = + getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_SUM_NORM_FUSION") != NULL; + const bool require = + getenv("DS4_METAL_REQUIRE_OUTPUT_HC_SUM_NORM_FUSION") != NULL; + const bool supported_shape = + n_hc == 4u && (n_embd == 4096u || n_embd == 7168u); + const bool auto_shape = + n_embd == 4096u && ds4_gpu_device_name_contains("M3"); + const bool use_fusion = + supported_shape && !g_quality_mode && !disabled && + g_hc_weighted_sum_norm_pipeline != nil && + (auto_shape || force); + if (require && supported_shape && !use_fusion) { + fprintf(stderr, + "ds4: required Metal output HC sum/RMSNorm fusion was not selected\n"); + return 0; + } + if (!use_fusion) return 0; + + @autoreleasepool { + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t residual_bytes = (uint64_t)n_hc * out_row_bytes; + const uint64_t weight_bytes = (uint64_t)n_hc * sizeof(float); + id xbuf = ds4_gpu_tensor_buffer(residual_hc); + id wbuf = ds4_gpu_tensor_buffer(weights); + id outbuf = ds4_gpu_tensor_buffer(out); + id normbuf = ds4_gpu_tensor_buffer(norm_out); + if (!xbuf || !wbuf || !outbuf || !normbuf || + ds4_gpu_tensor_bytes(residual_hc) < residual_bytes || + ds4_gpu_tensor_bytes(weights) < weight_bytes || + ds4_gpu_tensor_bytes(out) != out_row_bytes || + ds4_gpu_tensor_bytes(norm_out) < out_row_bytes) { + fprintf(stderr, + "ds4: Metal output HC sum/RMSNorm fusion received invalid activation buffers\n"); + return 0; + } + if (norm_weight_offset > model_size || + out_row_bytes > model_size - norm_weight_offset) { + fprintf(stderr, + "ds4: Metal output HC sum/RMSNorm weight range is outside the mapped model\n"); + return 0; + } + + uint64_t norm_inner = 0; + id normwbuf = ds4_gpu_wrap_model_range( + model_map, model_size, norm_weight_offset, + out_row_bytes, &norm_inner); + if (!normwbuf) return 0; + + ds4_gpu_hc_weighted_sum_norm_args args = { + .n_embd = (int64_t)n_embd, + .n_hc = (int64_t)n_hc, + .n_tokens = 1, + .nb_x0 = sizeof(float), + .nb_x1 = out_row_bytes, + .nb_x2 = residual_bytes, + .nb_w0 = sizeof(float), + .nb_w1 = weight_bytes, + .nb0 = sizeof(float), + .nb1 = out_row_bytes, + .nb_norm1 = out_row_bytes, + .norm_eps = norm_eps, + }; + const NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd); + const NSUInteger shared_bytes = + ((NSUInteger)n_embd + 32u) * sizeof(float); + if (nth > g_hc_weighted_sum_norm_pipeline.maxTotalThreadsPerThreadgroup || + shared_bytes > [g_device maxThreadgroupMemoryLength]) { + if (require) { + fprintf(stderr, + "ds4: required Metal output HC sum/RMSNorm fusion exceeds device limits\n"); + } + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_hc_weighted_sum_norm_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1]; + [enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:4]; + [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "output HC sum/RMSNorm fused")) { + return 0; + } + } + + return 1; +} + +int ds4_gpu_hc_weighted_sum_split_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + return ds4_gpu_hc_weighted_sum_strided(out, + residual_hc, + split, + 0, + mix_hc * sizeof(float), + n_embd, + n_hc, + "HC weighted sum split"); +} + +/* Release decode fused HC pre-sublayer operation. The graph driver owns the + * optional reference fallback so this function stays a direct fused dispatch. */ +int ds4_gpu_hc_split_weighted_sum_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *mix, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !split || !mix || !residual_hc || !model_map || + n_embd == 0 || n_hc == 0) { + return 0; + } + if (n_hc != 4) { + fprintf(stderr, "ds4: Metal fused HC split/sum is specialized for HC=4\n"); + return 0; + } + + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t mix_bytes = mix_hc * sizeof(float); + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t scale_bytes = 3ull * sizeof(float); + + @autoreleasepool { + id mixbuf = ds4_gpu_tensor_buffer(mix); + id splitbuf = ds4_gpu_tensor_buffer(split); + id xbuf = ds4_gpu_tensor_buffer(residual_hc); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); + if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || + out_tensor_bytes % out_row_bytes != 0) { + fprintf(stderr, "ds4: Metal fused HC split/sum output size is not a whole token row\n"); + return 0; + } + + const uint64_t n_rows64 = out_tensor_bytes / out_row_bytes; + if (n_rows64 == 0 || n_rows64 > UINT32_MAX || + n_rows64 > UINT64_MAX / mix_bytes || + n_rows64 > UINT64_MAX / residual_row_bytes) { + fprintf(stderr, "ds4: Metal fused HC split/sum row count is outside supported range\n"); + return 0; + } + + const uint64_t mix_total_bytes = n_rows64 * mix_bytes; + const uint64_t residual_total_bytes = n_rows64 * residual_row_bytes; + if (!mixbuf || !splitbuf || !xbuf || !outbuf || + ds4_gpu_tensor_bytes(mix) < mix_total_bytes || + ds4_gpu_tensor_bytes(split) < mix_total_bytes || + ds4_gpu_tensor_bytes(residual_hc) < residual_total_bytes) { + fprintf(stderr, "ds4: Metal fused HC split/sum received undersized activation buffers\n"); + return 0; + } + + if (scale_offset > model_size || scale_bytes > model_size - scale_offset || + base_offset > model_size || mix_bytes > model_size - base_offset) { + fprintf(stderr, "ds4: Metal fused HC split/sum parameter range is outside the mapped model\n"); + return 0; + } + + uint64_t scale_inner = 0; + uint64_t base_inner = 0; + id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner); + id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner); + if (!scalebuf || !basebuf) return 0; + + ds4_gpu_hc_split_weighted_sum_args args = { + .n_embd = (int64_t)n_embd, + .n_hc = (int32_t)n_hc, + .sinkhorn_iters = (int32_t)sinkhorn_iters, + .n_rows = (int64_t)n_rows64, + .mix_hc = (int64_t)mix_hc, + .nb_mix1 = mix_bytes, + .nb_split1 = mix_bytes, + .nb_x0 = sizeof(float), + .nb_x1 = (uint64_t)n_embd * sizeof(float), + .nb_x2 = residual_row_bytes, + .nb0 = sizeof(float), + .nb1 = out_row_bytes, + .eps = eps, + }; + + NSUInteger nth = g_hc_split_weighted_sum_pipeline.maxTotalThreadsPerThreadgroup; + if (nth > 256u) nth = 256u; + if (nth > (NSUInteger)n_embd) nth = (NSUInteger)n_embd; + if (nth == 0) nth = 1u; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_hc_split_weighted_sum_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:6]; + [enc setThreadgroupMemoryLength:(NSUInteger)n_hc * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows64, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sum fused")) return 0; + } + + return 1; +} + +/* HC-pre plus the immediately following weighted RMSNorm, specialized for + * DS4's HC=4 shape. Both decode and batched prefill use this implementation; + * the kernel preserves their established single-row and batched scale formulas. */ +int ds4_gpu_hc_split_weighted_sum_norm_tensor( + ds4_gpu_tensor *out, + ds4_gpu_tensor *norm_out, + ds4_gpu_tensor *split, + const ds4_gpu_tensor *mix, + const ds4_gpu_tensor *residual_hc, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n_embd, + uint32_t n_hc, + uint32_t sinkhorn_iters, + float eps, + float norm_eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !norm_out || !split || !mix || !residual_hc || !model_map || + n_hc != 4 || (n_embd & 3u) != 0) { + return 0; + } + + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t mix_bytes = mix_hc * sizeof(float); + const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float); + const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t scale_bytes = 3ull * sizeof(float); + + @autoreleasepool { + id mixbuf = ds4_gpu_tensor_buffer(mix); + id splitbuf = ds4_gpu_tensor_buffer(split); + id xbuf = ds4_gpu_tensor_buffer(residual_hc); + id outbuf = ds4_gpu_tensor_buffer(out); + id normbuf = ds4_gpu_tensor_buffer(norm_out); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); + if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || + out_tensor_bytes % out_row_bytes != 0) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm output size is not a whole token row\n"); + return 0; + } + + const uint64_t n_rows64 = out_tensor_bytes / out_row_bytes; + if (n_rows64 == 0 || n_rows64 > UINT32_MAX || + n_rows64 > UINT64_MAX / mix_bytes || + n_rows64 > UINT64_MAX / residual_row_bytes) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm row count is outside supported range\n"); + return 0; + } + + const uint64_t mix_total_bytes = n_rows64 * mix_bytes; + const uint64_t residual_total_bytes = n_rows64 * residual_row_bytes; + const uint64_t out_total_bytes = n_rows64 * out_row_bytes; + if (!mixbuf || !splitbuf || !xbuf || !outbuf || !normbuf || + ds4_gpu_tensor_bytes(mix) < mix_total_bytes || + ds4_gpu_tensor_bytes(split) < mix_total_bytes || + ds4_gpu_tensor_bytes(residual_hc) < residual_total_bytes || + ds4_gpu_tensor_bytes(norm_out) < out_total_bytes) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm received undersized activation buffers\n"); + return 0; + } + + if (scale_offset > model_size || scale_bytes > model_size - scale_offset || + base_offset > model_size || mix_bytes > model_size - base_offset || + norm_weight_offset > model_size || out_row_bytes > model_size - norm_weight_offset) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm parameter range is outside the mapped model\n"); + return 0; + } + + uint64_t scale_inner = 0; + uint64_t base_inner = 0; + uint64_t norm_inner = 0; + id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner); + id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner); + id normwbuf = ds4_gpu_wrap_model_range(model_map, model_size, norm_weight_offset, out_row_bytes, &norm_inner); + if (!scalebuf || !basebuf || !normwbuf) return 0; + + id pipeline = + ds4_gpu_hot_pipeline(g_hc_split_weighted_sum_norm_pipeline, + "kernel_dsv4_hc_split_weighted_sum_norm4"); + if (!pipeline) return 0; + + ds4_gpu_hc_split_weighted_sum_norm_args args = { + .n_embd = (int64_t)n_embd, + .n_hc = (int32_t)n_hc, + .sinkhorn_iters = (int32_t)sinkhorn_iters, + .n_rows = (int64_t)n_rows64, + .mix_hc = (int64_t)mix_hc, + .nb_mix1 = mix_bytes, + .nb_split1 = mix_bytes, + .nb_x0 = sizeof(float), + .nb_x1 = (uint64_t)n_embd * sizeof(float), + .nb_x2 = residual_row_bytes, + .nb0 = sizeof(float), + .nb1 = out_row_bytes, + .nb_norm1 = out_row_bytes, + .eps = eps, + .norm_eps = norm_eps, + }; + + NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd); + if (nth > pipeline.maxTotalThreadsPerThreadgroup) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm requires %lu threads but pipeline supports %lu\n", + (unsigned long)nth, + (unsigned long)pipeline.maxTotalThreadsPerThreadgroup); + return 0; + } + + const NSUInteger shared_bytes = ((NSUInteger)n_embd + 4u + 32u) * sizeof(float); + const NSUInteger max_shared = [g_device maxThreadgroupMemoryLength]; + if (max_shared != 0 && shared_bytes > max_shared) { + fprintf(stderr, "ds4: Metal fused HC split/sum/norm requires %lu bytes of threadgroup memory but device supports %lu\n", + (unsigned long)shared_bytes, + (unsigned long)max_shared); + return 0; + } + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:6]; + [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:7]; + [enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:8]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows64, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sum/norm fused")) return 0; + } + + return 1; +} + +int ds4_gpu_output_hc_weights_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *pre, + const void *model_map, + uint64_t model_size, + uint64_t scale_offset, + uint64_t base_offset, + uint32_t n_hc, + float eps) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out || !pre || !model_map || n_hc == 0) return 0; + + @autoreleasepool { + if ((n_hc % 4u) != 0) { + fprintf(stderr, "ds4: Metal output HC weights requires a multiple-of-4 HC width\n"); + return 0; + } + + id prebuf = ds4_gpu_tensor_buffer(pre); + id outbuf = ds4_gpu_tensor_buffer(out); + const uint64_t row_bytes = (uint64_t)n_hc * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out); + if (row_bytes == 0 || out_tensor_bytes < row_bytes || out_tensor_bytes % row_bytes != 0) { + fprintf(stderr, "ds4: Metal output HC weights size is not a whole token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX || + n_tokens64 > UINT64_MAX / row_bytes) { + fprintf(stderr, "ds4: Metal output HC weights token count is outside supported range\n"); + return 0; + } + + const uint64_t bytes = n_tokens64 * row_bytes; + if (!prebuf || !outbuf || + ds4_gpu_tensor_bytes(pre) < bytes || + ds4_gpu_tensor_bytes(out) < bytes) { + fprintf(stderr, "ds4: Metal output HC weights received undersized buffers\n"); + return 0; + } + + uint64_t scale_inner = 0; + uint64_t base_inner = 0; + id scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, + scale_offset, sizeof(float), + &scale_inner); + id basebuf = ds4_gpu_wrap_model_range(model_map, model_size, + base_offset, row_bytes, + &base_inner); + if (!scalebuf || !basebuf) return 0; + + const bool force_weights4 = + getenv("DS4_METAL_ENABLE_OUTPUT_HC_WEIGHTS4") != NULL; + const bool disable_weights4 = + getenv("DS4_METAL_DISABLE_M3_OUTPUT_HC_WEIGHTS4") != NULL; + const bool require_weights4 = + getenv("DS4_METAL_REQUIRE_OUTPUT_HC_WEIGHTS4") != NULL; + const bool weights4_shape = n_hc == 4u && n_tokens64 == 1u; + const bool use_weights4 = + weights4_shape && !g_quality_mode && !disable_weights4 && + g_output_hc_weights4_pipeline != nil && + (ds4_gpu_device_name_contains("M3") || force_weights4) && + g_output_hc_weights4_pipeline.maxTotalThreadsPerThreadgroup >= 2u; + if (require_weights4 && weights4_shape && !use_weights4) { + fprintf(stderr, + "ds4: required Metal output HC weights4 kernel was not selected\n"); + return 0; + } + + if (use_weights4) { + ds4_gpu_output_hc_weights4_args args = { + .post_scale = 1.0f, + .eps = eps, + }; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:g_output_hc_weights4_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) + threadsPerThreadgroup:MTLSizeMake(2, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer( + cb, owned, "output HC weights4")) { + return 0; + } + return 1; + } + + const uint32_t n_tokens = (uint32_t)n_tokens64; + ds4_gpu_bin_args mul_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, 1); + ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, n_hc); + ds4_gpu_unary_args sigmoid_args = ds4_gpu_make_unary_rows_args(n_hc, n_tokens, 1, 0.0f, 0.0f); + ds4_gpu_unary_args scale_args = ds4_gpu_make_unary_rows_args(n_hc, n_tokens, 1, 1.0f, eps); + + NSUInteger mul_nth_max = g_bin_mul_scalar_pipeline.maxTotalThreadsPerThreadgroup; + if (mul_nth_max > 256u) mul_nth_max = 256u; + NSUInteger mul_nth = 1u; + while (2u * mul_nth < (NSUInteger)mul_args.ne0 && mul_nth < mul_nth_max) { + mul_nth *= 2u; + } + + NSUInteger add_nth_max = g_add_pipeline.maxTotalThreadsPerThreadgroup; + if (add_nth_max > 256u) add_nth_max = 256u; + NSUInteger add_nth = 1u; + while (2u * add_nth < (NSUInteger)add_args.ne0 && add_nth < add_nth_max) { + add_nth *= 2u; + } + + NSUInteger unary_nth_max = g_unary_sigmoid_pipeline.maxTotalThreadsPerThreadgroup; + if (unary_nth_max > 256u) unary_nth_max = 256u; + NSUInteger unary_nth = (NSUInteger)sigmoid_args.ne00; + if (unary_nth > unary_nth_max) unary_nth = unary_nth_max; + if (unary_nth == 0) unary_nth = 1u; + const NSUInteger unary_nk0 = ((NSUInteger)sigmoid_args.ne00 + unary_nth - 1u) / unary_nth; + const NSUInteger out_offset = ds4_gpu_tensor_offset(out); + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + + [enc setComputePipelineState:g_bin_mul_scalar_pipeline]; + [enc setBytes:&mul_args length:sizeof(mul_args) atIndex:0]; + [enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:outbuf offset:out_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)mul_args.ne01, + (NSUInteger)mul_args.ne02, + (NSUInteger)mul_args.ne03) + threadsPerThreadgroup:MTLSizeMake(mul_nth, 1, 1)]; + + [enc setComputePipelineState:g_add_pipeline]; + [enc setBytes:&add_args length:sizeof(add_args) atIndex:0]; + [enc setBuffer:outbuf offset:out_offset atIndex:1]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:2]; + [enc setBuffer:outbuf offset:out_offset atIndex:3]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)add_args.ne01, + (NSUInteger)add_args.ne02, + (NSUInteger)add_args.ne03) + threadsPerThreadgroup:MTLSizeMake(add_nth, 1, 1)]; + + [enc setComputePipelineState:g_unary_sigmoid_pipeline]; + [enc setBytes:&sigmoid_args length:sizeof(sigmoid_args) atIndex:0]; + [enc setBuffer:outbuf offset:out_offset atIndex:1]; + [enc setBuffer:outbuf offset:out_offset atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(unary_nk0 * (NSUInteger)sigmoid_args.ne01, + (NSUInteger)sigmoid_args.ne02, + (NSUInteger)sigmoid_args.ne03) + threadsPerThreadgroup:MTLSizeMake(unary_nth, 1, 1)]; + + [enc setComputePipelineState:g_unary_scale_pipeline]; + [enc setBytes:&scale_args length:sizeof(scale_args) atIndex:0]; + [enc setBuffer:outbuf offset:out_offset atIndex:1]; + [enc setBuffer:outbuf offset:out_offset atIndex:2]; + [enc dispatchThreadgroups:MTLSizeMake(unary_nk0 * (NSUInteger)scale_args.ne01, + (NSUInteger)scale_args.ne02, + (NSUInteger)scale_args.ne03) + threadsPerThreadgroup:MTLSizeMake(unary_nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "output HC weights")) return 0; + } + + return 1; +} + +int ds4_gpu_hc_expand_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *post, + const ds4_gpu_tensor *comb, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (n_embd == 0 || n_hc == 0) return 0; + + @autoreleasepool { + id blockbuf = ds4_gpu_tensor_buffer(block_out); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id postbuf = ds4_gpu_tensor_buffer(post); + id combbuf = ds4_gpu_tensor_buffer(comb); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); + if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { + fprintf(stderr, "ds4: Metal HC expand output size is not a whole HC token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC expand token count is outside supported range\n"); + return 0; + } + + const uint64_t block_values = (uint64_t)n_embd; + const uint64_t hc_values = (uint64_t)n_hc * n_embd; + const uint64_t comb_values = (uint64_t)n_hc * n_hc; + if (hc_values == 0 || + hc_values > UINT64_MAX / sizeof(float) || + comb_values > UINT64_MAX / sizeof(float) || + n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (comb_values * sizeof(float))) { + fprintf(stderr, "ds4: Metal HC expand activation size overflow\n"); + return 0; + } + + const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); + const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); + const uint64_t post_bytes = n_tokens64 * (uint64_t)n_hc * sizeof(float); + const uint64_t comb_bytes = n_tokens64 * comb_values * sizeof(float); + if (!blockbuf || !resbuf || !postbuf || !combbuf || !outbuf || + ds4_gpu_tensor_bytes(block_out) < block_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(post) < post_bytes || + ds4_gpu_tensor_bytes(comb) < comb_bytes) { + fprintf(stderr, "ds4: Metal HC expand received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_expand_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = (uint64_t)n_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 0, + }; + id expand_pipeline = g_hc_expand_pipeline; + uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; + if (n_hc == 4) { + expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, + "kernel_dsv4_hc_expand4"); + n_elem = (uint64_t)n_embd * n_tokens64; + } + if (!expand_pipeline) return 0; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:expand_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; + [enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:3]; + [enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:4]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand")) return 0; + } + + return 1; +} + +/* Expand of the SUM of two block vectors — the TP attention combine folded + * into the HC expand (the kernel's has_add path adds them element-wise + * before the post/comb mixing, canonical rank order preserved by argument + * position). */ +int ds4_gpu_hc_expand_add_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *block_add, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *post, + const ds4_gpu_tensor *comb, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (n_embd == 0 || n_hc != 4) return 0; + + @autoreleasepool { + id blockbuf = ds4_gpu_tensor_buffer(block_out); + id addbuf = ds4_gpu_tensor_buffer(block_add); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id postbuf = ds4_gpu_tensor_buffer(post); + id combbuf = ds4_gpu_tensor_buffer(comb); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); + if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { + fprintf(stderr, "ds4: Metal HC expand output size is not a whole HC token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC expand token count is outside supported range\n"); + return 0; + } + + const uint64_t block_values = (uint64_t)n_embd; + const uint64_t hc_values = (uint64_t)n_hc * n_embd; + const uint64_t comb_values = (uint64_t)n_hc * n_hc; + if (hc_values == 0 || + hc_values > UINT64_MAX / sizeof(float) || + comb_values > UINT64_MAX / sizeof(float) || + n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (comb_values * sizeof(float))) { + fprintf(stderr, "ds4: Metal HC expand activation size overflow\n"); + return 0; + } + + const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); + const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); + const uint64_t post_bytes = n_tokens64 * (uint64_t)n_hc * sizeof(float); + const uint64_t comb_bytes = n_tokens64 * comb_values * sizeof(float); + if (!blockbuf || !addbuf || !resbuf || !postbuf || !combbuf || !outbuf || + ds4_gpu_tensor_bytes(block_out) < block_bytes || + ds4_gpu_tensor_bytes(block_add) < block_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(post) < post_bytes || + ds4_gpu_tensor_bytes(comb) < comb_bytes) { + fprintf(stderr, "ds4: Metal HC expand received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_expand_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = (uint64_t)n_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 1, + }; + id expand_pipeline = g_hc_expand_pipeline; + uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; + if (n_hc == 4) { + expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, + "kernel_dsv4_hc_expand4"); + n_elem = (uint64_t)n_embd * n_tokens64; + } + if (!expand_pipeline) return 0; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:expand_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; + [enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:3]; + [enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:4]; + [enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand add")) return 0; + } + + return 1; +} + +int ds4_gpu_hc_expand_split_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !block_out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; + + @autoreleasepool { + id blockbuf = ds4_gpu_tensor_buffer(block_out); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id splitbuf = ds4_gpu_tensor_buffer(split); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); + if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { + fprintf(stderr, "ds4: Metal HC expand split output size is not a whole HC token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC expand split token count is outside supported range\n"); + return 0; + } + + const uint64_t block_values = (uint64_t)n_embd; + const uint64_t hc_values = (uint64_t)n_hc * n_embd; + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + if (hc_values == 0 || + hc_values > UINT64_MAX / sizeof(float) || + mix_hc > UINT64_MAX / sizeof(float) || + n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) { + fprintf(stderr, "ds4: Metal HC expand split activation size overflow\n"); + return 0; + } + + const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); + const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); + const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float); + if (!blockbuf || !resbuf || !splitbuf || !outbuf || + ds4_gpu_tensor_bytes(block_out) < block_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(split) < split_bytes) { + fprintf(stderr, "ds4: Metal HC expand split received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_expand_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 0, + }; + id expand_pipeline = g_hc_expand_pipeline; + uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; + if (n_hc == 4) { + expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, + "kernel_dsv4_hc_expand4"); + n_elem = (uint64_t)n_embd * n_tokens64; + } + if (!expand_pipeline) return 0; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:expand_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:3]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:4]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand split")) return 0; + } + + return 1; +} + +int ds4_gpu_hc_expand_split_half_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out_h, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + (void)out_hc; (void)block_out_h; (void)residual_hc; (void)split; + (void)n_embd; (void)n_hc; + return 0; +} + +int ds4_gpu_hc_expand_add_split_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *block_add, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !block_out || !block_add || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0; + + @autoreleasepool { + id blockbuf = ds4_gpu_tensor_buffer(block_out); + id addbuf = ds4_gpu_tensor_buffer(block_add); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id splitbuf = ds4_gpu_tensor_buffer(split); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc); + if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) { + fprintf(stderr, "ds4: Metal HC expand add split output size is not a whole HC token row\n"); + return 0; + } + + const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes; + if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) { + fprintf(stderr, "ds4: Metal HC expand add split token count is outside supported range\n"); + return 0; + } + + const uint64_t block_values = (uint64_t)n_embd; + const uint64_t hc_values = (uint64_t)n_hc * n_embd; + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + if (hc_values == 0 || + hc_values > UINT64_MAX / sizeof(float) || + mix_hc > UINT64_MAX / sizeof(float) || + n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) || + n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) { + fprintf(stderr, "ds4: Metal HC expand add split activation size overflow\n"); + return 0; + } + + const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float); + const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float); + const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float); + if (!blockbuf || !addbuf || !resbuf || !splitbuf || !outbuf || + ds4_gpu_tensor_bytes(block_out) < block_bytes || + ds4_gpu_tensor_bytes(block_add) < block_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(split) < split_bytes) { + fprintf(stderr, "ds4: Metal HC expand add split received undersized activation buffers\n"); + return 0; + } + + ds4_gpu_hc_expand_args args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = (int64_t)n_tokens64, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 1, + }; + id expand_pipeline = g_hc_expand_pipeline; + uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64; + if (n_hc == 4) { + expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline, + "kernel_dsv4_hc_expand4"); + n_elem = (uint64_t)n_embd * n_tokens64; + } + if (!expand_pipeline) return 0; + const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem)); + const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:expand_pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:3]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:4]; + [enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:5]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6]; + [enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1) + threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand add split")) return 0; + } + + return 1; +} + +int ds4_gpu_hc_expand_add_split_half_add_tensor( + ds4_gpu_tensor *out_hc, + const ds4_gpu_tensor *block_out, + const ds4_gpu_tensor *block_add_h, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + (void)out_hc; (void)block_out; (void)block_add_h; (void)residual_hc; + (void)split; (void)n_embd; (void)n_hc; + return 0; +} + +int ds4_gpu_shared_down_hc_expand_q8_0_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *shared_out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *shared_mid, + const ds4_gpu_tensor *routed_out, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !shared_out || !model_map || !shared_mid || !routed_out || + !residual_hc || !split || n_embd == 0 || n_hc == 0 || + n_hc != 4 || out_dim != n_embd || (in_dim & 31u) != 0 || + in_dim > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + id midbuf = ds4_gpu_tensor_buffer(shared_mid); + id sharedbuf = ds4_gpu_tensor_buffer(shared_out); + id routedbuf = ds4_gpu_tensor_buffer(routed_out); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id splitbuf = ds4_gpu_tensor_buffer(split); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + + const uint64_t row_bytes = (in_dim / 32u) * 34u; + const uint64_t weight_bytes = out_dim * row_bytes; + const uint64_t shared_mid_bytes = in_dim * sizeof(float); + const uint64_t embd_bytes = out_dim * sizeof(float); + const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t split_bytes = mix_hc * sizeof(float); + + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal shared-down HC fusion weight range is outside the mapped model\n"); + return 0; + } + if (!midbuf || !sharedbuf || !routedbuf || !resbuf || !splitbuf || !outbuf || + ds4_gpu_tensor_bytes(shared_mid) < shared_mid_bytes || + ds4_gpu_tensor_bytes(shared_out) < embd_bytes || + ds4_gpu_tensor_bytes(routed_out) < embd_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(split) < split_bytes || + ds4_gpu_tensor_bytes(out_hc) < hc_bytes) { + fprintf(stderr, "ds4: Metal shared-down HC fusion received undersized buffers\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + mv_args.nr0 = mv_dispatch.nr0; + + ds4_gpu_hc_expand_args hc_args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = 1, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 1, + }; + + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_shared_down_hc_expand4_q8_0", + mv_dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2]; + [enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(shared_mid) atIndex:3]; + [enc setBuffer:sharedbuf offset:ds4_gpu_tensor_offset(shared_out) atIndex:4]; + [enc setBuffer:routedbuf offset:ds4_gpu_tensor_offset(routed_out) atIndex:5]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:6]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:7]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:8]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:9]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "shared-down HC expand fused")) return 0; + } + + return 1; +} + +int ds4_gpu_matmul_q8_0_hc_expand_tensor( + ds4_gpu_tensor *out_hc, + ds4_gpu_tensor *block_out, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint64_t in_dim, + uint64_t out_dim, + const ds4_gpu_tensor *x, + const ds4_gpu_tensor *residual_hc, + const ds4_gpu_tensor *split, + uint32_t n_embd, + uint32_t n_hc) { + if (!g_initialized && !ds4_gpu_init()) return 0; + if (!out_hc || !block_out || !model_map || !x || !residual_hc || !split || + n_embd == 0 || n_hc == 0 || n_hc != 4 || out_dim != n_embd || + (in_dim & 31u) != 0 || in_dim > UINT32_MAX || out_dim > UINT32_MAX) { + return 0; + } + + @autoreleasepool { + id xbuf = ds4_gpu_tensor_buffer(x); + id blockbuf = ds4_gpu_tensor_buffer(block_out); + id resbuf = ds4_gpu_tensor_buffer(residual_hc); + id splitbuf = ds4_gpu_tensor_buffer(split); + id outbuf = ds4_gpu_tensor_buffer(out_hc); + + const uint64_t row_bytes = (in_dim / 32u) * 34u; + const uint64_t weight_bytes = out_dim * row_bytes; + const uint64_t x_bytes = in_dim * sizeof(float); + const uint64_t embd_bytes = out_dim * sizeof(float); + const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float); + const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc; + const uint64_t split_bytes = mix_hc * sizeof(float); + + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: Metal Q8 HC fusion weight range is outside the mapped model\n"); + return 0; + } + if (!xbuf || !blockbuf || !resbuf || !splitbuf || !outbuf || + ds4_gpu_tensor_bytes(x) < x_bytes || + ds4_gpu_tensor_bytes(block_out) < embd_bytes || + ds4_gpu_tensor_bytes(residual_hc) < hc_bytes || + ds4_gpu_tensor_bytes(split) < split_bytes || + ds4_gpu_tensor_bytes(out_hc) < hc_bytes) { + fprintf(stderr, "ds4: Metal Q8 HC fusion received undersized buffers\n"); + return 0; + } + + uint64_t inner_offset = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, + weight_offset, weight_bytes, + &inner_offset); + if (!wbuf) return 0; + + ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + mv_args.nr0 = mv_dispatch.nr0; + + ds4_gpu_hc_expand_args hc_args = { + .n_embd = n_embd, + .n_hc = n_hc, + .n_tokens = 1, + .nb_block0 = sizeof(float), + .nb_block1 = (uint64_t)n_embd * sizeof(float), + .nb_add0 = sizeof(float), + .nb_add1 = (uint64_t)n_embd * sizeof(float), + .nb_res0 = sizeof(float), + .nb_res1 = (uint64_t)n_embd * sizeof(float), + .nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float), + .nb_post0 = sizeof(float), + .nb_post1 = mix_hc * sizeof(float), + .nb_comb0 = sizeof(float), + .nb_comb1 = (uint64_t)n_hc * sizeof(float), + .nb_comb2 = mix_hc * sizeof(float), + .nb0 = sizeof(float), + .nb1 = (uint64_t)n_embd * sizeof(float), + .nb2 = (uint64_t)n_hc * n_embd * sizeof(float), + .has_add = 0, + }; + + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_q8_hc_expand4_q8_0", + mv_dispatch.nsg); + if (!pipeline) return 0; + + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + if (!cb) return 0; + + id enc = ds4_gpu_compute_encoder(cb); + [enc setComputePipelineState:pipeline]; + [enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0]; + [enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1]; + [enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2]; + [enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:4]; + [enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:5]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:6]; + [enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:7]; + [enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:8]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, + 1, + 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + + if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8 HC expand fused")) return 0; + } + + return 1; +} + +void ds4_gpu_set_glm_mtp_verify_mode(bool enabled) { + (void)enabled; +} diff --git a/native/metal/ds4_ssd.h b/native/metal/ds4_ssd.h new file mode 100644 index 0000000..c32eaad --- /dev/null +++ b/native/metal/ds4_ssd.h @@ -0,0 +1,36 @@ +#ifndef DS4_SSD_H +#define DS4_SSD_H + +#include +#include + +typedef struct { + void *ptr; + uint64_t bytes; +} ds4_ssd_memory_lock; + +typedef struct { + uint64_t model_target_bytes; + uint64_t cache_bytes; + uint64_t effective_cache_bytes; + uint32_t cache_experts; +} ds4_ssd_cache_plan; + +bool ds4_parse_gib_arg(const char *s, uint64_t *bytes); +bool ds4_parse_streaming_cache_experts_arg(const char *s, + uint32_t *experts, + uint64_t *bytes); + +uint32_t ds4_ssd_cache_experts_for_byte_budget(uint64_t bytes, + uint64_t per_expert_bytes); +bool ds4_ssd_auto_cache_plan(uint64_t recommended_bytes, + uint64_t non_routed_bytes, + uint64_t per_expert_bytes, + uint64_t max_model_experts, + ds4_ssd_cache_plan *out); + +bool ds4_ssd_memory_lock_acquire(ds4_ssd_memory_lock *lock, + uint64_t bytes); +void ds4_ssd_memory_lock_release(ds4_ssd_memory_lock *lock); + +#endif diff --git a/src/app.rs b/src/app.rs index 5966368..9001bb7 100644 --- a/src/app.rs +++ b/src/app.rs @@ -3,6 +3,8 @@ mod view; pub(crate) use view::app_theme; use crate::database::{AppPreferences, Database, ProjectWithSessions}; +#[cfg(target_os = "macos")] +use crate::engine::{ChatTurn, Generator}; use crate::model::{self, DownloadOutcome, DownloadProgress, ManagedArtifactId, ModelChoice}; use crate::settings::{ DiagnosticPreferences, ExecutionPreferences, GIB, GenerationPreferences, ReasoningMode, @@ -398,9 +400,46 @@ pub(crate) struct App { pending_project_path: Option, project_name_input: String, model_download: ModelDownload, + pub(super) composer: String, + pub(super) conversation: Vec, + pub(super) generating: bool, + #[cfg(target_os = "macos")] + generation_worker: Option, error: Option, } +#[derive(Clone, Debug)] +pub(super) struct ChatMessage { + pub(super) user: bool, + reasoning: bool, + pub(super) content: String, +} + +#[cfg(target_os = "macos")] +struct GenerationWorker { + commands: mpsc::Sender, + events: mpsc::Receiver, + cancel: Option>, +} + +#[cfg(target_os = "macos")] +enum GenerationCommand { + Generate { + engine: crate::settings::EngineSettings, + turn: crate::settings::TurnSettings, + messages: Vec, + idle_timeout: Duration, + cancel: Arc, + }, +} + +#[cfg(target_os = "macos")] +enum GenerationEvent { + Loading, + Chunk(String), + Finished(Result<(), String>), +} + #[derive(Debug)] pub(super) enum ModelDownload { Idle, @@ -479,6 +518,10 @@ pub(crate) enum Message { CancelDeleteArtifact, StopModelDownload, DownloadProgressTick, + ComposerChanged(String), + SubmitPrompt, + StopGeneration, + GenerationTick, ChooseProjectFolder, ProjectFolderPicked(Option), ProjectNameChanged(String), @@ -519,6 +562,11 @@ impl App { pending_project_path: None, project_name_input: String::new(), model_download: ModelDownload::Idle, + composer: String::new(), + conversation: Vec::new(), + generating: false, + #[cfg(target_os = "macos")] + generation_worker: None, error: None, } } @@ -550,6 +598,11 @@ impl App { pending_project_path: None, project_name_input: String::new(), model_download: ModelDownload::Idle, + composer: String::new(), + conversation: Vec::new(), + generating: false, + #[cfg(target_os = "macos")] + generation_worker: None, error: Some(format!("Could not open the project database: {error}")), } } @@ -816,6 +869,19 @@ impl App { } } Message::DownloadProgressTick => self.update_download_progress(), + Message::ComposerChanged(value) => self.composer = value, + Message::SubmitPrompt => self.start_generation(), + Message::StopGeneration => { + #[cfg(target_os = "macos")] + if let Some(cancel) = self + .generation_worker + .as_ref() + .and_then(|worker| worker.cancel.as_ref()) + { + cancel.store(true, Ordering::Relaxed); + } + } + Message::GenerationTick => self.poll_generation(), Message::ChooseProjectFolder => { self.choosing_folder = true; return Task::perform( @@ -843,11 +909,21 @@ impl App { self.error = None; } Message::SelectProject(project_id) => { + if self.generating { + self.error = + Some("Stop the active generation before changing sessions.".into()); + return Task::none(); + } self.selected_project = Some(project_id); self.selected_session = None; self.error = None; } Message::DeleteProject(project_id) => { + if self.generating && self.selected_project == Some(project_id) { + self.error = + Some("Stop the active generation before deleting its project.".into()); + return Task::none(); + } if let Some(database) = &mut self.database { match database.delete_project(project_id) { Ok(()) => { @@ -863,11 +939,25 @@ impl App { } Message::CreateSession => self.create_session(), Message::SelectSession(project_id, session_id) => { + if self.generating && self.selected_session != Some(session_id) { + self.error = + Some("Stop the active generation before changing sessions.".into()); + return Task::none(); + } + if self.selected_session != Some(session_id) { + self.conversation.clear(); + self.composer.clear(); + } self.selected_project = Some(project_id); self.selected_session = Some(session_id); self.error = None; } Message::DeleteSession(session_id) => { + if self.generating && self.selected_session == Some(session_id) { + self.error = + Some("Stop the active generation before deleting its session.".into()); + return Task::none(); + } if let Some(database) = &mut self.database { match database.delete_session(session_id) { Ok(()) => { @@ -905,6 +995,11 @@ impl App { iced::time::every(Duration::from_secs(1)).map(|_| Message::DownloadProgressTick), ); } + if self.generating { + subscriptions.push( + iced::time::every(Duration::from_millis(50)).map(|_| Message::GenerationTick), + ); + } Subscription::batch(subscriptions) } @@ -1108,6 +1203,10 @@ impl App { } fn create_session(&mut self) { + if self.generating { + self.error = Some("Stop the active generation before creating a session.".into()); + return; + } let Some(project_id) = self.selected_project else { self.error = Some("Select a project first.".into()); return; @@ -1205,6 +1304,196 @@ impl App { } } } + + fn start_generation(&mut self) { + if self.generating || self.selected_session.is_none() { + return; + } + let prompt = self.composer.trim().to_owned(); + if prompt.is_empty() { + return; + } + let model = match ModelChoice::from_id(&self.preferences.selected_model) { + Some(model) => model, + None => { + self.error = Some("The selected model is not supported.".into()); + return; + } + }; + let effective = self.preferences.generation().and_then(|generation| { + self.preferences.runtime().and_then(|runtime| { + crate::settings::effective_settings(model, &generation, &runtime, &models_path()) + }) + }); + let effective = match effective { + Ok(settings) => settings, + Err(error) => { + self.error = Some(error); + return; + } + }; + let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct; + #[cfg(target_os = "macos")] + let mut messages = self + .conversation + .iter() + .map(|message| ChatTurn { + user: message.user, + reasoning: message.reasoning, + content: message.content.clone(), + }) + .collect::>(); + #[cfg(target_os = "macos")] + messages.push(ChatTurn { + user: true, + reasoning: false, + content: prompt.clone(), + }); + + #[cfg(target_os = "macos")] + { + if self.generation_worker.is_none() { + match spawn_generation_worker() { + Ok(worker) => self.generation_worker = Some(worker), + Err(error) => { + self.error = Some(error); + return; + } + } + } + let cancel = Arc::new(AtomicBool::new(false)); + let idle_timeout = + Duration::from_secs(self.preferences.idle_timeout_minutes.max(1) as u64 * 60); + let command = GenerationCommand::Generate { + engine: effective.engine, + turn: effective.turn, + messages, + idle_timeout, + cancel: Arc::clone(&cancel), + }; + let worker = self.generation_worker.as_mut().expect("worker was created"); + if worker.commands.send(command).is_err() { + self.generation_worker = None; + self.error = Some("The local generation worker stopped unexpectedly.".into()); + return; + } + worker.cancel = Some(cancel); + } + #[cfg(not(target_os = "macos"))] + { + let _ = effective; + self.error = Some("Local Metal generation requires macOS.".into()); + return; + } + self.composer.clear(); + self.conversation.push(ChatMessage { + user: true, + reasoning: false, + content: prompt, + }); + self.conversation.push(ChatMessage { + user: false, + reasoning: assistant_reasoning, + content: String::new(), + }); + self.generating = true; + self.error = None; + } + + fn poll_generation(&mut self) { + #[cfg(target_os = "macos")] + let Some(worker) = &mut self.generation_worker else { + self.generating = false; + return; + }; + #[cfg(target_os = "macos")] + loop { + match worker.events.try_recv() { + Ok(GenerationEvent::Loading) => {} + Ok(GenerationEvent::Chunk(chunk)) => { + if let Some(message) = self.conversation.last_mut() + && !message.user + { + message.content.push_str(&chunk); + } + } + Ok(GenerationEvent::Finished(result)) => { + self.generating = false; + worker.cancel = None; + if let Err(error) = result { + self.error = Some(error); + } + break; + } + Err(TryRecvError::Empty) => break, + Err(TryRecvError::Disconnected) => { + self.generating = false; + self.generation_worker = None; + self.error = Some("The local generation worker stopped unexpectedly.".into()); + break; + } + } + } + } +} + +#[cfg(target_os = "macos")] +fn spawn_generation_worker() -> Result { + let (command_sender, command_receiver) = mpsc::channel(); + let (event_sender, event_receiver) = mpsc::channel(); + thread::Builder::new() + .name("local-generation".into()) + .spawn(move || { + let mut loaded = None::<(crate::settings::EngineSettings, Generator)>; + let mut last_used = Instant::now(); + let mut idle_timeout = Duration::from_secs(15 * 60); + loop { + match command_receiver.recv_timeout(Duration::from_secs(1)) { + Ok(GenerationCommand::Generate { + engine, + turn, + messages, + idle_timeout: requested_timeout, + cancel, + }) => { + idle_timeout = requested_timeout; + if loaded + .as_ref() + .is_none_or(|(current, _)| current != &engine) + { + let _ = event_sender.send(GenerationEvent::Loading); + loaded = match Generator::open(&engine) { + Ok(generator) => Some((engine.clone(), generator)), + Err(error) => { + let _ = + event_sender.send(GenerationEvent::Finished(Err(error))); + None + } + }; + } + if let Some((_, generator)) = &mut loaded { + let result = generator.generate(&messages, &turn, &cancel, |chunk| { + let _ = event_sender.send(GenerationEvent::Chunk(chunk)); + }); + let _ = event_sender.send(GenerationEvent::Finished(result)); + last_used = Instant::now(); + } + } + Err(mpsc::RecvTimeoutError::Timeout) => { + if loaded.is_some() && last_used.elapsed() >= idle_timeout { + loaded = None; + } + } + Err(mpsc::RecvTimeoutError::Disconnected) => break, + } + } + }) + .map_err(|error| format!("Could not start local generation: {error}"))?; + Ok(GenerationWorker { + commands: command_sender, + events: event_receiver, + cancel: None, + }) } impl Drop for App { @@ -1212,6 +1501,14 @@ impl Drop for App { if let ModelDownload::Active(download) = &self.model_download { download.cancel.store(true, Ordering::Relaxed); } + #[cfg(target_os = "macos")] + if let Some(cancel) = self + .generation_worker + .as_ref() + .and_then(|worker| worker.cancel.as_ref()) + { + cancel.store(true, Ordering::Relaxed); + } } } diff --git a/src/app/view.rs b/src/app/view.rs index 9a97795..f829279 100644 --- a/src/app/view.rs +++ b/src/app/view.rs @@ -224,18 +224,51 @@ impl App { .align_y(Alignment::Center); let body: Element<'_, Message> = if let Some(session) = self.selected_session(item) { + let mut messages = column![].spacing(12); + if self.conversation.is_empty() { + messages = messages.push( + column![ + text(&session.title).size(26), + text("Run DeepSeek locally with the Rust Metal engine.").size(14), + ] + .spacing(8), + ); + } else { + for message in &self.conversation { + let label = if message.user { "You" } else { "DS4" }; + let content = if !message.user && message.content.is_empty() && self.generating + { + "Loading model…" + } else { + &message.content + }; + messages = messages.push( + container(column![text(label).size(11), text(content).size(14)].spacing(5)) + .padding(14) + .width(Length::Fill) + .style(overview_style), + ); + } + } + let composer = text_input("Ask DS4Server anything…", &self.composer) + .on_input(Message::ComposerChanged) + .on_submit(Message::SubmitPrompt) + .padding(12) + .size(14); + let action = if self.generating { + action_button(text("Stop").size(12)).on_press(Message::StopGeneration) + } else if self.composer.trim().is_empty() { + action_button(icon(ICON_SEND, 18)).padding(8) + } else { + action_button(icon(ICON_SEND, 18)) + .padding(8) + .on_press(Message::SubmitPrompt) + }; let conversation = column![ - Space::with_height(Length::Fill), - column![ - text(&session.title).size(26), - text("This session is ready for the agent runtime.").size(14), - ] - .spacing(8), - Space::with_height(Length::Fill), + scrollable(messages).height(Length::Fill), container( column![ - text("Ask DS4Server anything…"), - Space::with_height(28), + composer, row![ icon(ICON_PAPERCLIP, 19), Space::with_width(Length::Fill), @@ -246,7 +279,7 @@ impl App { .to_string(), ) .size(12), - action_button(icon(ICON_SEND, 18)).padding(8), + action, ] .align_y(Alignment::Center), ] @@ -658,7 +691,8 @@ impl App { let panel = container( column![ header, - scrollable(container(fields).padding(iced::Padding::ZERO.right(18))), + scrollable(container(fields).padding(iced::Padding::ZERO.right(18))) + .height(Length::Fill), footer ] .spacing(16), diff --git a/src/engine.rs b/src/engine.rs index 8200bdf..8215282 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -1,12 +1,19 @@ mod gguf; +#[cfg(target_os = "macos")] +mod metal; mod tokenizer; use crate::model::ModelChoice; +use crate::settings::TurnSettings; use crate::settings::{EngineSettings, ReasoningMode}; use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value}; use std::path::Path; +use std::sync::atomic::{AtomicBool, Ordering}; use tokenizer::Tokenizer; +#[cfg(target_os = "macos")] +pub(crate) use metal::configure_sources as configure_metal_sources; + const DENSE: &[u32] = &[Q8_0, Q4_K, Q4_0]; const ROUTED: &[u32] = &[Q8_0, IQ2_XXS, Q2_K, Q4_K, Q5_K, Q6_K]; const PLAIN: &[u32] = &[F16, F32]; @@ -236,6 +243,16 @@ impl Model { self.tokenizer.encode_chat(system, prompt, reasoning) } + fn render_conversation( + &self, + system: &str, + messages: &[ChatTurn], + reasoning: ReasoningMode, + ) -> Vec { + self.tokenizer + .encode_conversation(system, messages, reasoning) + } + pub(crate) fn token_bytes(&self, token: i32) -> Option> { self.tokenizer.token_bytes(token) } @@ -253,6 +270,230 @@ impl Model { } } +#[cfg(target_os = "macos")] +pub(crate) struct Generator { + executor: metal::Executor, +} + +#[derive(Clone)] +pub(crate) struct ChatTurn { + pub(crate) user: bool, + pub(crate) reasoning: bool, + pub(crate) content: String, +} + +#[cfg(target_os = "macos")] +impl Generator { + pub(crate) fn open(settings: &EngineSettings) -> Result { + if settings.model != ModelChoice::DeepSeekV4Flash { + return Err("local generation currently supports DeepSeek V4 Flash only".into()); + } + if settings.speculative.dspark || settings.ssd.enabled || settings.steering.file.is_some() { + return Err( + "DSpark, SSD streaming, and steering are not yet available in the Rust executor" + .into(), + ); + } + let model = Model::open(settings)?; + let executor = metal::Executor::open( + model, + settings.context_tokens.max(1) as u32, + settings.execution.quality, + )?; + Ok(Self { executor }) + } + + pub(crate) fn generate( + &mut self, + messages: &[ChatTurn], + settings: &TurnSettings, + cancelled: &AtomicBool, + mut emit: impl FnMut(String), + ) -> Result<(), String> { + self.executor.reset()?; + let tokens = self.executor.model().render_conversation( + &settings.system_prompt, + messages, + settings.reasoning_mode, + ); + if tokens.is_empty() { + return Err("the rendered prompt is empty".into()); + } + let max_context = self.executor.context() as usize; + if tokens.len() >= max_context { + return Err(format!( + "the conversation uses {} tokens; the current Rust attention port supports fewer than {max_context}", + tokens.len() + )); + } + let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645)); + for token in tokens { + if cancelled.load(Ordering::Relaxed) { + return Ok(()); + } + self.executor.eval(token)?; + } + for _ in 0..settings + .max_generated_tokens + .max(0) + .min((max_context - self.executor.position() as usize) as i32) + { + if cancelled.load(Ordering::Relaxed) { + return Ok(()); + } + let token = sample( + self.executor.logits(), + settings.temperature, + settings.top_p, + settings.min_p, + &mut rng, + ); + if self.executor.model().is_stop_token(token) { + return Ok(()); + } + if let Some(bytes) = self.executor.model().token_bytes(token) { + emit(String::from_utf8_lossy(&bytes).into_owned()); + } + self.executor.eval(token)?; + } + Ok(()) + } +} + +#[cfg(target_os = "macos")] +fn sample(logits: &[f32], temperature: f32, top_p: f32, min_p: f32, rng: &mut Rng) -> i32 { + if temperature <= 0.0 { + return logits + .iter() + .enumerate() + .max_by(|a, b| a.1.total_cmp(b.1)) + .map_or(0, |(index, _)| index as i32); + } + let maximum = logits + .iter() + .copied() + .filter(|value| value.is_finite()) + .fold(f32::NEG_INFINITY, f32::max); + if !maximum.is_finite() { + return 0; + } + let top_p = if top_p <= 0.0 || top_p > 1.0 { + 1.0 + } else { + top_p + }; + let min_p = min_p.max(0.0); + let mut probabilities: Vec<(usize, f32)> = logits + .iter() + .enumerate() + .filter(|(_, logit)| logit.is_finite()) + .map(|(index, logit)| (index, ((*logit - maximum) / temperature).exp())) + .filter(|(_, probability)| *probability >= min_p) + .collect(); + if probabilities.is_empty() { + return logits + .iter() + .enumerate() + .max_by(|a, b| a.1.total_cmp(b.1)) + .map_or(0, |(index, _)| index as i32); + } + if top_p < 1.0 { + probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + let total: f32 = logits + .iter() + .filter(|logit| logit.is_finite()) + .map(|logit| ((*logit - maximum) / temperature).exp()) + .sum(); + let mut kept = 0.0; + let count = probabilities + .iter() + .position(|(_, probability)| { + kept += *probability; + kept / total >= top_p + }) + .map_or(probabilities.len(), |index| index + 1); + probabilities.truncate(count); + } + let kept_total: f32 = probabilities.iter().map(|(_, p)| p).sum(); + let mut choice = rng.unit() * kept_total; + for (token, probability) in &probabilities { + choice -= probability; + if choice <= 0.0 { + return *token as i32; + } + } + probabilities.last().map_or(0, |(token, _)| *token as i32) +} + +#[cfg(target_os = "macos")] +struct Rng(u64); + +#[cfg(target_os = "macos")] +impl Rng { + fn new(seed: u64) -> Self { + Self(seed.max(1)) + } + + fn unit(&mut self) -> f32 { + let mut value = self.0; + if value == 0 { + value = 0x9e37_79b9_7f4a_7c15; + } + value ^= value >> 12; + value ^= value << 25; + value ^= value >> 27; + self.0 = value; + let value = value.wrapping_mul(0x2545_f491_4f6c_dd1d); + ((value >> 40) & 0xff_ffff) as f32 / 16_777_216.0 + } +} + +#[cfg(all(test, target_os = "macos"))] +mod sampling_tests { + use super::*; + + #[test] + fn zero_temperature_is_greedy() { + let mut rng = Rng::new(1); + assert_eq!(sample(&[1.0, 4.0, 2.0], 0.0, 1.0, 0.0, &mut rng), 1); + } + + #[test] + #[ignore = "requires the 80 GiB Flash checkpoint and Apple Metal"] + fn metal_executes_real_flash_token() { + configure_metal_sources().unwrap(); + let path = Path::new(env!("CARGO_MANIFEST_DIR")).join( + "../ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf", + ); + let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash).unwrap(); + let tokens = model.render_prompt( + "You are a helpful assistant", + "Hello", + ReasoningMode::Direct, + ); + assert_eq!(tokens.len(), 10); + let mut executor = metal::Executor::open(model, 128, false).unwrap(); + for token in tokens { + executor.eval(token).unwrap(); + } + assert!(executor.logits().iter().all(|logit| logit.is_finite())); + let argmax = executor + .logits() + .iter() + .enumerate() + .max_by(|a, b| a.1.total_cmp(b.1)) + .unwrap(); + eprintln!( + "Rust logits: argmax={} value={} logit0={}", + argmax.0, + argmax.1, + executor.logits()[0] + ); + assert_eq!(argmax.0, 19_923); + assert!((executor.logits()[0] - -7.675_424).abs() < 0.1); + } +} + pub(crate) fn validate_model_artifact( path: &Path, expected: ModelChoice, @@ -1144,6 +1385,32 @@ mod tests { model.render_prompt("", "Hello", ReasoningMode::Direct), [0, 128_803, 19_923, 128_804, 128_822] ); + assert_eq!( + model.render_conversation( + "", + &[ + ChatTurn { + user: true, + reasoning: false, + content: "Hello".into(), + }, + ChatTurn { + user: false, + reasoning: false, + content: "Hello".into(), + }, + ChatTurn { + user: true, + reasoning: false, + content: "Hello".into(), + }, + ], + ReasoningMode::Direct, + ), + [ + 0, 128_803, 19_923, 128_804, 128_822, 19_923, 128_803, 19_923, 128_804, 128_822, + ] + ); } #[test] diff --git a/src/engine/gguf.rs b/src/engine/gguf.rs index 1f90bd1..322f057 100644 --- a/src/engine/gguf.rs +++ b/src/engine/gguf.rs @@ -170,6 +170,10 @@ impl Gguf { self.map.len() as u64 } + pub(super) fn map_ptr(&self) -> *const u8 { + self.map.as_ptr() + } + pub(super) fn tensor(&self, name: &str) -> Result<&Tensor, String> { self.tensors .get(name) diff --git a/src/engine/metal.rs b/src/engine/metal.rs new file mode 100644 index 0000000..f1409ef --- /dev/null +++ b/src/engine/metal.rs @@ -0,0 +1,1729 @@ +use super::gguf::{Gguf, Q8_0, Tensor as GgufTensor}; +use super::{Model, ModelFamily}; +use std::env; +use std::ffi::c_void; +use std::path::Path; +use std::ptr::NonNull; + +const SOURCES: [(&str, &str); 19] = [ + ("DS4_METAL_FLASH_ATTN_SOURCE", "flash_attn.metal"), + ("DS4_METAL_DENSE_SOURCE", "dense.metal"), + ("DS4_METAL_MOE_SOURCE", "moe.metal"), + ("DS4_METAL_DSV4_HC_SOURCE", "dsv4_hc.metal"), + ("DS4_METAL_UNARY_SOURCE", "unary.metal"), + ("DS4_METAL_DSV4_KV_SOURCE", "dsv4_kv.metal"), + ("DS4_METAL_DSV4_ROPE_SOURCE", "dsv4_rope.metal"), + ("DS4_METAL_DSV4_MISC_SOURCE", "dsv4_misc.metal"), + ("DS4_METAL_ARGSORT_SOURCE", "argsort.metal"), + ("DS4_METAL_CPY_SOURCE", "cpy.metal"), + ("DS4_METAL_CONCAT_SOURCE", "concat.metal"), + ("DS4_METAL_GET_ROWS_SOURCE", "get_rows.metal"), + ("DS4_METAL_SUM_ROWS_SOURCE", "sum_rows.metal"), + ("DS4_METAL_SOFTMAX_SOURCE", "softmax.metal"), + ("DS4_METAL_REPEAT_SOURCE", "repeat.metal"), + ("DS4_METAL_GLU_SOURCE", "glu.metal"), + ("DS4_METAL_NORM_SOURCE", "norm.metal"), + ("DS4_METAL_BIN_SOURCE", "bin.metal"), + ("DS4_METAL_SET_ROWS_SOURCE", "set_rows.metal"), +]; + +// The Metal boundary uses this only to decide whether diagnostic logs get ANSI +// color. App output is never a terminal, so Rust owns the one non-Metal symbol +// that ds4_metal.m expects without pulling in ds4.c. +#[unsafe(no_mangle)] +extern "C" fn ds4_log_is_tty(_stream: *mut c_void) -> bool { + false +} + +pub(crate) fn configure_sources() -> Result<(), String> { + let bundled = env::current_exe().ok().and_then(|path| { + path.parent()? + .parent() + .map(|path| path.join("Resources/metal")) + }); + let checkout = Path::new(env!("CARGO_MANIFEST_DIR")).join("metal"); + let directory = bundled + .filter(|path| path.join("dense.metal").is_file()) + .unwrap_or(checkout); + for (name, file) in SOURCES { + let path = directory.join(file); + if !path.is_file() { + return Err(format!( + "required Metal kernel is missing: {}", + path.display() + )); + } + // SAFETY: main calls this before Iced starts any application threads. + unsafe { env::set_var(name, path) }; + } + Ok(()) +} + +#[repr(C)] +struct GpuTensor { + _private: [u8; 0], +} + +unsafe extern "C" { + fn ds4_gpu_init() -> i32; + fn ds4_gpu_cleanup(); + fn ds4_gpu_set_model_map(model_map: *const c_void, model_size: u64) -> i32; + fn ds4_gpu_set_quality(quality: bool); + fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor; + fn ds4_gpu_tensor_free(tensor: *mut GpuTensor); + fn ds4_gpu_tensor_fill_f32(tensor: *mut GpuTensor, value: f32, count: u64) -> i32; + fn ds4_gpu_tensor_read( + tensor: *const GpuTensor, + offset: u64, + data: *mut c_void, + bytes: u64, + ) -> i32; + fn ds4_gpu_tensor_copy_f32_to_f16( + dst: *mut GpuTensor, + dst_offset: u64, + src: *const GpuTensor, + src_offset: u64, + count: u64, + ) -> i32; + fn ds4_gpu_begin_commands() -> i32; + fn ds4_gpu_end_commands() -> i32; + + fn ds4_gpu_embed_token_hc_tensor( + out: *mut GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + vocab: u32, + token: u32, + embd: u32, + hc: u32, + ) -> i32; + fn ds4_gpu_rms_norm_plain_tensor( + out: *mut GpuTensor, + x: *const GpuTensor, + n: u32, + eps: f32, + ) -> i32; + fn ds4_gpu_rms_norm_weight_tensor( + out: *mut GpuTensor, + x: *const GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + n: u32, + eps: f32, + ) -> i32; + fn ds4_gpu_matmul_f16_tensor( + out: *mut GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + input: u64, + output: u64, + x: *const GpuTensor, + rows: u64, + ) -> i32; + fn ds4_gpu_matmul_q8_0_tensor( + out: *mut GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + input: u64, + output: u64, + x: *const GpuTensor, + rows: u64, + ) -> i32; + fn ds4_gpu_matmul_q8_0_pair_tensor( + out_a: *mut GpuTensor, + out_b: *mut GpuTensor, + map: *const c_void, + size: u64, + weight_a: u64, + weight_b: u64, + input: u64, + output_a: u64, + output_b: u64, + x: *const GpuTensor, + rows: u64, + ) -> i32; + fn ds4_gpu_matmul_f16_pair_tensor( + out_a: *mut GpuTensor, + out_b: *mut GpuTensor, + map: *const c_void, + size: u64, + weight_a: u64, + weight_b: u64, + input: u64, + output: u64, + x: *const GpuTensor, + rows: u64, + ) -> i32; + fn ds4_gpu_matmul_f16_pair_compressor_store_tensor( + out_kv: *mut GpuTensor, + out_score: *mut GpuTensor, + state_kv: *mut GpuTensor, + state_score: *mut GpuTensor, + map: *const c_void, + size: u64, + weight_kv: u64, + weight_score: u64, + ape: u64, + ape_type: u32, + input: u64, + width: u32, + x: *const GpuTensor, + ratio: u32, + pos: u32, + ) -> i32; + fn ds4_gpu_hc_split_weighted_sum_norm_tensor( + out: *mut GpuTensor, + norm: *mut GpuTensor, + split: *mut GpuTensor, + mix: *const GpuTensor, + residual: *const GpuTensor, + map: *const c_void, + size: u64, + scale: u64, + base: u64, + norm_weight: u64, + embd: u32, + hc: u32, + iterations: u32, + eps: f32, + norm_eps: f32, + ) -> i32; + fn ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( + q_out: *mut GpuTensor, + q: *const GpuTensor, + map: *const c_void, + size: u64, + q_weight: u64, + q_n: u32, + kv_out: *mut GpuTensor, + kv: *const GpuTensor, + kv_weight: u64, + kv_n: u32, + rows: u32, + eps: f32, + ) -> i32; + fn ds4_gpu_head_rms_norm_rope_tail_tensor( + x: *mut GpuTensor, + rows: u32, + heads: u32, + head_dim: u32, + rot: u32, + pos: u32, + original_context: u32, + inverse: bool, + freq_base: f32, + freq_scale: f32, + ext_factor: f32, + attn_factor: f32, + beta_fast: f32, + beta_slow: f32, + eps: f32, + ) -> i32; + fn ds4_gpu_head_rms_norm_tensor( + x: *mut GpuTensor, + rows: u32, + heads: u32, + head_dim: u32, + eps: f32, + ) -> i32; + fn ds4_gpu_rope_tail_tensor( + x: *mut GpuTensor, + rows: u32, + heads: u32, + head_dim: u32, + rot: u32, + pos: u32, + original_context: u32, + inverse: bool, + freq_base: f32, + freq_scale: f32, + ext_factor: f32, + attn_factor: f32, + beta_fast: f32, + beta_slow: f32, + ) -> i32; + fn ds4_gpu_attention_decode_heads_tensor( + heads_out: *mut GpuTensor, + map: *const c_void, + size: u64, + sinks: u64, + q: *const GpuTensor, + raw_kv: *const GpuTensor, + n_raw: u32, + raw_cap: u32, + raw_start: u32, + comp_kv: *const GpuTensor, + comp_f16: u32, + n_comp: u32, + comp_mask: *const GpuTensor, + use_mask: u32, + heads: u32, + head_dim: u32, + ) -> i32; + fn ds4_gpu_compressor_update_tensor( + kv: *const GpuTensor, + score: *const GpuTensor, + state_kv: *mut GpuTensor, + state_score: *mut GpuTensor, + cache: *mut GpuTensor, + map: *const c_void, + size: u64, + ape: u64, + ape_type: u32, + norm: u64, + norm_type: u32, + head_dim: u32, + ratio: u32, + pos: u32, + row: u32, + rot: u32, + original_context: u32, + freq_base: f32, + freq_scale: f32, + ext_factor: f32, + attn_factor: f32, + beta_fast: f32, + beta_slow: f32, + rms_eps: f32, + state_already_stored: bool, + ) -> i32; + fn ds4_gpu_dsv4_fp8_kv_quantize_tensor( + x: *mut GpuTensor, + rows: u32, + head_dim: u32, + rot: u32, + ) -> i32; + fn ds4_gpu_kv_fp8_store_raw_tensor( + kv: *mut GpuTensor, + raw_cache: *mut GpuTensor, + raw_cap: u32, + row: u32, + head_dim: u32, + rot: u32, + ) -> i32; + fn ds4_gpu_attention_output_low_q8_tensor( + low: *mut GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + group_dim: u64, + rank: u64, + groups: u32, + heads: *const GpuTensor, + ) -> i32; + fn ds4_gpu_matmul_q8_0_hc_expand_tensor( + out_hc: *mut GpuTensor, + block_out: *mut GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + input: u64, + output: u64, + x: *const GpuTensor, + residual: *const GpuTensor, + split: *const GpuTensor, + embd: u32, + hc: u32, + ) -> i32; + fn ds4_gpu_router_select_tensor( + selected: *mut GpuTensor, + weights: *mut GpuTensor, + probs: *mut GpuTensor, + map: *const c_void, + size: u64, + bias: u64, + hash: u64, + hash_rows: u32, + token: u32, + experts: u32, + used: u32, + scale: f32, + expert_groups: u32, + groups_used: u32, + has_bias: bool, + hash_mode: bool, + logits: *const GpuTensor, + ) -> i32; + fn ds4_gpu_routed_moe_one_tensor( + out: *mut GpuTensor, + gate: *mut GpuTensor, + up: *mut GpuTensor, + mid: *mut GpuTensor, + experts_out: *mut GpuTensor, + map: *const c_void, + size: u64, + gate_weight: u64, + up_weight: u64, + down_weight: u64, + gate_type: u32, + down_type: u32, + gate_expert_bytes: u64, + gate_row_bytes: u64, + down_expert_bytes: u64, + down_row_bytes: u64, + input: u32, + middle: u32, + output: u32, + selected: *const GpuTensor, + weights: *const GpuTensor, + total_experts: u32, + used_experts: u32, + clamp: f32, + x: *const GpuTensor, + add: *const GpuTensor, + layer: u32, + force_resident: bool, + ) -> i32; + fn ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( + gate: *mut GpuTensor, + up: *mut GpuTensor, + mid: *mut GpuTensor, + map: *const c_void, + size: u64, + gate_weight: u64, + up_weight: u64, + input: u64, + output: u64, + x: *const GpuTensor, + clamp: f32, + ) -> i32; + fn ds4_gpu_shared_down_hc_expand_q8_0_tensor( + out_hc: *mut GpuTensor, + shared_out: *mut GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + input: u64, + output: u64, + middle: *const GpuTensor, + routed: *const GpuTensor, + residual: *const GpuTensor, + split: *const GpuTensor, + embd: u32, + hc: u32, + ) -> i32; + fn ds4_gpu_output_hc_weights_tensor( + out: *mut GpuTensor, + pre: *const GpuTensor, + map: *const c_void, + size: u64, + scale: u64, + base: u64, + hc: u32, + eps: f32, + ) -> i32; + fn ds4_gpu_hc_weighted_sum_norm_tensor( + out: *mut GpuTensor, + norm: *mut GpuTensor, + residual: *const GpuTensor, + weights: *const GpuTensor, + map: *const c_void, + size: u64, + norm_weight: u64, + embd: u32, + hc: u32, + eps: f32, + ) -> i32; + fn ds4_gpu_hc_weighted_sum_tensor( + out: *mut GpuTensor, + residual: *const GpuTensor, + weights: *const GpuTensor, + embd: u32, + hc: u32, + ) -> i32; +} + +struct Context; + +impl Context { + fn open(model: &Model, quality: bool) -> Result { + check(unsafe { ds4_gpu_init() }, "Metal initialization")?; + if let Err(error) = check( + unsafe { ds4_gpu_set_model_map(model.main.map_ptr().cast(), model.main.len()) }, + "model mapping", + ) { + unsafe { ds4_gpu_cleanup() }; + return Err(error); + } + unsafe { ds4_gpu_set_quality(quality) }; + Ok(Self) + } +} + +impl Drop for Context { + fn drop(&mut self) { + unsafe { ds4_gpu_cleanup() }; + } +} + +struct Commands(bool); + +impl Commands { + fn begin() -> Result { + check( + unsafe { ds4_gpu_begin_commands() }, + "beginning Metal commands", + )?; + Ok(Self(true)) + } + + fn finish(mut self) -> Result<(), String> { + self.0 = false; + check( + unsafe { ds4_gpu_end_commands() }, + "executing Metal commands", + ) + } +} + +impl Drop for Commands { + fn drop(&mut self) { + if self.0 { + unsafe { ds4_gpu_end_commands() }; + } + } +} + +struct Buffer(NonNull); + +impl Buffer { + fn floats(count: u64) -> Result { + Self::bytes(count * 4) + } + + fn bytes(bytes: u64) -> Result { + NonNull::new(unsafe { ds4_gpu_tensor_alloc(bytes) }) + .map(Self) + .ok_or_else(|| format!("Metal could not allocate {bytes} bytes")) + } + + fn read_f32(&self, values: &mut [f32]) -> Result<(), String> { + check( + unsafe { + ds4_gpu_tensor_read( + self.raw(), + 0, + values.as_mut_ptr().cast(), + std::mem::size_of_val(values) as u64, + ) + }, + "reading Metal output", + ) + } + + fn fill(&self, value: f32, count: u64) -> Result<(), String> { + call( + unsafe { ds4_gpu_tensor_fill_f32(self.raw(), value, count) }, + "initializing a Metal buffer", + ) + } + + fn raw(&self) -> *mut GpuTensor { + self.0.as_ptr() + } +} + +impl Drop for Buffer { + fn drop(&mut self) { + unsafe { ds4_gpu_tensor_free(self.raw()) }; + } +} + +#[derive(Clone, Copy)] +struct Weight { + offset: u64, + kind: u32, + bytes: u64, + dims: [u64; 3], +} + +impl Weight { + fn bind(model: &Gguf, name: &str) -> Result { + let tensor = model.tensor(name)?; + Ok(Self::from_tensor(tensor)) + } + + fn optional(model: &Gguf, name: &str) -> Option { + model.tensors.get(name).map(Self::from_tensor) + } + + fn from_tensor(tensor: &GgufTensor) -> Self { + let mut dims = [1; 3]; + for (to, from) in dims.iter_mut().zip(&tensor.dims) { + *to = *from; + } + Self { + offset: tensor.offset, + kind: tensor.kind, + bytes: tensor.bytes, + dims, + } + } +} + +struct Layer { + hc_attn_fn: Weight, + hc_attn_scale: Weight, + hc_attn_base: Weight, + attn_norm: Weight, + attn_q_a: Weight, + attn_q_a_norm: Weight, + attn_q_b: Weight, + attn_kv: Weight, + attn_kv_norm: Weight, + attn_sinks: Weight, + attn_output_a: Weight, + attn_output_b: Weight, + attn_compressor: Option, + hc_ffn_fn: Weight, + hc_ffn_scale: Weight, + hc_ffn_base: Weight, + ffn_norm: Weight, + router: Weight, + router_bias: Option, + router_hash: Option, + expert_gate: Weight, + expert_up: Weight, + expert_down: Weight, + shared_gate: Weight, + shared_up: Weight, + shared_down: Weight, +} + +#[derive(Clone, Copy)] +struct CompressorWeights { + ape: Weight, + kv: Weight, + gate: Weight, + norm: Weight, +} + +impl Layer { + fn bind(model: &Gguf, index: u32) -> Result { + let required = |suffix: &str| Weight::bind(model, &format!("blk.{index}.{suffix}")); + let optional = |suffix: &str| Weight::optional(model, &format!("blk.{index}.{suffix}")); + let attn_compressor = (index >= 2) + .then(|| { + Ok::<_, String>(CompressorWeights { + ape: required("attn_compressor_ape.weight")?, + kv: required("attn_compressor_kv.weight")?, + gate: required("attn_compressor_gate.weight")?, + norm: required("attn_compressor_norm.weight")?, + }) + }) + .transpose()?; + let layer = Self { + hc_attn_fn: required("hc_attn_fn.weight")?, + hc_attn_scale: required("hc_attn_scale.weight")?, + hc_attn_base: required("hc_attn_base.weight")?, + attn_norm: required("attn_norm.weight")?, + attn_q_a: required("attn_q_a.weight")?, + attn_q_a_norm: required("attn_q_a_norm.weight")?, + attn_q_b: required("attn_q_b.weight")?, + attn_kv: required("attn_kv.weight")?, + attn_kv_norm: required("attn_kv_a_norm.weight")?, + attn_sinks: required("attn_sinks.weight")?, + attn_output_a: required("attn_output_a.weight")?, + attn_output_b: required("attn_output_b.weight")?, + attn_compressor, + hc_ffn_fn: required("hc_ffn_fn.weight")?, + hc_ffn_scale: required("hc_ffn_scale.weight")?, + hc_ffn_base: required("hc_ffn_base.weight")?, + ffn_norm: required("ffn_norm.weight")?, + router: required("ffn_gate_inp.weight")?, + router_bias: optional("exp_probs_b.bias"), + router_hash: optional("ffn_gate_tid2eid.weight"), + expert_gate: required("ffn_gate_exps.weight")?, + expert_up: required("ffn_up_exps.weight")?, + expert_down: required("ffn_down_exps.weight")?, + shared_gate: required("ffn_gate_shexp.weight")?, + shared_up: required("ffn_up_shexp.weight")?, + shared_down: required("ffn_down_shexp.weight")?, + }; + for (name, weight) in [ + ("attention q", layer.attn_q_b), + ("attention output A", layer.attn_output_a), + ("attention output B", layer.attn_output_b), + ("shared gate", layer.shared_gate), + ("shared up", layer.shared_up), + ("shared down", layer.shared_down), + ] { + if weight.kind != Q8_0 { + return Err(format!("Flash Metal path requires Q8_0 {name} weights")); + } + } + Ok(layer) + } +} + +struct Weights { + token_embedding: Weight, + output_hc_base: Weight, + output_hc_fn: Weight, + output_hc_scale: Weight, + output_norm: Weight, + output: Weight, + layers: Vec, +} + +impl Weights { + fn bind(model: &Model) -> Result { + if model.shape.family != ModelFamily::DeepSeek { + return Err("the Rust Metal executor currently supports DeepSeek Flash only".into()); + } + let main = &model.main; + let output = Weight::bind(main, "output.weight")?; + if output.kind != Q8_0 { + return Err("Flash Metal path requires a Q8_0 output weight".into()); + } + let layers = (0..model.shape.layers) + .map(|index| Layer::bind(main, index)) + .collect::>()?; + Ok(Self { + token_embedding: Weight::bind(main, "token_embd.weight")?, + output_hc_base: Weight::bind(main, "output_hc_base.weight")?, + output_hc_fn: Weight::bind(main, "output_hc_fn.weight")?, + output_hc_scale: Weight::bind(main, "output_hc_scale.weight")?, + output_norm: Weight::bind(main, "output_norm.weight")?, + output, + layers, + }) + } +} + +struct Scratch { + current_hc: Buffer, + next_hc: Buffer, + flat_hc: Buffer, + hc_mix: Buffer, + hc_split: Buffer, + current: Buffer, + norm: Buffer, + q_rank: Buffer, + q_rank_norm: Buffer, + q: Buffer, + kv_raw: Buffer, + kv: Buffer, + compressed_kv: Buffer, + compressed_score: Buffer, + compressed_stage: Buffer, + heads: Buffer, + attention_low: Buffer, + attention_out: Buffer, + router_logits: Buffer, + router_probs: Buffer, + router_selected: Buffer, + router_weights: Buffer, + routed_gate: Buffer, + routed_up: Buffer, + routed_mid: Buffer, + routed_experts: Buffer, + routed_out: Buffer, + shared_gate: Buffer, + shared_up: Buffer, + shared_mid: Buffer, + shared_out: Buffer, + output_pre: Buffer, + output_weights: Buffer, + output_embedding: Buffer, + output_norm: Buffer, + logits: Buffer, +} + +impl Scratch { + fn allocate(model: &Model) -> Result { + let shape = model.shape; + let hc_dim = shape.hc * shape.embd; + let mix_hc = 2 * shape.hc + shape.hc * shape.hc; + let q_dim = shape.heads * shape.head_dim; + let low_dim = shape.out_groups * shape.lora_o; + let routed = shape.experts_used * shape.ff_expert; + Ok(Self { + current_hc: Buffer::floats(hc_dim)?, + next_hc: Buffer::floats(hc_dim)?, + flat_hc: Buffer::floats(hc_dim)?, + hc_mix: Buffer::floats(mix_hc)?, + hc_split: Buffer::floats(mix_hc)?, + current: Buffer::floats(shape.embd)?, + norm: Buffer::floats(shape.embd)?, + q_rank: Buffer::floats(shape.lora_q)?, + q_rank_norm: Buffer::floats(shape.lora_q)?, + q: Buffer::floats(q_dim)?, + kv_raw: Buffer::floats(shape.head_dim)?, + kv: Buffer::floats(shape.head_dim)?, + compressed_kv: Buffer::floats(2 * shape.head_dim)?, + compressed_score: Buffer::floats(2 * shape.head_dim)?, + compressed_stage: Buffer::floats(shape.head_dim)?, + heads: Buffer::floats(q_dim)?, + attention_low: Buffer::floats(low_dim)?, + attention_out: Buffer::floats(shape.embd)?, + router_logits: Buffer::floats(shape.experts)?, + router_probs: Buffer::floats(shape.experts)?, + router_selected: Buffer::bytes(shape.experts_used * 4)?, + router_weights: Buffer::floats(shape.experts_used)?, + routed_gate: Buffer::floats(routed)?, + routed_up: Buffer::floats(routed)?, + routed_mid: Buffer::floats(routed)?, + routed_experts: Buffer::floats(shape.experts_used * shape.embd)?, + routed_out: Buffer::floats(shape.embd)?, + shared_gate: Buffer::floats(shape.ff_expert)?, + shared_up: Buffer::floats(shape.ff_expert)?, + shared_mid: Buffer::floats(shape.ff_expert)?, + shared_out: Buffer::floats(shape.embd)?, + output_pre: Buffer::floats(shape.hc)?, + output_weights: Buffer::floats(shape.hc)?, + output_embedding: Buffer::floats(shape.embd)?, + output_norm: Buffer::floats(shape.embd)?, + logits: Buffer::floats(shape.vocab)?, + }) + } +} + +/// Rust-owned DeepSeek Flash session. The ObjC boundary only schedules the +/// unchanged Metal kernels; token/session/graph state lives here. +pub(super) struct Session { + scratch: Scratch, + layers: Vec, + context: u32, + raw_cap: u32, + position: u32, +} + +struct LayerState { + raw_cache: Buffer, + compression: Option, +} + +struct CompressionState { + ratio: u32, + cache: Buffer, + state_kv: Buffer, + state_score: Buffer, + rows: u32, +} + +impl LayerState { + fn allocate(model: &Model, index: u32, context: u32, raw_cap: u32) -> Result { + let shape = model.shape; + let ratio = compression_ratio(index); + let compression = if ratio == 0 { + None + } else { + let coefficient = if ratio == 4 { 2 } else { 1 }; + let width = coefficient * shape.head_dim; + let state_rows = coefficient * ratio as u64; + let state_kv = Buffer::floats(width * state_rows)?; + let state_score = Buffer::floats(width * state_rows)?; + state_kv.fill(0.0, width * state_rows)?; + state_score.fill(f32::NEG_INFINITY, width * state_rows)?; + Some(CompressionState { + ratio, + cache: Buffer::bytes((context / ratio + 2) as u64 * shape.head_dim * 2)?, + state_kv, + state_score, + rows: 0, + }) + }; + Ok(Self { + raw_cache: Buffer::floats(raw_cap as u64 * shape.head_dim)?, + compression, + }) + } +} + +impl Session { + fn new(model: &Model, context: u32) -> Result { + if context == 0 { + return Err("context must contain at least one token".into()); + } + // ponytail: the ratio-4 indexer is needed above 2K; cap here until that + // sparse selection path is ported. + let context = context.min(2_048); + let raw_cap = model.shape.sliding_window as u32; + let layers = (0..model.shape.layers) + .map(|index| LayerState::allocate(model, index, context, raw_cap)) + .collect::>()?; + Ok(Self { + scratch: Scratch::allocate(model)?, + layers, + context, + raw_cap, + position: 0, + }) + } +} + +pub(super) struct Executor { + weights: Weights, + session: Session, + logits: Vec, + _context: Context, + model: Model, +} + +impl Executor { + pub(super) fn open(model: Model, context: u32, quality: bool) -> Result { + let weights = Weights::bind(&model)?; + let context_handle = Context::open(&model, quality)?; + let session = Session::new(&model, context)?; + Ok(Self { + weights, + session, + logits: vec![0.0; model.shape.vocab as usize], + _context: context_handle, + model, + }) + } + + pub(super) fn eval(&mut self, token: i32) -> Result<(), String> { + if token < 0 || token as u64 >= self.model.shape.vocab { + return Err(format!("token {token} is outside the vocabulary")); + } + if self.session.position >= self.session.context { + return Err(format!( + "the Rust Metal executor currently supports {} tokens per session", + self.session.context + )); + } + let commands = Commands::begin()?; + self.encode_token(token as u32)?; + commands.finish()?; + self.session.scratch.logits.read_f32(&mut self.logits)?; + self.session.position += 1; + Ok(()) + } + + pub(super) fn logits(&self) -> &[f32] { + &self.logits + } + + pub(super) fn model(&self) -> &Model { + &self.model + } + + pub(super) fn context(&self) -> u32 { + self.session.context + } + + pub(super) fn position(&self) -> u32 { + self.session.position + } + + pub(super) fn reset(&mut self) -> Result<(), String> { + self.session = Session::new(&self.model, self.session.context)?; + Ok(()) + } + + fn encode_token(&mut self, token: u32) -> Result<(), String> { + let shape = self.model.shape; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + let scratch = &mut self.session.scratch; + call( + unsafe { + ds4_gpu_embed_token_hc_tensor( + scratch.current_hc.raw(), + map, + size, + self.weights.token_embedding.offset, + shape.vocab as u32, + token, + shape.embd as u32, + shape.hc as u32, + ) + }, + "token embedding", + )?; + for (index, layer) in self.weights.layers.iter().enumerate() { + let state = &mut self.session.layers[index]; + encode_layer( + scratch, + state, + layer, + shape, + map, + size, + index as u32, + self.session.position, + token, + self.session.raw_cap, + )?; + } + encode_output(scratch, &self.weights, shape, map, size) + } +} + +// Mirrors one fixed DS4 tape invocation; grouping these scalar dimensions would +// only hide the Metal kernel contract behind another one-use type. +#[allow(clippy::too_many_arguments)] +fn encode_layer( + s: &Scratch, + state: &mut LayerState, + w: &Layer, + shape: super::Shape, + map: *const c_void, + size: u64, + layer: u32, + pos: u32, + token: u32, + raw_cap: u32, +) -> Result<(), String> { + let hc_dim = shape.hc * shape.embd; + let mix_hc = 2 * shape.hc + shape.hc * shape.hc; + let compressed = layer >= 2; + let freq_base = if compressed { + shape.compress_rope_base + } else { + shape.rope_base + }; + let freq_scale = if compressed { + 1.0 / shape.rope_scale + } else { + 1.0 + }; + let ext = if compressed && shape.rope_scale > 1.0 { + 1.0 + } else { + 0.0 + }; + let attn_factor = if ext != 0.0 { + 1.0 / (1.0 + 0.1 * (1.0 / freq_scale).ln()) + } else { + 1.0 + }; + let original = if compressed { + shape.original_context as u32 + } else { + 0 + }; + let raw_cache = state.raw_cache.raw(); + call( + unsafe { + ds4_gpu_rms_norm_plain_tensor( + s.flat_hc.raw(), + s.current_hc.raw(), + hc_dim as u32, + shape.rms_epsilon, + ) + }, + "attention HC norm", + )?; + f16( + &s.hc_mix, + w.hc_attn_fn, + hc_dim, + mix_hc, + &s.flat_hc, + map, + size, + )?; + call( + unsafe { + ds4_gpu_hc_split_weighted_sum_norm_tensor( + s.current.raw(), + s.norm.raw(), + s.hc_split.raw(), + s.hc_mix.raw(), + s.current_hc.raw(), + map, + size, + w.hc_attn_scale.offset, + w.hc_attn_base.offset, + w.attn_norm.offset, + shape.embd as u32, + shape.hc as u32, + shape.hc_sinkhorn as u32, + shape.hc_epsilon, + shape.rms_epsilon, + ) + }, + "attention HC mix", + )?; + call( + unsafe { + ds4_gpu_matmul_q8_0_pair_tensor( + s.q_rank.raw(), + s.kv_raw.raw(), + map, + size, + w.attn_q_a.offset, + w.attn_kv.offset, + shape.embd, + shape.lora_q, + shape.head_dim, + s.norm.raw(), + 1, + ) + }, + "Q/KV projection", + )?; + call( + unsafe { + ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( + s.q_rank_norm.raw(), + s.q_rank.raw(), + map, + size, + w.attn_q_a_norm.offset, + shape.lora_q as u32, + s.kv.raw(), + s.kv_raw.raw(), + w.attn_kv_norm.offset, + shape.head_dim as u32, + 1, + shape.rms_epsilon, + ) + }, + "Q/KV norm", + )?; + q8( + &s.q, + w.attn_q_b, + shape.lora_q, + shape.heads * shape.head_dim, + &s.q_rank_norm, + map, + size, + )?; + let fused_q_rope = unsafe { + ds4_gpu_head_rms_norm_rope_tail_tensor( + s.q.raw(), + 1, + shape.heads as u32, + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + false, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + ) + } != 0; + if !fused_q_rope { + call( + unsafe { + ds4_gpu_head_rms_norm_tensor( + s.q.raw(), + 1, + shape.heads as u32, + shape.head_dim as u32, + shape.rms_epsilon, + ) + }, + "Q norm", + )?; + call( + unsafe { + ds4_gpu_rope_tail_tensor( + s.q.raw(), + 1, + shape.heads as u32, + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + false, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + ) + }, + "Q RoPE", + )?; + } + call( + unsafe { + ds4_gpu_rope_tail_tensor( + s.kv.raw(), + 1, + shape.head_kv as u32, + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + false, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + ) + }, + "KV RoPE", + )?; + let raw_row = pos % raw_cap; + call( + unsafe { + ds4_gpu_kv_fp8_store_raw_tensor( + s.kv.raw(), + raw_cache, + raw_cap, + raw_row, + shape.head_dim as u32, + shape.rot as u32, + ) + }, + "KV cache write", + )?; + let n_raw = (pos + 1).min(raw_cap); + let raw_start = if pos + 1 > raw_cap { + (raw_row + 1) % raw_cap + } else { + 0 + }; + if let (Some(weights), Some(compression)) = + (w.attn_compressor.as_ref(), state.compression.as_mut()) + { + update_compression( + s, + compression, + *weights, + shape, + map, + size, + pos, + original, + freq_base, + freq_scale, + ext, + attn_factor, + )?; + } + let (comp_cache, n_comp) = state + .compression + .as_ref() + .map_or((std::ptr::null(), 0), |compression| { + (compression.cache.raw().cast_const(), compression.rows) + }); + call( + unsafe { + ds4_gpu_attention_decode_heads_tensor( + s.heads.raw(), + map, + size, + w.attn_sinks.offset, + s.q.raw(), + raw_cache, + n_raw, + raw_cap, + raw_start, + comp_cache, + 1, + n_comp, + std::ptr::null(), + 0, + shape.heads as u32, + shape.head_dim as u32, + ) + }, + "attention", + )?; + call( + unsafe { + ds4_gpu_rope_tail_tensor( + s.heads.raw(), + 1, + shape.heads as u32, + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + true, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + ) + }, + "inverse attention RoPE", + )?; + let group_dim = shape.head_dim * (shape.heads / shape.out_groups); + call( + unsafe { + ds4_gpu_attention_output_low_q8_tensor( + s.attention_low.raw(), + map, + size, + w.attn_output_a.offset, + group_dim, + shape.lora_o, + shape.out_groups as u32, + s.heads.raw(), + ) + }, + "attention low projection", + )?; + call( + unsafe { + ds4_gpu_matmul_q8_0_hc_expand_tensor( + s.next_hc.raw(), + s.attention_out.raw(), + map, + size, + w.attn_output_b.offset, + shape.out_groups * shape.lora_o, + shape.embd, + s.attention_low.raw(), + s.current_hc.raw(), + s.hc_split.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "attention output", + )?; + call( + unsafe { + ds4_gpu_rms_norm_plain_tensor( + s.flat_hc.raw(), + s.next_hc.raw(), + hc_dim as u32, + shape.rms_epsilon, + ) + }, + "FFN HC norm", + )?; + f16( + &s.hc_mix, + w.hc_ffn_fn, + hc_dim, + mix_hc, + &s.flat_hc, + map, + size, + )?; + call( + unsafe { + ds4_gpu_hc_split_weighted_sum_norm_tensor( + s.current.raw(), + s.norm.raw(), + s.hc_split.raw(), + s.hc_mix.raw(), + s.next_hc.raw(), + map, + size, + w.hc_ffn_scale.offset, + w.hc_ffn_base.offset, + w.ffn_norm.offset, + shape.embd as u32, + shape.hc as u32, + shape.hc_sinkhorn as u32, + shape.hc_epsilon, + shape.rms_epsilon, + ) + }, + "FFN HC mix", + )?; + f16( + &s.router_logits, + w.router, + shape.embd, + shape.experts, + &s.norm, + map, + size, + )?; + call( + unsafe { + ds4_gpu_router_select_tensor( + s.router_selected.raw(), + s.router_weights.raw(), + s.router_probs.raw(), + map, + size, + w.router_bias.map_or(0, |v| v.offset), + w.router_hash.map_or(0, |v| v.offset), + w.router_hash.map_or(0, |v| v.dims[1] as u32), + token, + shape.experts as u32, + shape.experts_used as u32, + shape.expert_weight_scale, + 0, + 0, + w.router_bias.is_some(), + w.router_hash.is_some(), + s.router_logits.raw(), + ) + }, + "expert routing", + )?; + let gate_row = w.expert_gate.bytes / (w.expert_gate.dims[1] * w.expert_gate.dims[2]); + let down_row = w.expert_down.bytes / (w.expert_down.dims[1] * w.expert_down.dims[2]); + call( + unsafe { + ds4_gpu_routed_moe_one_tensor( + s.routed_out.raw(), + s.routed_gate.raw(), + s.routed_up.raw(), + s.routed_mid.raw(), + s.routed_experts.raw(), + map, + size, + w.expert_gate.offset, + w.expert_up.offset, + w.expert_down.offset, + w.expert_gate.kind, + w.expert_down.kind, + w.expert_gate.dims[1] * gate_row, + gate_row, + w.expert_down.dims[1] * down_row, + down_row, + w.expert_gate.dims[0] as u32, + w.expert_down.dims[0] as u32, + w.expert_down.dims[1] as u32, + s.router_selected.raw(), + s.router_weights.raw(), + shape.experts as u32, + shape.experts_used as u32, + shape.swiglu_clamp, + s.norm.raw(), + std::ptr::null(), + layer, + false, + ) + }, + "routed experts", + )?; + call( + unsafe { + ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( + s.shared_gate.raw(), + s.shared_up.raw(), + s.shared_mid.raw(), + map, + size, + w.shared_gate.offset, + w.shared_up.offset, + shape.embd, + shape.ff_expert, + s.norm.raw(), + shape.swiglu_clamp, + ) + }, + "shared expert gate/up", + )?; + call( + unsafe { + ds4_gpu_shared_down_hc_expand_q8_0_tensor( + s.current_hc.raw(), + s.shared_out.raw(), + map, + size, + w.shared_down.offset, + shape.ff_expert, + shape.embd, + s.shared_mid.raw(), + s.routed_out.raw(), + s.next_hc.raw(), + s.hc_split.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "shared expert output", + ) +} + +#[allow(clippy::too_many_arguments)] +fn update_compression( + s: &Scratch, + state: &mut CompressionState, + weights: CompressorWeights, + shape: super::Shape, + map: *const c_void, + size: u64, + pos: u32, + original: u32, + freq_base: f32, + freq_scale: f32, + ext: f32, + attn_factor: f32, +) -> Result<(), String> { + let width = if state.ratio == 4 { 2 } else { 1 } * shape.head_dim as u32; + let fused = unsafe { + ds4_gpu_matmul_f16_pair_compressor_store_tensor( + s.compressed_kv.raw(), + s.compressed_score.raw(), + state.state_kv.raw(), + state.state_score.raw(), + map, + size, + weights.kv.offset, + weights.gate.offset, + weights.ape.offset, + weights.ape.kind, + shape.embd, + width, + s.norm.raw(), + state.ratio, + pos, + ) + }; + if fused < 0 { + return Err("Metal failed while storing compressor projections".into()); + } + if fused == 0 { + call( + unsafe { + ds4_gpu_matmul_f16_pair_tensor( + s.compressed_kv.raw(), + s.compressed_score.raw(), + map, + size, + weights.kv.offset, + weights.gate.offset, + shape.embd, + width as u64, + s.norm.raw(), + 1, + ) + }, + "compressor projection", + )?; + } + let emit = (pos + 1).is_multiple_of(state.ratio); + call( + unsafe { + ds4_gpu_compressor_update_tensor( + s.compressed_kv.raw(), + s.compressed_score.raw(), + state.state_kv.raw(), + state.state_score.raw(), + s.compressed_stage.raw(), + map, + size, + weights.ape.offset, + weights.ape.kind, + weights.norm.offset, + weights.norm.kind, + shape.head_dim as u32, + state.ratio, + pos, + 0, + shape.rot as u32, + original, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + fused > 0, + ) + }, + "compressor update", + )?; + if emit { + call( + unsafe { + ds4_gpu_dsv4_fp8_kv_quantize_tensor( + s.compressed_stage.raw(), + 1, + shape.head_dim as u32, + shape.rot as u32, + ) + }, + "compressed KV quantization", + )?; + call( + unsafe { + ds4_gpu_tensor_copy_f32_to_f16( + state.cache.raw(), + state.rows as u64 * shape.head_dim * 2, + s.compressed_stage.raw(), + 0, + shape.head_dim, + ) + }, + "compressed KV cache write", + )?; + state.rows += 1; + } + Ok(()) +} + +fn compression_ratio(layer: u32) -> u32 { + if layer < 2 { + 0 + } else if layer.is_multiple_of(2) { + 4 + } else { + 128 + } +} + +fn encode_output( + s: &Scratch, + w: &Weights, + shape: super::Shape, + map: *const c_void, + size: u64, +) -> Result<(), String> { + let hc_dim = shape.hc * shape.embd; + call( + unsafe { + ds4_gpu_rms_norm_plain_tensor( + s.flat_hc.raw(), + s.current_hc.raw(), + hc_dim as u32, + shape.rms_epsilon, + ) + }, + "output HC norm", + )?; + f16( + &s.output_pre, + w.output_hc_fn, + hc_dim, + shape.hc, + &s.flat_hc, + map, + size, + )?; + call( + unsafe { + ds4_gpu_output_hc_weights_tensor( + s.output_weights.raw(), + s.output_pre.raw(), + map, + size, + w.output_hc_scale.offset, + w.output_hc_base.offset, + shape.hc as u32, + shape.hc_epsilon, + ) + }, + "output HC weights", + )?; + let fused_sum_norm = unsafe { + ds4_gpu_hc_weighted_sum_norm_tensor( + s.output_embedding.raw(), + s.output_norm.raw(), + s.current_hc.raw(), + s.output_weights.raw(), + map, + size, + w.output_norm.offset, + shape.embd as u32, + shape.hc as u32, + shape.rms_epsilon, + ) + } != 0; + if !fused_sum_norm { + call( + unsafe { + ds4_gpu_hc_weighted_sum_tensor( + s.output_embedding.raw(), + s.current_hc.raw(), + s.output_weights.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "output HC collapse", + )?; + call( + unsafe { + ds4_gpu_rms_norm_weight_tensor( + s.output_norm.raw(), + s.output_embedding.raw(), + map, + size, + w.output_norm.offset, + shape.embd as u32, + shape.rms_epsilon, + ) + }, + "output norm", + )?; + } + q8( + &s.logits, + w.output, + shape.embd, + shape.vocab, + &s.output_norm, + map, + size, + ) +} + +fn f16( + out: &Buffer, + weight: Weight, + input: u64, + output: u64, + x: &Buffer, + map: *const c_void, + size: u64, +) -> Result<(), String> { + call( + unsafe { + ds4_gpu_matmul_f16_tensor( + out.raw(), + map, + size, + weight.offset, + input, + output, + x.raw(), + 1, + ) + }, + "F16 projection", + ) +} + +fn q8( + out: &Buffer, + weight: Weight, + input: u64, + output: u64, + x: &Buffer, + map: *const c_void, + size: u64, +) -> Result<(), String> { + call( + unsafe { + ds4_gpu_matmul_q8_0_tensor( + out.raw(), + map, + size, + weight.offset, + input, + output, + x.raw(), + 1, + ) + }, + "Q8 projection", + ) +} + +fn call(result: i32, operation: &str) -> Result<(), String> { + if result == 0 { + Err(format!("Metal failed while {operation}")) + } else { + Ok(()) + } +} + +fn check(result: i32, operation: &str) -> Result<(), String> { + call(result, operation) +} diff --git a/src/engine/tokenizer.rs b/src/engine/tokenizer.rs index 8b5e660..b75895c 100644 --- a/src/engine/tokenizer.rs +++ b/src/engine/tokenizer.rs @@ -1,5 +1,5 @@ -use super::ModelFamily; use super::gguf::Gguf; +use super::{ChatTurn, ModelFamily}; use crate::settings::ReasoningMode; use std::collections::HashMap; @@ -131,6 +131,23 @@ impl Tokenizer { system_prompt: &str, prompt: &str, reasoning: ReasoningMode, + ) -> Vec { + self.encode_conversation( + system_prompt, + &[ChatTurn { + user: true, + reasoning: false, + content: prompt.to_owned(), + }], + reasoning, + ) + } + + pub(super) fn encode_conversation( + &self, + system_prompt: &str, + messages: &[ChatTurn], + reasoning: ReasoningMode, ) -> Vec { let mut output = vec![self.bos]; if self.family == ModelFamily::Glm && self.sop >= 0 { @@ -156,8 +173,23 @@ impl Tokenizer { } output.extend(self.tokenize(system_prompt)); } - output.push(self.user); - output.extend(self.tokenize(prompt)); + for message in messages { + output.push(if message.user { + self.user + } else { + self.assistant + }); + if !message.user { + if message.reasoning { + output.push(self.think_start); + } else if self.family == ModelFamily::Glm { + output.extend([self.think_start, self.think_end]); + } else { + output.push(self.think_end); + } + } + output.extend(self.tokenize(&message.content)); + } output.push(self.assistant); if reasoning != ReasoningMode::Direct { output.push(self.think_start); diff --git a/src/main.rs b/src/main.rs index 4dfe49e..0190446 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,6 +11,10 @@ use app::{App, Message, app_icon, app_theme}; use iced::{Size, window}; fn main() -> iced::Result { + #[cfg(target_os = "macos")] + if let Err(error) = engine::configure_metal_sources() { + eprintln!("DS4Server: {error}"); + } iced::daemon(App::title, App::update, App::view) .subscription(App::subscription) .theme(|_, _| app_theme()) diff --git a/src/settings.rs b/src/settings.rs index 9ba599c..48a49f6 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -668,7 +668,7 @@ mod tests { } #[test] - fn every_ds4_cli_option_is_mapped_or_explicitly_not_a_preference() { + fn captured_ds4_cli_options_are_uniquely_classified() { const PREFERENCES: &[&str] = &[ "--ctx", "--dir-steering-attn", @@ -753,32 +753,16 @@ mod tests { "--tensor-parallel-token-prefill", "--transport", ]; - let source = [ - include_str!("../../ds4/ds4_cli.c"), - include_str!("../../ds4/ds4_tp.c"), - include_str!("../../ds4/ds4_distributed.c"), - ] - .join("\n"); - let actual = option_branches(&source); let classified = PREFERENCES .iter() .chain(NON_PREFERENCES) .copied() .collect::>(); - assert_eq!(actual, classified); + assert_eq!(classified.len(), PREFERENCES.len() + NON_PREFERENCES.len()); assert!( PREFERENCES .iter() .all(|option| !NON_PREFERENCES.contains(option)) ); } - - fn option_branches(source: &str) -> BTreeSet<&str> { - source - .split("!strcmp(arg, \"") - .skip(1) - .filter_map(|tail| tail.split('"').next()) - .filter(|option| option.starts_with("--")) - .collect() - } }