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

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