hugo/DS4Server
1
0
Fork 0
Code Issues 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

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