hugo/DS4Server
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases Wiki Activity

feat: first cut at actual token generation and model loading

Browse Source
This commit is contained in:
Georg Bauer
2026-07-24 17:44:41 +02:00
parent 616b550849
commit ff984bb96a
38 changed files with 66885 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

@@ -1133,6 +1133,7 @@ dependencies = [
name = "ds4-server"
version = "0.1.0"
dependencies = [
"cc",
"diesel",
"diesel_migrations",
"iced",

5
Cargo.toml
View File

@@ -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"

56
PLAN.md
View File

@@ -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.

12
README.md
View File

@@ -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

18
build.rs Normal file
View File

@@ -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");
}

275
metal/argsort.metal Normal file
View File

@@ -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<ds4_sort_order order>
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<DS4_SORT_ORDER_DESC>;
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<ds4_sort_order order>
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<DS4_SORT_ORDER_DESC>;

217
metal/bin.metal Normal file
View File

@@ -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 <typename T0, typename T1, typename T>
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<float, float, float>) 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<float, float, float>;
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];
}

62
metal/concat.metal Normal file
View File

@@ -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;
}
}

247
metal/cpy.metal Normal file
View File

@@ -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<typename T0, typename T1>
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<float, float>) 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<float, float>;
template [[host_name("kernel_cpy_f32_f16")]] kernel kernel_cpy_t kernel_cpy_t_t<float, half>;
template [[host_name("kernel_cpy_f16_f32")]] kernel kernel_cpy_t kernel_cpy_t_t<half, float>;
template [[host_name("kernel_cpy_f16_f16")]] kernel kernel_cpy_t kernel_cpy_t_t<half, half>;
// 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;
}
}

2389
metal/dense.metal Normal file
View File

File diff suppressed because it is too large Load Diff

1017
metal/dsv4_hc.metal Normal file
View File

File diff suppressed because it is too large Load Diff

369
metal/dsv4_kv.metal Normal file
View File

@@ -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;
}

6157
metal/dsv4_misc.metal Normal file
View File

File diff suppressed because it is too large Load Diff

385
metal/dsv4_rope.metal Normal file
View File

@@ -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<int>(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;
}
}

1441
metal/flash_attn.metal Normal file
View File

File diff suppressed because it is too large Load Diff

186
metal/get_rows.metal Normal file
View File

@@ -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<typename T0, typename T>
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<float, float>) 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<float, float>;
template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f<half, float>;
template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
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;
}
}
}

63
metal/glu.metal Normal file
View File

@@ -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;
}

7858
metal/moe.metal Normal file
View File

File diff suppressed because it is too large Load Diff

243
metal/norm.metal Normal file
View File

@@ -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 <typename T, short F>
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<float4, 1>) 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<float4, 1>;
template [[host_name("kernel_rms_norm_mul_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 2>;
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];
}
}

52
metal/repeat.metal Normal file
View File

@@ -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<typename T>
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<float>) 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<float>;

55
metal/set_rows.metal Normal file
View File

@@ -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<typename T, typename TI>
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<float, int64_t>) 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<float, int32_t>;

241
metal/softmax.metal Normal file
View File

@@ -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<typename T>
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<typename T>
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<float>) kernel_soft_max_t;
typedef decltype(kernel_soft_max_4<float4>) 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<float>;
template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<float4>;

102
metal/sum_rows.metal Normal file
View File

@@ -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 <typename T0, typename T>
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<float4, T0>::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<float, float>) 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<float, float>;

312
metal/unary.metal Normal file
View File

@@ -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<typename T>
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<typename T> T elu_approx(T x);
template<> inline float elu_approx<float>(float x) {
return (x > 0.f) ? x : (exp(x) - 1);
}
template<> inline float4 elu_approx<float4>(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 <typename T0, typename T, typename TC>
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<float, float, float>) 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<float, float, float>;
template [[host_name("kernel_unary_f32_f32_4")]] kernel kernel_unary_t kernel_unary_impl<float4, float4, float4>;
template [[host_name("kernel_unary_f16_f16")]] kernel kernel_unary_t kernel_unary_impl<half, half, float>;

22
native/metal/LICENSE Normal file
View File

@@ -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.

15
native/metal/README.md Normal file
View File

@@ -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.

477
native/metal/ds4.h Normal file
View File

@@ -0,0 +1,477 @@
#ifndef DS4_H
#define DS4_H
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#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

2603
native/metal/ds4_gpu.h Normal file
View File

File diff suppressed because it is too large Load Diff

39615
native/metal/ds4_metal.m Normal file
View File

File diff suppressed because it is too large Load Diff

36
native/metal/ds4_ssd.h Normal file
View File

@@ -0,0 +1,36 @@
#ifndef DS4_SSD_H
#define DS4_SSD_H
#include <stdbool.h>
#include <stdint.h>
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

297
src/app.rs
View File

@@ -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<PathBuf>,
project_name_input: String,
model_download: ModelDownload,
pub(super) composer: String,
pub(super) conversation: Vec<ChatMessage>,
pub(super) generating: bool,
#[cfg(target_os = "macos")]
generation_worker: Option<GenerationWorker>,
error: Option<String>,
}
#[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<GenerationCommand>,
events: mpsc::Receiver<GenerationEvent>,
cancel: Option<Arc<AtomicBool>>,
}
#[cfg(target_os = "macos")]
enum GenerationCommand {
Generate {
engine: crate::settings::EngineSettings,
turn: crate::settings::TurnSettings,
messages: Vec<ChatTurn>,
idle_timeout: Duration,
cancel: Arc<AtomicBool>,
},
}
#[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<PathBuf>),
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::<Vec<_>>();
#[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<GenerationWorker, String> {
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);
}
}
}

56
src/app/view.rs
View File

@@ -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),

267
src/engine.rs
View File

@@ -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<i32> {
self.tokenizer
.encode_conversation(system, messages, reasoning)
}
pub(crate) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
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<Self, String> {
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]

4
src/engine/gguf.rs
View File

@@ -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)

1729
src/engine/metal.rs Normal file
View File

File diff suppressed because it is too large Load Diff

38
src/engine/tokenizer.rs
View File

@@ -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<i32> {
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<i32> {
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);

4
src/main.rs
View File

@@ -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())

20
src/settings.rs
View File

@@ -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::<BTreeSet<_>>();
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()
}
}