From c60faf602a6e3a73e40b80e5adf1855ab443078a Mon Sep 17 00:00:00 2001 From: iCreil Date: Tue, 7 Jul 2026 11:10:53 +0200 Subject: [PATCH 1/3] Engine: share prefill/decode work buffers across sessions Split ds4_gpu_graph into per-session state (KV caches, compressor and indexer frontiers, MTP verification state) and a new engine-owned ds4_gpu_scratch holding the transient work buffers: one-token decode tensors, per-layer stage tensors, MTP draft tensors and the batched prefill buffers. Everything in the scratch is written and consumed inside a single sync/eval call, so sessions sharing an engine can share one scratch as long as they are driven from one thread at a time -- which is exactly the serialization the server's single graph worker already provides. This is what makes a second live session affordable: the prefill batch buffers dominate a session's non-KV footprint (~1.8 MiB per prefill-chunk token), and they are now allocated once per engine instead of once per session. The scratch grows on demand to the union of the attached sessions' capacities and is freed with the engine. No behavior change with a single session: tensors, sizes and the execution order are unchanged, only their ownership moved. Add tests/two_sessions_smoke: drives two sessions of one engine interleaved (one greedy token each, alternating on one thread) and checks both streams match isolated single-session runs token by token. Verified on CUDA (RTX PRO 6000, DeepSeek V4 Flash q2): MATCH for both sessions, and isolated outputs identical to the pre-refactor build. --- Makefile | 6 + ds4.c | 2344 +++++++++++++++++++----------------- ds4.h | 6 + tests/two_sessions_smoke.c | 144 +++ 4 files changed, 1391 insertions(+), 1109 deletions(-) create mode 100644 tests/two_sessions_smoke.c diff --git a/Makefile b/Makefile index 9711dc1a4..36b5a2738 100644 --- a/Makefile +++ b/Makefile @@ -217,6 +217,12 @@ ds4_rocm.o: ds4_rocm.cu ds4_gpu.h ds4_iq2_tables_cuda.inc $(ROCM_SRCS) tests/cuda_long_context_smoke: tests/cuda_long_context_smoke.o ds4_cuda.o $(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS) +tests/two_sessions_smoke.o: tests/two_sessions_smoke.c ds4.h + $(CC) $(CFLAGS) -I. -c -o $@ tests/two_sessions_smoke.c + +tests/two_sessions_smoke: tests/two_sessions_smoke.o $(CORE_OBJS) + $(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS) + ds4_test: ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) ifeq ($(UNAME_S),Darwin) $(CC) $(CFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(METAL_LDLIBS) diff --git a/ds4.c b/ds4.c index 640511eb0..26ac783ac 100644 --- a/ds4.c +++ b/ds4.c @@ -10309,10 +10309,22 @@ static void print_vec_stats(const char *name, const float *x, uint64_t n) { enum { DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS = 64 }; +/* Shared per-engine work buffers. Everything in here is written and + * consumed inside a single ds4_session_sync()/ds4_session_eval() call: + * no content survives across calls, so sessions that share one engine can + * also share one scratch as long as they are driven from one thread at a + * time (the same serialization the server's single graph worker already + * provides). Persistent KV caches, compressor/indexer frontiers and MTP + * verification state stay in ds4_gpu_graph: they belong to one session's + * timeline. Sharing the scratch is what makes a second live session + * affordable: the prefill batch buffers dominate a session's non-KV + * footprint (~1.8 MiB per prefill-chunk token). */ typedef struct { - /* One-token decode tensors. These stay allocated for the life of a - * session; a generated token enters as an embedding in cur_hc and leaves as - * logits after all 43 layers update their raw/compressed/indexer caches. */ + uint32_t cap_prefill; /* prefill_cap the buffers were sized for */ + uint32_t cap_comp; /* comp_cap ceiling across sessions */ + uint32_t cap_attn_comp_stage; + bool mtp_ready; /* MTP work tensors are allocated */ + int refcount; /* live sessions attached to this scratch */ ds4_gpu_tensor *cur_hc; ds4_gpu_tensor *flat_hc; ds4_gpu_tensor *hc_mix; @@ -10327,51 +10339,7 @@ typedef struct { ds4_gpu_tensor *q; ds4_gpu_tensor *kv_raw; ds4_gpu_tensor *kv; - - /* Persistent KV state. Raw KV is a sliding-window ring per layer. Ratio-4 - * layers also keep an indexer-compressed cache; ratio-128 layers keep only - * the attention-compressed cache. The small state tensors are compressor - * frontiers for the next compressed row, so they must be snapshotted with - * the row counters whenever a checkpoint is saved or partially rewound. */ - ds4_gpu_tensor *layer_raw_cache[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_attn_comp_cache[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_attn_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_attn_state_score[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_index_comp_cache[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_index_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *layer_index_state_score[DS4_MAX_LAYER]; - - /* Speculative decoding scratch. MTP is allowed to mutate graph state only - * if the target verifier can either commit it or restore the saved - * frontiers. The prefix1 buffers are the cheap partial-accept state for the - * common N=2 case. */ - ds4_gpu_tensor *spec_attn_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_attn_state_score[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_index_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_index_state_score[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_prefix1_attn_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_prefix1_attn_state_score[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_prefix1_index_state_kv[DS4_MAX_LAYER]; - ds4_gpu_tensor *spec_prefix1_index_state_score[DS4_MAX_LAYER]; ds4_gpu_tensor *spec_logits; - uint32_t layer_n_comp[DS4_MAX_LAYER]; - uint32_t layer_n_index_comp[DS4_MAX_LAYER]; - uint32_t spec_prefix1_n_comp[DS4_MAX_LAYER]; - uint32_t spec_prefix1_n_index_comp[DS4_MAX_LAYER]; - bool spec_capture_prefix1; - uint32_t raw_cap; - /* Maximum compressed-row capacity across layers. Shared work buffers use - * this worst-case size because ratio-4 indexer layers can still reach it. */ - uint32_t comp_cap; - /* Persistent compressed caches are per layer, so size them from the actual - * layer compression ratio instead of pessimistically using the ratio-4 cap - * for every ratio-128 layer. */ - uint32_t layer_comp_cap[DS4_MAX_LAYER]; - uint32_t attn_comp_stage_cap; - - /* Per-layer work tensors. They are reused in place by every layer instead - * of allocating a generic graph arena. This is why the code is verbose but - * predictable: each pointer names an actual DS4 stage. */ ds4_gpu_tensor *comp_kv_cur; ds4_gpu_tensor *comp_sc_cur; ds4_gpu_tensor *attn_comp_stage; @@ -10406,10 +10374,6 @@ typedef struct { ds4_gpu_tensor *output_embd; ds4_gpu_tensor *output_norm; ds4_gpu_tensor *logits; - - /* Optional MTP model state. It has its own raw cache because the drafter - * runs on speculative future tokens; target KV state is updated only after - * verification accepts draft tokens. */ ds4_gpu_tensor *mtp_embed; ds4_gpu_tensor *mtp_enorm; ds4_gpu_tensor *mtp_eproj; @@ -10419,15 +10383,6 @@ typedef struct { ds4_gpu_tensor *mtp_input_hc; ds4_gpu_tensor *mtp_state_hc; ds4_gpu_tensor *mtp_next_hc; - ds4_gpu_tensor *mtp_raw_cache; - uint32_t mtp_n_raw; - uint32_t prefill_cap; - uint32_t raw_window; - - /* Batched prefill tensors. Prefill is layer-major: a chunk of prompt - * tokens moves through layer 0, then layer 1, and so on, updating the same - * persistent caches used by decode. Keeping this separate from decode - * avoids a slow loop of one-token graph steps for long prompts. */ ds4_gpu_tensor *prefill_tokens; ds4_gpu_tensor *batch_cur_hc; ds4_gpu_tensor *batch_next_hc; @@ -10464,6 +10419,79 @@ typedef struct { ds4_gpu_tensor *batch_router_weights; ds4_gpu_tensor *prefill_seed_router_selected; uint32_t prefill_seed_tokens; + ds4_gpu_tensor *batch_routed_gate; + ds4_gpu_tensor *batch_routed_up; + ds4_gpu_tensor *batch_routed_mid; + ds4_gpu_tensor *batch_routed_down; + ds4_gpu_tensor *batch_routed_out; + bool batch_routed_mid_is_f16; + bool materialize_ffn_out; + float *cpu_router_norm; +} ds4_gpu_scratch; + +typedef struct { + /* One-token decode tensors. These stay allocated for the life of a + * session; a generated token enters as an embedding in cur_hc and leaves as + * logits after all 43 layers update their raw/compressed/indexer caches. */ + + /* Persistent KV state. Raw KV is a sliding-window ring per layer. Ratio-4 + * layers also keep an indexer-compressed cache; ratio-128 layers keep only + * the attention-compressed cache. The small state tensors are compressor + * frontiers for the next compressed row, so they must be snapshotted with + * the row counters whenever a checkpoint is saved or partially rewound. */ + ds4_gpu_tensor *layer_raw_cache[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_attn_comp_cache[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_attn_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_attn_state_score[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_index_comp_cache[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_index_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *layer_index_state_score[DS4_MAX_LAYER]; + + /* Speculative decoding scratch. MTP is allowed to mutate graph state only + * if the target verifier can either commit it or restore the saved + * frontiers. The prefix1 buffers are the cheap partial-accept state for the + * common N=2 case. */ + ds4_gpu_tensor *spec_attn_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_attn_state_score[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_index_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_index_state_score[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_prefix1_attn_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_prefix1_attn_state_score[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_prefix1_index_state_kv[DS4_MAX_LAYER]; + ds4_gpu_tensor *spec_prefix1_index_state_score[DS4_MAX_LAYER]; + uint32_t layer_n_comp[DS4_MAX_LAYER]; + uint32_t layer_n_index_comp[DS4_MAX_LAYER]; + uint32_t spec_prefix1_n_comp[DS4_MAX_LAYER]; + uint32_t spec_prefix1_n_index_comp[DS4_MAX_LAYER]; + bool spec_capture_prefix1; + /* Engine-owned shared work buffers (see ds4_gpu_scratch). */ + ds4_gpu_scratch *scr; + uint32_t raw_cap; + /* Maximum compressed-row capacity across layers. Shared work buffers use + * this worst-case size because ratio-4 indexer layers can still reach it. */ + uint32_t comp_cap; + /* Persistent compressed caches are per layer, so size them from the actual + * layer compression ratio instead of pessimistically using the ratio-4 cap + * for every ratio-128 layer. */ + uint32_t layer_comp_cap[DS4_MAX_LAYER]; + uint32_t attn_comp_stage_cap; + + /* Per-layer work tensors. They are reused in place by every layer instead + * of allocating a generic graph arena. This is why the code is verbose but + * predictable: each pointer names an actual DS4 stage. */ + + /* Optional MTP model state. It has its own raw cache because the drafter + * runs on speculative future tokens; target KV state is updated only after + * verification accepts draft tokens. */ + ds4_gpu_tensor *mtp_raw_cache; + uint32_t mtp_n_raw; + uint32_t prefill_cap; + uint32_t raw_window; + + /* Batched prefill tensors. Prefill is layer-major: a chunk of prompt + * tokens moves through layer 0, then layer 1, and so on, updating the same + * persistent caches used by decode. Keeping this separate from decode + * avoids a slow loop of one-token graph steps for long prompts. */ uint64_t prefill_selected_profile_rows; uint64_t prefill_selected_profile_unique; uint64_t prefill_selected_profile_selected_bytes; @@ -10471,14 +10499,7 @@ typedef struct { uint32_t prefill_selected_profile_layers; uint32_t prefill_selected_profile_min_unique; uint32_t prefill_selected_profile_max_unique; - ds4_gpu_tensor *batch_routed_gate; - ds4_gpu_tensor *batch_routed_up; - ds4_gpu_tensor *batch_routed_mid; - ds4_gpu_tensor *batch_routed_down; - ds4_gpu_tensor *batch_routed_out; - bool batch_routed_mid_is_f16; ds4_gpu_tensor *batch_ffn_out; - bool materialize_ffn_out; ds4_gpu_tensor *directional_steering_dirs; float directional_steering_attn_scale; float directional_steering_ffn_scale; @@ -10491,7 +10512,6 @@ typedef struct { bool ssd_streaming_cold; bool streaming_static_decode_map_current; bool mtp_enabled; - float *cpu_router_norm; } ds4_gpu_graph; static bool graph_power_throttle_enabled(const ds4_gpu_graph *g) { @@ -10533,93 +10553,12 @@ static void graph_power_note_decode_token(ds4_gpu_graph *g, double elapsed_sec) /* Release every Metal tensor owned by the whole-model graph runtime. */ static void metal_graph_free(ds4_gpu_graph *g) { + /* The shared scratch is engine-owned: just detach from it. The last + * engine close frees it. */ + if (g->scr && g->scr->refcount > 0) g->scr->refcount--; ds4_gpu_tensor_free(g->directional_steering_dirs); ds4_gpu_tensor_free(g->batch_ffn_out); - ds4_gpu_tensor_free(g->batch_routed_out); - ds4_gpu_tensor_free(g->batch_routed_down); - ds4_gpu_tensor_free(g->batch_routed_mid); - ds4_gpu_tensor_free(g->batch_routed_up); - ds4_gpu_tensor_free(g->batch_routed_gate); - ds4_gpu_tensor_free(g->batch_router_weights); - ds4_gpu_tensor_free(g->prefill_seed_router_selected); - ds4_gpu_tensor_free(g->batch_router_selected); - ds4_gpu_tensor_free(g->batch_router_probs); - ds4_gpu_tensor_free(g->batch_router_logits); - ds4_gpu_tensor_free(g->batch_shared_out); - ds4_gpu_tensor_free(g->batch_shared_mid); - ds4_gpu_tensor_free(g->batch_shared_up); - ds4_gpu_tensor_free(g->batch_shared_gate); - ds4_gpu_tensor_free(g->batch_ffn_norm); - ds4_gpu_tensor_free(g->batch_ffn_cur); - ds4_gpu_tensor_free(g->batch_after_attn_hc); - ds4_gpu_tensor_free(g->batch_low_tmp); - ds4_gpu_tensor_free(g->batch_group_tmp); - ds4_gpu_tensor_free(g->batch_attn_out); - ds4_gpu_tensor_free(g->batch_attn_low); - ds4_gpu_tensor_free(g->batch_heads); - ds4_gpu_tensor_free(g->batch_indexer_weights); - ds4_gpu_tensor_free(g->batch_indexer_q); - ds4_gpu_tensor_free(g->batch_comp_sc); - ds4_gpu_tensor_free(g->batch_comp_kv); - ds4_gpu_tensor_free(g->batch_kv); - ds4_gpu_tensor_free(g->batch_kv_raw); - ds4_gpu_tensor_free(g->batch_q_half); - ds4_gpu_tensor_free(g->batch_q); - ds4_gpu_tensor_free(g->batch_qr_norm); - ds4_gpu_tensor_free(g->batch_qr); - ds4_gpu_tensor_free(g->batch_attn_norm); - ds4_gpu_tensor_free(g->batch_attn_cur); - ds4_gpu_tensor_free(g->batch_hc_split); - ds4_gpu_tensor_free(g->batch_hc_mix); - ds4_gpu_tensor_free(g->batch_flat_hc); - ds4_gpu_tensor_free(g->batch_next_hc); - ds4_gpu_tensor_free(g->batch_cur_hc); - ds4_gpu_tensor_free(g->prefill_tokens); - ds4_gpu_tensor_free(g->logits); ds4_gpu_tensor_free(g->mtp_raw_cache); - ds4_gpu_tensor_free(g->mtp_next_hc); - ds4_gpu_tensor_free(g->mtp_state_hc); - ds4_gpu_tensor_free(g->mtp_input_hc); - ds4_gpu_tensor_free(g->mtp_hproj_hc); - ds4_gpu_tensor_free(g->mtp_hnorm_hc); - ds4_gpu_tensor_free(g->mtp_eproj_hc); - ds4_gpu_tensor_free(g->mtp_eproj); - ds4_gpu_tensor_free(g->mtp_enorm); - ds4_gpu_tensor_free(g->mtp_embed); - ds4_gpu_tensor_free(g->spec_logits); - ds4_gpu_tensor_free(g->output_norm); - ds4_gpu_tensor_free(g->output_embd); - ds4_gpu_tensor_free(g->output_weights); - ds4_gpu_tensor_free(g->output_pre); - ds4_gpu_tensor_free(g->after_ffn_hc); - ds4_gpu_tensor_free(g->ffn_out); - ds4_gpu_tensor_free(g->routed_out); - ds4_gpu_tensor_free(g->routed_down); - ds4_gpu_tensor_free(g->routed_mid); - ds4_gpu_tensor_free(g->routed_up); - ds4_gpu_tensor_free(g->routed_gate); - ds4_gpu_tensor_free(g->router_weights); - ds4_gpu_tensor_free(g->router_selected); - ds4_gpu_tensor_free(g->router_probs); - ds4_gpu_tensor_free(g->router_logits); - ds4_gpu_tensor_free(g->shared_out); - ds4_gpu_tensor_free(g->shared_mid); - ds4_gpu_tensor_free(g->shared_up); - ds4_gpu_tensor_free(g->shared_gate); - ds4_gpu_tensor_free(g->ffn_norm); - ds4_gpu_tensor_free(g->ffn_cur); - ds4_gpu_tensor_free(g->after_attn_hc); - ds4_gpu_tensor_free(g->attn_out); - ds4_gpu_tensor_free(g->attn_low); - ds4_gpu_tensor_free(g->heads); - ds4_gpu_tensor_free(g->comp_sc_cur); - ds4_gpu_tensor_free(g->comp_kv_cur); - ds4_gpu_tensor_free(g->attn_comp_stage); - ds4_gpu_tensor_free(g->comp_mask); - ds4_gpu_tensor_free(g->comp_selected); - ds4_gpu_tensor_free(g->indexer_scores); - ds4_gpu_tensor_free(g->indexer_weights); - ds4_gpu_tensor_free(g->indexer_q); for (uint32_t il = 0; il < DS4_N_LAYER; il++) { ds4_gpu_tensor_free(g->layer_raw_cache[il]); } @@ -10651,21 +10590,6 @@ static void metal_graph_free(ds4_gpu_graph *g) { ds4_gpu_tensor_free(g->spec_prefix1_index_state_kv[il]); ds4_gpu_tensor_free(g->spec_prefix1_index_state_score[il]); } - ds4_gpu_tensor_free(g->kv); - ds4_gpu_tensor_free(g->kv_raw); - ds4_gpu_tensor_free(g->q); - ds4_gpu_tensor_free(g->qr_norm); - ds4_gpu_tensor_free(g->qr); - ds4_gpu_tensor_free(g->attn_norm); - ds4_gpu_tensor_free(g->attn_cur); - ds4_gpu_tensor_free(g->hc_comb); - ds4_gpu_tensor_free(g->hc_post); - ds4_gpu_tensor_free(g->hc_pre); - ds4_gpu_tensor_free(g->hc_split); - ds4_gpu_tensor_free(g->hc_mix); - ds4_gpu_tensor_free(g->flat_hc); - ds4_gpu_tensor_free(g->cur_hc); - free(g->cpu_router_norm); memset(g, 0, sizeof(*g)); } @@ -10930,15 +10854,15 @@ static void metal_graph_debug_dump_i32_tensor( static bool metal_graph_needs_ffn_out(const ds4_gpu_graph *g, uint32_t il, uint32_t pos) { return metal_graph_directional_steering_ffn_enabled(g) || - g->materialize_ffn_out || + g->scr->materialize_ffn_out || metal_graph_debug_wants("ffn_out", il, pos); } static bool metal_graph_ensure_ffn_out(ds4_gpu_graph *g) { - if (!g->ffn_out) { - g->ffn_out = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + if (!g->scr->ffn_out) { + g->scr->ffn_out = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); } - return g->ffn_out != NULL; + return g->scr->ffn_out != NULL; } static bool metal_graph_ensure_batch_ffn_out(ds4_gpu_graph *g) { @@ -10952,10 +10876,333 @@ static bool metal_graph_ensure_batch_ffn_out(ds4_gpu_graph *g) { * Metal Release Graph Allocation. * ========================================================================= */ +static void metal_scratch_free(ds4_gpu_scratch *scr) { + if (!scr) return; + const int refs = scr->refcount; + ds4_gpu_tensor_free(scr->batch_routed_out); + ds4_gpu_tensor_free(scr->batch_routed_down); + ds4_gpu_tensor_free(scr->batch_routed_mid); + ds4_gpu_tensor_free(scr->batch_routed_up); + ds4_gpu_tensor_free(scr->batch_routed_gate); + ds4_gpu_tensor_free(scr->batch_router_weights); + ds4_gpu_tensor_free(scr->prefill_seed_router_selected); + ds4_gpu_tensor_free(scr->batch_router_selected); + ds4_gpu_tensor_free(scr->batch_router_probs); + ds4_gpu_tensor_free(scr->batch_router_logits); + ds4_gpu_tensor_free(scr->batch_shared_out); + ds4_gpu_tensor_free(scr->batch_shared_mid); + ds4_gpu_tensor_free(scr->batch_shared_up); + ds4_gpu_tensor_free(scr->batch_shared_gate); + ds4_gpu_tensor_free(scr->batch_ffn_norm); + ds4_gpu_tensor_free(scr->batch_ffn_cur); + ds4_gpu_tensor_free(scr->batch_after_attn_hc); + ds4_gpu_tensor_free(scr->batch_low_tmp); + ds4_gpu_tensor_free(scr->batch_group_tmp); + ds4_gpu_tensor_free(scr->batch_attn_out); + ds4_gpu_tensor_free(scr->batch_attn_low); + ds4_gpu_tensor_free(scr->batch_heads); + ds4_gpu_tensor_free(scr->batch_indexer_weights); + ds4_gpu_tensor_free(scr->batch_indexer_q); + ds4_gpu_tensor_free(scr->batch_comp_sc); + ds4_gpu_tensor_free(scr->batch_comp_kv); + ds4_gpu_tensor_free(scr->batch_kv); + ds4_gpu_tensor_free(scr->batch_kv_raw); + ds4_gpu_tensor_free(scr->batch_q_half); + ds4_gpu_tensor_free(scr->batch_q); + ds4_gpu_tensor_free(scr->batch_qr_norm); + ds4_gpu_tensor_free(scr->batch_qr); + ds4_gpu_tensor_free(scr->batch_attn_norm); + ds4_gpu_tensor_free(scr->batch_attn_cur); + ds4_gpu_tensor_free(scr->batch_hc_split); + ds4_gpu_tensor_free(scr->batch_hc_mix); + ds4_gpu_tensor_free(scr->batch_flat_hc); + ds4_gpu_tensor_free(scr->batch_next_hc); + ds4_gpu_tensor_free(scr->batch_cur_hc); + ds4_gpu_tensor_free(scr->prefill_tokens); + ds4_gpu_tensor_free(scr->logits); + ds4_gpu_tensor_free(scr->mtp_next_hc); + ds4_gpu_tensor_free(scr->mtp_state_hc); + ds4_gpu_tensor_free(scr->mtp_input_hc); + ds4_gpu_tensor_free(scr->mtp_hproj_hc); + ds4_gpu_tensor_free(scr->mtp_hnorm_hc); + ds4_gpu_tensor_free(scr->mtp_eproj_hc); + ds4_gpu_tensor_free(scr->mtp_eproj); + ds4_gpu_tensor_free(scr->mtp_enorm); + ds4_gpu_tensor_free(scr->mtp_embed); + ds4_gpu_tensor_free(scr->spec_logits); + ds4_gpu_tensor_free(scr->output_norm); + ds4_gpu_tensor_free(scr->output_embd); + ds4_gpu_tensor_free(scr->output_weights); + ds4_gpu_tensor_free(scr->output_pre); + ds4_gpu_tensor_free(scr->after_ffn_hc); + ds4_gpu_tensor_free(scr->ffn_out); + ds4_gpu_tensor_free(scr->routed_out); + ds4_gpu_tensor_free(scr->routed_down); + ds4_gpu_tensor_free(scr->routed_mid); + ds4_gpu_tensor_free(scr->routed_up); + ds4_gpu_tensor_free(scr->routed_gate); + ds4_gpu_tensor_free(scr->router_weights); + ds4_gpu_tensor_free(scr->router_selected); + ds4_gpu_tensor_free(scr->router_probs); + ds4_gpu_tensor_free(scr->router_logits); + ds4_gpu_tensor_free(scr->shared_out); + ds4_gpu_tensor_free(scr->shared_mid); + ds4_gpu_tensor_free(scr->shared_up); + ds4_gpu_tensor_free(scr->shared_gate); + ds4_gpu_tensor_free(scr->ffn_norm); + ds4_gpu_tensor_free(scr->ffn_cur); + ds4_gpu_tensor_free(scr->after_attn_hc); + ds4_gpu_tensor_free(scr->attn_out); + ds4_gpu_tensor_free(scr->attn_low); + ds4_gpu_tensor_free(scr->heads); + ds4_gpu_tensor_free(scr->comp_sc_cur); + ds4_gpu_tensor_free(scr->comp_kv_cur); + ds4_gpu_tensor_free(scr->attn_comp_stage); + ds4_gpu_tensor_free(scr->comp_mask); + ds4_gpu_tensor_free(scr->comp_selected); + ds4_gpu_tensor_free(scr->indexer_scores); + ds4_gpu_tensor_free(scr->indexer_weights); + ds4_gpu_tensor_free(scr->indexer_q); + ds4_gpu_tensor_free(scr->kv); + ds4_gpu_tensor_free(scr->kv_raw); + ds4_gpu_tensor_free(scr->q); + ds4_gpu_tensor_free(scr->qr_norm); + ds4_gpu_tensor_free(scr->qr); + ds4_gpu_tensor_free(scr->attn_norm); + ds4_gpu_tensor_free(scr->attn_cur); + ds4_gpu_tensor_free(scr->hc_comb); + ds4_gpu_tensor_free(scr->hc_post); + ds4_gpu_tensor_free(scr->hc_pre); + ds4_gpu_tensor_free(scr->hc_split); + ds4_gpu_tensor_free(scr->hc_mix); + ds4_gpu_tensor_free(scr->flat_hc); + ds4_gpu_tensor_free(scr->cur_hc); + free(scr->cpu_router_norm); + memset(scr, 0, sizeof(*scr)); + scr->refcount = refs; +} + +/* Allocate the shared work buffers for the requested capacities. Sizing + * mirrors what the graph allocator used to do when every session owned a + * private copy of these tensors. */ +static bool metal_scratch_alloc( + ds4_gpu_scratch *scr, + const ds4_weights *weights, + const ds4_layer_weights *layer, + uint32_t comp_cap, + uint32_t attn_comp_stage_cap, + uint32_t prefill_cap, + bool enable_mtp) { + const int refs = scr->refcount; + memset(scr, 0, sizeof(*scr)); + scr->refcount = refs; + scr->cap_prefill = prefill_cap; + scr->cap_comp = comp_cap; + scr->cap_attn_comp_stage = attn_comp_stage_cap; + scr->mtp_ready = enable_mtp; + + const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; + const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; + const uint64_t q_rank = layer->attn_q_a->dim[1]; + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; + const uint64_t group_dim = (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); + const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; + const uint64_t routed_mid_dim = layer->ffn_gate_exps->dim[1]; + const uint64_t vocab_dim = weights->output ? weights->output->dim[1] : DS4_N_VOCAB; + const uint64_t comp_width_max = 2ull * (DS4_N_HEAD_DIM > DS4_N_INDEXER_HEAD_DIM + ? DS4_N_HEAD_DIM + : DS4_N_INDEXER_HEAD_DIM); + const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; + const uint64_t pc = prefill_cap; + + scr->cur_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + scr->flat_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + scr->hc_mix = ds4_gpu_tensor_alloc(mix_hc * sizeof(float)); + scr->hc_split = ds4_gpu_tensor_alloc(mix_hc * sizeof(float)); + scr->hc_pre = ds4_gpu_tensor_view(scr->hc_split, 0, (uint64_t)DS4_N_HC * sizeof(float)); + scr->hc_post = ds4_gpu_tensor_view(scr->hc_split, + (uint64_t)DS4_N_HC * sizeof(float), + (uint64_t)DS4_N_HC * sizeof(float)); + scr->hc_comb = ds4_gpu_tensor_view(scr->hc_split, + 2ull * DS4_N_HC * sizeof(float), + (uint64_t)DS4_N_HC * DS4_N_HC * sizeof(float)); + scr->attn_cur = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->attn_norm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->qr = ds4_gpu_tensor_alloc(q_rank * sizeof(float)); + scr->qr_norm = ds4_gpu_tensor_alloc(q_rank * sizeof(float)); + scr->q = ds4_gpu_tensor_alloc(q_dim * sizeof(float)); + scr->kv_raw = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD_DIM * sizeof(float)); + scr->kv = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD_DIM * sizeof(float)); + scr->comp_kv_cur = ds4_gpu_tensor_alloc(comp_width_max * sizeof(float)); + scr->comp_sc_cur = ds4_gpu_tensor_alloc(comp_width_max * sizeof(float)); + scr->indexer_q = ds4_gpu_tensor_alloc(indexer_q_dim * sizeof(float)); + scr->indexer_weights = ds4_gpu_tensor_alloc((uint64_t)DS4_N_INDEXER_HEAD * sizeof(float)); + scr->indexer_scores = ds4_gpu_tensor_alloc((uint64_t)comp_cap * pc * sizeof(float)); + scr->comp_mask = ds4_gpu_tensor_alloc((uint64_t)comp_cap * pc * sizeof(float)); + scr->comp_selected = ds4_gpu_tensor_alloc((uint64_t)(DS4_N_INDEXER_TOP_K ? DS4_N_INDEXER_TOP_K : 1u) * + pc * sizeof(uint32_t)); + scr->heads = ds4_gpu_tensor_alloc(q_dim * sizeof(float)); + scr->attn_low = ds4_gpu_tensor_alloc(low_dim * sizeof(float)); + scr->attn_out = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->after_attn_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + scr->ffn_cur = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->ffn_norm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->cpu_router_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(scr->cpu_router_norm[0])); + scr->shared_gate = ds4_gpu_tensor_alloc(shared_dim * sizeof(float)); + scr->shared_up = ds4_gpu_tensor_alloc(shared_dim * sizeof(float)); + scr->shared_mid = ds4_gpu_tensor_alloc(shared_dim * sizeof(float)); + scr->shared_out = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->router_logits = ds4_gpu_tensor_alloc(DS4_N_EXPERT * sizeof(float)); + scr->router_probs = ds4_gpu_tensor_alloc(DS4_N_EXPERT * sizeof(float)); + scr->router_selected = ds4_gpu_tensor_alloc(DS4_N_EXPERT_USED * sizeof(int)); + scr->router_weights = ds4_gpu_tensor_alloc(DS4_N_EXPERT_USED * sizeof(float)); + scr->routed_gate = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + scr->routed_up = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + scr->routed_mid = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + scr->routed_down = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float)); + scr->routed_out = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->after_ffn_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + scr->output_pre = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HC * sizeof(float)); + scr->output_weights = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HC * sizeof(float)); + scr->output_embd = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->output_norm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->logits = ds4_gpu_tensor_alloc(vocab_dim * sizeof(float)); + scr->prefill_tokens = ds4_gpu_tensor_alloc(pc * sizeof(int32_t)); + scr->batch_cur_hc = ds4_gpu_tensor_alloc(pc * hc_dim * sizeof(float)); + scr->batch_next_hc = ds4_gpu_tensor_alloc(pc * hc_dim * sizeof(float)); + scr->batch_flat_hc = ds4_gpu_tensor_alloc(pc * hc_dim * sizeof(float)); + scr->batch_hc_mix = ds4_gpu_tensor_alloc(pc * mix_hc * sizeof(float)); + scr->batch_hc_split = ds4_gpu_tensor_alloc(pc * mix_hc * sizeof(float)); + scr->batch_attn_cur = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); + scr->batch_attn_norm = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); + scr->batch_qr = ds4_gpu_tensor_alloc(pc * q_rank * sizeof(float)); + scr->batch_qr_norm = ds4_gpu_tensor_alloc(pc * q_rank * sizeof(float)); + scr->batch_q = ds4_gpu_tensor_alloc(pc * q_dim * sizeof(float)); + scr->batch_q_half = ds4_gpu_tensor_alloc(pc * q_dim * sizeof(uint16_t)); + scr->batch_kv_raw = ds4_gpu_tensor_alloc(pc * DS4_N_HEAD_DIM * sizeof(float)); + scr->batch_kv = ds4_gpu_tensor_alloc(pc * DS4_N_HEAD_DIM * sizeof(float)); + scr->batch_comp_kv = ds4_gpu_tensor_alloc(pc * comp_width_max * sizeof(float)); + scr->batch_comp_sc = ds4_gpu_tensor_alloc(pc * comp_width_max * sizeof(float)); + scr->batch_indexer_q = ds4_gpu_tensor_alloc(pc * indexer_q_dim * sizeof(float)); + scr->batch_indexer_weights = ds4_gpu_tensor_alloc(pc * DS4_N_INDEXER_HEAD * sizeof(float)); + scr->batch_heads = ds4_gpu_tensor_alloc(pc * q_dim * sizeof(float)); + scr->batch_attn_low = ds4_gpu_tensor_alloc(pc * low_dim * sizeof(float)); + scr->batch_attn_out = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); + scr->batch_group_tmp = ds4_gpu_tensor_alloc(pc * group_dim * sizeof(float)); + scr->batch_low_tmp = ds4_gpu_tensor_alloc(pc * DS4_N_LORA_O * sizeof(float)); + scr->batch_after_attn_hc = ds4_gpu_tensor_alloc(pc * hc_dim * sizeof(float)); + scr->batch_ffn_cur = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); + scr->batch_ffn_norm = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); + scr->batch_shared_gate = ds4_gpu_tensor_alloc(pc * shared_dim * sizeof(float)); + scr->batch_shared_up = ds4_gpu_tensor_alloc(pc * shared_dim * sizeof(float)); + scr->batch_shared_mid = ds4_gpu_tensor_alloc(pc * shared_dim * sizeof(float)); + scr->batch_shared_out = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); + scr->batch_router_logits = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT * sizeof(float)); + scr->batch_router_probs = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT * sizeof(float)); + scr->batch_router_selected = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * sizeof(int)); + scr->batch_router_weights = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * sizeof(float)); + scr->prefill_seed_router_selected = + ds4_gpu_tensor_alloc((uint64_t)DS4_N_LAYER * + DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * + DS4_N_EXPERT_USED * + sizeof(int32_t)); + scr->batch_routed_gate = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + scr->batch_routed_up = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + scr->batch_routed_mid = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); + scr->batch_routed_down = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float)); + scr->batch_routed_out = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); + if (DS4_GPU_ATTN_COMP_CACHE_F16) { + scr->attn_comp_stage = ds4_gpu_tensor_alloc((uint64_t)attn_comp_stage_cap * + DS4_N_HEAD_DIM * sizeof(float)); + } + if (enable_mtp) { + scr->mtp_embed = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->mtp_enorm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->mtp_eproj = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + scr->mtp_eproj_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + scr->mtp_hnorm_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + scr->mtp_hproj_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + scr->mtp_input_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + scr->mtp_state_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + scr->mtp_next_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + scr->spec_logits = ds4_gpu_tensor_alloc((uint64_t)16 * DS4_N_VOCAB * sizeof(float)); + } + + const bool ok = scr->cpu_router_norm && + (!DS4_GPU_ATTN_COMP_CACHE_F16 || scr->attn_comp_stage) && + scr->cur_hc && scr->flat_hc && scr->hc_mix && + scr->hc_split && scr->hc_pre && scr->hc_post && + scr->hc_comb && scr->attn_cur && scr->attn_norm && + scr->qr && scr->qr_norm && scr->q && + scr->kv_raw && scr->kv && scr->comp_kv_cur && + scr->comp_sc_cur && scr->indexer_q && scr->indexer_weights && + scr->indexer_scores && scr->comp_mask && scr->comp_selected && + scr->heads && scr->attn_low && scr->attn_out && + scr->after_attn_hc && scr->ffn_cur && scr->ffn_norm && + scr->shared_gate && scr->shared_up && scr->shared_mid && + scr->shared_out && scr->router_logits && scr->router_probs && + scr->router_selected && scr->router_weights && scr->routed_gate && + scr->routed_up && scr->routed_mid && scr->routed_down && + scr->routed_out && scr->after_ffn_hc && scr->output_pre && + scr->output_weights && scr->output_embd && scr->output_norm && + scr->logits && scr->prefill_tokens && scr->batch_cur_hc && + scr->batch_next_hc && scr->batch_flat_hc && scr->batch_hc_mix && + scr->batch_hc_split && scr->batch_attn_cur && scr->batch_attn_norm && + scr->batch_qr && scr->batch_qr_norm && scr->batch_q && + scr->batch_q_half && scr->batch_kv_raw && scr->batch_kv && + scr->batch_comp_kv && scr->batch_comp_sc && scr->batch_indexer_q && + scr->batch_indexer_weights && scr->batch_heads && scr->batch_attn_low && + scr->batch_attn_out && scr->batch_group_tmp && scr->batch_low_tmp && + scr->batch_after_attn_hc && scr->batch_ffn_cur && scr->batch_ffn_norm && + scr->batch_shared_gate && scr->batch_shared_up && scr->batch_shared_mid && + scr->batch_shared_out && scr->batch_router_logits && scr->batch_router_probs && + scr->batch_router_selected && scr->batch_router_weights && scr->prefill_seed_router_selected && + scr->batch_routed_gate && scr->batch_routed_up && scr->batch_routed_mid && + scr->batch_routed_down && scr->batch_routed_out && + (!enable_mtp || + (scr->mtp_embed && scr->mtp_enorm && scr->mtp_eproj && + scr->mtp_eproj_hc && scr->mtp_hnorm_hc && scr->mtp_hproj_hc && + scr->mtp_input_hc && scr->mtp_state_hc && scr->mtp_next_hc && + scr->spec_logits)); + if (!ok) metal_scratch_free(scr); + return ok; +} + +/* Reuse the current scratch when it is already big enough, otherwise grow it + * to the union of the old and requested capacities. Sessions sharing an + * engine are driven from one thread at a time, so no live kernel can be using + * these buffers while a new session is being created. */ +static bool metal_scratch_ensure( + ds4_gpu_scratch *scr, + const ds4_weights *weights, + const ds4_layer_weights *layer, + uint32_t comp_cap, + uint32_t attn_comp_stage_cap, + uint32_t prefill_cap, + bool enable_mtp) { + if (scr->cur_hc && + scr->cap_prefill >= prefill_cap && + scr->cap_comp >= comp_cap && + scr->cap_attn_comp_stage >= attn_comp_stage_cap && + (!enable_mtp || scr->mtp_ready)) { + return true; + } + if (comp_cap < scr->cap_comp) comp_cap = scr->cap_comp; + if (attn_comp_stage_cap < scr->cap_attn_comp_stage) + attn_comp_stage_cap = scr->cap_attn_comp_stage; + if (prefill_cap < scr->cap_prefill) prefill_cap = scr->cap_prefill; + enable_mtp = enable_mtp || scr->mtp_ready; + metal_scratch_free(scr); + return metal_scratch_alloc(scr, weights, layer, comp_cap, + attn_comp_stage_cap, prefill_cap, enable_mtp); +} + /* Allocate the Metal graph state for a chosen raw-cache capacity. The model * weights are not copied here; tensors reference the mapped GGUF. */ static bool metal_graph_alloc_raw_cap( ds4_gpu_graph *g, + ds4_gpu_scratch *scr, const ds4_weights *weights, const ds4_layer_weights *layer, uint32_t raw_cap, @@ -10963,6 +11210,7 @@ static bool metal_graph_alloc_raw_cap( uint32_t prefill_cap, bool enable_mtp) { memset(g, 0, sizeof(*g)); + g->scr = scr; g->mtp_enabled = enable_mtp; if (raw_cap == 0) raw_cap = 1; if (ctx_size == 0) ctx_size = raw_cap; @@ -10998,20 +11246,15 @@ static bool metal_graph_alloc_raw_cap( } } - const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - const uint64_t mix_hc = 2ull * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; - const uint64_t q_rank = layer->attn_q_a->dim[1]; - const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; - const uint64_t low_dim = (uint64_t)DS4_N_OUT_GROUP * DS4_N_LORA_O; - const uint64_t group_dim = (uint64_t)DS4_N_HEAD_DIM * (DS4_N_HEAD / DS4_N_OUT_GROUP); - const uint64_t shared_dim = layer->ffn_gate_shexp->dim[1]; - const uint64_t routed_mid_dim = layer->ffn_gate_exps->dim[1]; - const uint64_t vocab_dim = weights->output ? weights->output->dim[1] : DS4_N_VOCAB; - const uint64_t comp_width_max = 2ull * (DS4_N_HEAD_DIM > DS4_N_INDEXER_HEAD_DIM - ? DS4_N_HEAD_DIM - : DS4_N_INDEXER_HEAD_DIM); - const uint64_t indexer_q_dim = (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM; - const uint64_t pc = prefill_cap; + /* Work buffers are shared across the engine's sessions; grow them if + * this session needs more capacity than any session before it. */ + if (!metal_scratch_ensure(scr, weights, layer, g->comp_cap, + g->attn_comp_stage_cap, prefill_cap, + enable_mtp)) { + return false; + } + scr->refcount++; + uint64_t kv_cache_bytes = 0; const uint64_t context_bytes = metal_graph_context_bytes_for_kv_policy(ctx_size, raw_cap, prefill_cap, &kv_cache_bytes); @@ -11035,24 +11278,6 @@ static bool metal_graph_alloc_raw_cap( (double)context_bytes / 1073741824.0); } - g->cur_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->flat_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->hc_mix = ds4_gpu_tensor_alloc(mix_hc * sizeof(float)); - g->hc_split = ds4_gpu_tensor_alloc(mix_hc * sizeof(float)); - g->hc_pre = ds4_gpu_tensor_view(g->hc_split, 0, (uint64_t)DS4_N_HC * sizeof(float)); - g->hc_post = ds4_gpu_tensor_view(g->hc_split, - (uint64_t)DS4_N_HC * sizeof(float), - (uint64_t)DS4_N_HC * sizeof(float)); - g->hc_comb = ds4_gpu_tensor_view(g->hc_split, - 2ull * DS4_N_HC * sizeof(float), - (uint64_t)DS4_N_HC * DS4_N_HC * sizeof(float)); - g->attn_cur = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->attn_norm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->qr = ds4_gpu_tensor_alloc(q_rank * sizeof(float)); - g->qr_norm = ds4_gpu_tensor_alloc(q_rank * sizeof(float)); - g->q = ds4_gpu_tensor_alloc(q_dim * sizeof(float)); - g->kv_raw = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD_DIM * sizeof(float)); - g->kv = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD_DIM * sizeof(float)); bool state_init_ok = true; for (uint32_t il = 0; il < DS4_N_LAYER; il++) { g->layer_raw_cache[il] = metal_graph_alloc_kv_cache_tensor( @@ -11109,44 +11334,6 @@ static bool metal_graph_alloc_raw_cap( } } } - g->comp_kv_cur = ds4_gpu_tensor_alloc(comp_width_max * sizeof(float)); - g->comp_sc_cur = ds4_gpu_tensor_alloc(comp_width_max * sizeof(float)); - if (DS4_GPU_ATTN_COMP_CACHE_F16) { - g->attn_comp_stage = ds4_gpu_tensor_alloc((uint64_t)g->attn_comp_stage_cap * - DS4_N_HEAD_DIM * sizeof(float)); - } - g->indexer_q = ds4_gpu_tensor_alloc(indexer_q_dim * sizeof(float)); - g->indexer_weights = ds4_gpu_tensor_alloc((uint64_t)DS4_N_INDEXER_HEAD * sizeof(float)); - g->indexer_scores = ds4_gpu_tensor_alloc((uint64_t)g->comp_cap * pc * sizeof(float)); - g->comp_mask = ds4_gpu_tensor_alloc((uint64_t)g->comp_cap * pc * sizeof(float)); - g->comp_selected = ds4_gpu_tensor_alloc((uint64_t)(DS4_N_INDEXER_TOP_K ? DS4_N_INDEXER_TOP_K : 1u) * - pc * sizeof(uint32_t)); - g->heads = ds4_gpu_tensor_alloc(q_dim * sizeof(float)); - g->attn_low = ds4_gpu_tensor_alloc(low_dim * sizeof(float)); - g->attn_out = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->after_attn_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->ffn_cur = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->ffn_norm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->cpu_router_norm = xmalloc((size_t)DS4_N_EMBD * sizeof(g->cpu_router_norm[0])); - g->shared_gate = ds4_gpu_tensor_alloc(shared_dim * sizeof(float)); - g->shared_up = ds4_gpu_tensor_alloc(shared_dim * sizeof(float)); - g->shared_mid = ds4_gpu_tensor_alloc(shared_dim * sizeof(float)); - g->shared_out = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->router_logits = ds4_gpu_tensor_alloc(DS4_N_EXPERT * sizeof(float)); - g->router_probs = ds4_gpu_tensor_alloc(DS4_N_EXPERT * sizeof(float)); - g->router_selected = ds4_gpu_tensor_alloc(DS4_N_EXPERT_USED * sizeof(int)); - g->router_weights = ds4_gpu_tensor_alloc(DS4_N_EXPERT_USED * sizeof(float)); - g->routed_gate = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->routed_up = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->routed_mid = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->routed_down = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float)); - g->routed_out = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->after_ffn_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->output_pre = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HC * sizeof(float)); - g->output_weights = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HC * sizeof(float)); - g->output_embd = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->output_norm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->logits = ds4_gpu_tensor_alloc(vocab_dim * sizeof(float)); /* * MTP is deliberately outside the normal graph footprint. A session that * does not opt in with --mtp must allocate and execute exactly the same @@ -11154,66 +11341,12 @@ static bool metal_graph_alloc_raw_cap( * and no MTP scratch hidden behind otherwise unused tensors. */ if (enable_mtp) { - g->mtp_embed = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->mtp_enorm = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->mtp_eproj = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); - g->mtp_eproj_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_hnorm_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_hproj_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_input_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_state_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); - g->mtp_next_hc = ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); g->mtp_raw_cache = metal_graph_alloc_kv_cache_tensor( managed_kv_cache, (uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float)); - g->spec_logits = ds4_gpu_tensor_alloc((uint64_t)16 * DS4_N_VOCAB * sizeof(float)); g->mtp_n_raw = 0; } - g->prefill_tokens = ds4_gpu_tensor_alloc(pc * sizeof(int32_t)); - g->batch_cur_hc = ds4_gpu_tensor_alloc(pc * hc_dim * sizeof(float)); - g->batch_next_hc = ds4_gpu_tensor_alloc(pc * hc_dim * sizeof(float)); - g->batch_flat_hc = ds4_gpu_tensor_alloc(pc * hc_dim * sizeof(float)); - g->batch_hc_mix = ds4_gpu_tensor_alloc(pc * mix_hc * sizeof(float)); - g->batch_hc_split = ds4_gpu_tensor_alloc(pc * mix_hc * sizeof(float)); - g->batch_attn_cur = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); - g->batch_attn_norm = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); - g->batch_qr = ds4_gpu_tensor_alloc(pc * q_rank * sizeof(float)); - g->batch_qr_norm = ds4_gpu_tensor_alloc(pc * q_rank * sizeof(float)); - g->batch_q = ds4_gpu_tensor_alloc(pc * q_dim * sizeof(float)); - g->batch_q_half = ds4_gpu_tensor_alloc(pc * q_dim * sizeof(uint16_t)); - g->batch_kv_raw = ds4_gpu_tensor_alloc(pc * DS4_N_HEAD_DIM * sizeof(float)); - g->batch_kv = ds4_gpu_tensor_alloc(pc * DS4_N_HEAD_DIM * sizeof(float)); - g->batch_comp_kv = ds4_gpu_tensor_alloc(pc * comp_width_max * sizeof(float)); - g->batch_comp_sc = ds4_gpu_tensor_alloc(pc * comp_width_max * sizeof(float)); - g->batch_indexer_q = ds4_gpu_tensor_alloc(pc * indexer_q_dim * sizeof(float)); - g->batch_indexer_weights = ds4_gpu_tensor_alloc(pc * DS4_N_INDEXER_HEAD * sizeof(float)); - g->batch_heads = ds4_gpu_tensor_alloc(pc * q_dim * sizeof(float)); - g->batch_attn_low = ds4_gpu_tensor_alloc(pc * low_dim * sizeof(float)); - g->batch_attn_out = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); - g->batch_group_tmp = ds4_gpu_tensor_alloc(pc * group_dim * sizeof(float)); - g->batch_low_tmp = ds4_gpu_tensor_alloc(pc * DS4_N_LORA_O * sizeof(float)); - g->batch_after_attn_hc = ds4_gpu_tensor_alloc(pc * hc_dim * sizeof(float)); - g->batch_ffn_cur = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); - g->batch_ffn_norm = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); - g->batch_shared_gate = ds4_gpu_tensor_alloc(pc * shared_dim * sizeof(float)); - g->batch_shared_up = ds4_gpu_tensor_alloc(pc * shared_dim * sizeof(float)); - g->batch_shared_mid = ds4_gpu_tensor_alloc(pc * shared_dim * sizeof(float)); - g->batch_shared_out = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); - g->batch_router_logits = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT * sizeof(float)); - g->batch_router_probs = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT * sizeof(float)); - g->batch_router_selected = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * sizeof(int)); - g->batch_router_weights = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * sizeof(float)); - g->prefill_seed_router_selected = - ds4_gpu_tensor_alloc((uint64_t)DS4_N_LAYER * - DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * - DS4_N_EXPERT_USED * - sizeof(int32_t)); - g->batch_routed_gate = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->batch_routed_up = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->batch_routed_mid = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * routed_mid_dim * sizeof(float)); - g->batch_routed_down = ds4_gpu_tensor_alloc(pc * DS4_N_EXPERT_USED * DS4_N_EMBD * sizeof(float)); - g->batch_routed_out = ds4_gpu_tensor_alloc(pc * DS4_N_EMBD * sizeof(float)); bool layer_cache_ok = true; for (uint32_t il = 0; layer_cache_ok && il < DS4_N_LAYER; il++) { @@ -11242,57 +11375,16 @@ static bool metal_graph_alloc_raw_cap( } const bool ok = state_init_ok && layer_cache_ok && - g->cur_hc && g->flat_hc && g->hc_mix && g->hc_split && - g->hc_pre && g->hc_post && g->hc_comb && - g->attn_cur && g->attn_norm && g->qr && g->qr_norm && - g->q && g->kv_raw && g->kv && - g->comp_kv_cur && g->comp_sc_cur && - (!DS4_GPU_ATTN_COMP_CACHE_F16 || g->attn_comp_stage) && - g->indexer_q && g->indexer_weights && g->indexer_scores && - g->comp_mask && g->comp_selected && - g->heads && g->attn_low && g->attn_out && - g->after_attn_hc && g->ffn_cur && g->ffn_norm && - g->shared_gate && g->shared_up && g->shared_mid && - g->shared_out && - g->router_logits && g->router_probs && g->router_selected && g->router_weights && - g->routed_gate && g->routed_up && g->routed_mid && - g->routed_down && g->routed_out && - g->after_ffn_hc && - g->output_pre && g->output_weights && g->output_embd && - g->output_norm && g->logits && - (!enable_mtp || - (g->mtp_embed && g->mtp_enorm && g->mtp_eproj && - g->mtp_eproj_hc && g->mtp_hnorm_hc && g->mtp_hproj_hc && - g->mtp_input_hc && g->mtp_state_hc && g->mtp_next_hc && - g->mtp_raw_cache && g->spec_logits)) && - g->prefill_tokens && - g->batch_cur_hc && g->batch_next_hc && g->batch_flat_hc && - g->batch_hc_mix && g->batch_hc_split && - g->batch_attn_cur && g->batch_attn_norm && - g->batch_qr && g->batch_qr_norm && g->batch_q && g->batch_q_half && - g->batch_kv_raw && g->batch_kv && - g->batch_comp_kv && g->batch_comp_sc && - g->batch_indexer_q && g->batch_indexer_weights && - g->batch_heads && g->batch_attn_low && g->batch_attn_out && - g->batch_group_tmp && g->batch_low_tmp && g->batch_after_attn_hc && - g->batch_ffn_cur && g->batch_ffn_norm && - g->batch_shared_gate && g->batch_shared_up && - g->batch_shared_mid && g->batch_shared_out && - g->batch_router_logits && g->batch_router_probs && - g->batch_router_selected && g->batch_router_weights && - g->prefill_seed_router_selected && - g->batch_routed_gate && g->batch_routed_up && - g->batch_routed_mid && g->batch_routed_down && - g->batch_routed_out; - if (!ok) metal_graph_free(g); + (!enable_mtp || g->mtp_raw_cache); return ok; } static bool metal_graph_alloc( ds4_gpu_graph *g, + ds4_gpu_scratch *scr, const ds4_weights *weights, const ds4_layer_weights *layer) { - return metal_graph_alloc_raw_cap(g, weights, layer, DS4_N_SWA, DS4_N_SWA, 1, false); + return metal_graph_alloc_raw_cap(g, scr, weights, layer, DS4_N_SWA, DS4_N_SWA, 1, false); } static bool metal_graph_install_model_spans( @@ -11899,7 +11991,7 @@ static bool rocm_graph_stream_seed_full_layer_selected( down_expert_bytes); if (ds4_gpu_stream_expert_cache_seed_from_layer_selected( &table, - g->batch_router_selected, + g->scr->batch_router_selected, n_tokens, rocm_graph_stream_prefill_full_layer_seed_tokens(), DS4_N_EXPERT_USED) == 0) { @@ -11947,7 +12039,7 @@ static bool metal_graph_stream_prefill_selected_profile_layer( uint32_t il, uint32_t n_tokens) { if (!metal_graph_stream_prefill_selected_profile_enabled(g)) return true; - if (!layer || !g->batch_router_selected || n_tokens == 0 || + if (!layer || !g->scr->batch_router_selected || n_tokens == 0 || DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { return false; @@ -11956,7 +12048,7 @@ static bool metal_graph_stream_prefill_selected_profile_layer( const uint64_t n_ids = (uint64_t)n_tokens * DS4_N_EXPERT_USED; if (n_ids > SIZE_MAX / sizeof(int32_t)) return false; int32_t *selected = xmalloc((size_t)n_ids * sizeof(selected[0])); - const bool read_ok = ds4_gpu_tensor_read(g->batch_router_selected, + const bool read_ok = ds4_gpu_tensor_read(g->scr->batch_router_selected, 0, selected, n_ids * sizeof(selected[0])) != 0; @@ -12335,7 +12427,7 @@ static bool metal_graph_stream_prefill_selected_pagein_start( job->madvise_only = madvise_only; if (!metal_graph_stream_prefill_selected_pagein_enabled(g) && !madvise_only) return true; - if (!model || !layer || !g->batch_router_selected || n_tokens == 0) { + if (!model || !layer || !g->scr->batch_router_selected || n_tokens == 0) { return false; } @@ -12344,7 +12436,7 @@ static bool metal_graph_stream_prefill_selected_pagein_start( int32_t *selected = xmalloc((size_t)n_ids * sizeof(selected[0])); const double t_read0 = job->profile ? now_sec() : 0.0; - bool ok = ds4_gpu_tensor_read(g->batch_router_selected, + bool ok = ds4_gpu_tensor_read(g->scr->batch_router_selected, 0, selected, n_ids * sizeof(selected[0])) != 0; @@ -12978,7 +13070,7 @@ static bool metal_graph_stream_readahead_selected_experts_from_gpu( uint64_t gate_expert_bytes, uint64_t down_expert_bytes) { if (!metal_graph_stream_prefill_selected_readahead_enabled(g)) return true; - if (!model || !layer || !g || !g->batch_router_selected || n_tokens == 0) { + if (!model || !layer || !g || !g->scr->batch_router_selected || n_tokens == 0) { return false; } if (sizeof(int) != sizeof(int32_t) || DS4_N_EXPERT > DS4_MAX_EXPERT) { @@ -12992,7 +13084,7 @@ static bool metal_graph_stream_readahead_selected_experts_from_gpu( const bool profile = getenv("DS4_METAL_STREAMING_PREFILL_SELECTED_READAHEAD_PROFILE") != NULL; const double t0 = profile ? now_sec() : 0.0; - bool ok = ds4_gpu_tensor_read(g->batch_router_selected, + bool ok = ds4_gpu_tensor_read(g->scr->batch_router_selected, 0, selected, n_ids * sizeof(selected[0])) != 0; @@ -13494,7 +13586,7 @@ static bool metal_graph_store_attn_comp_stage( uint32_t rows) { if (!g || il >= DS4_N_LAYER) return false; if (rows == 0) return true; - if (!g->layer_attn_comp_cache[il] || !g->attn_comp_stage) return false; + if (!g->layer_attn_comp_cache[il] || !g->scr->attn_comp_stage) return false; if (rows > g->attn_comp_stage_cap || first_row > g->layer_comp_cap[il] || rows > g->layer_comp_cap[il] - first_row) { return false; @@ -13506,14 +13598,14 @@ static bool metal_graph_store_attn_comp_stage( if (DS4_GPU_ATTN_COMP_CACHE_F16) { return ds4_gpu_tensor_copy_f32_to_f16(g->layer_attn_comp_cache[il], dst_offset, - g->attn_comp_stage, + g->scr->attn_comp_stage, 0, count) != 0; } return ds4_gpu_tensor_copy(g->layer_attn_comp_cache[il], dst_offset, - g->attn_comp_stage, + g->scr->attn_comp_stage, 0, count * sizeof(float)) != 0; } @@ -13522,7 +13614,7 @@ static ds4_gpu_tensor *metal_graph_attn_comp_update_target( ds4_gpu_graph *g, uint32_t il) { return DS4_GPU_ATTN_COMP_CACHE_F16 - ? g->attn_comp_stage + ? g->scr->attn_comp_stage : g->layer_attn_comp_cache[il]; } @@ -13544,7 +13636,7 @@ static ds4_gpu_tensor *metal_graph_attn_comp_row_view( uint32_t il, uint32_t row) { if (DS4_GPU_ATTN_COMP_CACHE_F16) { - return ds4_gpu_tensor_view(g->attn_comp_stage, + return ds4_gpu_tensor_view(g->scr->attn_comp_stage, 0, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)); } @@ -13558,7 +13650,7 @@ static ds4_gpu_tensor *metal_graph_attn_comp_prefill_target( uint32_t il, uint32_t first_row, uint32_t rows) { - if (DS4_GPU_ATTN_COMP_CACHE_F16) return g->attn_comp_stage; + if (DS4_GPU_ATTN_COMP_CACHE_F16) return g->scr->attn_comp_stage; const uint32_t view_rows = rows ? rows : 1u; return ds4_gpu_tensor_view(g->layer_attn_comp_cache[il], (uint64_t)first_row * DS4_N_HEAD_DIM * sizeof(float), @@ -14080,10 +14172,10 @@ static bool metal_graph_decode_cpu_router( const double t0 = profile ? now_sec() : 0.0; if (ds4_gpu_end_commands() == 0) return false; const double t_sync = profile ? now_sec() : 0.0; - if (ds4_gpu_tensor_read(g->ffn_norm, + if (ds4_gpu_tensor_read(g->scr->ffn_norm, 0, - g->cpu_router_norm, - (uint64_t)DS4_N_EMBD * sizeof(g->cpu_router_norm[0])) == 0) { + g->scr->cpu_router_norm, + (uint64_t)DS4_N_EMBD * sizeof(g->scr->cpu_router_norm[0])) == 0) { return false; } const double t_read = profile ? now_sec() : 0.0; @@ -14094,7 +14186,7 @@ static bool metal_graph_decode_cpu_router( int32_t selected_i32[DS4_MAX_EXPERT_USED]; float weights[DS4_MAX_EXPERT_USED]; - matvec_any(logits, model, layer->ffn_gate_inp, g->cpu_router_norm); + matvec_any(logits, model, layer->ffn_gate_inp, g->scr->cpu_router_norm); for (uint32_t i = 0; i < DS4_N_EXPERT; i++) { probs[i] = sqrtf(softplus_stable(logits[i])); } @@ -14109,19 +14201,19 @@ static bool metal_graph_decode_cpu_router( } const double t_cpu = profile ? now_sec() : 0.0; - if (ds4_gpu_tensor_write(g->router_logits, + if (ds4_gpu_tensor_write(g->scr->router_logits, 0, logits, (uint64_t)DS4_N_EXPERT * sizeof(logits[0])) == 0 || - ds4_gpu_tensor_write(g->router_probs, + ds4_gpu_tensor_write(g->scr->router_probs, 0, probs, (uint64_t)DS4_N_EXPERT * sizeof(probs[0])) == 0 || - ds4_gpu_tensor_write(g->router_selected, + ds4_gpu_tensor_write(g->scr->router_selected, 0, selected_i32, (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_i32[0])) == 0 || - ds4_gpu_tensor_write(g->router_weights, + ds4_gpu_tensor_write(g->scr->router_weights, 0, weights, (uint64_t)DS4_N_EXPERT_USED * sizeof(weights[0])) == 0) { @@ -14161,7 +14253,7 @@ static bool metal_graph_decode_selected_readahead_override( uint32_t il, uint64_t gate_expert_bytes, uint64_t down_expert_bytes) { - if (!g || !model || !layer || !g->router_selected || + if (!g || !model || !layer || !g->scr->router_selected || DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { return false; @@ -14174,7 +14266,7 @@ static bool metal_graph_decode_selected_readahead_override( const double t_sync = profile ? now_sec() : 0.0; int32_t selected_ids[DS4_MAX_EXPERT_USED] = {0}; - if (ds4_gpu_tensor_read(g->router_selected, + if (ds4_gpu_tensor_read(g->scr->router_selected, 0, selected_ids, (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_ids[0])) == 0) { @@ -14271,7 +14363,7 @@ static bool metal_graph_decode_cuda_selected_load( #if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) if (!metal_graph_decode_cuda_selected_slots_expected(g, layer) || !model || - !g->router_selected || + !g->scr->router_selected || DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || DS4_N_EXPERT_USED == 0 || @@ -14287,7 +14379,7 @@ static bool metal_graph_decode_cuda_selected_load( const double t_sync = profile ? now_sec() : 0.0; int32_t selected_ids[DS4_MAX_EXPERT_USED] = {0}; - bool ok = ds4_gpu_tensor_read(g->router_selected, + bool ok = ds4_gpu_tensor_read(g->scr->router_selected, 0, selected_ids, (uint64_t)DS4_N_EXPERT_USED * @@ -14344,7 +14436,7 @@ static bool metal_graph_cuda_stream_prefill_batch_selected_load( #if !defined(DS4_ROCM_BUILD) && !defined(DS4_NO_GPU) && !defined(__APPLE__) if (!metal_graph_decode_cuda_selected_slots_expected(g, layer) || !model || - !g->batch_router_selected || + !g->scr->batch_router_selected || n_tokens <= 1 || DS4_N_EXPERT == 0 || DS4_N_EXPERT_USED == 0 || @@ -14370,7 +14462,7 @@ static bool metal_graph_cuda_stream_prefill_batch_selected_load( const double t_sync = profile ? now_sec() : 0.0; int32_t *selected_ids = xmalloc((size_t)n_ids64 * sizeof(selected_ids[0])); - bool ok = ds4_gpu_tensor_read(g->batch_router_selected, + bool ok = ds4_gpu_tensor_read(g->scr->batch_router_selected, 0, selected_ids, n_ids64 * sizeof(selected_ids[0])) != 0; @@ -14447,7 +14539,7 @@ static void metal_graph_selected_async_load_run( metal_graph_selected_async_load *job) { job->ok = false; - if (!job->g || !job->model || !job->layer || !job->g->router_selected || + if (!job->g || !job->model || !job->layer || !job->g->scr->router_selected || DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { return; @@ -14455,7 +14547,7 @@ static void metal_graph_selected_async_load_run( if (job->event_value != 0) { #ifdef DS4_ROCM_BUILD if (ds4_gpu_tensor_read_after_selected_event( - job->g->router_selected, + job->g->scr->router_selected, 0, job->selected_ids, (uint64_t)DS4_N_EXPERT_USED * @@ -14469,7 +14561,7 @@ static void metal_graph_selected_async_load_run( "selected-id async expert load") == 0) { return; } - if (ds4_gpu_tensor_read(job->g->router_selected, + if (ds4_gpu_tensor_read(job->g->scr->router_selected, 0, job->selected_ids, (uint64_t)DS4_N_EXPERT_USED * @@ -14635,7 +14727,7 @@ static void rocm_graph_batch_selected_async_load_run( rocm_graph_batch_selected_async_load *job) { job->ok = false; if (!job->g || !job->model || !job->layer || - !job->g->batch_router_selected || !job->selected_ids || + !job->g->scr->batch_router_selected || !job->selected_ids || job->n_tokens <= 1 || DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || DS4_N_EXPERT_USED == 0 || DS4_N_EXPERT_USED > DS4_MAX_EXPERT_USED) { @@ -14648,7 +14740,7 @@ static void rocm_graph_batch_selected_async_load_run( const uint64_t n_ids = (uint64_t)job->n_tokens * DS4_N_EXPERT_USED; if (n_ids > SIZE_MAX / sizeof(job->selected_ids[0])) return; if (ds4_gpu_tensor_read_after_selected_event( - job->g->batch_router_selected, + job->g->scr->batch_router_selected, 0, job->selected_ids, n_ids * sizeof(job->selected_ids[0]), @@ -14801,7 +14893,7 @@ static bool metal_graph_profile_router_selection( uint32_t il, uint32_t pos) { if (!g_expert_profile.active) return true; - if (!g || !layer || !g->router_selected || !g->router_weights) return false; + if (!g || !layer || !g->scr->router_selected || !g->scr->router_weights) return false; if (ds4_gpu_end_commands() == 0) { fprintf(stderr, @@ -14812,11 +14904,11 @@ static bool metal_graph_profile_router_selection( int32_t selected[DS4_MAX_EXPERT_USED] = {0}; float weights[DS4_MAX_EXPERT_USED] = {0}; const bool read_ok = - ds4_gpu_tensor_read(g->router_selected, + ds4_gpu_tensor_read(g->scr->router_selected, 0, selected, (uint64_t)DS4_N_EXPERT_USED * sizeof(selected[0])) != 0 && - ds4_gpu_tensor_read(g->router_weights, + ds4_gpu_tensor_read(g->scr->router_weights, 0, weights, (uint64_t)DS4_N_EXPERT_USED * sizeof(weights[0])) != 0; @@ -14881,19 +14973,19 @@ static bool metal_graph_encode_decode_layer( ok = metal_graph_layer_stage_profile_boundary("decode", (name), il, pos, 1, &decode_stage_t0); \ } \ } while (0) - if (ok) ok = ds4_gpu_rms_norm_plain_tensor(g->flat_hc, g->cur_hc, (uint32_t)hc_dim, DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(g->hc_mix, model, layer->hc_attn_fn, - hc_dim, mix_hc, g->flat_hc, 1); + if (ok) ok = ds4_gpu_rms_norm_plain_tensor(g->scr->flat_hc, g->scr->cur_hc, (uint32_t)hc_dim, DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(g->scr->hc_mix, model, layer->hc_attn_fn, + hc_dim, mix_hc, g->scr->flat_hc, 1); const bool fuse_hc_norm = DS4_N_HC == 4 && !metal_graph_use_reference_hc_decode() && !metal_graph_use_reference_hc_norm_decode(); if (ok && fuse_hc_norm) { - ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(g->attn_cur, - g->attn_norm, - g->hc_split, - g->hc_mix, - g->cur_hc, + ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(g->scr->attn_cur, + g->scr->attn_norm, + g->scr->hc_split, + g->scr->hc_mix, + g->scr->cur_hc, model->map, model->size, layer->hc_attn_scale->abs_offset, @@ -14906,10 +14998,10 @@ static bool metal_graph_encode_decode_layer( DS4_RMS_EPS) != 0; if (ok) { ok = metal_graph_check_hc_norm_fusion("attn", - g->attn_cur, - g->attn_norm, - g->hc_mix, - g->cur_hc, + g->scr->attn_cur, + g->scr->attn_norm, + g->scr->hc_mix, + g->scr->cur_hc, model, layer->hc_attn_scale->abs_offset, layer->hc_attn_base->abs_offset, @@ -14918,36 +15010,36 @@ static bool metal_graph_encode_decode_layer( pos); } } else if (ok) { - ok = metal_graph_decode_hc_pre(g->attn_cur, - g->hc_split, - g->hc_mix, - g->cur_hc, + ok = metal_graph_decode_hc_pre(g->scr->attn_cur, + g->scr->hc_split, + g->scr->hc_mix, + g->scr->cur_hc, model, layer->hc_attn_scale->abs_offset, layer->hc_attn_base->abs_offset); } DS4_METAL_PROFILE_DECODE_STAGE("attn_hc_pre"); if (ok) { - metal_graph_debug_dump_tensor("hc_attn_pre_mixes", g->hc_mix, mix_hc, il, pos); - metal_graph_debug_dump_tensor("hc_attn_pre_weights", g->hc_pre, DS4_N_HC, il, pos); - metal_graph_debug_dump_tensor("hc_attn_pre_post_weights", g->hc_post, DS4_N_HC, il, pos); - metal_graph_debug_dump_tensor("hc_attn_pre_comb", g->hc_comb, (uint64_t)DS4_N_HC * DS4_N_HC, il, pos); + metal_graph_debug_dump_tensor("hc_attn_pre_mixes", g->scr->hc_mix, mix_hc, il, pos); + metal_graph_debug_dump_tensor("hc_attn_pre_weights", g->scr->hc_pre, DS4_N_HC, il, pos); + metal_graph_debug_dump_tensor("hc_attn_pre_post_weights", g->scr->hc_post, DS4_N_HC, il, pos); + metal_graph_debug_dump_tensor("hc_attn_pre_comb", g->scr->hc_comb, (uint64_t)DS4_N_HC * DS4_N_HC, il, pos); } if (ok) { - metal_graph_debug_dump_tensor("hc_attn_pre", g->attn_cur, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("hc_attn_pre", g->scr->attn_cur, DS4_N_EMBD, il, pos); } - if (ok && !fuse_hc_norm) ok = ds4_gpu_rms_norm_weight_tensor(g->attn_norm, g->attn_cur, + if (ok && !fuse_hc_norm) ok = ds4_gpu_rms_norm_weight_tensor(g->scr->attn_norm, g->scr->attn_cur, model->map, model->size, layer->attn_norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; DS4_METAL_PROFILE_DECODE_STAGE("attn_norm"); if (ok) { - metal_graph_debug_dump_tensor("attn_norm", g->attn_norm, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("attn_norm", g->scr->attn_norm, DS4_N_EMBD, il, pos); } bool qkv_pair_projected = false; if (ok && qkv_rms_fused) { - qkv_pair_projected = ds4_gpu_matmul_q8_0_pair_tensor(g->qr, - g->kv_raw, + qkv_pair_projected = ds4_gpu_matmul_q8_0_pair_tensor(g->scr->qr, + g->scr->kv_raw, model->map, model->size, layer->attn_q_a->abs_offset, @@ -14955,64 +15047,64 @@ static bool metal_graph_encode_decode_layer( DS4_N_EMBD, q_rank, DS4_N_HEAD_DIM, - g->attn_norm, + g->scr->attn_norm, 1) != 0; } - if (ok && !qkv_pair_projected) ok = ds4_gpu_matmul_q8_0_tensor(g->qr, + if (ok && !qkv_pair_projected) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->qr, model->map, model->size, layer->attn_q_a->abs_offset, DS4_N_EMBD, q_rank, - g->attn_norm, + g->scr->attn_norm, 1) != 0; if (ok) { - metal_graph_debug_dump_tensor("q_lora", g->qr, q_rank, il, pos); + metal_graph_debug_dump_tensor("q_lora", g->scr->qr, q_rank, il, pos); } if (qkv_rms_fused) { - if (ok && !qkv_pair_projected) ok = ds4_gpu_matmul_q8_0_tensor(g->kv_raw, model->map, model->size, + if (ok && !qkv_pair_projected) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->kv_raw, model->map, model->size, layer->attn_kv->abs_offset, DS4_N_EMBD, DS4_N_HEAD_DIM, - g->attn_norm, 1) != 0; + g->scr->attn_norm, 1) != 0; if (ok) { - metal_graph_debug_dump_tensor("KVraw", g->kv_raw, DS4_N_HEAD_DIM, il, pos); + metal_graph_debug_dump_tensor("KVraw", g->scr->kv_raw, DS4_N_HEAD_DIM, il, pos); } - if (ok) ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(g->qr_norm, - g->qr, + if (ok) ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(g->scr->qr_norm, + g->scr->qr, model->map, model->size, layer->attn_q_a_norm->abs_offset, (uint32_t)q_rank, - g->kv, - g->kv_raw, + g->scr->kv, + g->scr->kv_raw, layer->attn_kv_a_norm->abs_offset, DS4_N_HEAD_DIM, 1, DS4_RMS_EPS) != 0; } else { - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->qr_norm, g->qr, + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->scr->qr_norm, g->scr->qr, model->map, model->size, layer->attn_q_a_norm->abs_offset, (uint32_t)q_rank, DS4_RMS_EPS) != 0; } if (ok) { - metal_graph_debug_dump_tensor("q_lora_norm", g->qr_norm, q_rank, il, pos); + metal_graph_debug_dump_tensor("q_lora_norm", g->scr->qr_norm, q_rank, il, pos); } if (qkv_rms_fused && ok) { - metal_graph_debug_dump_tensor("KVnorm", g->kv, DS4_N_HEAD_DIM, il, pos); + metal_graph_debug_dump_tensor("KVnorm", g->scr->kv, DS4_N_HEAD_DIM, il, pos); } - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->q, model->map, model->size, + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->q, model->map, model->size, layer->attn_q_b->abs_offset, q_rank, q_dim, - g->qr_norm, 1) != 0; + g->scr->qr_norm, 1) != 0; if (ok) { - metal_graph_debug_dump_tensor("Qraw", g->q, q_dim, il, pos); + metal_graph_debug_dump_tensor("Qraw", g->scr->q, q_dim, il, pos); } const bool decode_q_norm_debug = metal_graph_debug_wants("Qnorm", il, pos); bool decode_q_norm_rope_fused = false; if (ok && !decode_q_norm_debug) { decode_q_norm_rope_fused = - ds4_gpu_head_rms_norm_rope_tail_tensor(g->q, + ds4_gpu_head_rms_norm_rope_tail_tensor(g->scr->q, 1, DS4_N_HEAD, DS4_N_HEAD_DIM, @@ -15029,11 +15121,11 @@ static bool metal_graph_encode_decode_layer( DS4_RMS_EPS) != 0; } if (!decode_q_norm_rope_fused) { - if (ok) ok = ds4_gpu_head_rms_norm_tensor(g->q, 1, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_head_rms_norm_tensor(g->scr->q, 1, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; if (ok) { - metal_graph_debug_dump_tensor("Qnorm", g->q, q_dim, il, pos); + metal_graph_debug_dump_tensor("Qnorm", g->scr->q, q_dim, il, pos); } - if (ok) ok = ds4_gpu_rope_tail_tensor(g->q, 1, DS4_N_HEAD, DS4_N_HEAD_DIM, + if (ok) ok = ds4_gpu_rope_tail_tensor(g->scr->q, 1, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, false, freq_base, freq_scale, ext_factor, attn_factor, @@ -15041,39 +15133,39 @@ static bool metal_graph_encode_decode_layer( } DS4_METAL_PROFILE_DECODE_STAGE("q_path"); if (ok) { - metal_graph_debug_dump_tensor("Qcur", g->q, q_dim, il, pos); + metal_graph_debug_dump_tensor("Qcur", g->scr->q, q_dim, il, pos); } if (!qkv_rms_fused) { - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->kv_raw, model->map, model->size, + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->kv_raw, model->map, model->size, layer->attn_kv->abs_offset, DS4_N_EMBD, DS4_N_HEAD_DIM, - g->attn_norm, 1) != 0; + g->scr->attn_norm, 1) != 0; if (ok) { - metal_graph_debug_dump_tensor("KVraw", g->kv_raw, DS4_N_HEAD_DIM, il, pos); + metal_graph_debug_dump_tensor("KVraw", g->scr->kv_raw, DS4_N_HEAD_DIM, il, pos); } - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->kv, g->kv_raw, + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->scr->kv, g->scr->kv_raw, model->map, model->size, layer->attn_kv_a_norm->abs_offset, DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; if (ok) { - metal_graph_debug_dump_tensor("KVnorm", g->kv, DS4_N_HEAD_DIM, il, pos); + metal_graph_debug_dump_tensor("KVnorm", g->scr->kv, DS4_N_HEAD_DIM, il, pos); } } - if (ok) ok = ds4_gpu_rope_tail_tensor(g->kv, 1, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, + if (ok) ok = ds4_gpu_rope_tail_tensor(g->scr->kv, 1, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_N_ROT, pos, compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, false, freq_base, freq_scale, ext_factor, attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; if (ok) { - metal_graph_debug_dump_tensor("KVrope", g->kv, DS4_N_HEAD_DIM, il, pos); + metal_graph_debug_dump_tensor("KVrope", g->scr->kv, DS4_N_HEAD_DIM, il, pos); } /* RoPE stays as the exact standalone kernel above. The decode fusion * starts after that, where FP8 KV quantization and raw-cache storage can * share one pass without changing the trigonometric path. */ - if (ok) ok = metal_graph_decode_kv_store(g->kv, raw_cache, raw_cap, raw_row); + if (ok) ok = metal_graph_decode_kv_store(g->scr->kv, raw_cache, raw_cap, raw_row); DS4_METAL_PROFILE_DECODE_STAGE("kv_path"); if (ok) { - metal_graph_debug_dump_tensor("KVcur", g->kv, DS4_N_HEAD_DIM, il, pos); + metal_graph_debug_dump_tensor("KVcur", g->scr->kv, DS4_N_HEAD_DIM, il, pos); } uint32_t n_comp = 0; @@ -15103,29 +15195,29 @@ static bool metal_graph_encode_decode_layer( ok = false; } if (ok && !metal_graph_use_reference_compressor_pair_proj()) { - ok = ds4_gpu_matmul_f16_pair_tensor(g->comp_kv_cur, - g->comp_sc_cur, + ok = ds4_gpu_matmul_f16_pair_tensor(g->scr->comp_kv_cur, + g->scr->comp_sc_cur, model->map, model->size, layer->attn_compressor_kv->abs_offset, layer->attn_compressor_gate->abs_offset, DS4_N_EMBD, comp_width, - g->attn_norm, + g->scr->attn_norm, 1) != 0; } else { - if (ok) ok = ds4_gpu_matmul_f16_tensor(g->comp_kv_cur, model->map, model->size, + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->comp_kv_cur, model->map, model->size, layer->attn_compressor_kv->abs_offset, DS4_N_EMBD, comp_width, - g->attn_norm, 1) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(g->comp_sc_cur, model->map, model->size, + g->scr->attn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->comp_sc_cur, model->map, model->size, layer->attn_compressor_gate->abs_offset, DS4_N_EMBD, comp_width, - g->attn_norm, 1) != 0; + g->scr->attn_norm, 1) != 0; } const uint32_t comp_row = g->layer_n_comp[il]; - if (ok) ok = ds4_gpu_compressor_update_tensor(g->comp_kv_cur, - g->comp_sc_cur, + if (ok) ok = ds4_gpu_compressor_update_tensor(g->scr->comp_kv_cur, + g->scr->comp_sc_cur, g->layer_attn_state_kv[il], g->layer_attn_state_score[il], metal_graph_attn_comp_update_target(g, il), @@ -15181,29 +15273,29 @@ static bool metal_graph_encode_decode_layer( ok = false; } if (ok && !metal_graph_use_reference_compressor_pair_proj()) { - ok = ds4_gpu_matmul_f16_pair_tensor(g->comp_kv_cur, - g->comp_sc_cur, + ok = ds4_gpu_matmul_f16_pair_tensor(g->scr->comp_kv_cur, + g->scr->comp_sc_cur, model->map, model->size, layer->indexer_compressor_kv->abs_offset, layer->indexer_compressor_gate->abs_offset, DS4_N_EMBD, index_width, - g->attn_norm, + g->scr->attn_norm, 1) != 0; } else { - if (ok) ok = ds4_gpu_matmul_f16_tensor(g->comp_kv_cur, model->map, model->size, + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->comp_kv_cur, model->map, model->size, layer->indexer_compressor_kv->abs_offset, DS4_N_EMBD, index_width, - g->attn_norm, 1) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(g->comp_sc_cur, model->map, model->size, + g->scr->attn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->comp_sc_cur, model->map, model->size, layer->indexer_compressor_gate->abs_offset, DS4_N_EMBD, index_width, - g->attn_norm, 1) != 0; + g->scr->attn_norm, 1) != 0; } const uint32_t index_row = g->layer_n_index_comp[il]; - if (ok) ok = ds4_gpu_compressor_update_tensor(g->comp_kv_cur, - g->comp_sc_cur, + if (ok) ok = ds4_gpu_compressor_update_tensor(g->scr->comp_kv_cur, + g->scr->comp_sc_cur, g->layer_index_state_kv[il], g->layer_index_state_score[il], g->layer_index_comp_cache[il], @@ -15261,14 +15353,14 @@ static bool metal_graph_encode_decode_layer( fprintf(stderr, "ds4: Metal graph indexer weight projection expects F16 weights\n"); ok = false; } - if (ok) ok = metal_graph_matmul_plain_tensor(g->indexer_q, + if (ok) ok = metal_graph_matmul_plain_tensor(g->scr->indexer_q, model, layer->indexer_attn_q_b, q_rank, indexer_q_dim, - g->qr_norm, + g->scr->qr_norm, 1); - if (ok) ok = ds4_gpu_rope_tail_tensor(g->indexer_q, 1, + if (ok) ok = ds4_gpu_rope_tail_tensor(g->scr->indexer_q, 1, DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM, DS4_N_ROT, @@ -15281,13 +15373,13 @@ static bool metal_graph_encode_decode_layer( attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(g->indexer_q, + if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(g->scr->indexer_q, DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(g->indexer_weights, model->map, model->size, + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->indexer_weights, model->map, model->size, layer->indexer_proj->abs_offset, DS4_N_EMBD, DS4_N_INDEXER_HEAD, - g->attn_norm, 1) != 0; + g->scr->attn_norm, 1) != 0; const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); if (ok && decode_index_stage_profile) { ok = metal_graph_indexer_stage_profile_boundary(NULL, @@ -15297,9 +15389,9 @@ static bool metal_graph_encode_decode_layer( g->layer_n_index_comp[il], &decode_index_stage_t0); } - if (ok) ok = ds4_gpu_indexer_score_one_tensor(g->indexer_scores, - g->indexer_q, - g->indexer_weights, + if (ok) ok = ds4_gpu_indexer_score_one_tensor(g->scr->indexer_scores, + g->scr->indexer_q, + g->scr->indexer_weights, g->layer_index_comp_cache[il], g->layer_n_index_comp[il], DS4_N_INDEXER_HEAD, @@ -15313,8 +15405,8 @@ static bool metal_graph_encode_decode_layer( g->layer_n_index_comp[il], &decode_index_stage_t0); } - if (ok) ok = ds4_gpu_indexer_topk_tensor(g->comp_selected, - g->indexer_scores, + if (ok) ok = ds4_gpu_indexer_topk_tensor(g->scr->comp_selected, + g->scr->indexer_scores, g->layer_n_index_comp[il], 1, DS4_N_INDEXER_TOP_K) != 0; @@ -15335,7 +15427,7 @@ static bool metal_graph_encode_decode_layer( * plus the selected compressed rows, matching prefill and * avoiding the long-context decode failure. */ if (ok) { - comp_selected = g->comp_selected; + comp_selected = g->scr->comp_selected; /* * Contract: the indexer top-k is fixed by the model config * and must remain the full 512 rows. Do not reduce this for @@ -15377,11 +15469,11 @@ static bool metal_graph_encode_decode_layer( const uint32_t raw_start = metal_graph_raw_start_for_span(g, pos, n_raw); if (n_comp != 0 && comp_selected != NULL && n_selected != 0) { ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor( - g->heads, + g->scr->heads, model->map, model->size, layer->attn_sinks->abs_offset, - g->q, + g->scr->q, raw_cache, g->layer_attn_comp_cache[il], metal_graph_attn_comp_cache_is_f16(), @@ -15406,10 +15498,10 @@ static bool metal_graph_encode_decode_layer( &decode_index_stage_t0); } } else { - ok = ds4_gpu_attention_decode_heads_tensor(g->heads, + ok = ds4_gpu_attention_decode_heads_tensor(g->scr->heads, model->map, model->size, layer->attn_sinks->abs_offset, - g->q, raw_cache, n_raw, + g->scr->q, raw_cache, n_raw, raw_cap, raw_start, n_comp ? comp_cache : NULL, @@ -15422,9 +15514,9 @@ static bool metal_graph_encode_decode_layer( } DS4_METAL_PROFILE_DECODE_STAGE("attention"); if (ok) { - metal_graph_debug_dump_tensor("kqv_out", g->heads, q_dim, il, pos); + metal_graph_debug_dump_tensor("kqv_out", g->scr->heads, q_dim, il, pos); } - if (ok) ok = ds4_gpu_rope_tail_tensor(g->heads, + if (ok) ok = ds4_gpu_rope_tail_tensor(g->scr->heads, 1, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_N_ROT, pos, compressed ? (uint32_t)DS4_ROPE_ORIG_CTX : 0, @@ -15436,74 +15528,74 @@ static bool metal_graph_encode_decode_layer( DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; if (ok) { - metal_graph_debug_dump_tensor("kqv_back", g->heads, q_dim, il, pos); + metal_graph_debug_dump_tensor("kqv_back", g->scr->heads, q_dim, il, pos); } const bool fuse_attn_out_hc = !metal_graph_directional_steering_attn_enabled(g) && !metal_graph_use_reference_attn_out_hc(); if (ok && fuse_attn_out_hc) { - ok = ds4_gpu_attention_output_low_q8_tensor(g->attn_low, + ok = ds4_gpu_attention_output_low_q8_tensor(g->scr->attn_low, model->map, model->size, layer->attn_output_a->abs_offset, group_dim, rank, n_groups, - g->heads) != 0; + g->scr->heads) != 0; if (ok) { - ok = ds4_gpu_matmul_q8_0_hc_expand_tensor(g->after_attn_hc, - g->attn_out, + ok = ds4_gpu_matmul_q8_0_hc_expand_tensor(g->scr->after_attn_hc, + g->scr->attn_out, model->map, model->size, layer->attn_output_b->abs_offset, (uint64_t)n_groups * rank, DS4_N_EMBD, - g->attn_low, - g->cur_hc, - g->hc_split, + g->scr->attn_low, + g->scr->cur_hc, + g->scr->hc_split, DS4_N_EMBD, DS4_N_HC) != 0; } } else if (ok) { - ok = ds4_gpu_attention_output_q8_batch_tensor(g->attn_out, - g->attn_low, - g->batch_group_tmp, - g->batch_low_tmp, + ok = ds4_gpu_attention_output_q8_batch_tensor(g->scr->attn_out, + g->scr->attn_low, + g->scr->batch_group_tmp, + g->scr->batch_low_tmp, model->map, model->size, layer->attn_output_a->abs_offset, layer->attn_output_b->abs_offset, group_dim, rank, n_groups, DS4_N_EMBD, - g->heads, 1) != 0; + g->scr->heads, 1) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("attn_output"); if (ok) { - metal_graph_debug_dump_tensor("attn_low", g->attn_low, (uint64_t)n_groups * rank, il, pos); + metal_graph_debug_dump_tensor("attn_low", g->scr->attn_low, (uint64_t)n_groups * rank, il, pos); } if (ok) { - metal_graph_debug_dump_tensor("attn_out", g->attn_out, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("attn_out", g->scr->attn_out, DS4_N_EMBD, il, pos); } if (ok && metal_graph_directional_steering_attn_enabled(g)) { - ok = metal_graph_apply_directional_steering_attn(g, g->attn_out, il, 1); + ok = metal_graph_apply_directional_steering_attn(g, g->scr->attn_out, il, 1); } if (ok && !fuse_attn_out_hc) { - ok = ds4_gpu_hc_expand_tensor(g->after_attn_hc, g->attn_out, g->cur_hc, - g->hc_post, g->hc_comb, DS4_N_EMBD, DS4_N_HC) != 0; + ok = ds4_gpu_hc_expand_tensor(g->scr->after_attn_hc, g->scr->attn_out, g->scr->cur_hc, + g->scr->hc_post, g->scr->hc_comb, DS4_N_EMBD, DS4_N_HC) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("attn_hc_post"); if (ok) { - metal_graph_debug_dump_tensor("hc_attn_post", g->after_attn_hc, hc_dim, il, pos); + metal_graph_debug_dump_tensor("hc_attn_post", g->scr->after_attn_hc, hc_dim, il, pos); } - if (ok) ok = ds4_gpu_rms_norm_plain_tensor(g->flat_hc, g->after_attn_hc, (uint32_t)hc_dim, DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(g->hc_mix, model, layer->hc_ffn_fn, - hc_dim, mix_hc, g->flat_hc, 1); + if (ok) ok = ds4_gpu_rms_norm_plain_tensor(g->scr->flat_hc, g->scr->after_attn_hc, (uint32_t)hc_dim, DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(g->scr->hc_mix, model, layer->hc_ffn_fn, + hc_dim, mix_hc, g->scr->flat_hc, 1); if (ok && fuse_hc_norm) { - ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(g->ffn_cur, - g->ffn_norm, - g->hc_split, - g->hc_mix, - g->after_attn_hc, + ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(g->scr->ffn_cur, + g->scr->ffn_norm, + g->scr->hc_split, + g->scr->hc_mix, + g->scr->after_attn_hc, model->map, model->size, layer->hc_ffn_scale->abs_offset, @@ -15516,10 +15608,10 @@ static bool metal_graph_encode_decode_layer( DS4_RMS_EPS) != 0; if (ok) { ok = metal_graph_check_hc_norm_fusion("ffn", - g->ffn_cur, - g->ffn_norm, - g->hc_mix, - g->after_attn_hc, + g->scr->ffn_cur, + g->scr->ffn_norm, + g->scr->hc_mix, + g->scr->after_attn_hc, model, layer->hc_ffn_scale->abs_offset, layer->hc_ffn_base->abs_offset, @@ -15528,31 +15620,31 @@ static bool metal_graph_encode_decode_layer( pos); } } else if (ok) { - ok = metal_graph_decode_hc_pre(g->ffn_cur, - g->hc_split, - g->hc_mix, - g->after_attn_hc, + ok = metal_graph_decode_hc_pre(g->scr->ffn_cur, + g->scr->hc_split, + g->scr->hc_mix, + g->scr->after_attn_hc, model, layer->hc_ffn_scale->abs_offset, layer->hc_ffn_base->abs_offset); } DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_pre"); if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_pre_mixes", g->hc_mix, mix_hc, il, pos); - metal_graph_debug_dump_tensor("hc_ffn_pre_weights", g->hc_pre, DS4_N_HC, il, pos); - metal_graph_debug_dump_tensor("hc_ffn_pre_post_weights", g->hc_post, DS4_N_HC, il, pos); - metal_graph_debug_dump_tensor("hc_ffn_pre_comb", g->hc_comb, (uint64_t)DS4_N_HC * DS4_N_HC, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_pre_mixes", g->scr->hc_mix, mix_hc, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_pre_weights", g->scr->hc_pre, DS4_N_HC, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_pre_post_weights", g->scr->hc_post, DS4_N_HC, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_pre_comb", g->scr->hc_comb, (uint64_t)DS4_N_HC * DS4_N_HC, il, pos); } if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_pre", g->ffn_cur, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_pre", g->scr->ffn_cur, DS4_N_EMBD, il, pos); } - if (ok && !fuse_hc_norm) ok = ds4_gpu_rms_norm_weight_tensor(g->ffn_norm, g->ffn_cur, + if (ok && !fuse_hc_norm) ok = ds4_gpu_rms_norm_weight_tensor(g->scr->ffn_norm, g->scr->ffn_cur, model->map, model->size, layer->ffn_norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; DS4_METAL_PROFILE_DECODE_STAGE("ffn_norm"); if (ok) { - metal_graph_debug_dump_tensor("ffn_norm", g->ffn_norm, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_norm", g->scr->ffn_norm, DS4_N_EMBD, il, pos); } const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); const uint64_t gate_expert_bytes = expert_mid_dim * gate_row_bytes; @@ -15561,9 +15653,9 @@ static bool metal_graph_encode_decode_layer( if (ok && metal_graph_decode_cpu_router_applicable(g, layer)) { ok = metal_graph_decode_cpu_router(g, model, layer, il, (uint32_t)token); } else { - if (ok) ok = metal_graph_matmul_plain_tensor(g->router_logits, model, layer->ffn_gate_inp, - DS4_N_EMBD, DS4_N_EXPERT, g->ffn_norm, 1); - if (ok) ok = ds4_gpu_router_select_tensor(g->router_selected, g->router_weights, g->router_probs, + if (ok) ok = metal_graph_matmul_plain_tensor(g->scr->router_logits, model, layer->ffn_gate_inp, + DS4_N_EMBD, DS4_N_EXPERT, g->scr->ffn_norm, 1); + if (ok) ok = ds4_gpu_router_select_tensor(g->scr->router_selected, g->scr->router_weights, g->scr->router_probs, model->map, model->size, layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, layer->ffn_gate_tid2eid ? layer->ffn_gate_tid2eid->abs_offset : 0, @@ -15576,7 +15668,7 @@ static bool metal_graph_encode_decode_layer( 0, layer->ffn_exp_probs_b != NULL, layer->ffn_gate_tid2eid != NULL, - g->router_logits) != 0; + g->scr->router_logits) != 0; if (ok) ok = metal_graph_decode_set_hash_selected_override(model, layer, il, @@ -15588,10 +15680,10 @@ static bool metal_graph_encode_decode_layer( DS4_METAL_PROFILE_DECODE_STAGE("router"); if (ok) ok = metal_graph_profile_router_selection(g, layer, il, pos); if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_logits", g->router_logits, DS4_N_EXPERT, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_probs", g->router_probs, DS4_N_EXPERT, il, pos); - metal_graph_debug_dump_i32_tensor("ffn_moe_topk", g->router_selected, DS4_N_EXPERT_USED, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_weights_scaled", g->router_weights, DS4_N_EXPERT_USED, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_logits", g->scr->router_logits, DS4_N_EXPERT, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_probs", g->scr->router_probs, DS4_N_EXPERT, il, pos); + metal_graph_debug_dump_i32_tensor("ffn_moe_topk", g->scr->router_selected, DS4_N_EXPERT_USED, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_weights_scaled", g->scr->router_weights, DS4_N_EXPERT_USED, il, pos); } const bool keep_ffn_out = metal_graph_needs_ffn_out(g, il, pos); #ifdef DS4_ROCM_BUILD @@ -15664,35 +15756,35 @@ static bool metal_graph_encode_decode_layer( down_expert_bytes); } if (ok && fuse_shared_gate_up) { - ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(g->shared_gate, - g->shared_up, - g->shared_mid, + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(g->scr->shared_gate, + g->scr->shared_up, + g->scr->shared_mid, model->map, model->size, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, - g->ffn_norm, + g->scr->ffn_norm, DS4_SWIGLU_CLAMP_EXP) != 0; } else if (ok) { - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->shared_gate, model->map, model->size, + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->shared_gate, model->map, model->size, layer->ffn_gate_shexp->abs_offset, DS4_N_EMBD, shared_dim, - g->ffn_norm, 1) != 0; - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->shared_up, model->map, model->size, + g->scr->ffn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->shared_up, model->map, model->size, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, - g->ffn_norm, 1) != 0; - if (ok) ok = ds4_gpu_swiglu_tensor(g->shared_mid, g->shared_gate, g->shared_up, + g->scr->ffn_norm, 1) != 0; + if (ok) ok = ds4_gpu_swiglu_tensor(g->scr->shared_mid, g->scr->shared_gate, g->scr->shared_up, shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); - if (ok) ok = ds4_gpu_routed_moe_one_tensor(g->routed_out, - g->routed_gate, - g->routed_up, - g->routed_mid, - g->routed_down, + if (ok) ok = ds4_gpu_routed_moe_one_tensor(g->scr->routed_out, + g->scr->routed_gate, + g->scr->routed_up, + g->scr->routed_mid, + g->scr->routed_down, model->map, model->size, layer->ffn_gate_exps->abs_offset, layer->ffn_up_exps->abs_offset, @@ -15704,76 +15796,76 @@ static bool metal_graph_encode_decode_layer( (uint32_t)expert_in_dim, (uint32_t)down_in_dim, (uint32_t)routed_out_dim, - g->router_selected, g->router_weights, + g->scr->router_selected, g->scr->router_weights, DS4_N_EXPERT, - DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, g->ffn_norm, + DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, g->scr->ffn_norm, il) != 0; DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", g->routed_gate, + metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", g->scr->routed_gate, (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_up_clamped", g->routed_up, + metal_graph_debug_dump_tensor("ffn_moe_up_clamped", g->scr->routed_up, (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", g->routed_mid, + metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", g->scr->routed_mid, (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_down", g->routed_down, + metal_graph_debug_dump_tensor("ffn_moe_down", g->scr->routed_down, (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_out", g->routed_out, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_out", g->scr->routed_out, DS4_N_EMBD, il, pos); } if (ok && fuse_shared_down_hc) { - ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(g->after_ffn_hc, - g->shared_out, + ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(g->scr->after_ffn_hc, + g->scr->shared_out, model->map, model->size, layer->ffn_down_shexp->abs_offset, shared_dim, DS4_N_EMBD, - g->shared_mid, - g->routed_out, - g->after_attn_hc, - g->hc_split, + g->scr->shared_mid, + g->scr->routed_out, + g->scr->after_attn_hc, + g->scr->hc_split, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok) { - ok = ds4_gpu_matmul_q8_0_tensor(g->shared_out, model->map, model->size, + ok = ds4_gpu_matmul_q8_0_tensor(g->scr->shared_out, model->map, model->size, layer->ffn_down_shexp->abs_offset, shared_dim, DS4_N_EMBD, - g->shared_mid, 1) != 0; + g->scr->shared_mid, 1) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); if (ok) { - metal_graph_debug_dump_tensor("ffn_shexp", g->shared_out, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_shexp", g->scr->shared_out, DS4_N_EMBD, il, pos); } if (ok && keep_ffn_out) { ok = metal_graph_ensure_ffn_out(g) && - ds4_gpu_add_tensor(g->ffn_out, g->shared_out, g->routed_out, DS4_N_EMBD) != 0; + ds4_gpu_add_tensor(g->scr->ffn_out, g->scr->shared_out, g->scr->routed_out, DS4_N_EMBD) != 0; } if (ok && keep_ffn_out) { - metal_graph_debug_dump_tensor("ffn_out", g->ffn_out, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_out", g->scr->ffn_out, DS4_N_EMBD, il, pos); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = metal_graph_apply_directional_steering_ffn(g, g->ffn_out, il, 1); + ok = metal_graph_apply_directional_steering_ffn(g, g->scr->ffn_out, il, 1); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = ds4_gpu_hc_expand_tensor(g->after_ffn_hc, - g->ffn_out, - g->after_attn_hc, - g->hc_post, - g->hc_comb, + ok = ds4_gpu_hc_expand_tensor(g->scr->after_ffn_hc, + g->scr->ffn_out, + g->scr->after_attn_hc, + g->scr->hc_post, + g->scr->hc_comb, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok && !fuse_shared_down_hc) { - ok = ds4_gpu_hc_expand_add_split_tensor(g->after_ffn_hc, - g->routed_out, - g->shared_out, - g->after_attn_hc, - g->hc_split, + ok = ds4_gpu_hc_expand_add_split_tensor(g->scr->after_ffn_hc, + g->scr->routed_out, + g->scr->shared_out, + g->scr->after_attn_hc, + g->scr->hc_split, DS4_N_EMBD, DS4_N_HC) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_post", g->after_ffn_hc, hc_dim, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_post", g->scr->after_ffn_hc, hc_dim, il, pos); } return ok; } @@ -15800,35 +15892,35 @@ static bool metal_graph_encode_decode_layer( ok = ds4_gpu_flush_commands() != 0; } if (ok && fuse_shared_gate_up) { - ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(g->shared_gate, - g->shared_up, - g->shared_mid, + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(g->scr->shared_gate, + g->scr->shared_up, + g->scr->shared_mid, model->map, model->size, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, - g->ffn_norm, + g->scr->ffn_norm, DS4_SWIGLU_CLAMP_EXP) != 0; } else if (ok) { - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->shared_gate, model->map, model->size, + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->shared_gate, model->map, model->size, layer->ffn_gate_shexp->abs_offset, DS4_N_EMBD, shared_dim, - g->ffn_norm, 1) != 0; - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->shared_up, model->map, model->size, + g->scr->ffn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->shared_up, model->map, model->size, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, - g->ffn_norm, 1) != 0; - if (ok) ok = ds4_gpu_swiglu_tensor(g->shared_mid, g->shared_gate, g->shared_up, + g->scr->ffn_norm, 1) != 0; + if (ok) ok = ds4_gpu_swiglu_tensor(g->scr->shared_mid, g->scr->shared_gate, g->scr->shared_up, shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); if (ok && !fuse_shared_down_hc) { - ok = ds4_gpu_matmul_q8_0_tensor(g->shared_out, model->map, model->size, + ok = ds4_gpu_matmul_q8_0_tensor(g->scr->shared_out, model->map, model->size, layer->ffn_down_shexp->abs_offset, shared_dim, DS4_N_EMBD, - g->shared_mid, 1) != 0; + g->scr->shared_mid, 1) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); if (async_load_started) { @@ -15842,7 +15934,7 @@ static bool metal_graph_encode_decode_layer( } if (ok && !async_load_started) { int32_t selected_ids[DS4_MAX_EXPERT_USED]; - ok = ds4_gpu_tensor_read(g->router_selected, + ok = ds4_gpu_tensor_read(g->scr->router_selected, 0, selected_ids, (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_ids[0])) != 0 && @@ -15861,11 +15953,11 @@ static bool metal_graph_encode_decode_layer( DS4_N_EXPERT_USED) != 0; } } - if (ok) ok = ds4_gpu_routed_moe_one_tensor(g->routed_out, - g->routed_gate, - g->routed_up, - g->routed_mid, - g->routed_down, + if (ok) ok = ds4_gpu_routed_moe_one_tensor(g->scr->routed_out, + g->scr->routed_gate, + g->scr->routed_up, + g->scr->routed_mid, + g->scr->routed_down, model->map, model->size, layer->ffn_gate_exps->abs_offset, layer->ffn_up_exps->abs_offset, @@ -15877,79 +15969,79 @@ static bool metal_graph_encode_decode_layer( (uint32_t)expert_in_dim, (uint32_t)down_in_dim, (uint32_t)routed_out_dim, - g->router_selected, g->router_weights, + g->scr->router_selected, g->scr->router_weights, DS4_N_EXPERT, - DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, g->ffn_norm, + DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, g->scr->ffn_norm, il) != 0; DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", g->routed_gate, + metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", g->scr->routed_gate, (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_up_clamped", g->routed_up, + metal_graph_debug_dump_tensor("ffn_moe_up_clamped", g->scr->routed_up, (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", g->routed_mid, + metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", g->scr->routed_mid, (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_down", g->routed_down, + metal_graph_debug_dump_tensor("ffn_moe_down", g->scr->routed_down, (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_out", g->routed_out, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_out", g->scr->routed_out, DS4_N_EMBD, il, pos); } if (ok && fuse_shared_down_hc) { - ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(g->after_ffn_hc, - g->shared_out, + ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(g->scr->after_ffn_hc, + g->scr->shared_out, model->map, model->size, layer->ffn_down_shexp->abs_offset, shared_dim, DS4_N_EMBD, - g->shared_mid, - g->routed_out, - g->after_attn_hc, - g->hc_split, + g->scr->shared_mid, + g->scr->routed_out, + g->scr->after_attn_hc, + g->scr->hc_split, DS4_N_EMBD, DS4_N_HC) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); if (ok) { - metal_graph_debug_dump_tensor("ffn_shexp", g->shared_out, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_shexp", g->scr->shared_out, DS4_N_EMBD, il, pos); } if (ok && keep_ffn_out) { ok = metal_graph_ensure_ffn_out(g) && - ds4_gpu_add_tensor(g->ffn_out, g->shared_out, g->routed_out, DS4_N_EMBD) != 0; + ds4_gpu_add_tensor(g->scr->ffn_out, g->scr->shared_out, g->scr->routed_out, DS4_N_EMBD) != 0; } if (ok && keep_ffn_out) { - metal_graph_debug_dump_tensor("ffn_out", g->ffn_out, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_out", g->scr->ffn_out, DS4_N_EMBD, il, pos); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = metal_graph_apply_directional_steering_ffn(g, g->ffn_out, il, 1); + ok = metal_graph_apply_directional_steering_ffn(g, g->scr->ffn_out, il, 1); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = ds4_gpu_hc_expand_tensor(g->after_ffn_hc, - g->ffn_out, - g->after_attn_hc, - g->hc_post, - g->hc_comb, + ok = ds4_gpu_hc_expand_tensor(g->scr->after_ffn_hc, + g->scr->ffn_out, + g->scr->after_attn_hc, + g->scr->hc_post, + g->scr->hc_comb, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok && !fuse_shared_down_hc) { - ok = ds4_gpu_hc_expand_add_split_tensor(g->after_ffn_hc, - g->routed_out, - g->shared_out, - g->after_attn_hc, - g->hc_split, + ok = ds4_gpu_hc_expand_add_split_tensor(g->scr->after_ffn_hc, + g->scr->routed_out, + g->scr->shared_out, + g->scr->after_attn_hc, + g->scr->hc_split, DS4_N_EMBD, DS4_N_HC) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_post", g->after_ffn_hc, hc_dim, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_post", g->scr->after_ffn_hc, hc_dim, il, pos); } return ok; } - if (ok) ok = ds4_gpu_routed_moe_one_tensor(g->routed_out, - g->routed_gate, - g->routed_up, - g->routed_mid, - g->routed_down, + if (ok) ok = ds4_gpu_routed_moe_one_tensor(g->scr->routed_out, + g->scr->routed_gate, + g->scr->routed_up, + g->scr->routed_mid, + g->scr->routed_down, model->map, model->size, layer->ffn_gate_exps->abs_offset, layer->ffn_up_exps->abs_offset, @@ -15961,108 +16053,108 @@ static bool metal_graph_encode_decode_layer( (uint32_t)expert_in_dim, (uint32_t)down_in_dim, (uint32_t)routed_out_dim, - g->router_selected, g->router_weights, + g->scr->router_selected, g->scr->router_weights, DS4_N_EXPERT, - DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, g->ffn_norm, + DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, g->scr->ffn_norm, il) != 0; DS4_METAL_PROFILE_DECODE_STAGE("routed_moe"); if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", g->routed_gate, + metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", g->scr->routed_gate, (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); - metal_graph_debug_dump_tensor("ffn_moe_up_clamped", g->routed_up, + metal_graph_debug_dump_tensor("ffn_moe_up_clamped", g->scr->routed_up, (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); } if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", g->routed_mid, + metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", g->scr->routed_mid, (uint64_t)DS4_N_EXPERT_USED * down_in_dim, il, pos); } if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_down", g->routed_down, + metal_graph_debug_dump_tensor("ffn_moe_down", g->scr->routed_down, (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); } if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_out", g->routed_out, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_moe_out", g->scr->routed_out, DS4_N_EMBD, il, pos); } if (ok && fuse_shared_gate_up) { - ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(g->shared_gate, - g->shared_up, - g->shared_mid, + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(g->scr->shared_gate, + g->scr->shared_up, + g->scr->shared_mid, model->map, model->size, layer->ffn_gate_shexp->abs_offset, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, - g->ffn_norm, + g->scr->ffn_norm, DS4_SWIGLU_CLAMP_EXP) != 0; } else { - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->shared_gate, model->map, model->size, + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->shared_gate, model->map, model->size, layer->ffn_gate_shexp->abs_offset, DS4_N_EMBD, shared_dim, - g->ffn_norm, 1) != 0; - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->shared_up, model->map, model->size, + g->scr->ffn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->shared_up, model->map, model->size, layer->ffn_up_shexp->abs_offset, DS4_N_EMBD, shared_dim, - g->ffn_norm, 1) != 0; - if (ok) ok = ds4_gpu_swiglu_tensor(g->shared_mid, g->shared_gate, g->shared_up, + g->scr->ffn_norm, 1) != 0; + if (ok) ok = ds4_gpu_swiglu_tensor(g->scr->shared_mid, g->scr->shared_gate, g->scr->shared_up, shared_dim, DS4_SWIGLU_CLAMP_EXP, 1.0f) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("shared_gate_up"); if (ok && fuse_shared_down_hc) { - ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(g->after_ffn_hc, - g->shared_out, + ok = ds4_gpu_shared_down_hc_expand_q8_0_tensor(g->scr->after_ffn_hc, + g->scr->shared_out, model->map, model->size, layer->ffn_down_shexp->abs_offset, shared_dim, DS4_N_EMBD, - g->shared_mid, - g->routed_out, - g->after_attn_hc, - g->hc_split, + g->scr->shared_mid, + g->scr->routed_out, + g->scr->after_attn_hc, + g->scr->hc_split, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok) { - ok = ds4_gpu_matmul_q8_0_tensor(g->shared_out, model->map, model->size, + ok = ds4_gpu_matmul_q8_0_tensor(g->scr->shared_out, model->map, model->size, layer->ffn_down_shexp->abs_offset, shared_dim, DS4_N_EMBD, - g->shared_mid, 1) != 0; + g->scr->shared_mid, 1) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("shared_down"); if (ok) { - metal_graph_debug_dump_tensor("ffn_shexp", g->shared_out, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_shexp", g->scr->shared_out, DS4_N_EMBD, il, pos); } if (ok && keep_ffn_out) { ok = metal_graph_ensure_ffn_out(g) && - ds4_gpu_add_tensor(g->ffn_out, g->shared_out, g->routed_out, DS4_N_EMBD) != 0; + ds4_gpu_add_tensor(g->scr->ffn_out, g->scr->shared_out, g->scr->routed_out, DS4_N_EMBD) != 0; } if (ok && keep_ffn_out) { - metal_graph_debug_dump_tensor("ffn_out", g->ffn_out, DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor("ffn_out", g->scr->ffn_out, DS4_N_EMBD, il, pos); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = metal_graph_apply_directional_steering_ffn(g, g->ffn_out, il, 1); + ok = metal_graph_apply_directional_steering_ffn(g, g->scr->ffn_out, il, 1); } if (ok && metal_graph_directional_steering_ffn_enabled(g)) { - ok = ds4_gpu_hc_expand_tensor(g->after_ffn_hc, - g->ffn_out, - g->after_attn_hc, - g->hc_post, - g->hc_comb, + ok = ds4_gpu_hc_expand_tensor(g->scr->after_ffn_hc, + g->scr->ffn_out, + g->scr->after_attn_hc, + g->scr->hc_post, + g->scr->hc_comb, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok && !fuse_shared_down_hc) { - ok = ds4_gpu_hc_expand_add_split_tensor(g->after_ffn_hc, - g->routed_out, - g->shared_out, - g->after_attn_hc, - g->hc_split, + ok = ds4_gpu_hc_expand_add_split_tensor(g->scr->after_ffn_hc, + g->scr->routed_out, + g->scr->shared_out, + g->scr->after_attn_hc, + g->scr->hc_split, DS4_N_EMBD, DS4_N_HC) != 0; } DS4_METAL_PROFILE_DECODE_STAGE("ffn_hc_post"); #undef DS4_METAL_PROFILE_DECODE_STAGE if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_post", g->after_ffn_hc, hc_dim, il, pos); + metal_graph_debug_dump_tensor("hc_ffn_post", g->scr->after_ffn_hc, hc_dim, il, pos); } return ok; } @@ -16074,20 +16166,20 @@ static bool metal_graph_encode_output_head( const ds4_weights *weights, uint64_t vocab_dim) { const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - bool ok = ds4_gpu_rms_norm_plain_tensor(g->flat_hc, g->cur_hc, (uint32_t)hc_dim, DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(g->output_pre, + bool ok = ds4_gpu_rms_norm_plain_tensor(g->scr->flat_hc, g->scr->cur_hc, (uint32_t)hc_dim, DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->output_pre, model->map, model->size, weights->output_hc_fn->abs_offset, hc_dim, DS4_N_HC, - g->flat_hc, + g->scr->flat_hc, 1) != 0; if (ok) { - metal_graph_debug_dump_tensor("result_hc_pre", g->output_pre, DS4_N_HC, DS4_N_LAYER, 0); + metal_graph_debug_dump_tensor("result_hc_pre", g->scr->output_pre, DS4_N_HC, DS4_N_LAYER, 0); } - if (ok) ok = ds4_gpu_output_hc_weights_tensor(g->output_weights, - g->output_pre, + if (ok) ok = ds4_gpu_output_hc_weights_tensor(g->scr->output_weights, + g->scr->output_pre, model->map, model->size, weights->output_hc_scale->abs_offset, @@ -16095,36 +16187,36 @@ static bool metal_graph_encode_output_head( DS4_N_HC, DS4_HC_EPS) != 0; if (ok) { - metal_graph_debug_dump_tensor("result_hc_weights", g->output_weights, DS4_N_HC, DS4_N_LAYER, 0); + metal_graph_debug_dump_tensor("result_hc_weights", g->scr->output_weights, DS4_N_HC, DS4_N_LAYER, 0); } - if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(g->output_embd, - g->cur_hc, - g->output_weights, + if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(g->scr->output_embd, + g->scr->cur_hc, + g->scr->output_weights, DS4_N_EMBD, DS4_N_HC) != 0; if (ok) { - metal_graph_debug_dump_tensor("result_hc", g->output_embd, DS4_N_EMBD, DS4_N_LAYER, 0); + metal_graph_debug_dump_tensor("result_hc", g->scr->output_embd, DS4_N_EMBD, DS4_N_LAYER, 0); } - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->output_norm, - g->output_embd, + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->scr->output_norm, + g->scr->output_embd, model->map, model->size, weights->output_norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; if (ok) { - metal_graph_debug_dump_tensor("result_norm", g->output_norm, DS4_N_EMBD, DS4_N_LAYER, 0); + metal_graph_debug_dump_tensor("result_norm", g->scr->output_norm, DS4_N_EMBD, DS4_N_LAYER, 0); } - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->logits, + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->logits, model->map, model->size, weights->output->abs_offset, DS4_N_EMBD, vocab_dim, - g->output_norm, + g->scr->output_norm, 1) != 0; if (ok) { - metal_graph_debug_dump_tensor("result_output", g->logits, vocab_dim, DS4_N_LAYER, 0); + metal_graph_debug_dump_tensor("result_output", g->scr->logits, vocab_dim, DS4_N_LAYER, 0); } return ok; } @@ -16143,7 +16235,7 @@ static bool metal_graph_encode_output_head_batch( const ds4_weights *weights, uint32_t n_tokens, uint64_t vocab_dim) { - if (n_tokens == 0 || n_tokens > g->prefill_cap || !g->spec_logits) return false; + if (n_tokens == 0 || n_tokens > g->prefill_cap || !g->scr->spec_logits) return false; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; ds4_gpu_tensor *output_pre = NULL; @@ -16153,25 +16245,25 @@ static bool metal_graph_encode_output_head_batch( ds4_gpu_tensor *logits = NULL; bool ok = true; - output_pre = ds4_gpu_tensor_view(g->batch_hc_mix, + output_pre = ds4_gpu_tensor_view(g->scr->batch_hc_mix, 0, (uint64_t)n_tokens * DS4_N_HC * sizeof(float)); - output_weights = ds4_gpu_tensor_view(g->batch_hc_split, + output_weights = ds4_gpu_tensor_view(g->scr->batch_hc_split, 0, (uint64_t)n_tokens * DS4_N_HC * sizeof(float)); - output_embd = ds4_gpu_tensor_view(g->batch_ffn_cur, + output_embd = ds4_gpu_tensor_view(g->scr->batch_ffn_cur, 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); - output_norm = ds4_gpu_tensor_view(g->batch_ffn_norm, + output_norm = ds4_gpu_tensor_view(g->scr->batch_ffn_norm, 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); - logits = ds4_gpu_tensor_view(g->spec_logits, + logits = ds4_gpu_tensor_view(g->scr->spec_logits, 0, (uint64_t)n_tokens * vocab_dim * sizeof(float)); ok = output_pre && output_weights && output_embd && output_norm && logits; - if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(g->batch_flat_hc, - g->batch_cur_hc, + if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(g->scr->batch_flat_hc, + g->scr->batch_cur_hc, (uint32_t)hc_dim, n_tokens, DS4_RMS_EPS) != 0; @@ -16181,7 +16273,7 @@ static bool metal_graph_encode_output_head_batch( weights->output_hc_fn->abs_offset, hc_dim, DS4_N_HC, - g->batch_flat_hc, + g->scr->batch_flat_hc, n_tokens) != 0; if (ok) ok = ds4_gpu_output_hc_weights_tensor(output_weights, output_pre, @@ -16192,7 +16284,7 @@ static bool metal_graph_encode_output_head_batch( DS4_N_HC, DS4_HC_EPS) != 0; if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(output_embd, - g->batch_cur_hc, + g->scr->batch_cur_hc, output_weights, DS4_N_EMBD, DS4_N_HC) != 0; @@ -16278,36 +16370,36 @@ static bool metal_graph_encode_output_head_mtp( const ds4_mtp_weights *mtp, uint64_t vocab_dim) { const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - bool ok = ds4_gpu_rms_norm_plain_tensor(g->flat_hc, g->cur_hc, (uint32_t)hc_dim, DS4_RMS_EPS) != 0; - if (ok) ok = metal_graph_matmul_plain_tensor(g->output_pre, mtp_model, mtp->hc_head_fn, - hc_dim, DS4_N_HC, g->flat_hc, 1); - if (ok) ok = ds4_gpu_output_hc_weights_tensor(g->output_weights, - g->output_pre, + bool ok = ds4_gpu_rms_norm_plain_tensor(g->scr->flat_hc, g->scr->cur_hc, (uint32_t)hc_dim, DS4_RMS_EPS) != 0; + if (ok) ok = metal_graph_matmul_plain_tensor(g->scr->output_pre, mtp_model, mtp->hc_head_fn, + hc_dim, DS4_N_HC, g->scr->flat_hc, 1); + if (ok) ok = ds4_gpu_output_hc_weights_tensor(g->scr->output_weights, + g->scr->output_pre, mtp_model->map, mtp_model->size, mtp->hc_head_scale->abs_offset, mtp->hc_head_base->abs_offset, DS4_N_HC, DS4_HC_EPS) != 0; - if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(g->output_embd, - g->cur_hc, - g->output_weights, + if (ok) ok = ds4_gpu_hc_weighted_sum_tensor(g->scr->output_embd, + g->scr->cur_hc, + g->scr->output_weights, DS4_N_EMBD, DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->output_norm, - g->output_embd, + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->scr->output_norm, + g->scr->output_embd, mtp_model->map, mtp_model->size, mtp->norm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->logits, + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->logits, base_model->map, base_model->size, base_weights->output->abs_offset, DS4_N_EMBD, vocab_dim, - g->output_norm, + g->scr->output_norm, 1) != 0; return ok; } @@ -16436,24 +16528,24 @@ static void metal_graph_trace_layer_stages( int gpu_selected[DS4_MAX_EXPERT_USED]; float gpu_expert_weight[DS4_MAX_EXPERT_USED]; - bool ok = ds4_gpu_tensor_read(g->attn_cur, 0, gpu_attn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->attn_norm, 0, gpu_attn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->q, 0, gpu_q, q_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->kv, 0, gpu_kv, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->attn_out, 0, gpu_attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->after_attn_hc, 0, gpu_after_attn_hc, hc_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->ffn_cur, 0, gpu_ffn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->ffn_norm, 0, gpu_ffn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->shared_gate, 0, gpu_shared_gate, shared_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->shared_up, 0, gpu_shared_up, shared_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->shared_mid, 0, gpu_shared_mid, shared_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->shared_out, 0, gpu_shared, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->router_selected, 0, gpu_selected, sizeof(gpu_selected)) != 0 && - ds4_gpu_tensor_read(g->router_weights, 0, gpu_expert_weight, sizeof(gpu_expert_weight)) != 0 && - ds4_gpu_tensor_read(g->routed_mid, 0, gpu_routed_mid_all, (uint64_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->routed_out, 0, gpu_routed, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->ffn_out, 0, gpu_ffn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g->cur_hc, 0, gpu_after_ffn_hc, hc_dim * sizeof(float)) != 0; + bool ok = ds4_gpu_tensor_read(g->scr->attn_cur, 0, gpu_attn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->attn_norm, 0, gpu_attn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->q, 0, gpu_q, q_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->kv, 0, gpu_kv, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->attn_out, 0, gpu_attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->after_attn_hc, 0, gpu_after_attn_hc, hc_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->ffn_cur, 0, gpu_ffn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->ffn_norm, 0, gpu_ffn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->shared_gate, 0, gpu_shared_gate, shared_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->shared_up, 0, gpu_shared_up, shared_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->shared_mid, 0, gpu_shared_mid, shared_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->shared_out, 0, gpu_shared, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->router_selected, 0, gpu_selected, sizeof(gpu_selected)) != 0 && + ds4_gpu_tensor_read(g->scr->router_weights, 0, gpu_expert_weight, sizeof(gpu_expert_weight)) != 0 && + ds4_gpu_tensor_read(g->scr->routed_mid, 0, gpu_routed_mid_all, (uint64_t)DS4_N_EXPERT_USED * down_in_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->routed_out, 0, gpu_routed, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->ffn_out, 0, gpu_ffn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g->scr->cur_hc, 0, gpu_after_ffn_hc, hc_dim * sizeof(float)) != 0; if (ok) { fprintf(stderr, @@ -16650,11 +16742,13 @@ static int metal_graph_decode_test( output_logits_one(cpu_logits, model, weights, cpu_after_ffn_hc); ds4_gpu_graph g; - bool ok = metal_graph_alloc(&g, weights, layer); + ds4_gpu_scratch scr; + memset(&scr, 0, sizeof(scr)); + bool ok = metal_graph_alloc(&g, &scr, weights, layer); g.quality = quality; - g.materialize_ffn_out = true; + if (ok) g.scr->materialize_ffn_out = true; if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_embed_token_hc_tensor(g.cur_hc, + if (ok) ok = ds4_gpu_embed_token_hc_tensor(g.scr->cur_hc, model->map, model->size, weights->token_embd->abs_offset, @@ -16673,31 +16767,31 @@ static int metal_graph_decode_test( 1, token); if (ok) { - ds4_gpu_tensor *embedded_hc = g.cur_hc; - g.cur_hc = g.after_ffn_hc; - g.after_ffn_hc = embedded_hc; + ds4_gpu_tensor *embedded_hc = g.scr->cur_hc; + g.scr->cur_hc = g.scr->after_ffn_hc; + g.scr->after_ffn_hc = embedded_hc; } if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); if (ok) ok = ds4_gpu_end_commands() != 0; if (ok) { - ok = ds4_gpu_tensor_read(g.after_ffn_hc, 0, gpu_hc, hc_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.attn_cur, 0, gpu_attn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.attn_norm, 0, gpu_attn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.q, 0, gpu_q, q_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.kv, 0, gpu_kv, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && + ok = ds4_gpu_tensor_read(g.scr->after_ffn_hc, 0, gpu_hc, hc_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->attn_cur, 0, gpu_attn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->attn_norm, 0, gpu_attn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->q, 0, gpu_q, q_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->kv, 0, gpu_kv, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && ds4_gpu_tensor_read(g.layer_raw_cache[0], 0, gpu_raw, (uint64_t)DS4_N_HEAD_DIM * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.attn_out, 0, gpu_attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.after_attn_hc, 0, gpu_after_attn_hc, hc_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.ffn_cur, 0, gpu_ffn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.ffn_norm, 0, gpu_ffn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.shared_out, 0, gpu_shared, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.router_selected, 0, gpu_selected, sizeof(gpu_selected)) != 0 && - ds4_gpu_tensor_read(g.router_weights, 0, gpu_expert_weight, sizeof(gpu_expert_weight)) != 0 && - ds4_gpu_tensor_read(g.routed_out, 0, gpu_routed, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.ffn_out, 0, gpu_ffn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.cur_hc, 0, gpu_after_ffn_hc, hc_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.logits, 0, gpu_logits, vocab_dim * sizeof(float)) != 0; + ds4_gpu_tensor_read(g.scr->attn_out, 0, gpu_attn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->after_attn_hc, 0, gpu_after_attn_hc, hc_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->ffn_cur, 0, gpu_ffn_cur, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->ffn_norm, 0, gpu_ffn_norm, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->shared_out, 0, gpu_shared, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->router_selected, 0, gpu_selected, sizeof(gpu_selected)) != 0 && + ds4_gpu_tensor_read(g.scr->router_weights, 0, gpu_expert_weight, sizeof(gpu_expert_weight)) != 0 && + ds4_gpu_tensor_read(g.scr->routed_out, 0, gpu_routed, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->ffn_out, 0, gpu_ffn_out, (uint64_t)DS4_N_EMBD * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->cur_hc, 0, gpu_after_ffn_hc, hc_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->logits, 0, gpu_logits, vocab_dim * sizeof(float)) != 0; } if (ok) { @@ -16736,6 +16830,8 @@ static int metal_graph_decode_test( } metal_graph_free(&g); + + metal_scratch_free(&scr); free(routed_midq); free(routed_xq); free(routed_mid_all); @@ -16800,11 +16896,13 @@ static int metal_graph_first_token_full_test( output_logits_one(cpu_logits, model, weights, cpu_hc); ds4_gpu_graph g; - bool ok = metal_graph_alloc(&g, weights, &weights->layer[0]); + ds4_gpu_scratch scr; + memset(&scr, 0, sizeof(scr)); + bool ok = metal_graph_alloc(&g, &scr, weights, &weights->layer[0]); g.quality = quality; const bool trace_layers = getenv("DS4_METAL_GRAPH_TRACE_LAYERS") != NULL; if (trace_layers && ok) { - g.materialize_ffn_out = true; + g.scr->materialize_ffn_out = true; const bool teacher_force = getenv("DS4_METAL_GRAPH_TEACHER_FORCE") != NULL; const char *stage_layer_env = getenv("DS4_METAL_GRAPH_TRACE_STAGE_LAYER"); const long stage_layer = stage_layer_env ? strtol(stage_layer_env, NULL, 10) : -1; @@ -16815,7 +16913,7 @@ static int metal_graph_first_token_full_test( embed_token_f16(model, weights, token, plain); hc_from_plain_embedding(cpu_cur, plain, DS4_N_EMBD, DS4_N_HC); ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_embed_token_hc_tensor(g.cur_hc, + if (ok) ok = ds4_gpu_embed_token_hc_tensor(g.scr->cur_hc, model->map, model->size, weights->token_embd->abs_offset, @@ -16827,18 +16925,18 @@ static int metal_graph_first_token_full_test( for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { if (teacher_force) { - ok = ds4_gpu_tensor_write(g.cur_hc, 0, cpu_cur, hc_dim * sizeof(float)) != 0; + ok = ds4_gpu_tensor_write(g.scr->cur_hc, 0, cpu_cur, hc_dim * sizeof(float)) != 0; } ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_decode_layer(&g, model, &weights->layer[il], il, 0, g.layer_raw_cache[il], g.raw_cap, 0, 1, token); - ds4_gpu_tensor *tmp = g.cur_hc; - g.cur_hc = g.after_ffn_hc; - g.after_ffn_hc = tmp; + ds4_gpu_tensor *tmp = g.scr->cur_hc; + g.scr->cur_hc = g.scr->after_ffn_hc; + g.scr->after_ffn_hc = tmp; if (ok) ok = ds4_gpu_end_commands() != 0; layer_forward_self_one(cpu_next, model, &weights->layer[il], cpu_cur, il, 0, token); - if (ok) ok = ds4_gpu_tensor_read(g.cur_hc, 0, gpu_hc, hc_dim * sizeof(float)) != 0; + if (ok) ok = ds4_gpu_tensor_read(g.scr->cur_hc, 0, gpu_hc, hc_dim * sizeof(float)) != 0; if (ok) { fprintf(stderr, "ds4: Metal full graph layer %u%s hc_max=%g hc_rms=%g\n", @@ -16864,7 +16962,7 @@ static int metal_graph_first_token_full_test( free(plain); } else { if (ok) ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_embed_token_hc_tensor(g.cur_hc, + if (ok) ok = ds4_gpu_embed_token_hc_tensor(g.scr->cur_hc, model->map, model->size, weights->token_embd->abs_offset, @@ -16877,9 +16975,9 @@ static int metal_graph_first_token_full_test( ok = metal_graph_encode_decode_layer(&g, model, &weights->layer[il], il, 0, g.layer_raw_cache[il], g.raw_cap, 0, 1, token); - ds4_gpu_tensor *tmp = g.cur_hc; - g.cur_hc = g.after_ffn_hc; - g.after_ffn_hc = tmp; + ds4_gpu_tensor *tmp = g.scr->cur_hc; + g.scr->cur_hc = g.scr->after_ffn_hc; + g.scr->after_ffn_hc = tmp; } if (ok) ok = metal_graph_encode_output_head(&g, model, weights, vocab_dim); @@ -16887,8 +16985,8 @@ static int metal_graph_first_token_full_test( } if (ok) { - ok = ds4_gpu_tensor_read(g.cur_hc, 0, gpu_hc, hc_dim * sizeof(float)) != 0 && - ds4_gpu_tensor_read(g.logits, 0, gpu_logits, vocab_dim * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g.scr->cur_hc, 0, gpu_hc, hc_dim * sizeof(float)) != 0 && + ds4_gpu_tensor_read(g.scr->logits, 0, gpu_logits, vocab_dim * sizeof(float)) != 0; } if (ok) { @@ -16912,6 +17010,8 @@ static int metal_graph_first_token_full_test( } metal_graph_free(&g); + + metal_scratch_free(&scr); free(gpu_logits); free(cpu_logits); free(gpu_hc); @@ -16958,7 +17058,7 @@ static bool metal_graph_encode_token_raw_swa( const uint32_t raw_row = pos % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); - bool ok = ds4_gpu_embed_token_hc_tensor(g->cur_hc, + bool ok = ds4_gpu_embed_token_hc_tensor(g->scr->cur_hc, model->map, model->size, weights->token_embd->abs_offset, @@ -16987,9 +17087,9 @@ static bool metal_graph_encode_token_raw_swa( raw_row, n_raw, token); - ds4_gpu_tensor *tmp = g->cur_hc; - g->cur_hc = g->after_ffn_hc; - g->after_ffn_hc = tmp; + ds4_gpu_tensor *tmp = g->scr->cur_hc; + g->scr->cur_hc = g->scr->after_ffn_hc; + g->scr->after_ffn_hc = tmp; if (ok && allow_split_flush && split_after_layers != 0 && il + 1u == split_after_layers) { ok = ds4_gpu_flush_commands() != 0; } @@ -17061,12 +17161,12 @@ static bool metal_graph_refresh_ratio4_compressor_state( * decisions in later chunks. */ ds4_gpu_tensor *tail_hc = ds4_gpu_tensor_view( - g->batch_attn_norm, + g->scr->batch_attn_norm, (uint64_t)(n_tokens - 4u) * DS4_N_EMBD * sizeof(float), 4ull * DS4_N_EMBD * sizeof(float)); bool ok = tail_hc != NULL; if (ok) { - ok = ds4_gpu_matmul_f16_tensor(g->batch_comp_kv, + ok = ds4_gpu_matmul_f16_tensor(g->scr->batch_comp_kv, model->map, model->size, kv_weight->abs_offset, @@ -17076,7 +17176,7 @@ static bool metal_graph_refresh_ratio4_compressor_state( 4) != 0; } if (ok) { - ok = ds4_gpu_matmul_f16_tensor(g->batch_comp_sc, + ok = ds4_gpu_matmul_f16_tensor(g->scr->batch_comp_sc, model->map, model->size, score_weight->abs_offset, @@ -17088,8 +17188,8 @@ static bool metal_graph_refresh_ratio4_compressor_state( if (ok) { ok = ds4_gpu_compressor_prefill_state_ratio4_tensor(state_kv, state_score, - g->batch_comp_kv, - g->batch_comp_sc, + g->scr->batch_comp_kv, + g->scr->batch_comp_sc, model->map, model->size, ape->abs_offset, @@ -17200,13 +17300,13 @@ static bool metal_graph_warmup_prefill_kernels( bool ok = ds4_gpu_begin_commands() != 0; if (ok) { - ok = ds4_gpu_matmul_f16_tensor(g->batch_hc_mix, + ok = ds4_gpu_matmul_f16_tensor(g->scr->batch_hc_mix, model->map, model->size, weights->layer[0].hc_attn_fn->abs_offset, hc_dim, mix_hc, - g->batch_flat_hc, + g->scr->batch_flat_hc, n_tokens) != 0; } if (ok) ok = ds4_gpu_end_commands() != 0; @@ -17356,19 +17456,19 @@ static bool metal_graph_encode_layer_attention_batch( uint32_t *index_counts = ratio == 4 ? xcalloc(n_tokens, sizeof(index_counts[0])) : NULL; const bool qkv_rms_fused = !metal_graph_use_reference_qkv_norm(); ds4_gpu_tensor *hc_mix_view = ds4_gpu_tensor_view( - g->batch_hc_mix, 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); + g->scr->batch_hc_mix, 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); ds4_gpu_tensor *hc_split_view = ds4_gpu_tensor_view( - g->batch_hc_split, 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); + g->scr->batch_hc_split, 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); ds4_gpu_tensor *attn_cur_view = ds4_gpu_tensor_view( - g->batch_attn_cur, 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); + g->scr->batch_attn_cur, 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); ds4_gpu_tensor *after_attn_hc_view = ds4_gpu_tensor_view( - g->batch_after_attn_hc, 0, (uint64_t)n_tokens * hc_dim * sizeof(float)); + g->scr->batch_after_attn_hc, 0, (uint64_t)n_tokens * hc_dim * sizeof(float)); bool ok = hc_mix_view && hc_split_view && attn_cur_view && after_attn_hc_view; const bool fuse_hc_norm = DS4_N_HC == 4 && !metal_graph_use_reference_hc_decode() && metal_graph_enable_batch_hc_norm_fusion(); - if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(g->batch_flat_hc, - g->batch_cur_hc, + if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(g->scr->batch_flat_hc, + g->scr->batch_cur_hc, (uint32_t)hc_dim, n_tokens, DS4_RMS_EPS) != 0; @@ -17378,7 +17478,7 @@ static bool metal_graph_encode_layer_attention_batch( layer->hc_attn_fn->abs_offset, hc_dim, mix_hc, - g->batch_flat_hc, + g->scr->batch_flat_hc, n_tokens) != 0; if (metal_graph_use_reference_hc_decode()) { if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, @@ -17391,16 +17491,16 @@ static bool metal_graph_encode_layer_attention_batch( DS4_N_HC_SINKHORN_ITER, DS4_HC_EPS) != 0; if (ok) ok = ds4_gpu_hc_weighted_sum_split_tensor(attn_cur_view, - g->batch_cur_hc, + g->scr->batch_cur_hc, hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } else if (fuse_hc_norm) { if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(attn_cur_view, - g->batch_attn_norm, + g->scr->batch_attn_norm, hc_split_view, hc_mix_view, - g->batch_cur_hc, + g->scr->batch_cur_hc, model->map, model->size, layer->hc_attn_scale->abs_offset, @@ -17415,7 +17515,7 @@ static bool metal_graph_encode_layer_attention_batch( if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(attn_cur_view, hc_split_view, hc_mix_view, - g->batch_cur_hc, + g->scr->batch_cur_hc, model->map, model->size, layer->hc_attn_scale->abs_offset, @@ -17426,13 +17526,13 @@ static bool metal_graph_encode_layer_attention_batch( DS4_HC_EPS) != 0; } if (ok) { - metal_graph_debug_dump_tensor("hc_attn_pre", g->batch_attn_cur, + metal_graph_debug_dump_tensor("hc_attn_pre", g->scr->batch_attn_cur, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } DS4_METAL_PROFILE_ATTN_STAGE("hc_pre"); if (ok && !fuse_hc_norm) { - ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_attn_norm, - g->batch_attn_cur, + ok = ds4_gpu_rms_norm_weight_rows_tensor(g->scr->batch_attn_norm, + g->scr->batch_attn_cur, model->map, model->size, layer->attn_norm->abs_offset, @@ -17441,7 +17541,7 @@ static bool metal_graph_encode_layer_attention_batch( DS4_RMS_EPS) != 0; } if (ok) { - metal_graph_debug_dump_tensor("attn_norm", g->batch_attn_norm, + metal_graph_debug_dump_tensor("attn_norm", g->scr->batch_attn_norm, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } DS4_METAL_PROFILE_ATTN_STAGE("norm"); @@ -17449,15 +17549,15 @@ static bool metal_graph_encode_layer_attention_batch( if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_q_a", il, pos0, - g->batch_qr, + g->scr->batch_qr, model, layer->attn_q_a, DS4_N_EMBD, q_rank, - g->batch_attn_norm, + g->scr->batch_attn_norm, n_tokens); if (ok) { - metal_graph_debug_dump_tensor("q_lora", g->batch_qr, + metal_graph_debug_dump_tensor("q_lora", g->scr->batch_qr, (uint64_t)n_tokens * q_rank, il, pos0); } DS4_METAL_PROFILE_Q_STAGE("q_a"); @@ -17465,32 +17565,32 @@ static bool metal_graph_encode_layer_attention_batch( if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_kv", il, pos0, - g->batch_kv_raw, + g->scr->batch_kv_raw, model, layer->attn_kv, DS4_N_EMBD, DS4_N_HEAD_DIM, - g->batch_attn_norm, + g->scr->batch_attn_norm, n_tokens); if (ok) { - metal_graph_debug_dump_tensor("KVraw", g->batch_kv_raw, + metal_graph_debug_dump_tensor("KVraw", g->scr->batch_kv_raw, (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); } - if (ok) ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(g->batch_qr_norm, - g->batch_qr, + if (ok) ok = ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(g->scr->batch_qr_norm, + g->scr->batch_qr, model->map, model->size, layer->attn_q_a_norm->abs_offset, (uint32_t)q_rank, - g->batch_kv, - g->batch_kv_raw, + g->scr->batch_kv, + g->scr->batch_kv_raw, layer->attn_kv_a_norm->abs_offset, DS4_N_HEAD_DIM, n_tokens, DS4_RMS_EPS) != 0; } else { - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_qr_norm, - g->batch_qr, + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->scr->batch_qr_norm, + g->scr->batch_qr, model->map, model->size, layer->attn_q_a_norm->abs_offset, @@ -17499,11 +17599,11 @@ static bool metal_graph_encode_layer_attention_batch( DS4_RMS_EPS) != 0; } if (ok) { - metal_graph_debug_dump_tensor("q_lora_norm", g->batch_qr_norm, + metal_graph_debug_dump_tensor("q_lora_norm", g->scr->batch_qr_norm, (uint64_t)n_tokens * q_rank, il, pos0); } if (qkv_rms_fused && ok) { - metal_graph_debug_dump_tensor("KVnorm", g->batch_kv, + metal_graph_debug_dump_tensor("KVnorm", g->scr->batch_kv, (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); } DS4_METAL_PROFILE_Q_STAGE("q_a_norm"); @@ -17512,14 +17612,14 @@ static bool metal_graph_encode_layer_attention_batch( metal_graph_debug_wants("Qnorm", il, pos0); bool q_b_f16_out = false; if (ok && !q_path_debug) { - q_b_f16_out = ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(g->batch_q, - g->batch_q_half, + q_b_f16_out = ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(g->scr->batch_q, + g->scr->batch_q_half, model->map, model->size, layer->attn_q_b->abs_offset, q_rank, q_dim, - g->batch_qr_norm, + g->scr->batch_qr_norm, n_tokens, DS4_N_HEAD, DS4_N_HEAD_DIM, @@ -17539,7 +17639,7 @@ static bool metal_graph_encode_layer_attention_batch( DS4_METAL_PROFILE_Q_STAGE("q_b"); DS4_METAL_PROFILE_Q_STAGE("head_norm"); if (ok) { - metal_graph_debug_dump_tensor("Qcur", g->batch_q, + metal_graph_debug_dump_tensor("Qcur", g->scr->batch_q, (uint64_t)n_tokens * q_dim, il, pos0); } DS4_METAL_PROFILE_Q_STAGE("rope"); @@ -17547,29 +17647,29 @@ static bool metal_graph_encode_layer_attention_batch( if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_q_b", il, pos0, - g->batch_q, + g->scr->batch_q, model, layer->attn_q_b, q_rank, q_dim, - g->batch_qr_norm, + g->scr->batch_qr_norm, n_tokens); if (ok) { - metal_graph_debug_dump_tensor("Qraw", g->batch_q, + metal_graph_debug_dump_tensor("Qraw", g->scr->batch_q, (uint64_t)n_tokens * q_dim, il, pos0); } DS4_METAL_PROFILE_Q_STAGE("q_b"); - if (ok) ok = ds4_gpu_head_rms_norm_tensor(g->batch_q, + if (ok) ok = ds4_gpu_head_rms_norm_tensor(g->scr->batch_q, n_tokens, DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; if (ok) { - metal_graph_debug_dump_tensor("Qnorm", g->batch_q, + metal_graph_debug_dump_tensor("Qnorm", g->scr->batch_q, (uint64_t)n_tokens * q_dim, il, pos0); } DS4_METAL_PROFILE_Q_STAGE("head_norm"); - if (ok) ok = ds4_gpu_rope_tail_tensor(g->batch_q, + if (ok) ok = ds4_gpu_rope_tail_tensor(g->scr->batch_q, n_tokens, DS4_N_HEAD, DS4_N_HEAD_DIM, @@ -17584,7 +17684,7 @@ static bool metal_graph_encode_layer_attention_batch( DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; if (ok) { - metal_graph_debug_dump_tensor("Qcur", g->batch_q, + metal_graph_debug_dump_tensor("Qcur", g->scr->batch_q, (uint64_t)n_tokens * q_dim, il, pos0); } DS4_METAL_PROFILE_Q_STAGE("rope"); @@ -17594,19 +17694,19 @@ static bool metal_graph_encode_layer_attention_batch( if (ok) ok = metal_graph_matmul_q8_0_named_tensor("attn_kv", il, pos0, - g->batch_kv_raw, + g->scr->batch_kv_raw, model, layer->attn_kv, DS4_N_EMBD, DS4_N_HEAD_DIM, - g->batch_attn_norm, + g->scr->batch_attn_norm, n_tokens); if (ok) { - metal_graph_debug_dump_tensor("KVraw", g->batch_kv_raw, + metal_graph_debug_dump_tensor("KVraw", g->scr->batch_kv_raw, (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); } - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_kv, - g->batch_kv_raw, + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->scr->batch_kv, + g->scr->batch_kv_raw, model->map, model->size, layer->attn_kv_a_norm->abs_offset, @@ -17614,11 +17714,11 @@ static bool metal_graph_encode_layer_attention_batch( n_tokens, DS4_RMS_EPS) != 0; if (ok) { - metal_graph_debug_dump_tensor("KVnorm", g->batch_kv, + metal_graph_debug_dump_tensor("KVnorm", g->scr->batch_kv, (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); } } - if (ok) ok = ds4_gpu_rope_tail_tensor(g->batch_kv, + if (ok) ok = ds4_gpu_rope_tail_tensor(g->scr->batch_kv, n_tokens, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, @@ -17633,15 +17733,15 @@ static bool metal_graph_encode_layer_attention_batch( DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; if (ok) { - metal_graph_debug_dump_tensor("KVrope", g->batch_kv, + metal_graph_debug_dump_tensor("KVrope", g->scr->batch_kv, (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); } - if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(g->batch_kv, + if (ok) ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(g->scr->batch_kv, n_tokens, DS4_N_HEAD_DIM, DS4_N_ROT) != 0; if (ok) { - metal_graph_debug_dump_tensor("KVcur", g->batch_kv, + metal_graph_debug_dump_tensor("KVcur", g->scr->batch_kv, (uint64_t)n_tokens * DS4_N_HEAD_DIM, il, pos0); } DS4_METAL_PROFILE_ATTN_STAGE("kv_path"); @@ -17654,7 +17754,7 @@ static bool metal_graph_encode_layer_attention_batch( * attention mask still enforces the 128-token logical window. */ if (ok && zero_prefix) ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], - g->batch_kv, + g->scr->batch_kv, g->raw_cap, pos0, n_tokens, @@ -17663,12 +17763,12 @@ static bool metal_graph_encode_layer_attention_batch( bool batch_attention_done = false; if (ok && raw_batch_attention) { - ok = ds4_gpu_attention_prefill_raw_heads_tensor(g->batch_heads, + ok = ds4_gpu_attention_prefill_raw_heads_tensor(g->scr->batch_heads, model->map, model->size, layer->attn_sinks->abs_offset, - g->batch_q, - g->batch_kv, + g->scr->batch_q, + g->scr->batch_kv, n_tokens, g->raw_window, DS4_N_HEAD, @@ -17689,7 +17789,7 @@ static bool metal_graph_encode_layer_attention_batch( pos0 + n_tokens - 1u, n_raw); ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], - g->batch_kv, + g->scr->batch_kv, g->raw_cap, pos0, n_tokens, @@ -17702,11 +17802,11 @@ static bool metal_graph_encode_layer_attention_batch( pos0); } if (ok) { - ok = ds4_gpu_attention_decode_raw_batch_heads_tensor(g->batch_heads, + ok = ds4_gpu_attention_decode_raw_batch_heads_tensor(g->scr->batch_heads, model->map, model->size, layer->attn_sinks->abs_offset, - g->batch_q, + g->scr->batch_q, g->layer_raw_cache[il], n_tokens, pos0, @@ -17728,30 +17828,30 @@ static bool metal_graph_encode_layer_attention_batch( ok = false; } if (ok) { - ok = ds4_gpu_matmul_f16_tensor(g->batch_comp_kv, + ok = ds4_gpu_matmul_f16_tensor(g->scr->batch_comp_kv, model->map, model->size, layer->attn_compressor_kv->abs_offset, DS4_N_EMBD, comp_width, - g->batch_attn_norm, + g->scr->batch_attn_norm, n_tokens) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(g->batch_comp_sc, + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->batch_comp_sc, model->map, model->size, layer->attn_compressor_gate->abs_offset, DS4_N_EMBD, comp_width, - g->batch_attn_norm, + g->scr->batch_attn_norm, n_tokens) != 0; } if (ok) metal_graph_debug_dump_tensor("attn_comp_kv_raw", - g->batch_comp_kv, + g->scr->batch_comp_kv, (uint64_t)comp_width * n_tokens, il, pos0); if (ok) metal_graph_debug_dump_tensor("attn_comp_score_raw", - g->batch_comp_sc, + g->scr->batch_comp_sc, (uint64_t)comp_width * n_tokens, il, pos0); @@ -17773,8 +17873,8 @@ static bool metal_graph_encode_layer_attention_batch( ds4_gpu_compressor_prefill_tensor(attn_comp_target, g->layer_attn_state_kv[il], g->layer_attn_state_score[il], - g->batch_comp_kv, - g->batch_comp_sc, + g->scr->batch_comp_kv, + g->scr->batch_comp_sc, model->map, model->size, layer->attn_compressor_ape->abs_offset, @@ -17857,8 +17957,8 @@ static bool metal_graph_encode_layer_attention_batch( attn_comp_target, g->layer_attn_state_kv[il], g->layer_attn_state_score[il], - g->batch_comp_kv, - g->batch_comp_sc, + g->scr->batch_comp_kv, + g->scr->batch_comp_sc, model->map, model->size, layer->attn_compressor_ape->abs_offset, @@ -17883,8 +17983,8 @@ static bool metal_graph_encode_layer_attention_batch( attn_comp_target, g->layer_attn_state_kv[il], g->layer_attn_state_score[il], - g->batch_comp_kv, - g->batch_comp_sc, + g->scr->batch_comp_kv, + g->scr->batch_comp_sc, model->map, model->size, layer->attn_compressor_ape->abs_offset, @@ -17955,8 +18055,8 @@ static bool metal_graph_encode_layer_attention_batch( ok = false; break; } - ds4_gpu_tensor *kv_view = metal_graph_tensor_row_view(g->batch_comp_kv, t, comp_width); - ds4_gpu_tensor *sc_view = metal_graph_tensor_row_view(g->batch_comp_sc, t, comp_width); + ds4_gpu_tensor *kv_view = metal_graph_tensor_row_view(g->scr->batch_comp_kv, t, comp_width); + ds4_gpu_tensor *sc_view = metal_graph_tensor_row_view(g->scr->batch_comp_sc, t, comp_width); const uint32_t comp_row = g->layer_n_comp[il]; ok = kv_view && sc_view && ds4_gpu_compressor_update_tensor(kv_view, @@ -18020,41 +18120,41 @@ static bool metal_graph_encode_layer_attention_batch( ok = false; } if (ok) { - ok = ds4_gpu_matmul_f16_tensor(g->batch_comp_kv, + ok = ds4_gpu_matmul_f16_tensor(g->scr->batch_comp_kv, model->map, model->size, layer->indexer_compressor_kv->abs_offset, DS4_N_EMBD, index_width, - g->batch_attn_norm, + g->scr->batch_attn_norm, n_tokens) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(g->batch_comp_sc, + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->batch_comp_sc, model->map, model->size, layer->indexer_compressor_gate->abs_offset, DS4_N_EMBD, index_width, - g->batch_attn_norm, + g->scr->batch_attn_norm, n_tokens) != 0; } if (ok) metal_graph_debug_dump_tensor("indexer_comp_kv_raw", - g->batch_comp_kv, + g->scr->batch_comp_kv, (uint64_t)index_width * n_tokens, il, pos0); if (ok) metal_graph_debug_dump_tensor("indexer_comp_score_raw", - g->batch_comp_sc, + g->scr->batch_comp_sc, (uint64_t)index_width * n_tokens, il, pos0); - if (ok) ok = metal_graph_matmul_plain_tensor(g->batch_indexer_q, + if (ok) ok = metal_graph_matmul_plain_tensor(g->scr->batch_indexer_q, model, layer->indexer_attn_q_b, q_rank, (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM, - g->batch_qr_norm, + g->scr->batch_qr_norm, n_tokens); - if (ok) ok = ds4_gpu_rope_tail_tensor(g->batch_indexer_q, + if (ok) ok = ds4_gpu_rope_tail_tensor(g->scr->batch_indexer_q, n_tokens, DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM, @@ -18068,16 +18168,16 @@ static bool metal_graph_encode_layer_attention_batch( attn_factor, DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; - if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(g->batch_indexer_q, + if (ok) ok = ds4_gpu_dsv4_indexer_qat_tensor(g->scr->batch_indexer_q, n_tokens * DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM) != 0; - if (ok) ok = ds4_gpu_matmul_f16_tensor(g->batch_indexer_weights, + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->batch_indexer_weights, model->map, model->size, layer->indexer_proj->abs_offset, DS4_N_EMBD, DS4_N_INDEXER_HEAD, - g->batch_attn_norm, + g->scr->batch_attn_norm, n_tokens) != 0; if (zero_prefix) { if (ok && n_comp > g->layer_comp_cap[il]) { @@ -18088,8 +18188,8 @@ static bool metal_graph_encode_layer_attention_batch( ok = ds4_gpu_compressor_prefill_tensor(g->layer_index_comp_cache[il], g->layer_index_state_kv[il], g->layer_index_state_score[il], - g->batch_comp_kv, - g->batch_comp_sc, + g->scr->batch_comp_kv, + g->scr->batch_comp_sc, model->map, model->size, layer->indexer_compressor_ape->abs_offset, @@ -18174,8 +18274,8 @@ static bool metal_graph_encode_layer_attention_batch( index_view, g->layer_index_state_kv[il], g->layer_index_state_score[il], - g->batch_comp_kv, - g->batch_comp_sc, + g->scr->batch_comp_kv, + g->scr->batch_comp_sc, model->map, model->size, layer->indexer_compressor_ape->abs_offset, @@ -18247,8 +18347,8 @@ static bool metal_graph_encode_layer_attention_batch( ok = false; break; } - ds4_gpu_tensor *kv_view = metal_graph_tensor_row_view(g->batch_comp_kv, t, index_width); - ds4_gpu_tensor *sc_view = metal_graph_tensor_row_view(g->batch_comp_sc, t, index_width); + ds4_gpu_tensor *kv_view = metal_graph_tensor_row_view(g->scr->batch_comp_kv, t, index_width); + ds4_gpu_tensor *sc_view = metal_graph_tensor_row_view(g->scr->batch_comp_sc, t, index_width); const uint32_t index_row = g->layer_n_index_comp[il]; ok = kv_view && sc_view && ds4_gpu_compressor_update_tensor(kv_view, @@ -18312,7 +18412,7 @@ static bool metal_graph_encode_layer_attention_batch( double index_stage_t0 = 0.0; ok = ds4_gpu_store_raw_kv_batch_tensor(g->layer_raw_cache[il], - g->batch_kv, + g->scr->batch_kv, g->raw_cap, pos0, n_tokens, @@ -18327,9 +18427,9 @@ static bool metal_graph_encode_layer_attention_batch( n_comp, &index_stage_t0); } - ok = ds4_gpu_indexer_scores_decode_batch_tensor(g->indexer_scores, - g->batch_indexer_q, - g->batch_indexer_weights, + ok = ds4_gpu_indexer_scores_decode_batch_tensor(g->scr->indexer_scores, + g->scr->batch_indexer_q, + g->scr->batch_indexer_weights, g->layer_index_comp_cache[il], n_comp, n_tokens, @@ -18348,14 +18448,14 @@ static bool metal_graph_encode_layer_attention_batch( } if (ok) { metal_graph_debug_dump_tensor("indexer_scores", - g->indexer_scores, + g->scr->indexer_scores, (uint64_t)n_comp * n_tokens, il, pos0); } if (ok) { - ok = ds4_gpu_indexer_topk_tensor(g->comp_selected, - g->indexer_scores, + ok = ds4_gpu_indexer_topk_tensor(g->scr->comp_selected, + g->scr->indexer_scores, n_comp, n_tokens, DS4_N_INDEXER_TOP_K) != 0; @@ -18369,7 +18469,7 @@ static bool metal_graph_encode_layer_attention_batch( } if (ok) { metal_graph_debug_dump_i32_tensor("indexer_topk", - g->comp_selected, + g->scr->comp_selected, (uint64_t)n_tokens * DS4_N_INDEXER_TOP_K, il, pos0); @@ -18382,15 +18482,15 @@ static bool metal_graph_encode_layer_attention_batch( } if (ok) { if (use_indexed_comp) { - ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(g->batch_heads, + ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(g->scr->batch_heads, model->map, model->size, layer->attn_sinks->abs_offset, - g->batch_q, + g->scr->batch_q, g->layer_raw_cache[il], g->layer_attn_comp_cache[il], metal_graph_attn_comp_cache_is_f16(), - g->comp_selected, + g->scr->comp_selected, n_tokens, pos0, n_raw, @@ -18411,15 +18511,15 @@ static bool metal_graph_encode_layer_attention_batch( &index_stage_t0); } } else { - ok = ds4_gpu_attention_decode_mixed_batch_heads_tensor(g->batch_heads, + ok = ds4_gpu_attention_decode_mixed_batch_heads_tensor(g->scr->batch_heads, model->map, model->size, layer->attn_sinks->abs_offset, - g->batch_q, + g->scr->batch_q, g->layer_raw_cache[il], g->layer_attn_comp_cache[il], metal_graph_attn_comp_cache_is_f16(), - use_comp_mask ? g->comp_mask : NULL, + use_comp_mask ? g->scr->comp_mask : NULL, use_comp_mask, n_tokens, pos0, @@ -18448,9 +18548,9 @@ static bool metal_graph_encode_layer_attention_batch( n_comp, &index_stage_t0); } - ok = ds4_gpu_indexer_scores_prefill_tensor(g->indexer_scores, - g->batch_indexer_q, - g->batch_indexer_weights, + ok = ds4_gpu_indexer_scores_prefill_tensor(g->scr->indexer_scores, + g->scr->batch_indexer_q, + g->scr->batch_indexer_weights, g->layer_index_comp_cache[il], n_comp, n_tokens, @@ -18468,14 +18568,14 @@ static bool metal_graph_encode_layer_attention_batch( } if (ok) { metal_graph_debug_dump_tensor("indexer_scores", - g->indexer_scores, + g->scr->indexer_scores, (uint64_t)n_comp * n_tokens, il, pos0); } if (ok) { - ok = ds4_gpu_indexer_topk_tensor(g->comp_selected, - g->indexer_scores, + ok = ds4_gpu_indexer_topk_tensor(g->scr->comp_selected, + g->scr->indexer_scores, n_comp, n_tokens, DS4_N_INDEXER_TOP_K) != 0; @@ -18489,22 +18589,22 @@ static bool metal_graph_encode_layer_attention_batch( } if (ok) { metal_graph_debug_dump_i32_tensor("indexer_topk", - g->comp_selected, + g->scr->comp_selected, (uint64_t)n_tokens * DS4_N_INDEXER_TOP_K, il, pos0); } } if (ok) { - ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(g->batch_heads, + ok = ds4_gpu_attention_indexed_mixed_batch_heads_tensor(g->scr->batch_heads, model->map, model->size, layer->attn_sinks->abs_offset, - g->batch_q, + g->scr->batch_q, g->layer_raw_cache[il], g->layer_attn_comp_cache[il], metal_graph_attn_comp_cache_is_f16(), - g->comp_selected, + g->scr->comp_selected, n_tokens, pos0, n_tokens, @@ -18528,12 +18628,12 @@ static bool metal_graph_encode_layer_attention_batch( if (ok) batch_attention_done = true; } if (ok && zero_prefix && !topk_prefill_needed && n_comp != 0) { - ok = ds4_gpu_attention_prefill_static_mixed_heads_tensor(g->batch_heads, + ok = ds4_gpu_attention_prefill_static_mixed_heads_tensor(g->scr->batch_heads, model->map, model->size, layer->attn_sinks->abs_offset, - g->batch_q, - g->batch_kv, + g->scr->batch_q, + g->scr->batch_kv, g->layer_attn_comp_cache[il], metal_graph_attn_comp_cache_is_f16(), n_tokens, @@ -18555,12 +18655,12 @@ static bool metal_graph_encode_layer_attention_batch( } if (raw_prefix_tokens != 0) { - ok = ds4_gpu_attention_prefill_raw_heads_tensor(g->batch_heads, + ok = ds4_gpu_attention_prefill_raw_heads_tensor(g->scr->batch_heads, model->map, model->size, layer->attn_sinks->abs_offset, - g->batch_q, - g->batch_kv, + g->scr->batch_q, + g->scr->batch_kv, raw_prefix_tokens, g->raw_window, DS4_N_HEAD, @@ -18579,11 +18679,11 @@ static bool metal_graph_encode_layer_attention_batch( if (ratio == 4 && cur_comp > DS4_N_INDEXER_TOP_K) { const float index_scale = 1.0f / sqrtf((float)(DS4_N_INDEXER_HEAD_DIM * DS4_N_INDEXER_HEAD)); ds4_gpu_tensor *indexer_q_view = metal_graph_tensor_row_view( - g->batch_indexer_q, t, (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM); + g->scr->batch_indexer_q, t, (uint64_t)DS4_N_INDEXER_HEAD * DS4_N_INDEXER_HEAD_DIM); ds4_gpu_tensor *indexer_w_view = metal_graph_tensor_row_view( - g->batch_indexer_weights, t, DS4_N_INDEXER_HEAD); + g->scr->batch_indexer_weights, t, DS4_N_INDEXER_HEAD); ok = indexer_q_view && indexer_w_view && - ds4_gpu_indexer_score_one_tensor(g->indexer_scores, + ds4_gpu_indexer_score_one_tensor(g->scr->indexer_scores, indexer_q_view, indexer_w_view, g->layer_index_comp_cache[il], @@ -18591,29 +18691,29 @@ static bool metal_graph_encode_layer_attention_batch( DS4_N_INDEXER_HEAD, DS4_N_INDEXER_HEAD_DIM, index_scale) != 0 && - ds4_gpu_indexer_topk_tensor(g->comp_selected, - g->indexer_scores, + ds4_gpu_indexer_topk_tensor(g->scr->comp_selected, + g->scr->indexer_scores, cur_index, 1, DS4_N_INDEXER_TOP_K) != 0 && - ds4_gpu_dsv4_topk_mask_tensor(g->comp_mask, - g->comp_selected, + ds4_gpu_dsv4_topk_mask_tensor(g->scr->comp_mask, + g->scr->comp_selected, cur_index, 1, DS4_N_INDEXER_TOP_K) != 0; ds4_gpu_tensor_free(indexer_w_view); ds4_gpu_tensor_free(indexer_q_view); if (ok) { - comp_mask = g->comp_mask; + comp_mask = g->scr->comp_mask; n_selected = DS4_N_INDEXER_TOP_K < cur_index ? DS4_N_INDEXER_TOP_K : cur_index; } } - ds4_gpu_tensor *q_view = metal_graph_tensor_row_view(g->batch_q, t, q_dim); - ds4_gpu_tensor *kv_cache_view = metal_graph_tensor_row_view(g->batch_kv, t, DS4_N_HEAD_DIM); - ds4_gpu_tensor *heads_view = metal_graph_tensor_row_view(g->batch_heads, t, q_dim); + ds4_gpu_tensor *q_view = metal_graph_tensor_row_view(g->scr->batch_q, t, q_dim); + ds4_gpu_tensor *kv_cache_view = metal_graph_tensor_row_view(g->scr->batch_kv, t, DS4_N_HEAD_DIM); + ds4_gpu_tensor *heads_view = metal_graph_tensor_row_view(g->scr->batch_heads, t, q_dim); ok = ok && q_view && kv_cache_view && heads_view; if (ok && !zero_prefix) { ok = ds4_gpu_store_raw_kv_tensor(g->layer_raw_cache[il], @@ -18631,7 +18731,7 @@ static bool metal_graph_encode_layer_attention_batch( g->layer_raw_cache[il], g->layer_attn_comp_cache[il], metal_graph_attn_comp_cache_is_f16(), - g->comp_selected, + g->scr->comp_selected, 1, pos, n_raw, @@ -18670,10 +18770,10 @@ static bool metal_graph_encode_layer_attention_batch( DS4_METAL_PROFILE_ATTN_STAGE("attention"); if (ok) { - metal_graph_debug_dump_tensor("kqv_out", g->batch_heads, + metal_graph_debug_dump_tensor("kqv_out", g->scr->batch_heads, (uint64_t)n_tokens * q_dim, il, pos0); } - if (ok) ok = ds4_gpu_rope_tail_tensor(g->batch_heads, + if (ok) ok = ds4_gpu_rope_tail_tensor(g->scr->batch_heads, n_tokens, DS4_N_HEAD, DS4_N_HEAD_DIM, @@ -18688,7 +18788,7 @@ static bool metal_graph_encode_layer_attention_batch( DS4_ROPE_YARN_BETA_FAST, DS4_ROPE_YARN_BETA_SLOW) != 0; if (ok) { - metal_graph_debug_dump_tensor("kqv_back", g->batch_heads, + metal_graph_debug_dump_tensor("kqv_back", g->scr->batch_heads, (uint64_t)n_tokens * q_dim, il, pos0); } DS4_METAL_PROFILE_ATTN_STAGE("inv_rope"); @@ -18699,8 +18799,8 @@ static bool metal_graph_encode_layer_attention_batch( if (ok && !attn_out_debug && !metal_graph_directional_steering_attn_enabled(g)) { - attn_out_f16 = ds4_gpu_attention_output_q8_batch_f16_tensor(g->batch_q_half, - g->batch_attn_low, + attn_out_f16 = ds4_gpu_attention_output_q8_batch_f16_tensor(g->scr->batch_q_half, + g->scr->batch_attn_low, model->map, model->size, layer->attn_output_a->abs_offset, @@ -18709,15 +18809,15 @@ static bool metal_graph_encode_layer_attention_batch( rank, n_groups, DS4_N_EMBD, - g->batch_heads, + g->scr->batch_heads, n_tokens) != 0; } if (!attn_out_f16) { if (ok) { - ok = ds4_gpu_attention_output_q8_batch_tensor(g->batch_attn_out, - g->batch_attn_low, - g->batch_group_tmp, - g->batch_low_tmp, + ok = ds4_gpu_attention_output_q8_batch_tensor(g->scr->batch_attn_out, + g->scr->batch_attn_low, + g->scr->batch_group_tmp, + g->scr->batch_low_tmp, model->map, model->size, layer->attn_output_a->abs_offset, @@ -18726,41 +18826,41 @@ static bool metal_graph_encode_layer_attention_batch( rank, n_groups, DS4_N_EMBD, - g->batch_heads, + g->scr->batch_heads, n_tokens) != 0; } if (ok) { - metal_graph_debug_dump_tensor("attn_low", g->batch_attn_low, + metal_graph_debug_dump_tensor("attn_low", g->scr->batch_attn_low, (uint64_t)n_tokens * n_groups * rank, il, pos0); } if (ok) { - metal_graph_debug_dump_tensor("attn_out", g->batch_attn_out, + metal_graph_debug_dump_tensor("attn_out", g->scr->batch_attn_out, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } } DS4_METAL_PROFILE_ATTN_STAGE("output_proj"); if (ok && !attn_out_f16 && metal_graph_directional_steering_attn_enabled(g)) { - ok = metal_graph_apply_directional_steering_attn(g, g->batch_attn_out, il, n_tokens); + ok = metal_graph_apply_directional_steering_attn(g, g->scr->batch_attn_out, il, n_tokens); } if (ok && attn_out_f16) { ok = ds4_gpu_hc_expand_split_half_tensor(after_attn_hc_view, - g->batch_q_half, - g->batch_cur_hc, + g->scr->batch_q_half, + g->scr->batch_cur_hc, hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok) { ok = ds4_gpu_hc_expand_split_tensor(after_attn_hc_view, - g->batch_attn_out, - g->batch_cur_hc, + g->scr->batch_attn_out, + g->scr->batch_cur_hc, hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } if (ok) { - metal_graph_debug_dump_tensor("hc_attn_post", g->batch_after_attn_hc, + metal_graph_debug_dump_tensor("hc_attn_post", g->scr->batch_after_attn_hc, (uint64_t)n_tokens * hc_dim, il, pos0); } DS4_METAL_PROFILE_ATTN_STAGE("hc_post"); @@ -18806,19 +18906,19 @@ static bool metal_graph_encode_layer_ffn_batch( } while (0) ds4_gpu_tensor *hc_mix_view = ds4_gpu_tensor_view( - g->batch_hc_mix, 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); + g->scr->batch_hc_mix, 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); ds4_gpu_tensor *hc_split_view = ds4_gpu_tensor_view( - g->batch_hc_split, 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); + g->scr->batch_hc_split, 0, (uint64_t)n_tokens * mix_hc * sizeof(float)); ds4_gpu_tensor *ffn_cur_view = ds4_gpu_tensor_view( - g->batch_ffn_cur, 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); + g->scr->batch_ffn_cur, 0, (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float)); ds4_gpu_tensor *next_hc_view = ds4_gpu_tensor_view( - g->batch_next_hc, 0, (uint64_t)n_tokens * hc_dim * sizeof(float)); + g->scr->batch_next_hc, 0, (uint64_t)n_tokens * hc_dim * sizeof(float)); bool ok = hc_mix_view && hc_split_view && ffn_cur_view && next_hc_view; const bool fuse_hc_norm = DS4_N_HC == 4 && !metal_graph_use_reference_hc_decode() && metal_graph_enable_batch_hc_norm_fusion(); - if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(g->batch_flat_hc, - g->batch_after_attn_hc, + if (ok) ok = ds4_gpu_rms_norm_plain_rows_tensor(g->scr->batch_flat_hc, + g->scr->batch_after_attn_hc, (uint32_t)hc_dim, n_tokens, DS4_RMS_EPS) != 0; @@ -18828,7 +18928,7 @@ static bool metal_graph_encode_layer_ffn_batch( layer->hc_ffn_fn->abs_offset, hc_dim, mix_hc, - g->batch_flat_hc, + g->scr->batch_flat_hc, n_tokens) != 0; if (metal_graph_use_reference_hc_decode()) { if (ok) ok = ds4_gpu_hc_split_sinkhorn_tensor(hc_split_view, @@ -18841,16 +18941,16 @@ static bool metal_graph_encode_layer_ffn_batch( DS4_N_HC_SINKHORN_ITER, DS4_HC_EPS) != 0; if (ok) ok = ds4_gpu_hc_weighted_sum_split_tensor(ffn_cur_view, - g->batch_after_attn_hc, + g->scr->batch_after_attn_hc, hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } else if (fuse_hc_norm) { if (ok) ok = ds4_gpu_hc_split_weighted_sum_norm_tensor(ffn_cur_view, - g->batch_ffn_norm, + g->scr->batch_ffn_norm, hc_split_view, hc_mix_view, - g->batch_after_attn_hc, + g->scr->batch_after_attn_hc, model->map, model->size, layer->hc_ffn_scale->abs_offset, @@ -18865,7 +18965,7 @@ static bool metal_graph_encode_layer_ffn_batch( if (ok) ok = ds4_gpu_hc_split_weighted_sum_tensor(ffn_cur_view, hc_split_view, hc_mix_view, - g->batch_after_attn_hc, + g->scr->batch_after_attn_hc, model->map, model->size, layer->hc_ffn_scale->abs_offset, @@ -18876,13 +18976,13 @@ static bool metal_graph_encode_layer_ffn_batch( DS4_HC_EPS) != 0; } if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_pre", g->batch_ffn_cur, + metal_graph_debug_dump_tensor("hc_ffn_pre", g->scr->batch_ffn_cur, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } DS4_METAL_PROFILE_FFN_STAGE("hc_pre"); if (ok && !fuse_hc_norm) { - ok = ds4_gpu_rms_norm_weight_rows_tensor(g->batch_ffn_norm, - g->batch_ffn_cur, + ok = ds4_gpu_rms_norm_weight_rows_tensor(g->scr->batch_ffn_norm, + g->scr->batch_ffn_cur, model->map, model->size, layer->ffn_norm->abs_offset, @@ -18891,22 +18991,22 @@ static bool metal_graph_encode_layer_ffn_batch( DS4_RMS_EPS) != 0; } if (ok) { - metal_graph_debug_dump_tensor("ffn_norm", g->batch_ffn_norm, + metal_graph_debug_dump_tensor("ffn_norm", g->scr->batch_ffn_norm, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } DS4_METAL_PROFILE_FFN_STAGE("norm"); - if (ok) ok = ds4_gpu_matmul_f16_tensor(g->batch_router_logits, + if (ok) ok = ds4_gpu_matmul_f16_tensor(g->scr->batch_router_logits, model->map, model->size, layer->ffn_gate_inp->abs_offset, DS4_N_EMBD, DS4_N_EXPERT, - g->batch_ffn_norm, + g->scr->batch_ffn_norm, n_tokens) != 0; - if (ok) ok = ds4_gpu_router_select_batch_tensor(g->batch_router_selected, - g->batch_router_weights, - g->batch_router_probs, + if (ok) ok = ds4_gpu_router_select_batch_tensor(g->scr->batch_router_selected, + g->scr->batch_router_weights, + g->scr->batch_router_probs, model->map, model->size, layer->ffn_exp_probs_b ? layer->ffn_exp_probs_b->abs_offset : 0, @@ -18916,20 +19016,20 @@ static bool metal_graph_encode_layer_ffn_batch( 0, layer->ffn_exp_probs_b != NULL, layer->ffn_gate_tid2eid != NULL, - g->batch_router_logits, - g->prefill_tokens, + g->scr->batch_router_logits, + g->scr->prefill_tokens, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE, n_tokens) != 0; if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_logits", g->batch_router_logits, + metal_graph_debug_dump_tensor("ffn_moe_logits", g->scr->batch_router_logits, (uint64_t)n_tokens * DS4_N_EXPERT, il, pos0); - metal_graph_debug_dump_tensor("ffn_moe_probs", g->batch_router_probs, + metal_graph_debug_dump_tensor("ffn_moe_probs", g->scr->batch_router_probs, (uint64_t)n_tokens * DS4_N_EXPERT, il, pos0); - metal_graph_debug_dump_i32_tensor("ffn_moe_topk", g->batch_router_selected, + metal_graph_debug_dump_i32_tensor("ffn_moe_topk", g->scr->batch_router_selected, (uint64_t)n_tokens * DS4_N_EXPERT_USED, il, pos0); - metal_graph_debug_dump_tensor("ffn_moe_weights_scaled", g->batch_router_weights, + metal_graph_debug_dump_tensor("ffn_moe_weights_scaled", g->scr->batch_router_weights, (uint64_t)n_tokens * DS4_N_EXPERT_USED, il, pos0); } DS4_METAL_PROFILE_FFN_STAGE("router"); @@ -19008,13 +19108,13 @@ static bool metal_graph_encode_layer_ffn_batch( #define DS4_METAL_TRY_SHARED_DOWN_F16() do { \ if (ok && !keep_ffn_out && !metal_graph_debug_wants("ffn_shexp", il, pos0)) { \ - shared_down_f16 = ds4_gpu_matmul_q8_0_f16_out_tensor(g->batch_q_half, \ + shared_down_f16 = ds4_gpu_matmul_q8_0_f16_out_tensor(g->scr->batch_q_half, \ model->map, \ model->size, \ layer->ffn_down_shexp->abs_offset, \ shared_dim, \ DS4_N_EMBD, \ - g->batch_shared_mid, \ + g->scr->batch_shared_mid, \ n_tokens) != 0; \ } \ } while (0) @@ -19023,27 +19123,27 @@ static bool metal_graph_encode_layer_ffn_batch( if (ok) ok = metal_graph_matmul_q8_0_named_tensor("shared_gate", \ il, \ pos0, \ - g->batch_shared_gate, \ + g->scr->batch_shared_gate, \ model, \ layer->ffn_gate_shexp, \ DS4_N_EMBD, \ shared_dim, \ - g->batch_ffn_norm, \ + g->scr->batch_ffn_norm, \ n_tokens); \ if (ok) ok = metal_graph_matmul_q8_0_named_tensor("shared_up", \ il, \ pos0, \ - g->batch_shared_up, \ + g->scr->batch_shared_up, \ model, \ layer->ffn_up_shexp, \ DS4_N_EMBD, \ shared_dim, \ - g->batch_ffn_norm, \ + g->scr->batch_ffn_norm, \ n_tokens); \ DS4_METAL_PROFILE_FFN_STAGE("shared_gate_up"); \ - if (ok) ok = ds4_gpu_swiglu_tensor(g->batch_shared_mid, \ - g->batch_shared_gate, \ - g->batch_shared_up, \ + if (ok) ok = ds4_gpu_swiglu_tensor(g->scr->batch_shared_mid, \ + g->scr->batch_shared_gate, \ + g->scr->batch_shared_up, \ (uint32_t)((uint64_t)n_tokens * shared_dim), \ DS4_SWIGLU_CLAMP_EXP, \ 1.0f) != 0; \ @@ -19051,16 +19151,16 @@ static bool metal_graph_encode_layer_ffn_batch( if (ok && !shared_down_f16) ok = metal_graph_matmul_q8_0_named_tensor("shared_down", \ il, \ pos0, \ - g->batch_shared_out, \ + g->scr->batch_shared_out, \ model, \ layer->ffn_down_shexp, \ shared_dim, \ DS4_N_EMBD, \ - g->batch_shared_mid, \ + g->scr->batch_shared_mid, \ n_tokens); \ DS4_METAL_PROFILE_FFN_STAGE("shared_down"); \ if (ok && !shared_down_f16) { \ - metal_graph_debug_dump_tensor("ffn_shexp", g->batch_shared_out, \ + metal_graph_debug_dump_tensor("ffn_shexp", g->scr->batch_shared_out, \ (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); \ } \ } while (0) @@ -19144,11 +19244,11 @@ static bool metal_graph_encode_layer_ffn_batch( #endif if (ok) { - ok = ds4_gpu_routed_moe_batch_tensor(g->batch_routed_out, - g->batch_routed_gate, - g->batch_routed_up, - g->batch_routed_mid, - g->batch_routed_down, + ok = ds4_gpu_routed_moe_batch_tensor(g->scr->batch_routed_out, + g->scr->batch_routed_gate, + g->scr->batch_routed_up, + g->scr->batch_routed_mid, + g->scr->batch_routed_down, model->map, model->size, layer->ffn_gate_exps->abs_offset, @@ -19163,38 +19263,38 @@ static bool metal_graph_encode_layer_ffn_batch( (uint32_t)expert_in_dim, (uint32_t)down_in_dim, (uint32_t)routed_out_dim, - g->batch_router_selected, - g->batch_router_weights, + g->scr->batch_router_selected, + g->scr->batch_router_weights, DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_SWIGLU_CLAMP_EXP, - g->batch_ffn_norm, + g->scr->batch_ffn_norm, il, n_tokens, - &g->batch_routed_mid_is_f16) != 0; + &g->scr->batch_routed_mid_is_f16) != 0; } if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", g->batch_routed_gate, + metal_graph_debug_dump_tensor("ffn_moe_gate_clamped", g->scr->batch_routed_gate, (uint64_t)n_tokens * DS4_N_EXPERT_USED * down_in_dim, il, pos0); - metal_graph_debug_dump_tensor("ffn_moe_up_clamped", g->batch_routed_up, + metal_graph_debug_dump_tensor("ffn_moe_up_clamped", g->scr->batch_routed_up, (uint64_t)n_tokens * DS4_N_EXPERT_USED * down_in_dim, il, pos0); } if (ok) { const uint64_t routed_mid_elems = (uint64_t)n_tokens * DS4_N_EXPERT_USED * down_in_dim; - if (g->batch_routed_mid_is_f16) { - metal_graph_debug_dump_f16_tensor("ffn_moe_weighted_swiglu", g->batch_routed_mid, + if (g->scr->batch_routed_mid_is_f16) { + metal_graph_debug_dump_f16_tensor("ffn_moe_weighted_swiglu", g->scr->batch_routed_mid, routed_mid_elems, il, pos0); } else { - metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", g->batch_routed_mid, + metal_graph_debug_dump_tensor("ffn_moe_weighted_swiglu", g->scr->batch_routed_mid, routed_mid_elems, il, pos0); } } if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_down", g->batch_routed_down, + metal_graph_debug_dump_tensor("ffn_moe_down", g->scr->batch_routed_down, (uint64_t)n_tokens * DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos0); } if (ok) { - metal_graph_debug_dump_tensor("ffn_moe_out", g->batch_routed_out, + metal_graph_debug_dump_tensor("ffn_moe_out", g->scr->batch_routed_out, (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } DS4_METAL_PROFILE_FFN_STAGE("routed_moe"); @@ -19207,8 +19307,8 @@ static bool metal_graph_encode_layer_ffn_batch( if (ok && keep_ffn_out) { ok = metal_graph_ensure_batch_ffn_out(g) && ds4_gpu_add_tensor(g->batch_ffn_out, - g->batch_shared_out, - g->batch_routed_out, + g->scr->batch_shared_out, + g->scr->batch_routed_out, (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; } if (ok && keep_ffn_out) { @@ -19221,31 +19321,31 @@ static bool metal_graph_encode_layer_ffn_batch( if (ok && metal_graph_directional_steering_ffn_enabled(g)) { ok = ds4_gpu_hc_expand_split_tensor(next_hc_view, g->batch_ffn_out, - g->batch_after_attn_hc, + g->scr->batch_after_attn_hc, hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok && shared_down_f16) { ok = ds4_gpu_hc_expand_add_split_half_add_tensor(next_hc_view, - g->batch_routed_out, - g->batch_q_half, - g->batch_after_attn_hc, + g->scr->batch_routed_out, + g->scr->batch_q_half, + g->scr->batch_after_attn_hc, hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } else if (ok) { ok = ds4_gpu_hc_expand_add_split_tensor(next_hc_view, - g->batch_routed_out, - g->batch_shared_out, - g->batch_after_attn_hc, + g->scr->batch_routed_out, + g->scr->batch_shared_out, + g->scr->batch_after_attn_hc, hc_split_view, DS4_N_EMBD, DS4_N_HC) != 0; } if (ok) { - metal_graph_debug_dump_tensor("hc_ffn_post", g->batch_next_hc, + metal_graph_debug_dump_tensor("hc_ffn_post", g->scr->batch_next_hc, (uint64_t)n_tokens * hc_dim, il, pos0); } DS4_METAL_PROFILE_FFN_STAGE("hc_post"); @@ -19276,9 +19376,9 @@ static bool metal_graph_encode_layer_batch( } } if (ok) { - ds4_gpu_tensor *tmp = g->batch_cur_hc; - g->batch_cur_hc = g->batch_next_hc; - g->batch_next_hc = tmp; + ds4_gpu_tensor *tmp = g->scr->batch_cur_hc; + g->scr->batch_cur_hc = g->scr->batch_next_hc; + g->scr->batch_next_hc = tmp; } return ok; } @@ -19321,7 +19421,7 @@ static bool metal_graph_eval_token_raw_swa_streaming( } if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok) { - ok = ds4_gpu_embed_token_hc_tensor(g->cur_hc, + ok = ds4_gpu_embed_token_hc_tensor(g->scr->cur_hc, model->map, model->size, weights->token_embd->abs_offset, @@ -19343,9 +19443,9 @@ static bool metal_graph_eval_token_raw_swa_streaming( n_raw, token); if (ok) { - ds4_gpu_tensor *tmp = g->cur_hc; - g->cur_hc = g->after_ffn_hc; - g->after_ffn_hc = tmp; + ds4_gpu_tensor *tmp = g->scr->cur_hc; + g->scr->cur_hc = g->scr->after_ffn_hc; + g->scr->after_ffn_hc = tmp; } } if (ok && logits) { @@ -19355,7 +19455,7 @@ static bool metal_graph_eval_token_raw_swa_streaming( if (ok) ok = ds4_gpu_end_commands() != 0; const double t_done = (profile || throttle) ? now_sec() : 0.0; if (ok && logits) { - ok = ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } const double t_read = (profile || throttle) ? now_sec() : 0.0; if (profile) { @@ -19409,9 +19509,9 @@ static bool metal_graph_eval_token_raw_swa_streaming( encoded_layer = true; } if (encoded_layer) { - ds4_gpu_tensor *tmp = g->cur_hc; - g->cur_hc = g->after_ffn_hc; - g->after_ffn_hc = tmp; + ds4_gpu_tensor *tmp = g->scr->cur_hc; + g->scr->cur_hc = g->scr->after_ffn_hc; + g->scr->after_ffn_hc = tmp; } const double tl_encoded = profile ? now_sec() : 0.0; if (ok) ok = ds4_gpu_end_commands() != 0; @@ -19430,7 +19530,7 @@ static bool metal_graph_eval_token_raw_swa_streaming( if (ok && logits) ok = ds4_gpu_end_commands() != 0; const double t_done = (profile || throttle) ? now_sec() : 0.0; if (ok && logits) { - ok = ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } const double t_read = (profile || throttle) ? now_sec() : 0.0; @@ -19480,7 +19580,7 @@ static bool metal_graph_eval_token_raw_swa( const double t_done = (profile || throttle) ? now_sec() : 0.0; if (ok && logits) { - ok = ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } const double t_read = (profile || throttle) ? now_sec() : 0.0; if (profile) { @@ -19661,8 +19761,8 @@ static bool metal_graph_capture_prefill_seed_router_selected( uint32_t k = metal_graph_streaming_prefill_cache_seed_k(g); if (k == 0) return true; if (k > n_tokens) k = n_tokens; - g->prefill_seed_tokens = k; - if (!g->prefill_seed_router_selected || !g->batch_router_selected || + g->scr->prefill_seed_tokens = k; + if (!g->scr->prefill_seed_router_selected || !g->scr->batch_router_selected || il >= DS4_N_LAYER || n_tokens == 0 || sizeof(int) != sizeof(int32_t)) { return false; } @@ -19673,9 +19773,9 @@ static bool metal_graph_capture_prefill_seed_router_selected( const uint64_t dst_off = (uint64_t)il * DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * DS4_N_EXPERT_USED * sizeof(int32_t); - return ds4_gpu_tensor_copy(g->prefill_seed_router_selected, + return ds4_gpu_tensor_copy(g->scr->prefill_seed_router_selected, dst_off, - g->batch_router_selected, + g->scr->batch_router_selected, src_off, bytes) != 0; } @@ -19684,9 +19784,9 @@ static bool metal_graph_seed_streaming_expert_cache_from_prefill( ds4_gpu_graph *g, const ds4_model *model, const ds4_weights *weights) { - const uint32_t seed_tokens = g ? g->prefill_seed_tokens : 0; + const uint32_t seed_tokens = g ? g->scr->prefill_seed_tokens : 0; if (!metal_graph_streaming_prefill_cache_seed_enabled(g)) return true; - if (!model || !weights || !g->prefill_seed_router_selected || seed_tokens == 0) { + if (!model || !weights || !g->scr->prefill_seed_router_selected || seed_tokens == 0) { return false; } @@ -19696,7 +19796,7 @@ static bool metal_graph_seed_streaming_expert_cache_from_prefill( const uint64_t bytes = (uint64_t)DS4_N_LAYER * DS4_STREAMING_PREFILL_CACHE_SEED_MAX_TOKENS * DS4_N_EXPERT_USED * sizeof(selected[0]); - if (ds4_gpu_tensor_read(g->prefill_seed_router_selected, + if (ds4_gpu_tensor_read(g->scr->prefill_seed_router_selected, 0, selected, bytes) == 0) { @@ -19904,14 +20004,14 @@ static bool metal_graph_eval_token_raw_swa_top( if (ok) ok = metal_graph_encode_token_raw_swa(g, model, weights, token, pos, true, true); if (ok) { - ok = ds4_gpu_argmax_tensor(g->comp_selected, - g->logits, + ok = ds4_gpu_argmax_tensor(g->scr->comp_selected, + g->scr->logits, DS4_N_VOCAB) != 0; } if (ok) ok = ds4_gpu_end_commands() != 0; - if (ok) ok = ds4_gpu_tensor_read(g->comp_selected, 0, top_id, sizeof(*top_id)) != 0; + if (ok) ok = ds4_gpu_tensor_read(g->scr->comp_selected, 0, top_id, sizeof(*top_id)) != 0; if (ok && logits) { - ok = ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (!ok) { if (ds4_gpu_synchronize() == 0) { @@ -19941,10 +20041,10 @@ static bool metal_graph_eval_mtp_draft_from_hc( if (n_raw > g->raw_window) n_raw = g->raw_window; if (n_raw > g->raw_cap) n_raw = g->raw_cap; - ds4_gpu_tensor *saved_cur = g->cur_hc; - ds4_gpu_tensor *saved_after = g->after_ffn_hc; + ds4_gpu_tensor *saved_cur = g->scr->cur_hc; + ds4_gpu_tensor *saved_after = g->scr->after_ffn_hc; bool ok = ds4_gpu_begin_commands() != 0; - if (ok) ok = ds4_gpu_embed_token_hc_tensor(g->mtp_embed, + if (ok) ok = ds4_gpu_embed_token_hc_tensor(g->scr->mtp_embed, base_model->map, base_model->size, base_weights->token_embd->abs_offset, @@ -19952,26 +20052,26 @@ static bool metal_graph_eval_mtp_draft_from_hc( (uint32_t)token, DS4_N_EMBD, 1) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->mtp_enorm, - g->mtp_embed, + if (ok) ok = ds4_gpu_rms_norm_weight_tensor(g->scr->mtp_enorm, + g->scr->mtp_embed, mtp_model->map, mtp_model->size, mtp->enorm->abs_offset, DS4_N_EMBD, DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->mtp_eproj, + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->mtp_eproj, mtp_model->map, mtp_model->size, mtp->e_proj->abs_offset, DS4_N_EMBD, DS4_N_EMBD, - g->mtp_enorm, + g->scr->mtp_enorm, 1) != 0; - if (ok) ok = ds4_gpu_repeat_hc_tensor(g->mtp_eproj_hc, - g->mtp_eproj, + if (ok) ok = ds4_gpu_repeat_hc_tensor(g->scr->mtp_eproj_hc, + g->scr->mtp_eproj, DS4_N_EMBD, DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->mtp_hnorm_hc, + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(g->scr->mtp_hnorm_hc, prev_hc, mtp_model->map, mtp_model->size, @@ -19979,21 +20079,21 @@ static bool metal_graph_eval_mtp_draft_from_hc( DS4_N_EMBD, DS4_N_HC, DS4_RMS_EPS) != 0; - if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->mtp_hproj_hc, + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(g->scr->mtp_hproj_hc, mtp_model->map, mtp_model->size, mtp->h_proj->abs_offset, DS4_N_EMBD, DS4_N_EMBD, - g->mtp_hnorm_hc, + g->scr->mtp_hnorm_hc, DS4_N_HC) != 0; - if (ok) ok = ds4_gpu_add_tensor(g->mtp_input_hc, - g->mtp_eproj_hc, - g->mtp_hproj_hc, + if (ok) ok = ds4_gpu_add_tensor(g->scr->mtp_input_hc, + g->scr->mtp_eproj_hc, + g->scr->mtp_hproj_hc, (uint32_t)hc_dim) != 0; if (ok) { - g->cur_hc = g->mtp_input_hc; - g->after_ffn_hc = out_hc; + g->scr->cur_hc = g->scr->mtp_input_hc; + g->scr->after_ffn_hc = out_hc; ok = metal_graph_encode_decode_layer(g, mtp_model, &mtp->block, @@ -20005,7 +20105,7 @@ static bool metal_graph_eval_mtp_draft_from_hc( n_raw, token); } - if (ok) g->cur_hc = out_hc; + if (ok) g->scr->cur_hc = out_hc; if (ok) ok = metal_graph_encode_output_head_mtp(g, base_model, base_weights, @@ -20013,25 +20113,25 @@ static bool metal_graph_eval_mtp_draft_from_hc( mtp, base_weights->output->dim[1]); if (ok && top_id) { - ok = ds4_gpu_argmax_tensor(g->comp_selected, - g->logits, + ok = ds4_gpu_argmax_tensor(g->scr->comp_selected, + g->scr->logits, DS4_N_VOCAB) != 0; } if (ok) ok = ds4_gpu_end_commands() != 0; - g->cur_hc = saved_cur; - g->after_ffn_hc = saved_after; + g->scr->cur_hc = saved_cur; + g->scr->after_ffn_hc = saved_after; if (ok && logits) { - ok = ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (ok && top_id) { - ok = ds4_gpu_tensor_read(g->comp_selected, 0, top_id, sizeof(*top_id)) != 0; + ok = ds4_gpu_tensor_read(g->scr->comp_selected, 0, top_id, sizeof(*top_id)) != 0; } if (ok && g->mtp_n_raw < g->raw_window) g->mtp_n_raw++; if (!ok) { (void)ds4_gpu_synchronize(); - g->cur_hc = saved_cur; - g->after_ffn_hc = saved_after; + g->scr->cur_hc = saved_cur; + g->scr->after_ffn_hc = saved_after; } return ok; } @@ -20051,8 +20151,8 @@ static bool metal_graph_eval_mtp_draft( base_weights, mtp_model, mtp, - g->cur_hc, - g->mtp_state_hc, + g->scr->cur_hc, + g->scr->mtp_state_hc, token, pos, logits, @@ -20138,14 +20238,14 @@ static bool imatrix_collect_layer_batch( const uint64_t norm_bytes = (uint64_t)n_tokens * DS4_N_EMBD * sizeof(float); const uint64_t mid_elems = (uint64_t)n_tokens * DS4_N_EXPERT_USED * DS4_N_FF_EXP; - const uint64_t mid_bytes = mid_elems * (g->batch_routed_mid_is_f16 ? sizeof(uint16_t) : sizeof(float)); + const uint64_t mid_bytes = mid_elems * (g->scr->batch_routed_mid_is_f16 ? sizeof(uint16_t) : sizeof(float)); const uint64_t sel_bytes = (uint64_t)n_tokens * DS4_N_EXPERT_USED * sizeof(int); - void *mid_dst = g->batch_routed_mid_is_f16 + void *mid_dst = g->scr->batch_routed_mid_is_f16 ? (void *)c->routed_mid_f16_buf : (void *)c->routed_mid_buf; - if (ds4_gpu_tensor_read(g->batch_ffn_norm, 0, c->ffn_norm_buf, norm_bytes) == 0 || - ds4_gpu_tensor_read(g->batch_routed_mid, 0, mid_dst, mid_bytes) == 0 || - ds4_gpu_tensor_read(g->batch_router_selected, 0, c->selected_buf, sel_bytes) == 0) + if (ds4_gpu_tensor_read(g->scr->batch_ffn_norm, 0, c->ffn_norm_buf, norm_bytes) == 0 || + ds4_gpu_tensor_read(g->scr->batch_routed_mid, 0, mid_dst, mid_bytes) == 0 || + ds4_gpu_tensor_read(g->scr->batch_router_selected, 0, c->selected_buf, sel_bytes) == 0) { return false; } @@ -20164,7 +20264,7 @@ static bool imatrix_collect_layer_batch( float *down = imatrix_down_ptr(c, il, (uint32_t)expert); const size_t mid_off = ((size_t)t * DS4_N_EXPERT_USED + slot) * DS4_N_FF_EXP; - if (g->batch_routed_mid_is_f16) { + if (g->scr->batch_routed_mid_is_f16) { const uint16_t *mid = c->routed_mid_f16_buf + mid_off; for (uint32_t i = 0; i < DS4_N_FF_EXP; i++) { const float v = f16_to_f32(mid[i]); @@ -20326,7 +20426,7 @@ static bool metal_graph_prefill_layer_major( if (display_progress) display_progress(display_progress_ud, "prefill_display", (int)start, prompt->len); - bool ok = metal_graph_upload_prompt_tokens(g->prefill_tokens, prompt, start, n_tokens); + bool ok = metal_graph_upload_prompt_tokens(g->scr->prefill_tokens, prompt, start, n_tokens); if (!ok) return false; #ifdef DS4_ROCM_BUILD @@ -20358,8 +20458,8 @@ static bool metal_graph_prefill_layer_major( double execute_s = 0.0; if (!split_commands) { - ok = metal_graph_upload_prompt_embeddings_hc(g->batch_cur_hc, - g->prefill_tokens, + ok = metal_graph_upload_prompt_embeddings_hc(g->scr->batch_cur_hc, + g->scr->prefill_tokens, model, weights, prompt, @@ -20396,22 +20496,22 @@ static bool metal_graph_prefill_layer_major( output_row = (uint32_t)v; } } - ds4_gpu_tensor *saved_cur = g->cur_hc; + ds4_gpu_tensor *saved_cur = g->scr->cur_hc; ds4_gpu_tensor *last_hc = NULL; if (ok && logits) { - last_hc = metal_graph_tensor_row_view(g->batch_cur_hc, output_row, hc_dim); + last_hc = metal_graph_tensor_row_view(g->scr->batch_cur_hc, output_row, hc_dim); ok = last_hc != NULL; } if (ok && logits) { - g->cur_hc = last_hc; + g->scr->cur_hc = last_hc; ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); - g->cur_hc = saved_cur; + g->scr->cur_hc = saved_cur; } const double t_encoded = profile ? now_sec() : 0.0; if (ok) ok = ds4_gpu_end_commands() != 0; const double t_done = profile ? now_sec() : 0.0; - g->cur_hc = saved_cur; + g->scr->cur_hc = saved_cur; if (last_hc) ds4_gpu_tensor_free(last_hc); if (!ok) { if (ds4_gpu_synchronize() == 0) { @@ -20422,7 +20522,7 @@ static bool metal_graph_prefill_layer_major( const double t_before_read = profile ? now_sec() : 0.0; if (logits) { - ok = ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (profile) { const double t_read = now_sec(); @@ -20505,8 +20605,8 @@ static bool metal_graph_prefill_layer_major( #endif double t_layer0 = (profile || throttle) ? now_sec() : 0.0; - ok = metal_graph_upload_prompt_embeddings_hc(g->batch_cur_hc, - g->prefill_tokens, + ok = metal_graph_upload_prompt_embeddings_hc(g->scr->batch_cur_hc, + g->scr->prefill_tokens, model, weights, prompt, @@ -20659,9 +20759,9 @@ static bool metal_graph_prefill_layer_major( fprintf(stderr, "ds4: gpu layer-major prefill layer %u ffn encode failed\n", il); } if (ok) { - ds4_gpu_tensor *tmp = g->batch_cur_hc; - g->batch_cur_hc = g->batch_next_hc; - g->batch_next_hc = tmp; + ds4_gpu_tensor *tmp = g->scr->batch_cur_hc; + g->scr->batch_cur_hc = g->scr->batch_next_hc; + g->scr->batch_next_hc = tmp; } if (ok) ok = metal_graph_capture_prefill_seed_router_selected(g, il, @@ -20828,12 +20928,12 @@ static bool metal_graph_prefill_layer_major( output_row = (uint32_t)v; } } - ds4_gpu_tensor *saved_cur = g->cur_hc; + ds4_gpu_tensor *saved_cur = g->scr->cur_hc; ds4_gpu_tensor *last_hc = NULL; const double t_head0 = profile ? now_sec() : 0.0; if (logits) { - last_hc = metal_graph_tensor_row_view(g->batch_cur_hc, + last_hc = metal_graph_tensor_row_view(g->scr->batch_cur_hc, output_row, hc_dim); ok = last_hc != NULL; @@ -20853,20 +20953,20 @@ static bool metal_graph_prefill_layer_major( } } if (ok && logits) { - g->cur_hc = last_hc; + g->scr->cur_hc = last_hc; ok = ds4_gpu_begin_commands() != 0; } if (ok && logits) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); const double t_head_encoded = profile ? now_sec() : 0.0; if (ok && logits) ok = ds4_gpu_end_commands() != 0; const double t_head_done = profile ? now_sec() : 0.0; - g->cur_hc = saved_cur; + g->scr->cur_hc = saved_cur; if (last_hc) ds4_gpu_tensor_free(last_hc); if (!ok) return false; const double t_before_read = profile ? now_sec() : 0.0; if (logits) { - ok = ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (profile) { const double t_read = now_sec(); @@ -21124,14 +21224,14 @@ static bool metal_graph_verify_suffix_tops( bool capture_prefix1, int *row_tops, float *row_logits) { - if (n_tokens == 0 || n_tokens > g->prefill_cap || !g->spec_logits) return false; + if (n_tokens == 0 || n_tokens > g->prefill_cap || !g->scr->spec_logits) return false; if (start > (uint32_t)prompt->len || n_tokens > (uint32_t)prompt->len - start) return false; const uint32_t top_rows = n_tokens > 1 ? n_tokens - 1 : 0; if (top_rows && !row_tops) return false; - bool ok = metal_graph_upload_prompt_tokens(g->prefill_tokens, prompt, start, n_tokens); - if (ok) ok = metal_graph_upload_prompt_embeddings_hc(g->batch_cur_hc, - g->prefill_tokens, + bool ok = metal_graph_upload_prompt_tokens(g->scr->prefill_tokens, prompt, start, n_tokens); + if (ok) ok = metal_graph_upload_prompt_embeddings_hc(g->scr->batch_cur_hc, + g->scr->prefill_tokens, model, weights, prompt, @@ -21167,15 +21267,15 @@ static bool metal_graph_verify_suffix_tops( /* Common K=2 verify case: top_k=1 over n_vocab → use the dedicated * argmax kernel (single-block tree-reduce) instead of the legacy * indexer_topk_kernel's single-thread O(n_vocab * top_k) fall-through. */ - ok = ds4_gpu_argmax_tensor(g->comp_selected, - g->spec_logits, + ok = ds4_gpu_argmax_tensor(g->scr->comp_selected, + g->scr->spec_logits, DS4_N_VOCAB) != 0; } else if (top_rows) { /* top-1 of each of the top_rows rows: n_tokens=top_rows, top_k=1. * The order is transposed vs the indexer-score callers; a swap * silently scores row 0's runner-ups instead of each row. */ - ok = ds4_gpu_indexer_topk_tensor(g->comp_selected, - g->spec_logits, + ok = ds4_gpu_indexer_topk_tensor(g->scr->comp_selected, + g->scr->spec_logits, DS4_N_VOCAB, top_rows, 1) != 0; @@ -21184,13 +21284,13 @@ static bool metal_graph_verify_suffix_tops( if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); if (ok && top_rows) { - ok = ds4_gpu_tensor_read(g->comp_selected, + ok = ds4_gpu_tensor_read(g->scr->comp_selected, 0, row_tops, (uint64_t)top_rows * sizeof(row_tops[0])) != 0; } if (ok && row_logits) { - ok = ds4_gpu_tensor_read(g->spec_logits, + ok = ds4_gpu_tensor_read(g->scr->spec_logits, 0, row_logits, (uint64_t)n_tokens * DS4_N_VOCAB * sizeof(row_logits[0])) != 0; @@ -21199,9 +21299,9 @@ static bool metal_graph_verify_suffix_tops( } static bool metal_graph_read_spec_logits_row(ds4_gpu_graph *g, uint32_t row, float *logits) { - if (!g || !g->spec_logits || !logits || row >= g->prefill_cap) return false; + if (!g || !g->scr->spec_logits || !logits || row >= g->prefill_cap) return false; const uint64_t row_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float); - return ds4_gpu_tensor_read(g->spec_logits, + return ds4_gpu_tensor_read(g->scr->spec_logits, (uint64_t)row * row_bytes, logits, row_bytes) != 0; @@ -21228,10 +21328,10 @@ static bool metal_graph_verify_decode2_exact( if (!g || !top0 || !logits1 || g->raw_cap == 0) return false; const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; - ds4_gpu_tensor *cur0 = metal_graph_tensor_row_view(g->batch_cur_hc, 0, hc_dim); - ds4_gpu_tensor *cur1 = metal_graph_tensor_row_view(g->batch_cur_hc, 1, hc_dim); - ds4_gpu_tensor *next0 = metal_graph_tensor_row_view(g->batch_next_hc, 0, hc_dim); - ds4_gpu_tensor *next1 = metal_graph_tensor_row_view(g->batch_next_hc, 1, hc_dim); + ds4_gpu_tensor *cur0 = metal_graph_tensor_row_view(g->scr->batch_cur_hc, 0, hc_dim); + ds4_gpu_tensor *cur1 = metal_graph_tensor_row_view(g->scr->batch_cur_hc, 1, hc_dim); + ds4_gpu_tensor *next0 = metal_graph_tensor_row_view(g->scr->batch_next_hc, 0, hc_dim); + ds4_gpu_tensor *next1 = metal_graph_tensor_row_view(g->scr->batch_next_hc, 1, hc_dim); bool ok = cur0 && cur1 && next0 && next1; if (ok) ok = ds4_gpu_embed_token_hc_tensor(cur0, @@ -21251,8 +21351,8 @@ static bool metal_graph_verify_decode2_exact( DS4_N_EMBD, DS4_N_HC) != 0; - ds4_gpu_tensor *saved_cur = g->cur_hc; - ds4_gpu_tensor *saved_after = g->after_ffn_hc; + ds4_gpu_tensor *saved_cur = g->scr->cur_hc; + ds4_gpu_tensor *saved_after = g->scr->after_ffn_hc; const bool saved_capture = g->spec_capture_prefix1; g->spec_capture_prefix1 = true; if (ok) ok = ds4_gpu_begin_commands() != 0; @@ -21260,8 +21360,8 @@ static bool metal_graph_verify_decode2_exact( const uint32_t pos0 = start; const uint32_t pos1 = start + 1u; - g->cur_hc = cur0; - g->after_ffn_hc = next0; + g->scr->cur_hc = cur0; + g->scr->after_ffn_hc = next0; ok = metal_graph_encode_decode_layer(g, model, &weights->layer[il], @@ -21277,8 +21377,8 @@ static bool metal_graph_verify_decode2_exact( metal_graph_capture_prefix1_index_state(g, il); if (!ok) break; - g->cur_hc = cur1; - g->after_ffn_hc = next1; + g->scr->cur_hc = cur1; + g->scr->after_ffn_hc = next1; ok = metal_graph_encode_decode_layer(g, model, &weights->layer[il], @@ -21297,22 +21397,22 @@ static bool metal_graph_verify_decode2_exact( if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); g->spec_capture_prefix1 = saved_capture; - g->cur_hc = saved_cur; - g->after_ffn_hc = saved_after; + g->scr->cur_hc = saved_cur; + g->scr->after_ffn_hc = saved_after; if (ok) { - g->cur_hc = cur0; + g->scr->cur_hc = cur0; ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); - if (ok) ok = ds4_gpu_argmax_tensor(g->comp_selected, - g->logits, + if (ok) ok = ds4_gpu_argmax_tensor(g->scr->comp_selected, + g->scr->logits, DS4_N_VOCAB) != 0; if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); - g->cur_hc = saved_cur; - if (ok) ok = ds4_gpu_tensor_read(g->comp_selected, 0, top0, sizeof(*top0)) != 0; + g->scr->cur_hc = saved_cur; + if (ok) ok = ds4_gpu_tensor_read(g->scr->comp_selected, 0, top0, sizeof(*top0)) != 0; if (ok && logits0) { - ok = ds4_gpu_tensor_read(g->logits, + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits0, (uint64_t)DS4_N_VOCAB * sizeof(logits0[0])) != 0; @@ -21320,21 +21420,21 @@ static bool metal_graph_verify_decode2_exact( } if (ok) { - g->cur_hc = cur1; + g->scr->cur_hc = cur1; ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); if (ok) ok = ds4_gpu_end_commands() != 0; else (void)ds4_gpu_synchronize(); - g->cur_hc = saved_cur; + g->scr->cur_hc = saved_cur; if (ok) { - ok = ds4_gpu_tensor_read(g->logits, + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits1, (uint64_t)DS4_N_VOCAB * sizeof(logits1[0])) != 0; } } - g->cur_hc = saved_cur; - g->after_ffn_hc = saved_after; + g->scr->cur_hc = saved_cur; + g->scr->after_ffn_hc = saved_after; g->spec_capture_prefix1 = saved_capture; ds4_gpu_tensor_free(next1); @@ -21509,10 +21609,13 @@ static int metal_graph_prompt_logits_test( const uint32_t raw_cap = metal_graph_raw_cap_for_context(ctx_size, (uint32_t)n_test); ds4_gpu_graph g; - bool ok = metal_graph_alloc_raw_cap(&g, weights, &weights->layer[0], + ds4_gpu_scratch scr; + memset(&scr, 0, sizeof(scr)); + bool ok = metal_graph_alloc_raw_cap(&g, &scr, weights, &weights->layer[0], raw_cap, (uint32_t)ctx_size, (uint32_t)n_test, false); if (!ok) { metal_graph_free(&g); + metal_scratch_free(&scr); fprintf(stderr, "ds4: failed to initialize Metal graph prompt test runtime\n"); return 1; } @@ -21664,6 +21767,7 @@ static int metal_graph_prompt_logits_test( free(oracle_logits); kv_cache_free(&cpu_cache); metal_graph_free(&g); + metal_scratch_free(&scr); return ok ? 0 : 1; } @@ -21811,6 +21915,12 @@ struct ds4_engine { ds4_vocab vocab; ds4_weights weights; ds4_mtp_weights mtp_weights; +#ifndef DS4_NO_GPU + /* Work buffers shared by every session of this engine (see + * ds4_gpu_scratch). Grown on demand by session creation, freed with the + * engine. */ + ds4_gpu_scratch gpu_scratch; +#endif ds4_backend backend; int mtp_draft_tokens; float mtp_margin; @@ -22955,7 +23065,9 @@ static int generate_metal_graph_raw_swa( prompt->len); } ds4_gpu_graph g; - bool ok = metal_graph_alloc_raw_cap(&g, weights, &weights->layer[0], + ds4_gpu_scratch scr; + memset(&scr, 0, sizeof(scr)); + bool ok = metal_graph_alloc_raw_cap(&g, &scr, weights, &weights->layer[0], raw_cap, (uint32_t)ctx_size, prefill_cap, false); if (!ok) { fprintf(stderr, "ds4: failed to allocate GPU graph runtime\n"); @@ -22971,6 +23083,7 @@ static int generate_metal_graph_raw_swa( directional_steering_attn, directional_steering_ffn)) { metal_graph_free(&g); + metal_scratch_free(&scr); return 1; } const bool memory_report = getenv("DS4_METAL_MEMORY_REPORT") != NULL; @@ -22999,6 +23112,7 @@ static int generate_metal_graph_raw_swa( if (!ok) { free(logits); metal_graph_free(&g); + metal_scratch_free(&scr); return 1; } const char *dump_prefill_logits = getenv("DS4_METAL_DUMP_PREFILL_LOGITS"); @@ -23006,6 +23120,7 @@ static int generate_metal_graph_raw_swa( if (!write_f32_binary_file(dump_prefill_logits, logits, DS4_N_VOCAB)) { free(logits); metal_graph_free(&g); + metal_scratch_free(&scr); return 1; } fprintf(stderr, "ds4: wrote GPU prefill logits to %s\n", dump_prefill_logits); @@ -23062,6 +23177,7 @@ static int generate_metal_graph_raw_swa( if (memory_report) ds4_gpu_print_memory_report("before graph free"); free(logits); metal_graph_free(&g); + metal_scratch_free(&scr); return ok ? 0 : 1; } #endif @@ -25022,7 +25138,9 @@ int ds4_engine_collect_imatrix(ds4_engine *e, const uint32_t raw_cap = metal_graph_raw_cap_for_context(ctx_size, prefill_cap); ds4_gpu_graph g; - bool ok = metal_graph_alloc_raw_cap(&g, weights, &weights->layer[0], + ds4_gpu_scratch scr; + memset(&scr, 0, sizeof(scr)); + bool ok = metal_graph_alloc_raw_cap(&g, &scr, weights, &weights->layer[0], raw_cap, (uint32_t)ctx_size, prefill_cap, false); if (!ok) { fprintf(stderr, "ds4: failed to allocate imatrix Metal graph runtime\n"); @@ -25039,6 +25157,7 @@ int ds4_engine_collect_imatrix(ds4_engine *e, if (!imatrix_collector_init(&collector, prefill_cap, dataset_path)) { fprintf(stderr, "ds4: failed to allocate imatrix collector\n"); metal_graph_free(&g); + metal_scratch_free(&scr); free(dataset); return 1; } @@ -25135,6 +25254,7 @@ int ds4_engine_collect_imatrix(ds4_engine *e, imatrix_collector_free(&collector); metal_graph_free(&g); + metal_scratch_free(&scr); free(dataset); return ok ? 0 : 1; #endif @@ -26027,6 +26147,12 @@ void ds4_engine_close(ds4_engine *e) { if (e->mtp_ready) model_close(&e->mtp_model); model_close(&e->model); #ifndef DS4_NO_GPU + if (e->gpu_scratch.refcount > 0) { + fprintf(stderr, + "ds4: warning: closing engine with %d live session(s)\n", + e->gpu_scratch.refcount); + } + metal_scratch_free(&e->gpu_scratch); ds4_gpu_cleanup(); #endif ds4_ssd_memory_lock_release(&e->simulated_memory); @@ -26071,7 +26197,7 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { free(s); return 1; } - if (!metal_graph_alloc_raw_cap(&s->graph, &e->weights, shape_layer, + if (!metal_graph_alloc_raw_cap(&s->graph, &e->gpu_scratch, &e->weights, shape_layer, raw_cap, (uint32_t)ctx_size, s->prefill_cap, e->mtp_ready)) { free(s); @@ -26247,7 +26373,7 @@ int ds4_session_eval_output_head_from_hc(ds4_session *s, return 1; #else ds4_gpu_graph *g = &s->graph; - bool ok = ds4_gpu_tensor_write(g->cur_hc, + bool ok = ds4_gpu_tensor_write(g->scr->cur_hc, 0, last_hc, hc_dim * sizeof(float)) != 0; @@ -26257,7 +26383,7 @@ int ds4_session_eval_output_head_from_hc(ds4_session *s, &e->weights, e->weights.output->dim[1]); if (ok) ok = ds4_gpu_end_commands() != 0; - if (ok) ok = ds4_gpu_tensor_read(g->logits, + if (ok) ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; @@ -26461,11 +26587,11 @@ int ds4_session_eval_layer_slice(ds4_session *s, ok = metal_graph_stream_map_token(&e->model, &e->weights); } if (input_hc) { - ok = ds4_gpu_tensor_write(g->cur_hc, 0, input_hc, hc_dim * sizeof(float)) != 0; + ok = ds4_gpu_tensor_write(g->scr->cur_hc, 0, input_hc, hc_dim * sizeof(float)) != 0; } if (ok) ok = ds4_gpu_begin_commands() != 0; if (ok && !input_hc) { - ok = ds4_gpu_embed_token_hc_tensor(g->cur_hc, + ok = ds4_gpu_embed_token_hc_tensor(g->scr->cur_hc, e->model.map, e->model.size, e->weights.token_embd->abs_offset, @@ -26495,9 +26621,9 @@ int ds4_session_eval_layer_slice(ds4_session *s, raw_row, n_raw, tokens[0]); - ds4_gpu_tensor *tmp = g->cur_hc; - g->cur_hc = g->after_ffn_hc; - g->after_ffn_hc = tmp; + ds4_gpu_tensor *tmp = g->scr->cur_hc; + g->scr->cur_hc = g->scr->after_ffn_hc; + g->scr->after_ffn_hc = tmp; } if (ok) ok = ds4_gpu_end_commands() != 0; } @@ -26520,9 +26646,9 @@ int ds4_session_eval_layer_slice(ds4_session *s, raw_row, n_raw, tokens[0]); - ds4_gpu_tensor *tmp = g->cur_hc; - g->cur_hc = g->after_ffn_hc; - g->after_ffn_hc = tmp; + ds4_gpu_tensor *tmp = g->scr->cur_hc; + g->scr->cur_hc = g->scr->after_ffn_hc; + g->scr->after_ffn_hc = tmp; encoded_layers++; if (ok && split_after_layers != 0 && @@ -26539,10 +26665,10 @@ int ds4_session_eval_layer_slice(ds4_session *s, } if (ok && !output_hc && !output_logits) ok = ds4_gpu_synchronize() != 0; if (ok && output_hc) { - ok = ds4_gpu_tensor_read(g->cur_hc, 0, output_hc, hc_dim * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g->scr->cur_hc, 0, output_hc, hc_dim * sizeof(float)) != 0; } if (ok && output_logits) { - ok = ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (!ok) { if (ds4_gpu_synchronize() == 0) { @@ -26569,12 +26695,12 @@ int ds4_session_eval_layer_slice(ds4_session *s, g->streaming_static_decode_map_current = false; ok = metal_graph_stream_map_token(&e->model, &e->weights); } - if (ok) ok = metal_graph_upload_prompt_tokens(g->prefill_tokens, &span, 0, n_tokens); + if (ok) ok = metal_graph_upload_prompt_tokens(g->scr->prefill_tokens, &span, 0, n_tokens); if (ok && input_hc) { - ok = ds4_gpu_tensor_write(g->batch_cur_hc, 0, input_hc, hc_bytes) != 0; + ok = ds4_gpu_tensor_write(g->scr->batch_cur_hc, 0, input_hc, hc_bytes) != 0; } else if (ok) { - ok = metal_graph_upload_prompt_embeddings_hc(g->batch_cur_hc, - g->prefill_tokens, + ok = metal_graph_upload_prompt_embeddings_hc(g->scr->batch_cur_hc, + g->scr->prefill_tokens, &e->model, &e->weights, &span, @@ -26618,31 +26744,31 @@ int ds4_session_eval_layer_slice(ds4_session *s, } } if (ok && output_logits) { - saved_cur = g->cur_hc; - last_hc = metal_graph_tensor_row_view(g->batch_cur_hc, n_tokens - 1u, hc_dim); + saved_cur = g->scr->cur_hc; + last_hc = metal_graph_tensor_row_view(g->scr->batch_cur_hc, n_tokens - 1u, hc_dim); ok = last_hc != NULL; if (ok && g->ssd_streaming) { g->streaming_static_decode_map_current = false; ok = metal_graph_stream_map_output(&e->model, &e->weights); } if (ok) { - g->cur_hc = last_hc; + g->scr->cur_hc = last_hc; if (g->ssd_streaming) ok = ds4_gpu_begin_commands() != 0; if (ok) ok = metal_graph_encode_output_head(g, &e->model, &e->weights, e->weights.output->dim[1]); if (ok && g->ssd_streaming) ok = ds4_gpu_end_commands() != 0; - g->cur_hc = saved_cur; + g->scr->cur_hc = saved_cur; } } if (ok && !g->ssd_streaming) ok = ds4_gpu_end_commands() != 0; - if (saved_cur) g->cur_hc = saved_cur; + if (saved_cur) g->scr->cur_hc = saved_cur; if (last_hc) ds4_gpu_tensor_free(last_hc); if (ok && !output_hc && !output_logits) ok = ds4_gpu_synchronize() != 0; if (ok && output_hc) { - ok = ds4_gpu_tensor_read(g->batch_cur_hc, 0, output_hc, hc_bytes) != 0; + ok = ds4_gpu_tensor_read(g->scr->batch_cur_hc, 0, output_hc, hc_bytes) != 0; } if (ok && output_logits) { - ok = ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + ok = ds4_gpu_tensor_read(g->scr->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; } if (!ok) { if (ds4_gpu_synchronize() == 0) { @@ -27262,8 +27388,8 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, } while (0) for (; draft_n < draft_cap; draft_n++) { - ds4_gpu_tensor *prev_hc = (draft_n & 1) ? s->graph.mtp_state_hc : s->graph.mtp_next_hc; - ds4_gpu_tensor *out_hc = (draft_n & 1) ? s->graph.mtp_next_hc : s->graph.mtp_state_hc; + ds4_gpu_tensor *prev_hc = (draft_n & 1) ? s->graph.scr->mtp_state_hc : s->graph.scr->mtp_next_hc; + ds4_gpu_tensor *out_hc = (draft_n & 1) ? s->graph.scr->mtp_next_hc : s->graph.scr->mtp_state_hc; int mtp_top = -1; if (!metal_graph_eval_mtp_draft_from_hc(&s->graph, &e->model, @@ -27724,7 +27850,7 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, if (drafts[i] == eos_token) break; } if (verified > 0 && !logits_on_host) { - if (ds4_gpu_tensor_read(s->graph.logits, + if (ds4_gpu_tensor_read(s->graph.scr->logits, 0, s->logits, (uint64_t)DS4_N_VOCAB * sizeof(s->logits[0])) == 0) diff --git a/ds4.h b/ds4.h index 9d040c92b..012ac3b21 100644 --- a/ds4.h +++ b/ds4.h @@ -217,6 +217,12 @@ int ds4_token_eos(ds4_engine *e); int ds4_token_user(ds4_engine *e); int ds4_token_assistant(ds4_engine *e); +/* Multiple sessions may share one engine: each session owns its live KV + * cache and timeline, while transient work buffers are engine-owned and + * shared. Sessions sharing an engine must therefore be driven from one + * thread at a time (create/sync/eval/sample/free); interleaving calls across + * sessions on that thread is fine, concurrent calls are not. This matches + * the server's single graph worker. */ 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); diff --git a/tests/two_sessions_smoke.c b/tests/two_sessions_smoke.c new file mode 100644 index 000000000..bb24346e3 --- /dev/null +++ b/tests/two_sessions_smoke.c @@ -0,0 +1,144 @@ +/* two_sessions_smoke.c — prove that two ds4_session objects sharing one + * ds4_engine produce the same greedy token streams when driven interleaved + * (one token each, alternating on the same thread) as when each session is + * driven alone. This is the ground assumption for multi-session serving: + * backend statics must not carry per-graph state across ds4_session_eval() + * calls of different sessions. + * + * Usage: tests/two_sessions_smoke [n_tokens] + */ +#include +#include +#include + +#include "../ds4.h" + +#define MAX_GEN 128 + +typedef struct { + int toks[MAX_GEN]; + int len; +} gen_result; + +static void die(const char *msg, const char *detail) { + fprintf(stderr, "two_sessions_smoke: %s%s%s\n", msg, + detail ? ": " : "", detail ? detail : ""); + exit(1); +} + +static void encode_prompt(ds4_engine *e, const char *system, const char *user, + ds4_tokens *out) { + memset(out, 0, sizeof(*out)); + ds4_encode_chat_prompt(e, system, user, DS4_THINK_NONE, out); +} + +/* One greedy step: argmax over current logits, then advance the session. + * Returns the sampled token, or -1 once the stream hit EOS. */ +static int greedy_step(ds4_session *s, int eos, char *err, size_t errlen) { + int tok = ds4_session_argmax(s); + if (tok < 0 || tok == eos) return -1; + if (ds4_session_eval(s, tok, err, errlen) != 0) die("eval failed", err); + return tok; +} + +static void run_isolated(ds4_engine *e, const ds4_tokens *prompt, int n, + int ctx_size, gen_result *out) { + ds4_session *s = NULL; + char err[256] = ""; + if (ds4_session_create(&s, e, ctx_size) != 0) die("session_create failed", NULL); + if (ds4_session_sync(s, prompt, err, sizeof(err)) != 0) die("sync failed", err); + int eos = ds4_token_eos(e); + out->len = 0; + for (int i = 0; i < n; i++) { + int tok = greedy_step(s, eos, err, sizeof(err)); + if (tok < 0) break; + out->toks[out->len++] = tok; + } + ds4_session_free(s); +} + +static void run_interleaved(ds4_engine *e, const ds4_tokens *pa, + const ds4_tokens *pb, int n, int ctx_size, + gen_result *ra, gen_result *rb) { + ds4_session *sa = NULL, *sb = NULL; + char err[256] = ""; + if (ds4_session_create(&sa, e, ctx_size) != 0) die("session_create A failed", NULL); + if (ds4_session_create(&sb, e, ctx_size) != 0) die("session_create B failed", NULL); + if (ds4_session_sync(sa, pa, err, sizeof(err)) != 0) die("sync A failed", err); + if (ds4_session_sync(sb, pb, err, sizeof(err)) != 0) die("sync B failed", err); + int eos = ds4_token_eos(e); + ra->len = rb->len = 0; + bool a_done = false, b_done = false; + for (int i = 0; i < n && (!a_done || !b_done); i++) { + if (!a_done) { + int tok = greedy_step(sa, eos, err, sizeof(err)); + if (tok < 0) a_done = true; else ra->toks[ra->len++] = tok; + } + if (!b_done) { + int tok = greedy_step(sb, eos, err, sizeof(err)); + if (tok < 0) b_done = true; else rb->toks[rb->len++] = tok; + } + } + ds4_session_free(sa); + ds4_session_free(sb); +} + +static void print_stream(ds4_engine *e, const char *label, const gen_result *r) { + printf("%s (%d tokens): ", label, r->len); + for (int i = 0; i < r->len; i++) { + size_t len = 0; + char *txt = ds4_token_text(e, r->toks[i], &len); + if (txt) fwrite(txt, 1, len, stdout); + } + printf("\n"); +} + +static bool same(const gen_result *x, const gen_result *y) { + if (x->len != y->len) return false; + return memcmp(x->toks, y->toks, sizeof(int) * (size_t)x->len) == 0; +} + +int main(int argc, char **argv) { + if (argc < 2) die("usage: two_sessions_smoke [n_tokens]", NULL); + int n = argc > 2 ? atoi(argv[2]) : 24; + if (n <= 0 || n > MAX_GEN) n = 24; + int ctx_size = 4096; + + ds4_engine_options opt; + memset(&opt, 0, sizeof(opt)); + opt.model_path = argv[1]; + opt.backend = DS4_BACKEND_CUDA; + opt.prefill_chunk = 256; + + ds4_engine *e = NULL; + if (ds4_engine_open(&e, &opt) != 0) die("engine_open failed", argv[1]); + + const char *sys = "You are a concise technical assistant."; + ds4_tokens pa, pb; + encode_prompt(e, sys, "Briefly explain what a compressor does in a refrigeration cycle.", &pa); + encode_prompt(e, sys, "List three common causes of high condensing pressure.", &pb); + + gen_result iso_a, iso_b, mix_a, mix_b; + printf("== isolated runs ==\n"); + run_isolated(e, &pa, n, ctx_size, &iso_a); + print_stream(e, "A/isolated", &iso_a); + run_isolated(e, &pb, n, ctx_size, &iso_b); + print_stream(e, "B/isolated", &iso_b); + + printf("== interleaved run ==\n"); + run_interleaved(e, &pa, &pb, n, ctx_size, &mix_a, &mix_b); + print_stream(e, "A/interleaved", &mix_a); + print_stream(e, "B/interleaved", &mix_b); + + bool ok_a = same(&iso_a, &mix_a); + bool ok_b = same(&iso_b, &mix_b); + printf("A: %s, B: %s\n", ok_a ? "MATCH" : "MISMATCH", ok_b ? "MATCH" : "MISMATCH"); + + ds4_tokens_free(&pa); + ds4_tokens_free(&pb); + ds4_engine_close(e); + + if (!ok_a || !ok_b) { printf("FAILED\n"); return 1; } + printf("OK\n"); + return 0; +} From 499ff452f66e3def2b28c5d0992b9d42da85e0de Mon Sep 17 00:00:00 2001 From: iCreil Date: Tue, 7 Jul 2026 11:21:34 +0200 Subject: [PATCH 2/3] Server: split generate_job into stepwise phase functions Move the request-scoped locals of generate_job() into a gen_state struct and split the body into job_begin / job_prefill / job_start_decode / job_decode_round_init / job_decode_step / job_finish. The decode_again tool-recovery label becomes a GEN_STEP_REDECODE result from job_finish. generate_job() is now a small driver that performs exactly the same calls in exactly the same order, so behavior is unchanged; this is groundwork for a scheduler that drives several jobs stepwise (one prefill chunk or a few decode tokens at a time) from the single graph worker instead of owning it until a job completes. Verified against the previous build on CUDA (DeepSeek V4 Flash q2): greedy chat completion is byte-identical in both stream and non-stream mode, and the SSE frame flow is intact. --- ds4_server.c | 1138 +++++++++++++++++++++++++++----------------------- 1 file changed, 626 insertions(+), 512 deletions(-) diff --git a/ds4_server.c b/ds4_server.c index 34a9d5084..53df363d4 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -9988,184 +9988,260 @@ static bool should_canonicalize_tool_checkpoint(const server *s, const tool_call * shorter than the full prompt, we prefill to that boundary, store it, and * immediately continue to the real prompt. The live graph therefore always * moves forward. */ -static void generate_job(server *s, job *j) { +/* --------------------------------------------------------------------------- + * Generation state machine. + * + * generate_job() used to be one long function whose stack frame carried a + * request through cache resolution, prefill, decode and the final response. + * That frame now lives in gen_state and each stage is its own function, so a + * future scheduler can drive several jobs stepwise (one prefill chunk or a + * few decode tokens at a time) instead of owning the worker until the job is + * done. With a single session the driver below performs exactly the same + * calls in exactly the same order as the old function body. + * ------------------------------------------------------------------------ */ + +typedef enum { + GEN_STEP_DONE = 0, /* response fully sent (or client gone) */ + GEN_STEP_REDECODE, /* tool-error recovery wants another decode round */ +} gen_step; + +typedef struct { char err[160]; - err[0] = '\0'; - const int old_pos = ds4_session_pos(s->session); - const int common = ds4_session_common_prefix(s->session, &j->req.prompt); - trace_cache_diag cache_diag = {0}; - trace_cache_capture(&cache_diag, ds4_session_tokens(s->session), - &j->req.prompt, old_pos, common); - ds4_tokens effective_prompt = {0}; - const ds4_tokens *prompt_for_sync = &j->req.prompt; - const bool responses_protocol = j->req.api == API_RESPONSES; - bool responses_live_continuation = false; - bool anthropic_live_continuation = false; - bool thinking_live_continuation = false; - const char *responses_live_match = NULL; - int responses_live_match_ids = 0; - int anthropic_live_match_ids = 0; + int old_pos; + int common; + trace_cache_diag cache_diag; + ds4_tokens effective_prompt; + const ds4_tokens *prompt_for_sync; + bool responses_protocol; + bool responses_live_continuation; + bool anthropic_live_continuation; + bool thinking_live_continuation; + const char *responses_live_match; + int responses_live_match_ids; + int anthropic_live_match_ids; + int cached; + const char *cache_source; + int disk_cached; + char *disk_cache_path; + uint8_t disk_cache_ext_flags; + int prompt_tokens; + double t0; + uint64_t trace_id; + char ctx_span[48]; + server_prefill_progress progress; + char req_flags[64]; + int cold_store_len; + int suppressed_continued_last; + char id[96]; + bool structured_stream; + anthropic_stream anthropic_live; + openai_stream openai_live; + responses_stream responses_live; + bool openai_live_chat; + bool responses_live_chat; + long responses_created_at; + bool dsml_recovery_attempted; + uint64_t rng; + /* Per-decode-round state, reset by job_decode_round_init(). */ + buf text; + size_t plain_stream_pos; + size_t stop_scan_from; + const char *finish; + int completion; + int max_tokens; + bool saw_tool_start; + bool saw_tool_end; + bool saw_orphan_tool_end; + size_t tool_scan_from; + int next_tool_progress; + int next_decode_log; + double decode_t0; + double last_decode_log_t; + int last_decode_log_completion; + thinking_state thinking; + bool thinking_gates_tool_markers; + bool tool_scan_waiting_for_think_close; + size_t think_recovery_scan_from; + bool think_tool_recovery_enabled; + dsml_decode_tracker dsml_tracker; +} gen_state; + +/* Resolve continuations and caches, account usage, register the prefill + * progress callback and plan the cold disk checkpoint. Returns false when + * the request was already answered (continuation state unavailable). */ +static bool job_begin(server *s, job *j, gen_state *gs) { + gs->err[0] = '\0'; + gs->old_pos = ds4_session_pos(s->session); + gs->common = ds4_session_common_prefix(s->session, &j->req.prompt); + trace_cache_capture(&gs->cache_diag, ds4_session_tokens(s->session), + &j->req.prompt, gs->old_pos, gs->common); + gs->prompt_for_sync = &j->req.prompt; + gs->responses_protocol = j->req.api == API_RESPONSES; + gs->responses_live_continuation = false; + gs->anthropic_live_continuation = false; + gs->thinking_live_continuation = false; + gs->responses_live_match = NULL; + gs->responses_live_match_ids = 0; + gs->anthropic_live_match_ids = 0; /* Responses gets the first chance to continue from live state. This is * the whole point of the API shape: a request that is bound to prior live * output by visible transcript or tool call ids does not need to prove an - * exact token-prefix match. Exact token/text/disk matching remains the + * exact token-prefix match. Exact token/gs->text/disk matching remains the * fallback when the live state is absent or no longer describes the * request. */ - int cached = responses_live_visible_prefix_prompt(s, &j->req, old_pos, - &effective_prompt); - const char *cache_source = cached > 0 ? "responses-visible" : "none"; - if (cached > 0) { - responses_live_match = "visible-prefix"; + gs->cached = responses_live_visible_prefix_prompt(s, &j->req, gs->old_pos, + &gs->effective_prompt); + gs->cache_source = gs->cached > 0 ? "responses-visible" : "none"; + if (gs->cached > 0) { + gs->responses_live_match = "visible-prefix"; if (responses_live_matches_request(s, &j->req.responses_live_call_ids, - old_pos)) + gs->old_pos)) { - responses_live_match_ids = j->req.responses_live_call_ids.len; + gs->responses_live_match_ids = j->req.responses_live_call_ids.len; } } - if (cached == 0) { - cached = responses_live_continuation_prompt(s, &j->req, old_pos, - &effective_prompt, - &responses_live_match_ids); - cache_source = cached > 0 ? "responses-tool-output" : "none"; - if (cached > 0) responses_live_match = "tool-output-ids"; + if (gs->cached == 0) { + gs->cached = responses_live_continuation_prompt(s, &j->req, gs->old_pos, + &gs->effective_prompt, + &gs->responses_live_match_ids); + gs->cache_source = gs->cached > 0 ? "responses-tool-output" : "none"; + if (gs->cached > 0) gs->responses_live_match = "tool-output-ids"; } - if (cached > 0) { - responses_live_continuation = true; - prompt_for_sync = &effective_prompt; + if (gs->cached > 0) { + gs->responses_live_continuation = true; + gs->prompt_for_sync = &gs->effective_prompt; } else { - cached = anthropic_live_continuation_prompt(s, &j->req, old_pos, - &effective_prompt, - &anthropic_live_match_ids); - if (cached > 0) { - anthropic_live_continuation = true; - cache_source = "anthropic-tool-output"; - prompt_for_sync = &effective_prompt; + gs->cached = anthropic_live_continuation_prompt(s, &j->req, gs->old_pos, + &gs->effective_prompt, + &gs->anthropic_live_match_ids); + if (gs->cached > 0) { + gs->anthropic_live_continuation = true; + gs->cache_source = "anthropic-tool-output"; + gs->prompt_for_sync = &gs->effective_prompt; } } - if (cached == 0 && responses_protocol && + if (gs->cached == 0 && gs->responses_protocol && j->req.responses_requires_live_tool_state) { /* The parser saw a valid live call_id, but by worker execution time the * live frontier no longer matches. Since the request did not replay * the prior assistant call, there is no stateless prefix to match and * no disk key to search by. */ - ds4_tokens_free(&effective_prompt); + ds4_tokens_free(&gs->effective_prompt); http_error(j->fd, s->enable_cors, 409, "Responses continuation state is not available; retry by replaying the full input history"); - return; - } else if (cached == 0 && j->req.api == API_ANTHROPIC && + return false; + } else if (gs->cached == 0 && j->req.api == API_ANTHROPIC && j->req.anthropic_requires_live_tool_state) { - ds4_tokens_free(&effective_prompt); + ds4_tokens_free(&gs->effective_prompt); http_error(j->fd, s->enable_cors, 409, "Anthropic continuation state is not available; retry by replaying the full messages history"); - return; - } else if (cached == 0) { - cached = common == old_pos && j->req.prompt.len >= old_pos ? common : 0; - cache_source = cached > 0 ? "memory-token" : "none"; + return false; + } else if (gs->cached == 0) { + gs->cached = gs->common == gs->old_pos && j->req.prompt.len >= gs->old_pos ? gs->common : 0; + gs->cache_source = gs->cached > 0 ? "memory-token" : "none"; } - if (cached == 0) { + if (gs->cached == 0) { int thinking_cached = - thinking_live_visible_prefix_prompt(s, &j->req, old_pos, - &effective_prompt); + thinking_live_visible_prefix_prompt(s, &j->req, gs->old_pos, + &gs->effective_prompt); if (thinking_cached > 0) { - cached = thinking_cached; - cache_source = "thinking-visible"; - thinking_live_continuation = true; - prompt_for_sync = &effective_prompt; + gs->cached = thinking_cached; + gs->cache_source = "thinking-visible"; + gs->thinking_live_continuation = true; + gs->prompt_for_sync = &gs->effective_prompt; } } - int disk_cached = 0; - char *disk_cache_path = NULL; - uint8_t disk_cache_ext_flags = 0; - if (cached == 0) { - int text_cached = live_text_prefix_prompt(s, &j->req, &effective_prompt); + gs->disk_cached = 0; + gs->disk_cache_path = NULL; + gs->disk_cache_ext_flags = 0; + if (gs->cached == 0) { + int text_cached = live_text_prefix_prompt(s, &j->req, &gs->effective_prompt); if (text_cached > 0) { - cached = text_cached; - cache_source = "memory-text"; - prompt_for_sync = &effective_prompt; + gs->cached = text_cached; + gs->cache_source = "memory-text"; + gs->prompt_for_sync = &gs->effective_prompt; } } - if (cached == 0 && old_pos > 0) { + if (gs->cached == 0 && gs->old_pos > 0) { server_log(DS4_LOG_WARNING, "ds4-server: live kv cache miss%s live=%d prompt=%d common=%d reason=%s", - responses_protocol ? " RESPPROTO" : "", - old_pos, j->req.prompt.len, common, - trace_cache_miss_reason(&cache_diag)); + gs->responses_protocol ? " RESPPROTO" : "", + gs->old_pos, j->req.prompt.len, gs->common, + trace_cache_miss_reason(&gs->cache_diag)); } - if (cached == 0) s->kv.continued_last_store_tokens = 0; - if (s->kv.enabled && cached == 0 && old_pos >= s->kv.opt.min_tokens) { + if (gs->cached == 0) s->kv.continued_last_store_tokens = 0; + if (s->kv.enabled && gs->cached == 0 && gs->old_pos >= s->kv.opt.min_tokens) { /* Loading a disk snapshot replaces the live Metal session. Persist the * current checkpoint first, otherwise a cache hit for an older prefix * would silently discard the newer conversation state. */ kv_cache_store_current(s, "evict"); } - if (cached == 0) { - disk_cached = kv_cache_try_load(s, &j->req, &effective_prompt, - &disk_cache_path, - &disk_cache_ext_flags); - if (disk_cached > 0) { - cached = disk_cached; - cache_source = "disk-text"; - prompt_for_sync = &effective_prompt; + if (gs->cached == 0) { + gs->disk_cached = kv_cache_try_load(s, &j->req, &gs->effective_prompt, + &gs->disk_cache_path, + &gs->disk_cache_ext_flags); + if (gs->disk_cached > 0) { + gs->cached = gs->disk_cached; + gs->cache_source = "disk-text"; + gs->prompt_for_sync = &gs->effective_prompt; } } const bool responses_reasoning_state_preserved = - cached > 0 && - ((!strcmp(cache_source, "responses-visible") || - !strcmp(cache_source, "responses-tool-output")) || - (!strcmp(cache_source, "disk-text") && - (disk_cache_ext_flags & KV_EXT_RESPONSES_VISIBLE))); + gs->cached > 0 && + ((!strcmp(gs->cache_source, "responses-visible") || + !strcmp(gs->cache_source, "responses-tool-output")) || + (!strcmp(gs->cache_source, "disk-text") && + (gs->disk_cache_ext_flags & KV_EXT_RESPONSES_VISIBLE))); const bool responses_visible_replay_without_reasoning = - responses_protocol && + gs->responses_protocol && j->req.responses_requires_live_reasoning && !responses_reasoning_state_preserved; - const int prompt_tokens = prompt_for_sync->len; + gs->prompt_tokens = gs->prompt_for_sync->len; /* OpenAI usage details: the reusable prefix is a cache read, while the * effective prompt suffix evaluated by ds4_session_sync() is written into * the live KV cache and can be reused by the next request. */ - j->req.cache_read_tokens = cached; - j->req.cache_write_tokens = prompt_tokens > cached ? prompt_tokens - cached : 0; - - const double t0 = now_sec(); - uint64_t trace_id = trace_begin(s, j, cached, prompt_tokens, &cache_diag, - cache_source, disk_cached, disk_cache_path); - char ctx_span[48]; - request_ctx_span(ctx_span, sizeof(ctx_span), cached, prompt_tokens); - server_prefill_progress progress = { + j->req.cache_read_tokens = gs->cached; + j->req.cache_write_tokens = gs->prompt_tokens > gs->cached ? gs->prompt_tokens - gs->cached : 0; + + gs->t0 = now_sec(); + gs->trace_id = trace_begin(s, j, gs->cached, gs->prompt_tokens, &gs->cache_diag, + gs->cache_source, gs->disk_cached, gs->disk_cache_path); + request_ctx_span(gs->ctx_span, sizeof(gs->ctx_span), gs->cached, gs->prompt_tokens); + gs->progress = (server_prefill_progress){ .srv = s, .kind = j->req.kind, - .prompt_tokens = prompt_tokens, - .cached_tokens = cached, + .prompt_tokens = gs->prompt_tokens, + .cached_tokens = gs->cached, .has_tools = j->req.has_tools, - .responses_protocol = responses_protocol, - .t0 = t0, + .responses_protocol = gs->responses_protocol, + .t0 = gs->t0, .fd = j->fd, .stream = j->req.stream, .enable_cors = s->enable_cors, }; - snprintf(progress.ctx, sizeof(progress.ctx), "%s", ctx_span); - char req_flags[64]; - log_flags(req_flags, sizeof(req_flags), responses_protocol, + snprintf(gs->progress.ctx, sizeof(gs->progress.ctx), "%s", gs->ctx_span); + log_flags(gs->req_flags, sizeof(gs->req_flags), gs->responses_protocol, j->req.has_tools, false, false, false); - if (responses_live_continuation) { + if (gs->responses_live_continuation) { server_log(DS4_LOG_PREFILL, "ds4-server: responses live continuation RESPPROTO match=%s ids=%d cached=%d prompt=%d", - responses_live_match ? responses_live_match : "unknown", - responses_live_match_ids, - cached, - prompt_tokens); - } else if (anthropic_live_continuation) { + gs->responses_live_match ? gs->responses_live_match : "unknown", + gs->responses_live_match_ids, + gs->cached, + gs->prompt_tokens); + } else if (gs->anthropic_live_continuation) { server_log(DS4_LOG_PREFILL, "ds4-server: anthropic live continuation match=tool-output-ids ids=%d cached=%d prompt=%d", - anthropic_live_match_ids, - cached, - prompt_tokens); - } else if (thinking_live_continuation) { + gs->anthropic_live_match_ids, + gs->cached, + gs->prompt_tokens); + } else if (gs->thinking_live_continuation) { server_log(DS4_LOG_PREFILL, "ds4-server: thinking live continuation match=visible-prefix cached=%d prompt=%d", - cached, - prompt_tokens); + gs->cached, + gs->prompt_tokens); } if (responses_visible_replay_without_reasoning) { /* The request replays a prior tool-call turn but omits the hidden @@ -10177,221 +10253,239 @@ static void generate_job(server *s, job *j) { * client asked us to prefill. */ server_log(DS4_LOG_WARNING, "ds4-server: responses replay RESPPROTO missing reasoning state; continuing from visible history source=%s cached=%d prompt=%d", - cache_source, - cached, - prompt_tokens); - trace_event(s, trace_id, + gs->cache_source, + gs->cached, + gs->prompt_tokens); + trace_event(s, gs->trace_id, "responses replay missing reasoning state; continuing from visible history source=%s cached=%d", - cache_source, cached); + gs->cache_source, gs->cached); } server_log(DS4_LOG_PREFILL, "ds4-server: %s ctx=%s%s%s prompt start", j->req.kind == REQ_CHAT ? "chat" : "completion", - ctx_span, - req_flags[0] ? " " : "", - req_flags); - ds4_session_set_progress(s->session, server_progress_cb, &progress); - ds4_session_set_display_progress(s->session, server_progress_cb, &progress); - - int cold_store_len = 0; - if (cached == 0 && + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags); + ds4_session_set_progress(s->session, server_progress_cb, &gs->progress); + ds4_session_set_display_progress(s->session, server_progress_cb, &gs->progress); + + gs->cold_store_len = 0; + if (gs->cached == 0 && s->kv.enabled && - prompt_for_sync->len >= s->kv.opt.min_tokens && + gs->prompt_for_sync->len >= s->kv.opt.min_tokens && s->kv.opt.cold_max_tokens > 0 && - prompt_for_sync->len <= s->kv.opt.cold_max_tokens) + gs->prompt_for_sync->len <= s->kv.opt.cold_max_tokens) { - const int anchor = kv_cache_chat_anchor_pos(&s->kv, prompt_for_sync, + const int anchor = kv_cache_chat_anchor_pos(&s->kv, gs->prompt_for_sync, ds4_token_user(s->engine), ds4_token_assistant(s->engine)); - cold_store_len = anchor >= s->kv.opt.min_tokens ? - anchor : kv_cache_store_len(&s->kv, prompt_for_sync->len); + gs->cold_store_len = anchor >= s->kv.opt.min_tokens ? + anchor : kv_cache_store_len(&s->kv, gs->prompt_for_sync->len); } - int suppressed_continued_last = -1; - if (cold_store_len >= s->kv.opt.min_tokens) { + gs->suppressed_continued_last = -1; + if (gs->cold_store_len >= s->kv.opt.min_tokens) { /* A cold checkpoint can land exactly on the continued-checkpoint - * frontier. The prefill progress callback would then write the same + * frontier. The prefill gs->progress callback would then write the same * prefix as "continued" while we are intentionally stopping there to * write it as "cold". Mark the frontier as already handled before the * sync reaches it; if the cold write fails, restore the old schedule so * a later continued write can still try. */ - suppressed_continued_last = - kv_cache_suppress_continued_store(&s->kv, cold_store_len); + gs->suppressed_continued_last = + kv_cache_suppress_continued_store(&s->kv, gs->cold_store_len); } + return true; +} +/* Run the cold-prefix sync (when a cold disk checkpoint is planned) and the + * main prompt sync. Returns false when prefill failed and the failure + * response was sent. */ +static bool job_prefill(server *s, job *j, gen_state *gs) { if (s->kv.enabled && - cold_store_len >= s->kv.opt.min_tokens && - cold_store_len < prompt_for_sync->len) + gs->cold_store_len >= s->kv.opt.min_tokens && + gs->cold_store_len < gs->prompt_for_sync->len) { ds4_tokens prefix = {0}; - tokens_copy_prefix(&prefix, prompt_for_sync, cold_store_len); - if (ds4_session_sync(s->session, &prefix, err, sizeof(err)) != 0) { + tokens_copy_prefix(&prefix, gs->prompt_for_sync, gs->cold_store_len); + if (ds4_session_sync(s->session, &prefix, gs->err, sizeof(gs->err)) != 0) { ds4_tokens_free(&prefix); - ds4_tokens_free(&effective_prompt); + ds4_tokens_free(&gs->effective_prompt); ds4_session_set_progress(s->session, NULL, NULL); ds4_session_set_display_progress(s->session, NULL, NULL); - kv_cache_restore_suppressed_continued(&s->kv, suppressed_continued_last, - cold_store_len); - kv_cache_discard_failed_disk_entry(s, disk_cache_path); - free(disk_cache_path); - trace_event(s, trace_id, "prefill failed: %s", err); - send_prefill_failure_response(s, j, &progress, ctx_span, req_flags, err); - return; + kv_cache_restore_suppressed_continued(&s->kv, gs->suppressed_continued_last, + gs->cold_store_len); + kv_cache_discard_failed_disk_entry(s, gs->disk_cache_path); + free(gs->disk_cache_path); + trace_event(s, gs->trace_id, "prefill failed: %s", gs->err); + send_prefill_failure_response(s, j, &gs->progress, gs->ctx_span, gs->req_flags, gs->err); + return false; } - if (kv_cache_store_live_prefix(s, prompt_for_sync, cold_store_len, "cold")) { - kv_cache_note_store(&s->kv, cold_store_len); - suppressed_continued_last = -1; + if (kv_cache_store_live_prefix(s, gs->prompt_for_sync, gs->cold_store_len, "cold")) { + kv_cache_note_store(&s->kv, gs->cold_store_len); + gs->suppressed_continued_last = -1; } else { - kv_cache_restore_suppressed_continued(&s->kv, suppressed_continued_last, - cold_store_len); - suppressed_continued_last = -1; + kv_cache_restore_suppressed_continued(&s->kv, gs->suppressed_continued_last, + gs->cold_store_len); + gs->suppressed_continued_last = -1; } ds4_tokens_free(&prefix); } - if (ds4_session_sync(s->session, prompt_for_sync, err, sizeof(err)) != 0) { - ds4_tokens_free(&effective_prompt); + if (ds4_session_sync(s->session, gs->prompt_for_sync, gs->err, sizeof(gs->err)) != 0) { + ds4_tokens_free(&gs->effective_prompt); ds4_session_set_progress(s->session, NULL, NULL); ds4_session_set_display_progress(s->session, NULL, NULL); - kv_cache_restore_suppressed_continued(&s->kv, suppressed_continued_last, - cold_store_len); - kv_cache_discard_failed_disk_entry(s, disk_cache_path); - free(disk_cache_path); - trace_event(s, trace_id, "prefill failed: %s", err); - send_prefill_failure_response(s, j, &progress, ctx_span, req_flags, err); - return; + kv_cache_restore_suppressed_continued(&s->kv, gs->suppressed_continued_last, + gs->cold_store_len); + kv_cache_discard_failed_disk_entry(s, gs->disk_cache_path); + free(gs->disk_cache_path); + trace_event(s, gs->trace_id, "prefill failed: %s", gs->err); + send_prefill_failure_response(s, j, &gs->progress, gs->ctx_span, gs->req_flags, gs->err); + return false; } - free(disk_cache_path); + return true; +} + +/* Post-prefill bookkeeping: clear stale live bindings, store checkpoints, + * emit stream headers and protocol preludes, seed the sampler. Returns + * false when the client is already gone. */ +static bool job_start_decode(server *s, job *j, gen_state *gs) { + free(gs->disk_cache_path); /* Once a non-live request wins, old protocol live bindings are stale. Keep * a binding only when this request explicitly continued from it. */ - if (!responses_live_continuation) responses_live_clear(s); - if (!anthropic_live_continuation) anthropic_live_clear(s); - if (!thinking_live_continuation) thinking_live_clear(s); + if (!gs->responses_live_continuation) responses_live_clear(s); + if (!gs->anthropic_live_continuation) anthropic_live_clear(s); + if (!gs->thinking_live_continuation) thinking_live_clear(s); ds4_session_set_progress(s->session, NULL, NULL); ds4_session_set_display_progress(s->session, NULL, NULL); kv_cache_maybe_store_continued(s); server_log(DS4_LOG_PREFILL, "ds4-server: %s ctx=%s%s%s prompt done %.3fs", j->req.kind == REQ_CHAT ? "chat" : "completion", - ctx_span, - req_flags[0] ? " " : "", - req_flags, - now_sec() - t0); - if (cold_store_len == prompt_for_sync->len) { - if (kv_cache_store_live_prefix(s, prompt_for_sync, cold_store_len, "cold")) { - kv_cache_note_store(&s->kv, cold_store_len); - suppressed_continued_last = -1; + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags, + now_sec() - gs->t0); + if (gs->cold_store_len == gs->prompt_for_sync->len) { + if (kv_cache_store_live_prefix(s, gs->prompt_for_sync, gs->cold_store_len, "cold")) { + kv_cache_note_store(&s->kv, gs->cold_store_len); + gs->suppressed_continued_last = -1; } else { - kv_cache_restore_suppressed_continued(&s->kv, suppressed_continued_last, - cold_store_len); + kv_cache_restore_suppressed_continued(&s->kv, gs->suppressed_continued_last, + gs->cold_store_len); } } - char id[96]; - snprintf(id, sizeof(id), "%s-%llu", + snprintf(gs->id, sizeof(gs->id), "%s-%llu", j->req.kind == REQ_CHAT ? "chatcmpl" : "cmpl", (unsigned long long)++s->seq); - bool structured_stream = request_uses_structured_stream(&j->req); - anthropic_stream anthropic_live = {0}; - openai_stream openai_live = {0}; - responses_stream responses_live = {0}; - const bool openai_live_chat = request_uses_openai_live_stream(&j->req); - const bool responses_live_chat = request_uses_responses_live_stream(&j->req); - long responses_created_at = (long)time(NULL); + gs->structured_stream = request_uses_structured_stream(&j->req); + gs->openai_live_chat = request_uses_openai_live_stream(&j->req); + gs->responses_live_chat = request_uses_responses_live_stream(&j->req); + gs->responses_created_at = (long)time(NULL); if (j->req.stream) { - if (progress.stream_failed) { + if (gs->progress.stream_failed) { server_log(DS4_LOG_GENERATION, "ds4-server: %s ctx=%s%s%s stream closed during prefill", j->req.kind == REQ_CHAT ? "chat" : "completion", - ctx_span, - req_flags[0] ? " " : "", - req_flags); - ds4_tokens_free(&effective_prompt); - return; + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags); + ds4_tokens_free(&gs->effective_prompt); + return false; } - /* The prefill progress callback may have already sent the SSE headers + /* The prefill gs->progress callback may have already sent the SSE headers * to keep the connection alive during a long prefill. Only emit them - * here when prefill never fired (e.g. fully cached prompt). */ - if (!progress.headers_sent && !sse_headers(j->fd, s->enable_cors)) { + * here when prefill never fired (e.g. fully gs->cached prompt). */ + if (!gs->progress.headers_sent && !sse_headers(j->fd, s->enable_cors)) { server_log(DS4_LOG_GENERATION, "ds4-server: %s ctx=%s%s%s sse headers failed", j->req.kind == REQ_CHAT ? "chat" : "completion", - ctx_span, - req_flags[0] ? " " : "", - req_flags); - ds4_tokens_free(&effective_prompt); - return; + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags); + ds4_tokens_free(&gs->effective_prompt); + return false; } - progress.headers_sent = true; + gs->progress.headers_sent = true; if (j->req.api == API_ANTHROPIC && - !anthropic_sse_start_live(j->fd, &j->req, id, - prompt_tokens, &anthropic_live)) { - server_log(DS4_LOG_GENERATION, "ds4-server: chat ctx=%s anthropic stream start failed", ctx_span); - ds4_tokens_free(&effective_prompt); - return; + !anthropic_sse_start_live(j->fd, &j->req, gs->id, + gs->prompt_tokens, &gs->anthropic_live)) { + server_log(DS4_LOG_GENERATION, "ds4-server: chat ctx=%s anthropic stream start failed", gs->ctx_span); + ds4_tokens_free(&gs->effective_prompt); + return false; } if (j->req.api == API_OPENAI && j->req.kind == REQ_CHAT && - !sse_chunk(j->fd, &j->req, id, NULL, NULL)) { - server_log(DS4_LOG_GENERATION, "ds4-server: chat ctx=%s openai role chunk failed", ctx_span); - ds4_tokens_free(&effective_prompt); - return; + !sse_chunk(j->fd, &j->req, gs->id, NULL, NULL)) { + server_log(DS4_LOG_GENERATION, "ds4-server: chat ctx=%s openai role chunk failed", gs->ctx_span); + ds4_tokens_free(&gs->effective_prompt); + return false; } - if (openai_live_chat) openai_stream_start(&j->req, &openai_live); - if (responses_live_chat) { - responses_stream_init(&j->req, &responses_live); - responses_live.active = true; - if (!responses_sse_created(j->fd, &j->req, &responses_live, responses_created_at)) { + if (gs->openai_live_chat) openai_stream_start(&j->req, &gs->openai_live); + if (gs->responses_live_chat) { + responses_stream_init(&j->req, &gs->responses_live); + gs->responses_live.active = true; + if (!responses_sse_created(j->fd, &j->req, &gs->responses_live, gs->responses_created_at)) { server_log(DS4_LOG_GENERATION, "ds4-server: chat ctx=%s%s%s responses created event failed", - ctx_span, - req_flags[0] ? " " : "", - req_flags); - responses_stream_free(&responses_live); - ds4_tokens_free(&effective_prompt); - return; + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags); + responses_stream_free(&gs->responses_live); + ds4_tokens_free(&gs->effective_prompt); + return false; } } } - - bool dsml_recovery_attempted = false; - uint64_t rng = j->req.seed ? j->req.seed : + gs->dsml_recovery_attempted = false; + gs->rng = j->req.seed ? j->req.seed : (((uint64_t)time(NULL) << 32) ^ ((uint64_t)s->seq << 1) ^ (uint64_t)(uintptr_t)j); -decode_again: - ; - buf text = {0}; - size_t plain_stream_pos = 0; - size_t stop_scan_from = 0; - const char *finish = "length"; - int completion = 0; - int max_tokens = j->req.max_tokens; + return true; +} + +/* Reset the per-round decode state. Runs before the first decode step and + * again after a tool-error recovery round (the old decode_again label). */ +static void job_decode_round_init(server *s, job *j, gen_state *gs) { + memset(&gs->text, 0, sizeof(gs->text)); + gs->plain_stream_pos = 0; + gs->stop_scan_from = 0; + gs->finish = "length"; + gs->completion = 0; + gs->max_tokens = j->req.max_tokens; int room = ds4_session_ctx(s->session) - ds4_session_pos(s->session); - bool saw_tool_start = false; - bool saw_tool_end = false; - bool saw_orphan_tool_end = false; - size_t tool_scan_from = 0; - int next_tool_progress = 128; - int next_decode_log = 50; - if (max_tokens < 0) max_tokens = 0; - if (max_tokens > room) max_tokens = room; - trace_event(s, trace_id, "prefill done; decode_max=%d ctx_room=%d", max_tokens, room); - const double decode_t0 = now_sec(); - double last_decode_log_t = decode_t0; - int last_decode_log_completion = 0; - thinking_state thinking = thinking_state_from_prompt(&j->req); - const bool thinking_gates_tool_markers = ds4_think_mode_enabled(j->req.think_mode); - bool tool_scan_waiting_for_think_close = - thinking_gates_tool_markers && thinking.inside; - size_t think_recovery_scan_from = 0; - const bool think_tool_recovery_enabled = + gs->saw_tool_start = false; + gs->saw_tool_end = false; + gs->saw_orphan_tool_end = false; + gs->tool_scan_from = 0; + gs->next_tool_progress = 128; + gs->next_decode_log = 50; + if (gs->max_tokens < 0) gs->max_tokens = 0; + if (gs->max_tokens > room) gs->max_tokens = room; + trace_event(s, gs->trace_id, "prefill done; decode_max=%d ctx_room=%d", gs->max_tokens, room); + gs->decode_t0 = now_sec(); + gs->last_decode_log_t = gs->decode_t0; + gs->last_decode_log_completion = 0; + gs->thinking = thinking_state_from_prompt(&j->req); + gs->thinking_gates_tool_markers = ds4_think_mode_enabled(j->req.think_mode); + gs->tool_scan_waiting_for_think_close = + gs->thinking_gates_tool_markers && gs->thinking.inside; + gs->think_recovery_scan_from = 0; + gs->think_tool_recovery_enabled = getenv("DS4_SERVER_DISABLE_THINK_TOOL_RECOVERY") == NULL; - dsml_decode_tracker dsml_tracker; - dsml_decode_tracker_init(&dsml_tracker); + dsml_decode_tracker_init(&gs->dsml_tracker); +} - while (!g_stop_requested && completion < max_tokens && - ds4_session_pos(s->session) < ds4_session_ctx(s->session)) { +/* One iteration of the decode loop: sample, advance the session (plain eval + * or MTP burst), stream the new text and scan for stops and tool markers. + * Returns true while more tokens should be generated in this round. */ +static bool job_decode_step(server *s, job *j, gen_state *gs) { + if (g_stop_requested || gs->completion >= gs->max_tokens || + ds4_session_pos(s->session) >= ds4_session_ctx(s->session)) { + return false; + } dsml_decode_state dsml_state = j->req.kind == REQ_CHAT && j->req.has_tools ? - dsml_tracker.decode : DSML_DECODE_OUTSIDE; + gs->dsml_tracker.decode : DSML_DECODE_OUTSIDE; const bool in_tool_call = dsml_decode_state_is_tool(dsml_state); - if (!(j->req.kind == REQ_CHAT && j->req.has_tools && (saw_tool_start || in_tool_call))) { + if (!(j->req.kind == REQ_CHAT && j->req.has_tools && (gs->saw_tool_start || in_tool_call))) { kv_cache_maybe_store_continued(s); } float temperature = j->req.temperature; @@ -10407,10 +10501,10 @@ static void generate_job(server *s, job *j) { if (in_tool_call && !dsml_decode_state_uses_payload_sampling(dsml_state)) { temperature = 0.0f; } - int token = ds4_session_sample(s->session, temperature, top_k, top_p, min_p, &rng); + int token = ds4_session_sample(s->session, temperature, top_k, top_p, min_p, &gs->rng); if (token == ds4_token_eos(s->engine)) { - finish = "stop"; - break; + gs->finish = "stop"; + return false; } int toks[17]; @@ -10421,98 +10515,98 @@ static void generate_job(server *s, job *j) { { ntok = ds4_session_eval_speculative_argmax(s->session, token, - max_tokens - completion, + gs->max_tokens - gs->completion, ds4_token_eos(s->engine), toks, (int)(sizeof(toks) / sizeof(toks[0])), - err, - sizeof(err)); + gs->err, + sizeof(gs->err)); if (ntok < 0) { - finish = "error"; - break; + gs->finish = "error"; + return false; } } else { - if (ds4_session_eval(s->session, token, err, sizeof(err)) != 0) { - finish = "error"; - break; + if (ds4_session_eval(s->session, token, gs->err, sizeof(gs->err)) != 0) { + gs->finish = "error"; + return false; } toks[0] = token; ntok = 1; } bool stop_decode = false; - for (int ti = 0; ti < ntok && completion < max_tokens; ti++) { + for (int ti = 0; ti < ntok && gs->completion < gs->max_tokens; ti++) { token = toks[ti]; if (token == ds4_token_eos(s->engine)) { - finish = "stop"; + gs->finish = "stop"; stop_decode = true; break; } size_t piece_len = 0; char *piece = ds4_token_text(s->engine, token, &piece_len); - completion++; + gs->completion++; - trace_piece(s, trace_id, piece, piece_len); - buf_append(&text, piece, piece_len); - thinking_state_feed(&thinking, piece, piece_len); + trace_piece(s, gs->trace_id, piece, piece_len); + buf_append(&gs->text, piece, piece_len); + thinking_state_feed(&gs->thinking, piece, piece_len); if (j->req.kind == REQ_CHAT && j->req.has_tools) { - dsml_decode_tracker_update(&dsml_tracker, text.ptr, text.len); + dsml_decode_tracker_update(&gs->dsml_tracker, gs->text.ptr, gs->text.len); } size_t stop_pos = 0, stop_len = 0; - bool hit_stop = stop_list_find_from(&j->req.stops, text.ptr, - stop_scan_from, + bool hit_stop = stop_list_find_from(&j->req.stops, gs->text.ptr, + gs->stop_scan_from, &stop_pos, &stop_len); size_t stream_len = hit_stop ? - stop_pos : stop_list_stream_safe_len(&j->req.stops, text.len); - if (stream_len > text.len) stream_len = text.len; - stream_len = utf8_stream_safe_len(text.ptr, plain_stream_pos, + stop_pos : stop_list_stream_safe_len(&j->req.stops, gs->text.len); + if (stream_len > gs->text.len) stream_len = gs->text.len; + stream_len = utf8_stream_safe_len(gs->text.ptr, gs->plain_stream_pos, stream_len, hit_stop); if (!hit_stop && j->req.stops.max_len > 1) { const size_t hold = j->req.stops.max_len - 1; - stop_scan_from = text.len > hold ? text.len - hold : 0; + gs->stop_scan_from = gs->text.len > hold ? gs->text.len - hold : 0; } - if (j->req.stream && !structured_stream && stream_len > plain_stream_pos) { - char *delta = xstrndup(text.ptr + plain_stream_pos, stream_len - plain_stream_pos); - bool ok = sse_chunk(j->fd, &j->req, id, delta, NULL); + if (j->req.stream && !gs->structured_stream && stream_len > gs->plain_stream_pos) { + char *delta = xstrndup(gs->text.ptr + gs->plain_stream_pos, stream_len - gs->plain_stream_pos); + bool ok = sse_chunk(j->fd, &j->req, gs->id, delta, NULL); free(delta); if (!ok) { - finish = "error"; - snprintf(err, sizeof(err), "client stream write failed"); + gs->finish = "error"; + snprintf(gs->err, sizeof(gs->err), "client stream write failed"); free(piece); stop_decode = true; break; } - plain_stream_pos = stream_len; + gs->plain_stream_pos = stream_len; } if (j->req.stream && j->req.api == API_ANTHROPIC && - !anthropic_sse_stream_update(j->fd, s, &j->req, id, - &anthropic_live, text.ptr, stream_len, + !anthropic_sse_stream_update(j->fd, s, &j->req, gs->id, + &gs->anthropic_live, gs->text.ptr, stream_len, false)) { - finish = "error"; - snprintf(err, sizeof(err), "client stream write failed"); + gs->finish = "error"; + snprintf(gs->err, sizeof(gs->err), "client stream write failed"); free(piece); stop_decode = true; break; } - if (openai_live_chat && - !openai_sse_stream_update(j->fd, s, &j->req, id, - &openai_live, text.ptr, stream_len, + if (gs->openai_live_chat && + !openai_sse_stream_update(j->fd, s, &j->req, gs->id, + &gs->openai_live, gs->text.ptr, stream_len, false)) { - finish = "error"; - snprintf(err, sizeof(err), "client stream write failed"); + gs->finish = "error"; + snprintf(gs->err, sizeof(gs->err), "client stream write failed"); free(piece); stop_decode = true; break; } - if (responses_live_chat && + if (gs->responses_live_chat && !responses_sse_stream_update(j->fd, &j->req, - &responses_live, text.ptr, stream_len, + &gs->responses_live, gs->text.ptr, stream_len, false)) { - finish = "error"; - snprintf(err, sizeof(err), "client stream write failed"); + gs->finish = "error"; + snprintf(gs->err, sizeof(gs->err), "client stream write failed"); free(piece); stop_decode = true; break; @@ -10520,21 +10614,21 @@ static void generate_job(server *s, job *j) { free(piece); if (j->req.kind == REQ_CHAT && j->req.has_tools) { - if (thinking_gates_tool_markers && thinking.inside) { + if (gs->thinking_gates_tool_markers && gs->thinking.inside) { /* A DSML block inside reasoning is not executable. This is * the live guard: do not let a quoted or mistaken marker in * stop decoding as a real tool call. A complete * stanza opening, however, almost always means the model - * forgot to close its thinking; recover by forcing the + * forgot to close its gs->thinking; recover by forcing the * close so the model restarts the call on the executable * side. */ - const int recovered = think_tool_recovery_enabled ? - chat_think_tool_recovery(s, &text, &thinking, - &think_recovery_scan_from, - &completion, max_tokens, - err, sizeof(err)) : 0; + const int recovered = gs->think_tool_recovery_enabled ? + chat_think_tool_recovery(s, &gs->text, &gs->thinking, + &gs->think_recovery_scan_from, + &gs->completion, gs->max_tokens, + gs->err, sizeof(gs->err)) : 0; if (recovered < 0) { - finish = "error"; + gs->finish = "error"; stop_decode = true; break; } @@ -10542,101 +10636,105 @@ static void generate_job(server *s, job *j) { server_log(DS4_LOG_WARNING, "ds4-server: chat ctx=%s%s%s tool call inside unclosed ; " "forced after %d generated tokens", - ctx_span, - req_flags[0] ? " " : "", - req_flags, - completion); - trace_event(s, trace_id, + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags, + gs->completion); + trace_event(s, gs->trace_id, "think tool recovery after %d generated tokens", - completion); - dsml_decode_tracker_update(&dsml_tracker, text.ptr, text.len); - tool_scan_waiting_for_think_close = true; + gs->completion); + dsml_decode_tracker_update(&gs->dsml_tracker, gs->text.ptr, gs->text.len); + gs->tool_scan_waiting_for_think_close = true; } else { - tool_scan_waiting_for_think_close = true; - tool_scan_from = text.len; + gs->tool_scan_waiting_for_think_close = true; + gs->tool_scan_from = gs->text.len; } } else { - if (tool_scan_waiting_for_think_close) { - const char *think_end = find_last_substr(text.ptr, ""); - tool_scan_from = think_end ? (size_t)((think_end + 8) - text.ptr) : text.len; - if (tool_scan_from > text.len) tool_scan_from = text.len; - tool_scan_waiting_for_think_close = false; + if (gs->tool_scan_waiting_for_think_close) { + const char *think_end = find_last_substr(gs->text.ptr, ""); + gs->tool_scan_from = think_end ? (size_t)((think_end + 8) - gs->text.ptr) : gs->text.len; + if (gs->tool_scan_from > gs->text.len) gs->tool_scan_from = gs->text.len; + gs->tool_scan_waiting_for_think_close = false; } - if (tool_scan_from > text.len) tool_scan_from = text.len; - const char *tool_scan = text.ptr ? text.ptr + tool_scan_from : ""; + if (gs->tool_scan_from > gs->text.len) gs->tool_scan_from = gs->text.len; + const char *tool_scan = gs->text.ptr ? gs->text.ptr + gs->tool_scan_from : ""; bool orphan_end = false; - bool old_start = saw_tool_start; - bool old_end = saw_tool_end; - observe_tool_markers(tool_scan, &saw_tool_start, &saw_tool_end, &orphan_end); - if (orphan_end && !saw_orphan_tool_end) { - saw_orphan_tool_end = true; + bool old_start = gs->saw_tool_start; + bool old_end = gs->saw_tool_end; + observe_tool_markers(tool_scan, &gs->saw_tool_start, &gs->saw_tool_end, &orphan_end); + if (orphan_end && !gs->saw_orphan_tool_end) { + gs->saw_orphan_tool_end = true; server_log(DS4_LOG_WARNING, "ds4-server: chat ctx=%s%s%s ignored orphan tool-call end marker after %d generated tokens", - ctx_span, - req_flags[0] ? " " : "", - req_flags, - completion); - trace_event(s, trace_id, + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags, + gs->completion); + trace_event(s, gs->trace_id, "ignored orphan tool-call end marker after %d generated tokens", - completion); + gs->completion); } - if (saw_tool_start && !old_start) { - trace_event(s, trace_id, "entered tool-call block after %d generated tokens", completion); + if (gs->saw_tool_start && !old_start) { + trace_event(s, gs->trace_id, "entered tool-call block after %d generated tokens", gs->completion); } - if (saw_tool_end && !old_end) { - trace_event(s, trace_id, "closed tool-call block after %d generated tokens", completion); + if (gs->saw_tool_end && !old_end) { + trace_event(s, gs->trace_id, "closed tool-call block after %d generated tokens", gs->completion); } const size_t marker_hold = 80; - size_t hold_from = text.len > marker_hold ? text.len - marker_hold : 0; - if (hold_from > tool_scan_from) tool_scan_from = hold_from; - if (s->trace && completion >= next_tool_progress) { - trace_event(s, trace_id, + size_t hold_from = gs->text.len > marker_hold ? gs->text.len - marker_hold : 0; + if (hold_from > gs->tool_scan_from) gs->tool_scan_from = hold_from; + if (s->trace && gs->completion >= gs->next_tool_progress) { + trace_event(s, gs->trace_id, "progress gen=%d dsml_start=%d dsml_end=%d", - completion, saw_tool_start ? 1 : 0, saw_tool_end ? 1 : 0); - next_tool_progress += 128; + gs->completion, gs->saw_tool_start ? 1 : 0, gs->saw_tool_end ? 1 : 0); + gs->next_tool_progress += 128; } } } - if (completion >= next_decode_log) { - log_decode_progress(j->req.kind, prompt_tokens, completion, - responses_protocol, + if (gs->completion >= gs->next_decode_log) { + log_decode_progress(j->req.kind, gs->prompt_tokens, gs->completion, + gs->responses_protocol, j->req.has_tools, - thinking.inside, - saw_tool_start, - saw_tool_end, - decode_t0, - &last_decode_log_t, - &last_decode_log_completion); - next_decode_log += 50; + gs->thinking.inside, + gs->saw_tool_start, + gs->saw_tool_end, + gs->decode_t0, + &gs->last_decode_log_t, + &gs->last_decode_log_completion); + gs->next_decode_log += 50; } if (hit_stop) { (void)stop_len; - finish = "stop"; - text.len = stop_pos; - text.ptr[text.len] = '\0'; + gs->finish = "stop"; + gs->text.len = stop_pos; + gs->text.ptr[gs->text.len] = '\0'; ds4_session_invalidate(s->session); stop_decode = true; break; } - if (j->req.kind == REQ_CHAT && j->req.has_tools && saw_tool_end) { - finish = "tool_calls"; + if (j->req.kind == REQ_CHAT && j->req.has_tools && gs->saw_tool_end) { + gs->finish = "tool_calls"; stop_decode = true; break; } } - if (stop_decode) break; - } + if (stop_decode) return false; + return true; +} - if (g_stop_requested && strcmp(finish, "error") != 0) { - finish = "error"; - snprintf(err, sizeof(err), "shutdown requested"); +/* Repair or parse the generated text, register live tool/thinking state, + * send the final response (or stream tail) and log the outcome. */ +static gen_step job_finish(server *s, job *j, gen_state *gs) { + if (g_stop_requested && strcmp(gs->finish, "error") != 0) { + gs->finish = "error"; + snprintf(gs->err, sizeof(gs->err), "shutdown requested"); } if (j->req.kind == REQ_CHAT && j->req.has_tools && - saw_tool_start && !saw_tool_end && strcmp(finish, "error") != 0) + gs->saw_tool_start && !gs->saw_tool_end && strcmp(gs->finish, "error") != 0) { /* Deterministically complete a simple truncation. Anything more than * missing closing tags stays model-owned: for non-streaming requests, @@ -10644,8 +10742,8 @@ static void generate_job(server *s, job *j) { * the model issue a fresh call. */ bool completed_truncation = false; buf repaired = {0}; - if (try_repair_dsml(text.ptr, text.len, &repaired)) { - /* Parse repaired text to verify it produces valid tool calls */ + if (try_repair_dsml(gs->text.ptr, gs->text.len, &repaired)) { + /* Parse repaired gs->text to verify it produces valid tool calls */ tool_calls test_calls = {0}; char *test_content = NULL; char *test_reasoning = NULL; @@ -10653,140 +10751,140 @@ static void generate_job(server *s, job *j) { free(test_content); free(test_reasoning); if (repair_ok && test_calls.len > 0) { - /* Repair succeeded - replace text with repaired version */ - free(text.ptr); - text.ptr = buf_take(&repaired); - text.len = strlen(text.ptr); - saw_tool_end = true; + /* Repair succeeded - replace gs->text with repaired version */ + free(gs->text.ptr); + gs->text.ptr = buf_take(&repaired); + gs->text.len = strlen(gs->text.ptr); + gs->saw_tool_end = true; completed_truncation = true; server_log(DS4_LOG_WARNING, "ds4-server: chat ctx=%s%s%s repaired unterminated tool call (%d calls recovered)", - ctx_span, - req_flags[0] ? " " : "", - req_flags, + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags, test_calls.len); - trace_event(s, trace_id, "repaired unterminated tool call (%d calls recovered)", test_calls.len); + trace_event(s, gs->trace_id, "repaired unterminated tool call (%d calls recovered)", test_calls.len); } tool_calls_free(&test_calls); } if (!completed_truncation) { - if (!j->req.stream && !dsml_recovery_attempted) { + if (!j->req.stream && !gs->dsml_recovery_attempted) { int recovery_tokens = 0; char recovery_err[160] = {0}; server_log(DS4_LOG_WARNING, "ds4-server: chat ctx=%s%s%s unterminated tool call; continuing with model-visible tool error", - ctx_span, - req_flags[0] ? " " : "", - req_flags); - trace_event(s, trace_id, + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags); + trace_event(s, gs->trace_id, "unterminated tool call; continuing with model-visible tool error"); - if (continue_after_invalid_dsml(s, &j->req, &thinking, + if (continue_after_invalid_dsml(s, &j->req, &gs->thinking, "unterminated tool call", &recovery_tokens, recovery_err, sizeof(recovery_err))) { - dsml_recovery_attempted = true; + gs->dsml_recovery_attempted = true; server_log(DS4_LOG_GENERATION, "ds4-server: chat ctx=%s%s%s tool-error continuation appended %d tokens", - ctx_span, - req_flags[0] ? " " : "", - req_flags, + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags, recovery_tokens); - trace_event(s, trace_id, + trace_event(s, gs->trace_id, "tool-error continuation appended %d tokens", recovery_tokens); buf_free(&repaired); - buf_free(&text); - goto decode_again; + buf_free(&gs->text); + return GEN_STEP_REDECODE; } - finish = "error"; - snprintf(err, sizeof(err), "invalid tool call recovery failed: %s", + gs->finish = "error"; + snprintf(gs->err, sizeof(gs->err), "invalid tool call recovery failed: %s", recovery_err[0] ? recovery_err : "unknown error"); } else { - finish = "error"; - snprintf(err, sizeof(err), "unterminated tool call"); + gs->finish = "error"; + snprintf(gs->err, sizeof(gs->err), "unterminated tool call"); } } buf_free(&repaired); } - if (completion > last_decode_log_completion) { - log_decode_progress(j->req.kind, prompt_tokens, completion, - responses_protocol, + if (gs->completion > gs->last_decode_log_completion) { + log_decode_progress(j->req.kind, gs->prompt_tokens, gs->completion, + gs->responses_protocol, j->req.has_tools, - thinking.inside, - saw_tool_start, - saw_tool_end, - decode_t0, - &last_decode_log_t, - &last_decode_log_completion); + gs->thinking.inside, + gs->saw_tool_start, + gs->saw_tool_end, + gs->decode_t0, + &gs->last_decode_log_t, + &gs->last_decode_log_completion); } - if (j->req.stream && !structured_stream && text.len > plain_stream_pos) { - char *tail = xstrndup(text.ptr + plain_stream_pos, text.len - plain_stream_pos); - if (!sse_chunk(j->fd, &j->req, id, tail, NULL)) finish = "error"; + if (j->req.stream && !gs->structured_stream && gs->text.len > gs->plain_stream_pos) { + char *tail = xstrndup(gs->text.ptr + gs->plain_stream_pos, gs->text.len - gs->plain_stream_pos); + if (!sse_chunk(j->fd, &j->req, gs->id, tail, NULL)) gs->finish = "error"; free(tail); } tool_calls parsed_calls = {0}; char *parsed_content = NULL; char *parsed_reasoning = NULL; - const char *final_finish = finish; + const char *final_finish = gs->finish; bool recovered_tool_parse_failure = false; if (j->req.kind == REQ_CHAT) { bool parsed_ok = parse_generated_message_for_response( - text.ptr ? text.ptr : "", + gs->text.ptr ? gs->text.ptr : "", j->req.has_tools, - saw_tool_start, + gs->saw_tool_start, ds4_think_mode_enabled(j->req.think_mode), &final_finish, - err, - sizeof(err), + gs->err, + sizeof(gs->err), &parsed_content, &parsed_reasoning, &parsed_calls, &recovered_tool_parse_failure); - if (!parsed_ok && recovered_tool_parse_failure && j->req.has_tools && saw_tool_start) { + if (!parsed_ok && recovered_tool_parse_failure && j->req.has_tools && gs->saw_tool_start) { /* parse_generated_message failed even though DSML was present. * Semantic repair is intentionally avoided: if the parser cannot * execute the block, feed the model a tool error and the protocol * reminder so it owns the corrected next action. */ - if (!j->req.stream && !dsml_recovery_attempted) { + if (!j->req.stream && !gs->dsml_recovery_attempted) { int recovery_tokens = 0; char recovery_err[160] = {0}; - const char *detail = err[0] ? err : "invalid tool call"; + const char *detail = gs->err[0] ? gs->err : "invalid tool call"; server_log(DS4_LOG_WARNING, "ds4-server: chat ctx=%s%s%s invalid tool call; continuing with model-visible tool error", - ctx_span, - req_flags[0] ? " " : "", - req_flags); - trace_event(s, trace_id, + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags); + trace_event(s, gs->trace_id, "invalid tool call; continuing with model-visible tool error"); - if (continue_after_invalid_dsml(s, &j->req, &thinking, + if (continue_after_invalid_dsml(s, &j->req, &gs->thinking, detail, &recovery_tokens, recovery_err, sizeof(recovery_err))) { - dsml_recovery_attempted = true; + gs->dsml_recovery_attempted = true; server_log(DS4_LOG_GENERATION, "ds4-server: chat ctx=%s%s%s tool-error continuation appended %d tokens", - ctx_span, - req_flags[0] ? " " : "", - req_flags, + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags, recovery_tokens); - trace_event(s, trace_id, + trace_event(s, gs->trace_id, "tool-error continuation appended %d tokens", recovery_tokens); free(parsed_content); free(parsed_reasoning); tool_calls_free(&parsed_calls); - buf_free(&text); - goto decode_again; + buf_free(&gs->text); + return GEN_STEP_REDECODE; } final_finish = "error"; - snprintf(err, sizeof(err), "invalid tool call recovery failed: %s", + snprintf(gs->err, sizeof(gs->err), "invalid tool call recovery failed: %s", recovery_err[0] ? recovery_err : "unknown error"); } if (!parsed_ok) { @@ -10794,7 +10892,7 @@ static void generate_job(server *s, job *j) { size_t dsml_snippet_len = 0; const char *dsml_start = NULL; const char *p; - for (p = text.ptr; p && (size_t)(p - text.ptr) < text.len - 20; p++) { + for (p = gs->text.ptr; p && (size_t)(p - gs->text.ptr) < gs->text.len - 20; p++) { if ((strncmp(p, DS4_TOOL_CALLS_START, strlen(DS4_TOOL_CALLS_START)) == 0) || (strncmp(p, DS4_TOOL_CALLS_START_SHORT, strlen(DS4_TOOL_CALLS_START_SHORT)) == 0) || (strncmp(p, "", 12) == 0)) { @@ -10803,38 +10901,38 @@ static void generate_job(server *s, job *j) { } } if (dsml_start) { - dsml_snippet_len = text.len - (dsml_start - text.ptr); + dsml_snippet_len = gs->text.len - (dsml_start - gs->text.ptr); if (dsml_snippet_len > 500) dsml_snippet_len = 500; } - /* Also log a snippet of the full text to see what the model output */ - size_t text_snippet_len = text.len > 300 ? 300 : text.len; + /* Also log a snippet of the full gs->text to see what the model output */ + size_t text_snippet_len = gs->text.len > 300 ? 300 : gs->text.len; server_log(DS4_LOG_WARNING, "ds4-server: chat ctx=%s%s%s invalid tool call returned as assistant text finish=%s [text_len=%zu saw_start=%d saw_end=%d text_snippet: %.*s]", - ctx_span, - req_flags[0] ? " " : "", - req_flags, + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags, final_finish, - text.len, - saw_tool_start, - saw_tool_end, + gs->text.len, + gs->saw_tool_start, + gs->saw_tool_end, (int)text_snippet_len, - text.ptr ? text.ptr : "(null)"); + gs->text.ptr ? gs->text.ptr : "(null)"); server_log(DS4_LOG_WARNING, "ds4-server: chat ctx=%s%s%s invalid tool call dsml_snippet: %.*s", - ctx_span, - req_flags[0] ? " " : "", - req_flags, + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags, (int)dsml_snippet_len, dsml_start ? dsml_start : "(none)"); - trace_event(s, trace_id, + trace_event(s, gs->trace_id, "invalid tool call returned as assistant text finish=%s", final_finish); } } if (parsed_calls.len) { - if (openai_live_chat) apply_openai_stream_tool_ids(&parsed_calls, &openai_live); + if (gs->openai_live_chat) apply_openai_stream_tool_ids(&parsed_calls, &gs->openai_live); if (j->req.api == API_ANTHROPIC && j->req.stream) - apply_anthropic_stream_tool_ids(&parsed_calls, &anthropic_live); + apply_anthropic_stream_tool_ids(&parsed_calls, &gs->anthropic_live); assign_tool_call_ids(s, &parsed_calls, j->req.api); tool_memory_remember(s, &parsed_calls); final_finish = "tool_calls"; @@ -10842,13 +10940,13 @@ static void generate_job(server *s, job *j) { responses_live_clear(s); } } - log_tool_calls_summary(ctx_span, &parsed_calls, - responses_protocol); + log_tool_calls_summary(gs->ctx_span, &parsed_calls, + gs->responses_protocol); - trace_finish(s, trace_id, &j->req, final_finish, completion, - saw_tool_start, saw_tool_end, - parsed_content ? parsed_content : (text.ptr ? text.ptr : ""), - parsed_reasoning, &parsed_calls, now_sec() - t0); + trace_finish(s, gs->trace_id, &j->req, final_finish, gs->completion, + gs->saw_tool_start, gs->saw_tool_end, + parsed_content ? parsed_content : (gs->text.ptr ? gs->text.ptr : ""), + parsed_reasoning, &parsed_calls, now_sec() - gs->t0); if (j->req.api == API_RESPONSES) { if (strcmp(final_finish, "error") && strcmp(final_finish, "length")) { @@ -10892,15 +10990,15 @@ static void generate_job(server *s, job *j) { * replaying those bytes keeps future prompts aligned without rebuilding * hidden reasoning. Responses deliberately skips this path because its * previous_response_id contract binds the next turn to live state. */ - canonicalize_tool_checkpoint(s, j, ctx_span, trace_id, + canonicalize_tool_checkpoint(s, j, gs->ctx_span, gs->trace_id, parsed_content ? parsed_content : "", parsed_reasoning, &parsed_calls); thinking_live_clear(s); } else if (parsed_calls.len) { thinking_live_clear(s); } else if (!parsed_calls.len && - should_remember_thinking_checkpoint(&j->req, &thinking, final_finish)) { - remember_thinking_checkpoint(s, j, ctx_span, trace_id, + should_remember_thinking_checkpoint(&j->req, &gs->thinking, final_finish)) { + remember_thinking_checkpoint(s, j, gs->ctx_span, gs->trace_id, parsed_content ? parsed_content : ""); } else if (!parsed_calls.len) { thinking_live_clear(s); @@ -10909,133 +11007,149 @@ static void generate_job(server *s, job *j) { if (j->req.stream) { bool response_ok = true; if (j->req.api == API_ANTHROPIC) { - response_ok = anthropic_sse_finish_live(j->fd, s, &j->req, id, &anthropic_live, - text.ptr ? text.ptr : "", text.len, - &parsed_calls, final_finish, completion); - } else if (openai_live_chat) { - response_ok = openai_sse_finish_live(j->fd, s, &j->req, id, &openai_live, - text.ptr ? text.ptr : "", text.len, + response_ok = anthropic_sse_finish_live(j->fd, s, &j->req, gs->id, &gs->anthropic_live, + gs->text.ptr ? gs->text.ptr : "", gs->text.len, + &parsed_calls, final_finish, gs->completion); + } else if (gs->openai_live_chat) { + response_ok = openai_sse_finish_live(j->fd, s, &j->req, gs->id, &gs->openai_live, + gs->text.ptr ? gs->text.ptr : "", gs->text.len, &parsed_calls, final_finish, - prompt_tokens, completion); - } else if (responses_live_chat) { - /* If parse recovered a malformed tool call back to plain text, + gs->prompt_tokens, gs->completion); + } else if (gs->responses_live_chat) { + /* If parse recovered a malformed tool call back to plain gs->text, * pass parsed_content so the streaming tail can be flushed; in - * the normal path parsed_content is the assistant text we already + * the normal path parsed_content is the assistant gs->text we already * streamed and the diff is empty. */ const char *recover = recovered_tool_parse_failure ? parsed_content : NULL; - response_ok = responses_sse_finish_live(j->fd, &j->req, &responses_live, - text.ptr ? text.ptr : "", text.len, + response_ok = responses_sse_finish_live(j->fd, &j->req, &gs->responses_live, + gs->text.ptr ? gs->text.ptr : "", gs->text.len, recover, &parsed_calls, final_finish, - prompt_tokens, completion, - responses_created_at); - } else if (structured_stream) { - response_ok = sse_chat_finish(j->fd, &j->req, id, - parsed_content ? parsed_content : (text.ptr ? text.ptr : ""), + gs->prompt_tokens, gs->completion, + gs->responses_created_at); + } else if (gs->structured_stream) { + response_ok = sse_chat_finish(j->fd, &j->req, gs->id, + parsed_content ? parsed_content : (gs->text.ptr ? gs->text.ptr : ""), parsed_reasoning, &parsed_calls, final_finish, - prompt_tokens, completion); + gs->prompt_tokens, gs->completion); } else { - response_ok = sse_chunk(j->fd, &j->req, id, NULL, final_finish) && - sse_done(j->fd, &j->req, id, prompt_tokens, completion); + response_ok = sse_chunk(j->fd, &j->req, gs->id, NULL, final_finish) && + sse_done(j->fd, &j->req, gs->id, gs->prompt_tokens, gs->completion); } if (!response_ok) { server_log(DS4_LOG_DEFAULT, "ds4-server: %s ctx=%s%s%s final stream failed", j->req.kind == REQ_CHAT ? "chat" : "completion", - ctx_span, - req_flags[0] ? " " : "", - req_flags); + gs->ctx_span, + gs->req_flags[0] ? " " : "", + gs->req_flags); } } else if (j->req.api == API_ANTHROPIC) { - anthropic_final_response(j->fd, s->enable_cors, &j->req, id, - parsed_content ? parsed_content : (text.ptr ? text.ptr : ""), + anthropic_final_response(j->fd, s->enable_cors, &j->req, gs->id, + parsed_content ? parsed_content : (gs->text.ptr ? gs->text.ptr : ""), parsed_reasoning, &parsed_calls, final_finish, - prompt_tokens, completion); + gs->prompt_tokens, gs->completion); } else if (j->req.api == API_RESPONSES) { - responses_final_response(j->fd, s->enable_cors, &j->req, id, - parsed_content ? parsed_content : (text.ptr ? text.ptr : ""), + responses_final_response(j->fd, s->enable_cors, &j->req, gs->id, + parsed_content ? parsed_content : (gs->text.ptr ? gs->text.ptr : ""), parsed_reasoning, &parsed_calls, final_finish, - prompt_tokens, completion); + gs->prompt_tokens, gs->completion); } else { - final_response(j->fd, s->enable_cors, &j->req, id, - parsed_content ? parsed_content : (text.ptr ? text.ptr : ""), + final_response(j->fd, s->enable_cors, &j->req, gs->id, + parsed_content ? parsed_content : (gs->text.ptr ? gs->text.ptr : ""), parsed_reasoning, &parsed_calls, final_finish, - prompt_tokens, completion); + gs->prompt_tokens, gs->completion); } if (j->req.kind == REQ_CHAT && j->req.has_tools) { char flags[80]; log_flags(flags, sizeof(flags), - responses_protocol, + gs->responses_protocol, true, - thinking.inside, - saw_tool_start, - saw_tool_end); - if (!strcmp(final_finish, "error") && err[0]) { + gs->thinking.inside, + gs->saw_tool_start, + gs->saw_tool_end); + if (!strcmp(final_finish, "error") && gs->err[0]) { server_log(DS4_LOG_GENERATION, "ds4-server: chat ctx=%s gen=%d%s%s finish=%s error=\"%s\" %.3fs", - ctx_span, - completion, + gs->ctx_span, + gs->completion, flags[0] ? " " : "", flags, final_finish, - err, - now_sec() - t0); + gs->err, + now_sec() - gs->t0); } else { server_log(DS4_LOG_GENERATION, "ds4-server: chat ctx=%s gen=%d%s%s finish=%s %.3fs", - ctx_span, - completion, + gs->ctx_span, + gs->completion, flags[0] ? " " : "", flags, final_finish, - now_sec() - t0); + now_sec() - gs->t0); } } else { char flags[80]; log_flags(flags, sizeof(flags), - responses_protocol, + gs->responses_protocol, j->req.has_tools, - thinking.inside, + gs->thinking.inside, false, false); - if (!strcmp(final_finish, "error") && err[0]) { + if (!strcmp(final_finish, "error") && gs->err[0]) { server_log(DS4_LOG_GENERATION, "ds4-server: %s ctx=%s gen=%d%s%s finish=%s error=\"%s\" %.3fs", j->req.kind == REQ_CHAT ? "chat" : "completion", - ctx_span, - completion, + gs->ctx_span, + gs->completion, flags[0] ? " " : "", flags, final_finish, - err, - now_sec() - t0); + gs->err, + now_sec() - gs->t0); } else { server_log(DS4_LOG_GENERATION, "ds4-server: %s ctx=%s gen=%d%s%s finish=%s %.3fs", j->req.kind == REQ_CHAT ? "chat" : "completion", - ctx_span, - completion, + gs->ctx_span, + gs->completion, flags[0] ? " " : "", flags, final_finish, - now_sec() - t0); + now_sec() - gs->t0); } } free(parsed_content); free(parsed_reasoning); tool_calls_free(&parsed_calls); - anthropic_stream_free(&anthropic_live); - openai_stream_free(&openai_live); - responses_stream_free(&responses_live); - buf_free(&text); - ds4_tokens_free(&effective_prompt); + anthropic_stream_free(&gs->anthropic_live); + openai_stream_free(&gs->openai_live); + responses_stream_free(&gs->responses_live); + buf_free(&gs->text); + ds4_tokens_free(&gs->effective_prompt); + return GEN_STEP_DONE; } +static void generate_job(server *s, job *j) { + gen_state gs; + memset(&gs, 0, sizeof(gs)); + if (!job_begin(s, j, &gs)) return; + if (!job_prefill(s, j, &gs)) return; + if (!job_start_decode(s, j, &gs)) return; + for (;;) { + job_decode_round_init(s, j, &gs); + while (job_decode_step(s, j, &gs)) { + } + if (job_finish(s, j, &gs) != GEN_STEP_REDECODE) break; + } +} + + static bool enqueue(server *s, job *j) { pthread_mutex_lock(&s->mu); if (s->stopping) { From 1471a367e2aab231185260cdfee59a73d6d1d893 Mon Sep 17 00:00:00 2001 From: iCreil Date: Tue, 7 Jul 2026 11:45:51 +0200 Subject: [PATCH 3/3] Server: add --parallel N interleaved sessions (default 1) Replace the single live session with an array of slots, each owning one ds4_session plus the protocol live state (Responses/Anthropic/thinking bindings) tied to that timeline. The graph worker becomes a scheduler: it hands queued jobs to idle slots (longest-common-prefix affinity, then LRU) and steps one slot at a time through the gen_state phases -- a whole begin/finish transition, one prefill slice, or a burst of decode tokens per turn. Prefill preemption reuses the cooperative cancel callback: syncs are interrupted at chunk boundaries once the slice deadline passes and another job is waiting, and resume from the session checkpoint on the slot's next turn. Decode interleaves round-robin between slots, so every connected client keeps streaming while another prompt prefills. The scheduler swaps the kvstore continued-checkpoint frontier in and out per slot; parser-side live call-id lookups scan every slot under the existing tool mutex. --parallel stays opt-in: the default is one slot, the yield callback never fires (no contention), and the phase sequence is exactly the old generate_job order -- verified byte-identical against the previous build on a greedy chat completion (stream and non-stream). N>1 requires resident experts and is refused with --ssd-streaming. Timeslices are tunable via DS4_SERVER_SCHED_PREFILL_MS (default 1500) and DS4_SERVER_SCHED_DECODE_TOKENS (default 6). Measured on CUDA (RTX PRO 6000 96GB, DeepSeek V4 Flash q2, --parallel 2 --ctx 32768 --prefill-chunk 1024, 94.9 GiB VRAM): with a 350-token generation in flight, a second client got its first token in 0.7 s and the first stream's worst inter-token gap during the newcomer's prefill was 0.15 s; a third concurrent request queued and completed once a slot freed. --- ds4_help.c | 1117 +++++++++++++++++++++++++------------------------- ds4_server.c | 621 +++++++++++++++++++--------- 2 files changed, 994 insertions(+), 744 deletions(-) diff --git a/ds4_help.c b/ds4_help.c index d32e088cf..84d7cfc09 100644 --- a/ds4_help.c +++ b/ds4_help.c @@ -1,559 +1,560 @@ -#include "ds4_help.h" - -#include -#include -#include -#include - -typedef struct { - const char *off; - const char *cyan; - const char *title; - const char *yellow; - const char *grey; - const char *red; - const char *white; - const char *bright; -} help_colors; - -static help_colors help_make_colors(FILE *fp) { - bool color = isatty(fileno(fp)); - help_colors c = {0}; - if (!color) return c; - c.off = "\x1b[0m"; - c.cyan = "\x1b[38;5;81m"; - c.title = "\x1b[1;38;5;250m"; - c.yellow = "\x1b[38;5;179m"; - c.grey = "\x1b[38;5;240m"; - c.red = "\x1b[38;5;203m"; - c.white = "\x1b[38;5;252m"; - c.bright = "\x1b[1;38;5;231m"; - return c; -} - -static void title(FILE *fp, const help_colors *c, const char *s) { - fprintf(fp, "%s%s%s\n", c->title ? c->title : "", s, c->off ? c->off : ""); -} - -static void title_red(FILE *fp, const help_colors *c, const char *s) { - fprintf(fp, "%s%s%s\n", c->red ? c->red : "", s, c->off ? c->off : ""); -} - -static bool option_name_has_switch(const char *name) { - bool word_start = true; - while (*name) { - if (word_start && (*name == '-' || *name == '/')) return true; - word_start = (*name == ' '); - name++; - } - return false; -} - -static void print_colored_option_name(FILE *fp, const help_colors *c, const char *name) { - bool has_switch = option_name_has_switch(name); - bool word_start = true; - while (*name) { - const char *start = name; - while (*name && *name != ' ') name++; - bool is_option = !has_switch || *start == '-' || *start == '/' || - (word_start && has_switch && *start != '['); - const char *color = is_option ? c->cyan : c->bright; - if (color) fputs(color, fp); - fwrite(start, 1, (size_t)(name - start), fp); - if (color && c->off) fputs(c->off, fp); - if (*name == ' ') { - fputc(*name++, fp); - word_start = false; - } - } -} - -static void opt(FILE *fp, const help_colors *c, const char *name, const char *desc) { - if (c->cyan) { - fputs(" ", fp); - print_colored_option_name(fp, c, name); - fprintf(fp, " %s|%s ", c->grey ? c->grey : "", c->grey ? c->off : ""); - fprintf(fp, "%s%s%s\n", c->white ? c->white : "", desc, - c->white ? c->off : ""); - return; - } - - const int col = 30; - int n = (int)strlen(name); - if (n > col) { - fprintf(fp, " %s\n %s\n", name, desc); - } else { - fprintf(fp, " %-30s %s\n", name, desc); - } -} - -static void para(FILE *fp, const help_colors *c, const char *s) { - fprintf(fp, "%s%s%s\n", - c->yellow ? c->yellow : "", s, c->yellow ? c->off : ""); -} - -static bool streq(const char *a, const char *b) { - return a && b && strcmp(a, b) == 0; -} - -static bool topic_is(const char *topic, const char *name) { - return topic && strcmp(topic, name) == 0; -} - -static const char *tool_name(ds4_help_tool tool) { - switch (tool) { - case DS4_HELP_DS4: return "ds4"; - case DS4_HELP_SERVER: return "ds4-server"; - case DS4_HELP_AGENT: return "ds4-agent"; - case DS4_HELP_BENCH: return "ds4-bench"; - case DS4_HELP_EVAL: return "ds4-eval"; - } - return "ds4"; -} - -static const char *tool_usage(ds4_help_tool tool) { - switch (tool) { - case DS4_HELP_DS4: - return "Usage: ds4 [(-p PROMPT | --prompt-file FILE)] [options]"; - case DS4_HELP_SERVER: - return "Usage: ds4-server [options]"; - case DS4_HELP_AGENT: - return "Usage: ds4-agent [options]"; - case DS4_HELP_BENCH: - return "Usage: ds4-bench (--prompt-file FILE | --chat-prompt-file FILE) [options]"; - case DS4_HELP_EVAL: - return "Usage: ds4-eval [options]"; - } - return "Usage: ds4 [options]"; -} - -static const char *tool_summary(ds4_help_tool tool) { - switch (tool) { - case DS4_HELP_DS4: - return "Chat with a local DwarfStar model, run one-shot prompts, inspect models, or coordinate distributed inference."; - case DS4_HELP_SERVER: - return "Serve one loaded DwarfStar model through OpenAI, Responses, Anthropic, and completion-compatible HTTP APIs."; - case DS4_HELP_AGENT: - return "Run the native terminal coding agent with live tools, session save/restore, and a responsive prompt while the model works."; - case DS4_HELP_BENCH: - return "Measure prefill, decode, context growth, and KV-cache size across repeatable context frontiers."; - case DS4_HELP_EVAL: - return "Run the built-in reasoning, math, science, and security evaluation harness with a live terminal UI."; - } - return ""; -} - -static void print_model_runtime(FILE *fp, const help_colors *c, - ds4_help_tool tool, bool full) { - title(fp, c, "Model And Runtime"); - opt(fp, c, "-m, --model FILE", "GGUF model path. Default: ds4flash.gguf"); -#ifdef DS4_ROCM_BUILD - opt(fp, c, "--metal | --rocm | --cpu", "Select the backend explicitly."); - opt(fp, c, "--backend NAME", "Backend name: metal, rocm, or cpu."); -#else - opt(fp, c, "--metal | --cuda | --cpu", "Select the backend explicitly."); - opt(fp, c, "--backend NAME", "Backend name: metal, cuda, or cpu."); -#endif - if (tool != DS4_HELP_BENCH) { +#include "ds4_help.h" + +#include +#include +#include +#include + +typedef struct { + const char *off; + const char *cyan; + const char *title; + const char *yellow; + const char *grey; + const char *red; + const char *white; + const char *bright; +} help_colors; + +static help_colors help_make_colors(FILE *fp) { + bool color = isatty(fileno(fp)); + help_colors c = {0}; + if (!color) return c; + c.off = "\x1b[0m"; + c.cyan = "\x1b[38;5;81m"; + c.title = "\x1b[1;38;5;250m"; + c.yellow = "\x1b[38;5;179m"; + c.grey = "\x1b[38;5;240m"; + c.red = "\x1b[38;5;203m"; + c.white = "\x1b[38;5;252m"; + c.bright = "\x1b[1;38;5;231m"; + return c; +} + +static void title(FILE *fp, const help_colors *c, const char *s) { + fprintf(fp, "%s%s%s\n", c->title ? c->title : "", s, c->off ? c->off : ""); +} + +static void title_red(FILE *fp, const help_colors *c, const char *s) { + fprintf(fp, "%s%s%s\n", c->red ? c->red : "", s, c->off ? c->off : ""); +} + +static bool option_name_has_switch(const char *name) { + bool word_start = true; + while (*name) { + if (word_start && (*name == '-' || *name == '/')) return true; + word_start = (*name == ' '); + name++; + } + return false; +} + +static void print_colored_option_name(FILE *fp, const help_colors *c, const char *name) { + bool has_switch = option_name_has_switch(name); + bool word_start = true; + while (*name) { + const char *start = name; + while (*name && *name != ' ') name++; + bool is_option = !has_switch || *start == '-' || *start == '/' || + (word_start && has_switch && *start != '['); + const char *color = is_option ? c->cyan : c->bright; + if (color) fputs(color, fp); + fwrite(start, 1, (size_t)(name - start), fp); + if (color && c->off) fputs(c->off, fp); + if (*name == ' ') { + fputc(*name++, fp); + word_start = false; + } + } +} + +static void opt(FILE *fp, const help_colors *c, const char *name, const char *desc) { + if (c->cyan) { + fputs(" ", fp); + print_colored_option_name(fp, c, name); + fprintf(fp, " %s|%s ", c->grey ? c->grey : "", c->grey ? c->off : ""); + fprintf(fp, "%s%s%s\n", c->white ? c->white : "", desc, + c->white ? c->off : ""); + return; + } + + const int col = 30; + int n = (int)strlen(name); + if (n > col) { + fprintf(fp, " %s\n %s\n", name, desc); + } else { + fprintf(fp, " %-30s %s\n", name, desc); + } +} + +static void para(FILE *fp, const help_colors *c, const char *s) { + fprintf(fp, "%s%s%s\n", + c->yellow ? c->yellow : "", s, c->yellow ? c->off : ""); +} + +static bool streq(const char *a, const char *b) { + return a && b && strcmp(a, b) == 0; +} + +static bool topic_is(const char *topic, const char *name) { + return topic && strcmp(topic, name) == 0; +} + +static const char *tool_name(ds4_help_tool tool) { + switch (tool) { + case DS4_HELP_DS4: return "ds4"; + case DS4_HELP_SERVER: return "ds4-server"; + case DS4_HELP_AGENT: return "ds4-agent"; + case DS4_HELP_BENCH: return "ds4-bench"; + case DS4_HELP_EVAL: return "ds4-eval"; + } + return "ds4"; +} + +static const char *tool_usage(ds4_help_tool tool) { + switch (tool) { + case DS4_HELP_DS4: + return "Usage: ds4 [(-p PROMPT | --prompt-file FILE)] [options]"; + case DS4_HELP_SERVER: + return "Usage: ds4-server [options]"; + case DS4_HELP_AGENT: + return "Usage: ds4-agent [options]"; + case DS4_HELP_BENCH: + return "Usage: ds4-bench (--prompt-file FILE | --chat-prompt-file FILE) [options]"; + case DS4_HELP_EVAL: + return "Usage: ds4-eval [options]"; + } + return "Usage: ds4 [options]"; +} + +static const char *tool_summary(ds4_help_tool tool) { + switch (tool) { + case DS4_HELP_DS4: + return "Chat with a local DwarfStar model, run one-shot prompts, inspect models, or coordinate distributed inference."; + case DS4_HELP_SERVER: + return "Serve one loaded DwarfStar model through OpenAI, Responses, Anthropic, and completion-compatible HTTP APIs."; + case DS4_HELP_AGENT: + return "Run the native terminal coding agent with live tools, session save/restore, and a responsive prompt while the model works."; + case DS4_HELP_BENCH: + return "Measure prefill, decode, context growth, and KV-cache size across repeatable context frontiers."; + case DS4_HELP_EVAL: + return "Run the built-in reasoning, math, science, and security evaluation harness with a live terminal UI."; + } + return ""; +} + +static void print_model_runtime(FILE *fp, const help_colors *c, + ds4_help_tool tool, bool full) { + title(fp, c, "Model And Runtime"); + opt(fp, c, "-m, --model FILE", "GGUF model path. Default: ds4flash.gguf"); +#ifdef DS4_ROCM_BUILD + opt(fp, c, "--metal | --rocm | --cpu", "Select the backend explicitly."); + opt(fp, c, "--backend NAME", "Backend name: metal, rocm, or cpu."); +#else + opt(fp, c, "--metal | --cuda | --cpu", "Select the backend explicitly."); + opt(fp, c, "--backend NAME", "Backend name: metal, cuda, or cpu."); +#endif + if (tool != DS4_HELP_BENCH) { opt(fp, c, "-c, --ctx N", "Allocated context tokens."); - } - if (tool == DS4_HELP_SERVER) { - opt(fp, c, "-n, --tokens N", "Default max output tokens when clients omit a limit."); - } - opt(fp, c, "-t, --threads N", "CPU helper threads for host-side/reference work."); - opt(fp, c, "--power N", "GPU duty-cycle target, 1..100. Default: 100"); - opt(fp, c, "--ssd-streaming", "Metal/CUDA/ROCm: opt in to SSD-backed model streaming instead of full residency."); - opt(fp, c, "--ssd-streaming-cold", "SSD streaming: skip default popularity-based expert-cache preload."); - opt(fp, c, "--ssd-streaming-cache-experts N|NGB", "SSD streaming: routed expert cache as expert count or GiB, e.g. 32GB. Metal/ROCm default: 80% working set minus non-routed weights; CUDA default: backend fixed cache."); - opt(fp, c, "--ssd-streaming-preload-experts N", "SSD streaming: upfront popularity preload count. Default: auto hot seed capped at 4096; use --ssd-streaming-cold to skip."); - opt(fp, c, "--simulate-used-memory NGB", "Diagnostic: lock N GiB before model load to simulate a smaller-memory machine."); - opt(fp, c, "--prefill-chunk N", "Metal graph prefill chunk size. Default: auto (PRO long prompts use 8192; others use 4096)."); - if (full) { - if (tool != DS4_HELP_BENCH) { - opt(fp, c, "--mtp FILE", "Optional MTP support GGUF used for draft-token probes."); - } - if (tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_SERVER) { - opt(fp, c, "--mtp-draft N", "Maximum autoregressive MTP draft tokens. Default: 1"); - opt(fp, c, "--mtp-margin F", "Verifier confidence margin for fast MTP acceptance. Default: 3"); - } - opt(fp, c, "--quality", "Prefer exact kernels where faster approximate paths exist."); - opt(fp, c, "--warm-weights", "Touch mapped tensor pages at startup to reduce first-use stalls."); - if (tool == DS4_HELP_DS4 || tool == DS4_HELP_BENCH) { - opt(fp, c, "--expert-profile FILE", "Metal-only: write routed expert locality/cache simulation JSON."); - } - } - fputc('\n', fp); -} - -static void print_sampling(FILE *fp, const help_colors *c, bool full) { - title(fp, c, "Prompt And Sampling"); - opt(fp, c, "-n, --tokens N", "Maximum generated tokens."); - opt(fp, c, "--temp F", "Sampling temperature. 0 is greedy/deterministic."); - opt(fp, c, "--top-p F", "Nucleus sampling probability."); - opt(fp, c, "--min-p F", "Keep tokens scoring at least F times the top token."); - opt(fp, c, "--seed N", "Sampling seed for reproducible non-greedy runs."); - opt(fp, c, "--think", "Use normal thinking mode."); - opt(fp, c, "--think-max", "Use Think Max when context is large enough."); - opt(fp, c, "--nothink", "Disable thinking and ask for direct replies."); - if (full) { - opt(fp, c, "-sys, --system TEXT", "System prompt. Empty string disables the default where supported."); - opt(fp, c, "-p, --prompt TEXT", "One-shot prompt text."); - opt(fp, c, "--prompt-file FILE", "Read one-shot prompt text from FILE."); - } - fputc('\n', fp); -} - -static void print_steering(FILE *fp, const help_colors *c) { - title(fp, c, "Directional Steering"); - opt(fp, c, "--dir-steering-file FILE", "Load one f32 direction vector per layer."); - opt(fp, c, "--dir-steering-ffn F", "Apply steering after FFN outputs. Default with file: 1"); - opt(fp, c, "--dir-steering-attn F", "Apply steering after attention outputs. Default: 0"); - fputc('\n', fp); -} - -static void print_distributed(FILE *fp, const help_colors *c) { - title(fp, c, "Distributed Inference"); - fputc('\n', fp); - para(fp, c, "Distributed mode runs one logical session across several machines by assigning contiguous model layer ranges to workers. Workers own their layer slice and KV-cache shard; the coordinator owns the prompt, sampling loop, and client/API flow. Start workers first, then start the coordinator. The coordinator waits for a complete route and streams hidden states through the workers."); - fputc('\n', fp); - opt(fp, c, "--role ROLE", "Distributed role: coordinator or worker."); - opt(fp, c, "--layers A:B", "Inclusive layer slice, e.g. 0:20 or 21:output."); - opt(fp, c, "--listen HOST PORT", "Coordinator listen address; workers may use it for their data listener."); - opt(fp, c, "--coordinator HOST PORT", "Coordinator address for --role worker."); - opt(fp, c, "--dist-prefill-chunk N", "Coordinator prefill pipeline chunk size. Default: session cap."); - opt(fp, c, "--dist-prefill-window N", "Max prefill chunks in flight. Default: workers+2, capped at 8."); - opt(fp, c, "--dist-activation-bits N", "Hidden-state transport width: 32, 16, or 8. Default: 32"); - opt(fp, c, "--dist-replay-check", "Diagnostic: reset and replay prompt, then compare logits."); - opt(fp, c, "--debug", "Print coordinator route/debug logs."); - fputc('\n', fp); -} - -static void print_cli_diagnostics(FILE *fp, const help_colors *c); - -static void print_cli_specific(FILE *fp, const help_colors *c, bool full) { - title(fp, c, "CLI Modes"); - opt(fp, c, "ds4", "Start the interactive prompt."); - opt(fp, c, "ds4 -p TEXT", "Run one prompt and exit."); - opt(fp, c, "ds4 --prompt-file FILE", "Run a long prompt from a file and exit."); - fputc('\n', fp); - if (full) { - print_cli_diagnostics(fp, c); - } -} - -static void print_cli_diagnostics(FILE *fp, const help_colors *c) { - title(fp, c, "Diagnostics And Data Collection"); - opt(fp, c, "--inspect", "Load the model and print a summary only."); - opt(fp, c, "--dump-tokens", "Tokenize the prompt exactly as written, then exit."); - opt(fp, c, "--dump-logits FILE", "Write full next-token logits as JSON."); - opt(fp, c, "--dump-logprobs FILE", "Write greedy continuation top-logprobs as JSON."); - opt(fp, c, "--logprobs-top-k N", "Alternatives stored by --dump-logprobs. Default: 20"); - opt(fp, c, "--expert-profile FILE", "Metal-only: write routed expert locality/cache simulation JSON."); - opt(fp, c, "--perplexity-file FILE", "Score raw text with teacher-forced NLL."); - opt(fp, c, "--imatrix-dataset FILE", "Rendered prompt dataset for imatrix collection."); - opt(fp, c, "--imatrix-out FILE", "Write llama-compatible routed-MoE imatrix .dat."); - opt(fp, c, "--imatrix-max-prompts N", "Stop imatrix collection after N prompts."); - opt(fp, c, "--imatrix-max-tokens N", "Stop imatrix collection after N prompt tokens."); - opt(fp, c, "--head-test", "Run the output HC/logits head after the native slice."); - opt(fp, c, "--first-token-test", "Run exact CPU whole-model pass for the first prompt token."); - opt(fp, c, "--metal-graph-test", "Compare first GPU-resident graph stages with CPU."); - opt(fp, c, "--metal-graph-full-test", "Run the GPU-resident self-token graph across all layers."); - opt(fp, c, "--metal-graph-prompt-test", "Compare CPU and GPU graph logits for the full prompt."); - fputc('\n', fp); -} - -static void print_cli_commands(FILE *fp, const help_colors *c) { - title_red(fp, c, "Interactive Commands"); - opt(fp, c, "/help", "Show interactive commands."); - opt(fp, c, "/think, /think-max, /nothink", "Switch thinking mode."); - opt(fp, c, "/ctx N", "Restart the interactive session with a new context size."); - opt(fp, c, "/power N", "Set GPU duty cycle percentage, 1..100."); - opt(fp, c, "/read FILE", "Read FILE and submit it as the next user message."); - opt(fp, c, "/quit, /exit", "Leave the prompt."); - opt(fp, c, "Ctrl+C", "Stop current generation and return to ds4>."); - fputc('\n', fp); -} - -static void print_agent_specific(FILE *fp, const help_colors *c) { - title(fp, c, "Agent Options"); - opt(fp, c, "-p, --prompt TEXT", "Submit an initial prompt after startup."); - opt(fp, c, "--non-interactive", "Run without TUI. With -p: one turn; without -p: repeated stdin prompts."); - opt(fp, c, "-sys, --system TEXT", "Extra system prompt. Empty disables extra text."); - opt(fp, c, "--trace FILE", "Write prompt, token, and DSML debug trace."); - opt(fp, c, "--chdir DIR", "Change working directory before loading runtime assets."); - fputc('\n', fp); -} - -static void print_agent_sessions(FILE *fp, const help_colors *c) { - title(fp, c, "Agent Runtime Commands"); - opt(fp, c, "/save", "Save the current session in ~/.ds4/kvcache."); - opt(fp, c, "/compact", "Compact the current session context now."); - opt(fp, c, "/list", "List saved sessions, sorted by recent update time."); - opt(fp, c, "/switch ID", "Load a saved session and show recent history."); - opt(fp, c, "/del ID", "Delete a saved session."); - opt(fp, c, "/strip ID", "Remove KV payload; the text history can be rebuilt later."); - opt(fp, c, "/history [N]", "Show N recent user turns from the current session."); - opt(fp, c, "/power N", "Set GPU duty cycle percentage, 1..100."); - opt(fp, c, "/new", "Start a fresh session from the system prompt."); - opt(fp, c, "/quit, /exit", "Exit."); - fputc('\n', fp); -} - -static void print_server_api(FILE *fp, const help_colors *c) { - title(fp, c, "HTTP API"); - opt(fp, c, "--host HOST", "Bind address. Default: 127.0.0.1"); - opt(fp, c, "--port N", "Bind port. Default: 8000"); - opt(fp, c, "--cors", "Add Access-Control-Allow-* headers for browser JS clients."); - opt(fp, c, "--trace FILE", "Write prompts, cache decisions, output, and tool calls."); - para(fp, c, "Endpoints: /v1/chat/completions, /v1/responses, /v1/completions, and /v1/messages."); - para(fp, c, "Model endpoint aliases include deepseek-v4-flash and deepseek-v4-pro; both serve the loaded GGUF."); - fputc('\n', fp); -} - -static void print_server_thinking(FILE *fp, const help_colors *c) { - title(fp, c, "Server Thinking Defaults"); - para(fp, c, "DeepSeek-compatible chat requests default to high-effort thinking."); - para(fp, c, "reasoning_effort=max or output_config.effort=max requests Think Max."); - para(fp, c, "Think Max requires --ctx >= 393216; smaller contexts use high."); - para(fp, c, "thinking={type:disabled}, think=false, or model=deepseek-chat selects non-thinking mode."); - para(fp, c, "In thinking mode, client sampling knobs are ignored like the official API."); - fputc('\n', fp); -} - -static void print_kv_cache(FILE *fp, const help_colors *c) { - title(fp, c, "Disk KV Cache"); - opt(fp, c, "--kv-disk-dir DIR", "Enable disk KV checkpoints in DIR."); - opt(fp, c, "--kv-disk-space-mb N", "Disk budget. Default when enabled: 4096"); - opt(fp, c, "--kv-cache-min-tokens N", "Do not save/load checkpoints shorter than N. Default: 512"); - opt(fp, c, "--kv-cache-cold-max-tokens N", "Save cold first prompts up to N tokens. 0 disables. Default: 30000"); - opt(fp, c, "--kv-cache-continued-interval-tokens N", "Save aligned continued frontiers. 0 disables. Default: 10000"); - opt(fp, c, "--kv-cache-boundary-trim-tokens N", "Trim tail tokens for cold boundary saves. Default: 32"); - opt(fp, c, "--kv-cache-boundary-align-tokens N", "Align cold boundary saves to this multiple. Default: 2048"); - opt(fp, c, "--kv-cache-reject-different-quant", "Reject checkpoints written with different routed-expert quantization."); - opt(fp, c, "--disable-exact-dsml-tool-replay", "Disable exact sampled DSML tool replay map."); - opt(fp, c, "--tool-memory-max-ids N", "Exact tool-call IDs kept in RAM. Default: 100000"); - fputc('\n', fp); -} - -static void print_bench_specific(FILE *fp, const help_colors *c) { - title(fp, c, "Benchmark Input"); - opt(fp, c, "--prompt-file FILE", "Raw benchmark text; token sequence is sliced at each frontier."); - opt(fp, c, "--chat-prompt-file FILE", "Render FILE as one no-thinking chat user message."); - opt(fp, c, "-sys, --system TEXT", "System prompt used only with --chat-prompt-file."); - fputc('\n', fp); - title(fp, c, "Benchmark Sweep"); - opt(fp, c, "--ctx-start N", "First measured frontier. Default: 2048"); - opt(fp, c, "--ctx-max N", "Last measured frontier. Default: 32768"); - opt(fp, c, "--ctx-alloc N", "Allocated context. Default: ctx-max + gen-tokens + 1"); - opt(fp, c, "--step-mul F", "Multiplicative step. Default: 1"); - opt(fp, c, "--step-incr N", "Linear step when --step-mul is 1. Default: 2048"); - opt(fp, c, "--gen-tokens N", "Greedy decode tokens per frontier. 0 for pure prefill. Default: 128"); - opt(fp, c, "--csv FILE", "Write CSV there instead of stdout."); - opt(fp, c, "--dump-frontier-logits-dir DIR", "Write one full-logit JSON file per frontier."); - fputc('\n', fp); -} - -static void print_eval_specific(FILE *fp, const help_colors *c) { - title(fp, c, "Evaluation"); - opt(fp, c, "-n, --tokens N", "Max generated tokens per question. Default: 16000"); - opt(fp, c, "--questions N", "Run only the first N embedded questions."); - opt(fp, c, "--case-sequence LIST", "Run 1-based case numbers in this comma-separated order."); - opt(fp, c, "--trace FILE", "Write questions, outputs, and grading decisions."); - opt(fp, c, "--regrade-trace FILE", "Regrade a prior trace without loading the model."); - opt(fp, c, "--soft-limit-reply-budget N", "Soft close thinking near the end of reply budget. Default: 1024"); - opt(fp, c, "--hard-limit-reply-budget N", "Force with N tokens left. Default: 512"); - opt(fp, c, "--soft-limit-think-close-rank N", "Soft-close when is in top N tokens. Default: 3"); - opt(fp, c, "--pause-ms N", "Pause after each result in the TTY UI. Default: 350"); - opt(fp, c, "--plain", "Disable split-screen ANSI UI."); - opt(fp, c, "--self-test-extractors", "Run answer-extractor self-tests and exit."); - fputc('\n', fp); -} - -static bool tool_has_topic(ds4_help_tool tool, const char *topic) { - if (!topic) return true; - if (streq(topic, "all")) return true; - if (streq(topic, "runtime") || streq(topic, "distributed")) return true; - if (streq(topic, "sampling")) - return tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_EVAL; - if (streq(topic, "steering")) - return tool == DS4_HELP_DS4 || tool == DS4_HELP_SERVER || tool == DS4_HELP_AGENT; - switch (tool) { - case DS4_HELP_DS4: - return streq(topic, "diagnostics") || streq(topic, "commands"); - case DS4_HELP_SERVER: - return streq(topic, "api") || streq(topic, "kv-cache") || streq(topic, "thinking"); - case DS4_HELP_AGENT: - return streq(topic, "sessions") || streq(topic, "commands") || streq(topic, "tools"); - case DS4_HELP_BENCH: - return streq(topic, "benchmark"); - case DS4_HELP_EVAL: - return streq(topic, "evaluation"); - } - return false; -} - -static void more_line(FILE *fp, const help_colors *c, const char *label, const char *topic) { - static const char *colors[] = { - "\x1b[38;5;81m", "\x1b[38;5;114m", "\x1b[38;5;179m", - "\x1b[38;5;141m", "\x1b[38;5;147m" - }; - static size_t idx; - const char *on = c->cyan ? colors[idx++ % (sizeof(colors) / sizeof(colors[0]))] : ""; - if (streq(label, "Interactive commands:") && c->red) on = c->red; - const char *off = c->off ? c->off : ""; - fprintf(fp, " %s%-26s%s --help %s\n", on, label, off, topic); -} - -static void print_more_info(FILE *fp, const help_colors *c, ds4_help_tool tool) { - title(fp, c, "More Info"); - more_line(fp, c, "Runtime full info:", "runtime"); - if (tool_has_topic(tool, "sampling")) - more_line(fp, c, "Sampling full info:", "sampling"); - more_line(fp, c, "Distributed inference:", "distributed"); - if (tool_has_topic(tool, "steering")) - more_line(fp, c, "Steering full info:", "steering"); - if (tool == DS4_HELP_DS4) { - more_line(fp, c, "Interactive commands:", "commands"); - more_line(fp, c, "Diagnostics:", "diagnostics"); - } else if (tool == DS4_HELP_SERVER) { - more_line(fp, c, "HTTP API:", "api"); - more_line(fp, c, "Disk KV cache:", "kv-cache"); - more_line(fp, c, "Thinking behavior:", "thinking"); - } else if (tool == DS4_HELP_AGENT) { - more_line(fp, c, "Agent sessions:", "sessions"); - more_line(fp, c, "Agent commands:", "commands"); - more_line(fp, c, "Agent tool system:", "tools"); - } else if (tool == DS4_HELP_BENCH) { - more_line(fp, c, "Benchmark sweep:", "benchmark"); - } else if (tool == DS4_HELP_EVAL) { - more_line(fp, c, "Evaluation options:", "evaluation"); - } - fputc('\n', fp); -} - -static void print_examples(FILE *fp, const help_colors *c, ds4_help_tool tool, const char *topic) { - title(fp, c, "Examples"); - if (topic_is(topic, "distributed")) { - opt(fp, c, "worker", "./ds4 --role worker --layers 21:output --coordinator 192.168.0.181 9000 -m ds4flash.gguf"); - opt(fp, c, "coordinator", "./ds4 --role coordinator --layers 0:20 --listen 0.0.0.0 9000 -p \"Hello\" -m ds4flash.gguf"); - } else if (topic_is(topic, "runtime")) { - if (tool == DS4_HELP_SERVER) { - opt(fp, c, "Metal API", "./ds4-server -m ds4flash.gguf --metal --ctx 100000"); - opt(fp, c, "quiet API", "./ds4-server --power 60 --host 127.0.0.1 --port 8000"); - } else if (tool == DS4_HELP_AGENT) { - opt(fp, c, "agent", "./ds4-agent -m ds4flash.gguf --ctx 100000"); - opt(fp, c, "quiet agent", "./ds4-agent --power 50"); - } else if (tool == DS4_HELP_BENCH) { - opt(fp, c, "bench", "./ds4-bench --prompt-file long.txt --ctx-max 32768"); - opt(fp, c, "quiet bench", "./ds4-bench --prompt-file long.txt --power 70"); - } else if (tool == DS4_HELP_EVAL) { - opt(fp, c, "eval", "./ds4-eval --questions 10 --ctx 100000"); - opt(fp, c, "CPU debug", "./ds4-eval --cpu --questions 1 --tokens 32"); - } else { - opt(fp, c, "Metal", "./ds4 -m ds4flash.gguf --metal -c 100000"); - opt(fp, c, "quiet thermals", "./ds4 -p \"Summarize README\" --power 50"); - } - } else if (topic_is(topic, "steering")) { - opt(fp, c, "steer FFN", "./ds4 -p \"Write tersely\" --dir-steering-file dir.bin --dir-steering-ffn 0.8"); - } else if (tool == DS4_HELP_SERVER || topic_is(topic, "api") || topic_is(topic, "kv-cache")) { - opt(fp, c, "local API", "./ds4-server --ctx 100000 --kv-disk-dir ~/.ds4/server-kv --kv-disk-space-mb 8192"); - opt(fp, c, "curl", "curl http://127.0.0.1:8000/v1/models"); - } else if (tool == DS4_HELP_AGENT || topic_is(topic, "sessions") || topic_is(topic, "tools")) { - opt(fp, c, "interactive", "./ds4-agent"); - opt(fp, c, "one shot", "./ds4-agent --non-interactive -p \"Create /tmp/hello.c\""); - } else if (tool == DS4_HELP_BENCH || topic_is(topic, "benchmark")) { - opt(fp, c, "csv", "./ds4-bench --prompt-file long.txt --ctx-max 32768 --csv speed.csv"); - opt(fp, c, "prefill only", "./ds4-bench --prompt-file long.txt --gen-tokens 0"); - } else if (tool == DS4_HELP_EVAL || topic_is(topic, "evaluation")) { - opt(fp, c, "first 10", "./ds4-eval --questions 10 --trace eval.trace"); - opt(fp, c, "plain", "./ds4-eval --plain --nothink --tokens 512"); - } else { - opt(fp, c, "chat", "./ds4"); - opt(fp, c, "one shot", "./ds4 -p \"Explain mmap in C\""); - opt(fp, c, "long prompt", "./ds4 --think-max --prompt-file prompt.txt --ctx 393216"); - } - fputc('\n', fp); -} - -static void print_topic(FILE *fp, const help_colors *c, ds4_help_tool tool, const char *topic) { - if (streq(topic, "all")) { - print_model_runtime(fp, c, tool, true); - if (tool_has_topic(tool, "sampling")) print_sampling(fp, c, true); - if (tool_has_topic(tool, "steering")) print_steering(fp, c); - print_distributed(fp, c); - if (tool == DS4_HELP_DS4) { - print_cli_specific(fp, c, true); - print_cli_commands(fp, c); - } else if (tool == DS4_HELP_SERVER) { - print_server_api(fp, c); - print_server_thinking(fp, c); - print_kv_cache(fp, c); - } else if (tool == DS4_HELP_AGENT) { - print_agent_specific(fp, c); - print_agent_sessions(fp, c); - } else if (tool == DS4_HELP_BENCH) { - print_bench_specific(fp, c); - } else if (tool == DS4_HELP_EVAL) { - print_eval_specific(fp, c); - } - return; - } - - if (streq(topic, "runtime")) print_model_runtime(fp, c, tool, true); - else if (streq(topic, "sampling")) print_sampling(fp, c, true); - else if (streq(topic, "steering")) print_steering(fp, c); - else if (streq(topic, "distributed")) print_distributed(fp, c); - else if (tool == DS4_HELP_DS4 && streq(topic, "diagnostics")) print_cli_diagnostics(fp, c); - else if (tool == DS4_HELP_DS4 && streq(topic, "commands")) print_cli_commands(fp, c); - else if (tool == DS4_HELP_SERVER && streq(topic, "api")) print_server_api(fp, c); - else if (tool == DS4_HELP_SERVER && streq(topic, "kv-cache")) print_kv_cache(fp, c); - else if (tool == DS4_HELP_SERVER && streq(topic, "thinking")) print_server_thinking(fp, c); - else if (tool == DS4_HELP_AGENT && streq(topic, "sessions")) print_agent_sessions(fp, c); - else if (tool == DS4_HELP_AGENT && streq(topic, "commands")) print_agent_sessions(fp, c); - else if (tool == DS4_HELP_AGENT && streq(topic, "tools")) { - title(fp, c, "Agent Tool System"); - para(fp, c, "The agent can read, search, write, edit, run bash, and browse through Chrome-backed web tools."); - para(fp, c, "Tool calls are emitted by the model as DSML and rendered live in the terminal."); - para(fp, c, "Edit uses exact old/new replacement; [upto] can bridge a unique head and tail for large anchored edits."); - fputc('\n', fp); - } else if (tool == DS4_HELP_BENCH && streq(topic, "benchmark")) print_bench_specific(fp, c); - else if (tool == DS4_HELP_EVAL && streq(topic, "evaluation")) print_eval_specific(fp, c); -} - -static void print_default(FILE *fp, const help_colors *c, ds4_help_tool tool) { - print_model_runtime(fp, c, tool, false); - - if (tool == DS4_HELP_DS4) { - print_cli_specific(fp, c, true); - print_sampling(fp, c, false); - } else if (tool == DS4_HELP_SERVER) { - print_server_api(fp, c); - print_kv_cache(fp, c); - } else if (tool == DS4_HELP_AGENT) { - print_agent_specific(fp, c); - print_agent_sessions(fp, c); - } else if (tool == DS4_HELP_BENCH) { - print_bench_specific(fp, c); - } else if (tool == DS4_HELP_EVAL) { - print_eval_specific(fp, c); - } -} - -void ds4_help_print(FILE *fp, ds4_help_tool tool, const char *topic) { - help_colors c = help_make_colors(fp); - if (topic && !tool_has_topic(tool, topic)) { - fprintf(fp, "%s: unknown help topic '%s'\n\n", tool_name(tool), topic); - topic = NULL; - } - - fprintf(fp, "%s%s%s\n", c.bright ? c.bright : "", tool_name(tool), c.off ? c.off : ""); - fprintf(fp, "%s\n\n", tool_summary(tool)); - fprintf(fp, "%s\n\n", tool_usage(tool)); - - if (topic) print_topic(fp, &c, tool, topic); - else { - print_default(fp, &c, tool); - print_more_info(fp, &c, tool); - } - print_examples(fp, &c, tool, topic); -} + opt(fp, c, "--parallel N", "Serve up to N requests concurrently by interleaving N sessions on the graph worker. Requires resident experts. Default: 1"); + } + if (tool == DS4_HELP_SERVER) { + opt(fp, c, "-n, --tokens N", "Default max output tokens when clients omit a limit."); + } + opt(fp, c, "-t, --threads N", "CPU helper threads for host-side/reference work."); + opt(fp, c, "--power N", "GPU duty-cycle target, 1..100. Default: 100"); + opt(fp, c, "--ssd-streaming", "Metal/CUDA/ROCm: opt in to SSD-backed model streaming instead of full residency."); + opt(fp, c, "--ssd-streaming-cold", "SSD streaming: skip default popularity-based expert-cache preload."); + opt(fp, c, "--ssd-streaming-cache-experts N|NGB", "SSD streaming: routed expert cache as expert count or GiB, e.g. 32GB. Metal/ROCm default: 80% working set minus non-routed weights; CUDA default: backend fixed cache."); + opt(fp, c, "--ssd-streaming-preload-experts N", "SSD streaming: upfront popularity preload count. Default: auto hot seed capped at 4096; use --ssd-streaming-cold to skip."); + opt(fp, c, "--simulate-used-memory NGB", "Diagnostic: lock N GiB before model load to simulate a smaller-memory machine."); + opt(fp, c, "--prefill-chunk N", "Metal graph prefill chunk size. Default: auto (PRO long prompts use 8192; others use 4096)."); + if (full) { + if (tool != DS4_HELP_BENCH) { + opt(fp, c, "--mtp FILE", "Optional MTP support GGUF used for draft-token probes."); + } + if (tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_SERVER) { + opt(fp, c, "--mtp-draft N", "Maximum autoregressive MTP draft tokens. Default: 1"); + opt(fp, c, "--mtp-margin F", "Verifier confidence margin for fast MTP acceptance. Default: 3"); + } + opt(fp, c, "--quality", "Prefer exact kernels where faster approximate paths exist."); + opt(fp, c, "--warm-weights", "Touch mapped tensor pages at startup to reduce first-use stalls."); + if (tool == DS4_HELP_DS4 || tool == DS4_HELP_BENCH) { + opt(fp, c, "--expert-profile FILE", "Metal-only: write routed expert locality/cache simulation JSON."); + } + } + fputc('\n', fp); +} + +static void print_sampling(FILE *fp, const help_colors *c, bool full) { + title(fp, c, "Prompt And Sampling"); + opt(fp, c, "-n, --tokens N", "Maximum generated tokens."); + opt(fp, c, "--temp F", "Sampling temperature. 0 is greedy/deterministic."); + opt(fp, c, "--top-p F", "Nucleus sampling probability."); + opt(fp, c, "--min-p F", "Keep tokens scoring at least F times the top token."); + opt(fp, c, "--seed N", "Sampling seed for reproducible non-greedy runs."); + opt(fp, c, "--think", "Use normal thinking mode."); + opt(fp, c, "--think-max", "Use Think Max when context is large enough."); + opt(fp, c, "--nothink", "Disable thinking and ask for direct replies."); + if (full) { + opt(fp, c, "-sys, --system TEXT", "System prompt. Empty string disables the default where supported."); + opt(fp, c, "-p, --prompt TEXT", "One-shot prompt text."); + opt(fp, c, "--prompt-file FILE", "Read one-shot prompt text from FILE."); + } + fputc('\n', fp); +} + +static void print_steering(FILE *fp, const help_colors *c) { + title(fp, c, "Directional Steering"); + opt(fp, c, "--dir-steering-file FILE", "Load one f32 direction vector per layer."); + opt(fp, c, "--dir-steering-ffn F", "Apply steering after FFN outputs. Default with file: 1"); + opt(fp, c, "--dir-steering-attn F", "Apply steering after attention outputs. Default: 0"); + fputc('\n', fp); +} + +static void print_distributed(FILE *fp, const help_colors *c) { + title(fp, c, "Distributed Inference"); + fputc('\n', fp); + para(fp, c, "Distributed mode runs one logical session across several machines by assigning contiguous model layer ranges to workers. Workers own their layer slice and KV-cache shard; the coordinator owns the prompt, sampling loop, and client/API flow. Start workers first, then start the coordinator. The coordinator waits for a complete route and streams hidden states through the workers."); + fputc('\n', fp); + opt(fp, c, "--role ROLE", "Distributed role: coordinator or worker."); + opt(fp, c, "--layers A:B", "Inclusive layer slice, e.g. 0:20 or 21:output."); + opt(fp, c, "--listen HOST PORT", "Coordinator listen address; workers may use it for their data listener."); + opt(fp, c, "--coordinator HOST PORT", "Coordinator address for --role worker."); + opt(fp, c, "--dist-prefill-chunk N", "Coordinator prefill pipeline chunk size. Default: session cap."); + opt(fp, c, "--dist-prefill-window N", "Max prefill chunks in flight. Default: workers+2, capped at 8."); + opt(fp, c, "--dist-activation-bits N", "Hidden-state transport width: 32, 16, or 8. Default: 32"); + opt(fp, c, "--dist-replay-check", "Diagnostic: reset and replay prompt, then compare logits."); + opt(fp, c, "--debug", "Print coordinator route/debug logs."); + fputc('\n', fp); +} + +static void print_cli_diagnostics(FILE *fp, const help_colors *c); + +static void print_cli_specific(FILE *fp, const help_colors *c, bool full) { + title(fp, c, "CLI Modes"); + opt(fp, c, "ds4", "Start the interactive prompt."); + opt(fp, c, "ds4 -p TEXT", "Run one prompt and exit."); + opt(fp, c, "ds4 --prompt-file FILE", "Run a long prompt from a file and exit."); + fputc('\n', fp); + if (full) { + print_cli_diagnostics(fp, c); + } +} + +static void print_cli_diagnostics(FILE *fp, const help_colors *c) { + title(fp, c, "Diagnostics And Data Collection"); + opt(fp, c, "--inspect", "Load the model and print a summary only."); + opt(fp, c, "--dump-tokens", "Tokenize the prompt exactly as written, then exit."); + opt(fp, c, "--dump-logits FILE", "Write full next-token logits as JSON."); + opt(fp, c, "--dump-logprobs FILE", "Write greedy continuation top-logprobs as JSON."); + opt(fp, c, "--logprobs-top-k N", "Alternatives stored by --dump-logprobs. Default: 20"); + opt(fp, c, "--expert-profile FILE", "Metal-only: write routed expert locality/cache simulation JSON."); + opt(fp, c, "--perplexity-file FILE", "Score raw text with teacher-forced NLL."); + opt(fp, c, "--imatrix-dataset FILE", "Rendered prompt dataset for imatrix collection."); + opt(fp, c, "--imatrix-out FILE", "Write llama-compatible routed-MoE imatrix .dat."); + opt(fp, c, "--imatrix-max-prompts N", "Stop imatrix collection after N prompts."); + opt(fp, c, "--imatrix-max-tokens N", "Stop imatrix collection after N prompt tokens."); + opt(fp, c, "--head-test", "Run the output HC/logits head after the native slice."); + opt(fp, c, "--first-token-test", "Run exact CPU whole-model pass for the first prompt token."); + opt(fp, c, "--metal-graph-test", "Compare first GPU-resident graph stages with CPU."); + opt(fp, c, "--metal-graph-full-test", "Run the GPU-resident self-token graph across all layers."); + opt(fp, c, "--metal-graph-prompt-test", "Compare CPU and GPU graph logits for the full prompt."); + fputc('\n', fp); +} + +static void print_cli_commands(FILE *fp, const help_colors *c) { + title_red(fp, c, "Interactive Commands"); + opt(fp, c, "/help", "Show interactive commands."); + opt(fp, c, "/think, /think-max, /nothink", "Switch thinking mode."); + opt(fp, c, "/ctx N", "Restart the interactive session with a new context size."); + opt(fp, c, "/power N", "Set GPU duty cycle percentage, 1..100."); + opt(fp, c, "/read FILE", "Read FILE and submit it as the next user message."); + opt(fp, c, "/quit, /exit", "Leave the prompt."); + opt(fp, c, "Ctrl+C", "Stop current generation and return to ds4>."); + fputc('\n', fp); +} + +static void print_agent_specific(FILE *fp, const help_colors *c) { + title(fp, c, "Agent Options"); + opt(fp, c, "-p, --prompt TEXT", "Submit an initial prompt after startup."); + opt(fp, c, "--non-interactive", "Run without TUI. With -p: one turn; without -p: repeated stdin prompts."); + opt(fp, c, "-sys, --system TEXT", "Extra system prompt. Empty disables extra text."); + opt(fp, c, "--trace FILE", "Write prompt, token, and DSML debug trace."); + opt(fp, c, "--chdir DIR", "Change working directory before loading runtime assets."); + fputc('\n', fp); +} + +static void print_agent_sessions(FILE *fp, const help_colors *c) { + title(fp, c, "Agent Runtime Commands"); + opt(fp, c, "/save", "Save the current session in ~/.ds4/kvcache."); + opt(fp, c, "/compact", "Compact the current session context now."); + opt(fp, c, "/list", "List saved sessions, sorted by recent update time."); + opt(fp, c, "/switch ID", "Load a saved session and show recent history."); + opt(fp, c, "/del ID", "Delete a saved session."); + opt(fp, c, "/strip ID", "Remove KV payload; the text history can be rebuilt later."); + opt(fp, c, "/history [N]", "Show N recent user turns from the current session."); + opt(fp, c, "/power N", "Set GPU duty cycle percentage, 1..100."); + opt(fp, c, "/new", "Start a fresh session from the system prompt."); + opt(fp, c, "/quit, /exit", "Exit."); + fputc('\n', fp); +} + +static void print_server_api(FILE *fp, const help_colors *c) { + title(fp, c, "HTTP API"); + opt(fp, c, "--host HOST", "Bind address. Default: 127.0.0.1"); + opt(fp, c, "--port N", "Bind port. Default: 8000"); + opt(fp, c, "--cors", "Add Access-Control-Allow-* headers for browser JS clients."); + opt(fp, c, "--trace FILE", "Write prompts, cache decisions, output, and tool calls."); + para(fp, c, "Endpoints: /v1/chat/completions, /v1/responses, /v1/completions, and /v1/messages."); + para(fp, c, "Model endpoint aliases include deepseek-v4-flash and deepseek-v4-pro; both serve the loaded GGUF."); + fputc('\n', fp); +} + +static void print_server_thinking(FILE *fp, const help_colors *c) { + title(fp, c, "Server Thinking Defaults"); + para(fp, c, "DeepSeek-compatible chat requests default to high-effort thinking."); + para(fp, c, "reasoning_effort=max or output_config.effort=max requests Think Max."); + para(fp, c, "Think Max requires --ctx >= 393216; smaller contexts use high."); + para(fp, c, "thinking={type:disabled}, think=false, or model=deepseek-chat selects non-thinking mode."); + para(fp, c, "In thinking mode, client sampling knobs are ignored like the official API."); + fputc('\n', fp); +} + +static void print_kv_cache(FILE *fp, const help_colors *c) { + title(fp, c, "Disk KV Cache"); + opt(fp, c, "--kv-disk-dir DIR", "Enable disk KV checkpoints in DIR."); + opt(fp, c, "--kv-disk-space-mb N", "Disk budget. Default when enabled: 4096"); + opt(fp, c, "--kv-cache-min-tokens N", "Do not save/load checkpoints shorter than N. Default: 512"); + opt(fp, c, "--kv-cache-cold-max-tokens N", "Save cold first prompts up to N tokens. 0 disables. Default: 30000"); + opt(fp, c, "--kv-cache-continued-interval-tokens N", "Save aligned continued frontiers. 0 disables. Default: 10000"); + opt(fp, c, "--kv-cache-boundary-trim-tokens N", "Trim tail tokens for cold boundary saves. Default: 32"); + opt(fp, c, "--kv-cache-boundary-align-tokens N", "Align cold boundary saves to this multiple. Default: 2048"); + opt(fp, c, "--kv-cache-reject-different-quant", "Reject checkpoints written with different routed-expert quantization."); + opt(fp, c, "--disable-exact-dsml-tool-replay", "Disable exact sampled DSML tool replay map."); + opt(fp, c, "--tool-memory-max-ids N", "Exact tool-call IDs kept in RAM. Default: 100000"); + fputc('\n', fp); +} + +static void print_bench_specific(FILE *fp, const help_colors *c) { + title(fp, c, "Benchmark Input"); + opt(fp, c, "--prompt-file FILE", "Raw benchmark text; token sequence is sliced at each frontier."); + opt(fp, c, "--chat-prompt-file FILE", "Render FILE as one no-thinking chat user message."); + opt(fp, c, "-sys, --system TEXT", "System prompt used only with --chat-prompt-file."); + fputc('\n', fp); + title(fp, c, "Benchmark Sweep"); + opt(fp, c, "--ctx-start N", "First measured frontier. Default: 2048"); + opt(fp, c, "--ctx-max N", "Last measured frontier. Default: 32768"); + opt(fp, c, "--ctx-alloc N", "Allocated context. Default: ctx-max + gen-tokens + 1"); + opt(fp, c, "--step-mul F", "Multiplicative step. Default: 1"); + opt(fp, c, "--step-incr N", "Linear step when --step-mul is 1. Default: 2048"); + opt(fp, c, "--gen-tokens N", "Greedy decode tokens per frontier. 0 for pure prefill. Default: 128"); + opt(fp, c, "--csv FILE", "Write CSV there instead of stdout."); + opt(fp, c, "--dump-frontier-logits-dir DIR", "Write one full-logit JSON file per frontier."); + fputc('\n', fp); +} + +static void print_eval_specific(FILE *fp, const help_colors *c) { + title(fp, c, "Evaluation"); + opt(fp, c, "-n, --tokens N", "Max generated tokens per question. Default: 16000"); + opt(fp, c, "--questions N", "Run only the first N embedded questions."); + opt(fp, c, "--case-sequence LIST", "Run 1-based case numbers in this comma-separated order."); + opt(fp, c, "--trace FILE", "Write questions, outputs, and grading decisions."); + opt(fp, c, "--regrade-trace FILE", "Regrade a prior trace without loading the model."); + opt(fp, c, "--soft-limit-reply-budget N", "Soft close thinking near the end of reply budget. Default: 1024"); + opt(fp, c, "--hard-limit-reply-budget N", "Force with N tokens left. Default: 512"); + opt(fp, c, "--soft-limit-think-close-rank N", "Soft-close when is in top N tokens. Default: 3"); + opt(fp, c, "--pause-ms N", "Pause after each result in the TTY UI. Default: 350"); + opt(fp, c, "--plain", "Disable split-screen ANSI UI."); + opt(fp, c, "--self-test-extractors", "Run answer-extractor self-tests and exit."); + fputc('\n', fp); +} + +static bool tool_has_topic(ds4_help_tool tool, const char *topic) { + if (!topic) return true; + if (streq(topic, "all")) return true; + if (streq(topic, "runtime") || streq(topic, "distributed")) return true; + if (streq(topic, "sampling")) + return tool == DS4_HELP_DS4 || tool == DS4_HELP_AGENT || tool == DS4_HELP_EVAL; + if (streq(topic, "steering")) + return tool == DS4_HELP_DS4 || tool == DS4_HELP_SERVER || tool == DS4_HELP_AGENT; + switch (tool) { + case DS4_HELP_DS4: + return streq(topic, "diagnostics") || streq(topic, "commands"); + case DS4_HELP_SERVER: + return streq(topic, "api") || streq(topic, "kv-cache") || streq(topic, "thinking"); + case DS4_HELP_AGENT: + return streq(topic, "sessions") || streq(topic, "commands") || streq(topic, "tools"); + case DS4_HELP_BENCH: + return streq(topic, "benchmark"); + case DS4_HELP_EVAL: + return streq(topic, "evaluation"); + } + return false; +} + +static void more_line(FILE *fp, const help_colors *c, const char *label, const char *topic) { + static const char *colors[] = { + "\x1b[38;5;81m", "\x1b[38;5;114m", "\x1b[38;5;179m", + "\x1b[38;5;141m", "\x1b[38;5;147m" + }; + static size_t idx; + const char *on = c->cyan ? colors[idx++ % (sizeof(colors) / sizeof(colors[0]))] : ""; + if (streq(label, "Interactive commands:") && c->red) on = c->red; + const char *off = c->off ? c->off : ""; + fprintf(fp, " %s%-26s%s --help %s\n", on, label, off, topic); +} + +static void print_more_info(FILE *fp, const help_colors *c, ds4_help_tool tool) { + title(fp, c, "More Info"); + more_line(fp, c, "Runtime full info:", "runtime"); + if (tool_has_topic(tool, "sampling")) + more_line(fp, c, "Sampling full info:", "sampling"); + more_line(fp, c, "Distributed inference:", "distributed"); + if (tool_has_topic(tool, "steering")) + more_line(fp, c, "Steering full info:", "steering"); + if (tool == DS4_HELP_DS4) { + more_line(fp, c, "Interactive commands:", "commands"); + more_line(fp, c, "Diagnostics:", "diagnostics"); + } else if (tool == DS4_HELP_SERVER) { + more_line(fp, c, "HTTP API:", "api"); + more_line(fp, c, "Disk KV cache:", "kv-cache"); + more_line(fp, c, "Thinking behavior:", "thinking"); + } else if (tool == DS4_HELP_AGENT) { + more_line(fp, c, "Agent sessions:", "sessions"); + more_line(fp, c, "Agent commands:", "commands"); + more_line(fp, c, "Agent tool system:", "tools"); + } else if (tool == DS4_HELP_BENCH) { + more_line(fp, c, "Benchmark sweep:", "benchmark"); + } else if (tool == DS4_HELP_EVAL) { + more_line(fp, c, "Evaluation options:", "evaluation"); + } + fputc('\n', fp); +} + +static void print_examples(FILE *fp, const help_colors *c, ds4_help_tool tool, const char *topic) { + title(fp, c, "Examples"); + if (topic_is(topic, "distributed")) { + opt(fp, c, "worker", "./ds4 --role worker --layers 21:output --coordinator 192.168.0.181 9000 -m ds4flash.gguf"); + opt(fp, c, "coordinator", "./ds4 --role coordinator --layers 0:20 --listen 0.0.0.0 9000 -p \"Hello\" -m ds4flash.gguf"); + } else if (topic_is(topic, "runtime")) { + if (tool == DS4_HELP_SERVER) { + opt(fp, c, "Metal API", "./ds4-server -m ds4flash.gguf --metal --ctx 100000"); + opt(fp, c, "quiet API", "./ds4-server --power 60 --host 127.0.0.1 --port 8000"); + } else if (tool == DS4_HELP_AGENT) { + opt(fp, c, "agent", "./ds4-agent -m ds4flash.gguf --ctx 100000"); + opt(fp, c, "quiet agent", "./ds4-agent --power 50"); + } else if (tool == DS4_HELP_BENCH) { + opt(fp, c, "bench", "./ds4-bench --prompt-file long.txt --ctx-max 32768"); + opt(fp, c, "quiet bench", "./ds4-bench --prompt-file long.txt --power 70"); + } else if (tool == DS4_HELP_EVAL) { + opt(fp, c, "eval", "./ds4-eval --questions 10 --ctx 100000"); + opt(fp, c, "CPU debug", "./ds4-eval --cpu --questions 1 --tokens 32"); + } else { + opt(fp, c, "Metal", "./ds4 -m ds4flash.gguf --metal -c 100000"); + opt(fp, c, "quiet thermals", "./ds4 -p \"Summarize README\" --power 50"); + } + } else if (topic_is(topic, "steering")) { + opt(fp, c, "steer FFN", "./ds4 -p \"Write tersely\" --dir-steering-file dir.bin --dir-steering-ffn 0.8"); + } else if (tool == DS4_HELP_SERVER || topic_is(topic, "api") || topic_is(topic, "kv-cache")) { + opt(fp, c, "local API", "./ds4-server --ctx 100000 --kv-disk-dir ~/.ds4/server-kv --kv-disk-space-mb 8192"); + opt(fp, c, "curl", "curl http://127.0.0.1:8000/v1/models"); + } else if (tool == DS4_HELP_AGENT || topic_is(topic, "sessions") || topic_is(topic, "tools")) { + opt(fp, c, "interactive", "./ds4-agent"); + opt(fp, c, "one shot", "./ds4-agent --non-interactive -p \"Create /tmp/hello.c\""); + } else if (tool == DS4_HELP_BENCH || topic_is(topic, "benchmark")) { + opt(fp, c, "csv", "./ds4-bench --prompt-file long.txt --ctx-max 32768 --csv speed.csv"); + opt(fp, c, "prefill only", "./ds4-bench --prompt-file long.txt --gen-tokens 0"); + } else if (tool == DS4_HELP_EVAL || topic_is(topic, "evaluation")) { + opt(fp, c, "first 10", "./ds4-eval --questions 10 --trace eval.trace"); + opt(fp, c, "plain", "./ds4-eval --plain --nothink --tokens 512"); + } else { + opt(fp, c, "chat", "./ds4"); + opt(fp, c, "one shot", "./ds4 -p \"Explain mmap in C\""); + opt(fp, c, "long prompt", "./ds4 --think-max --prompt-file prompt.txt --ctx 393216"); + } + fputc('\n', fp); +} + +static void print_topic(FILE *fp, const help_colors *c, ds4_help_tool tool, const char *topic) { + if (streq(topic, "all")) { + print_model_runtime(fp, c, tool, true); + if (tool_has_topic(tool, "sampling")) print_sampling(fp, c, true); + if (tool_has_topic(tool, "steering")) print_steering(fp, c); + print_distributed(fp, c); + if (tool == DS4_HELP_DS4) { + print_cli_specific(fp, c, true); + print_cli_commands(fp, c); + } else if (tool == DS4_HELP_SERVER) { + print_server_api(fp, c); + print_server_thinking(fp, c); + print_kv_cache(fp, c); + } else if (tool == DS4_HELP_AGENT) { + print_agent_specific(fp, c); + print_agent_sessions(fp, c); + } else if (tool == DS4_HELP_BENCH) { + print_bench_specific(fp, c); + } else if (tool == DS4_HELP_EVAL) { + print_eval_specific(fp, c); + } + return; + } + + if (streq(topic, "runtime")) print_model_runtime(fp, c, tool, true); + else if (streq(topic, "sampling")) print_sampling(fp, c, true); + else if (streq(topic, "steering")) print_steering(fp, c); + else if (streq(topic, "distributed")) print_distributed(fp, c); + else if (tool == DS4_HELP_DS4 && streq(topic, "diagnostics")) print_cli_diagnostics(fp, c); + else if (tool == DS4_HELP_DS4 && streq(topic, "commands")) print_cli_commands(fp, c); + else if (tool == DS4_HELP_SERVER && streq(topic, "api")) print_server_api(fp, c); + else if (tool == DS4_HELP_SERVER && streq(topic, "kv-cache")) print_kv_cache(fp, c); + else if (tool == DS4_HELP_SERVER && streq(topic, "thinking")) print_server_thinking(fp, c); + else if (tool == DS4_HELP_AGENT && streq(topic, "sessions")) print_agent_sessions(fp, c); + else if (tool == DS4_HELP_AGENT && streq(topic, "commands")) print_agent_sessions(fp, c); + else if (tool == DS4_HELP_AGENT && streq(topic, "tools")) { + title(fp, c, "Agent Tool System"); + para(fp, c, "The agent can read, search, write, edit, run bash, and browse through Chrome-backed web tools."); + para(fp, c, "Tool calls are emitted by the model as DSML and rendered live in the terminal."); + para(fp, c, "Edit uses exact old/new replacement; [upto] can bridge a unique head and tail for large anchored edits."); + fputc('\n', fp); + } else if (tool == DS4_HELP_BENCH && streq(topic, "benchmark")) print_bench_specific(fp, c); + else if (tool == DS4_HELP_EVAL && streq(topic, "evaluation")) print_eval_specific(fp, c); +} + +static void print_default(FILE *fp, const help_colors *c, ds4_help_tool tool) { + print_model_runtime(fp, c, tool, false); + + if (tool == DS4_HELP_DS4) { + print_cli_specific(fp, c, true); + print_sampling(fp, c, false); + } else if (tool == DS4_HELP_SERVER) { + print_server_api(fp, c); + print_kv_cache(fp, c); + } else if (tool == DS4_HELP_AGENT) { + print_agent_specific(fp, c); + print_agent_sessions(fp, c); + } else if (tool == DS4_HELP_BENCH) { + print_bench_specific(fp, c); + } else if (tool == DS4_HELP_EVAL) { + print_eval_specific(fp, c); + } +} + +void ds4_help_print(FILE *fp, ds4_help_tool tool, const char *topic) { + help_colors c = help_make_colors(fp); + if (topic && !tool_has_topic(tool, topic)) { + fprintf(fp, "%s: unknown help topic '%s'\n\n", tool_name(tool), topic); + topic = NULL; + } + + fprintf(fp, "%s%s%s\n", c.bright ? c.bright : "", tool_name(tool), c.off ? c.off : ""); + fprintf(fp, "%s\n\n", tool_summary(tool)); + fprintf(fp, "%s\n\n", tool_usage(tool)); + + if (topic) print_topic(fp, &c, tool, topic); + else { + print_default(fp, &c, tool); + print_more_info(fp, &c, tool); + } + print_examples(fp, &c, tool, topic); +} diff --git a/ds4_server.c b/ds4_server.c index 53df363d4..9eecd6458 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -519,6 +519,8 @@ static void random_tool_id(char *dst, size_t dstlen, api_style api) { } typedef struct server server; +typedef struct server_slot server_slot; +typedef struct gen_state gen_state; typedef struct { char *id; @@ -7710,15 +7712,51 @@ typedef struct { static bool id_list_contains(const stop_list *ids, const char *id); static void id_list_push_unique(stop_list *ids, const char *id); +typedef enum { + SLOT_IDLE = 0, + SLOT_BEGIN, + SLOT_PREFILL, + SLOT_START_DECODE, + SLOT_DECODE, + SLOT_FINISH, +} slot_phase; + +struct server_slot { + server *srv; + int id; + ds4_session *session; + live_tool_state responses_live; + live_tool_state anthropic_live; + visible_live_state thinking_live; + /* Per-slot copy of s->kv.continued_last_store_tokens: that field tracks + * the continued-checkpoint frontier of one live session, so the scheduler + * swaps it in and out around each turn instead of teaching the kvstore + * about slots. */ + int kv_continued_last_store_tokens; + job *job; + slot_phase phase; + gen_state *gs; + uint64_t last_step_seq; + double prefill_deadline; + bool yield_requested; +}; + struct server { ds4_engine *engine; - ds4_session *session; + /* One slot per live session (--parallel N, default 1). Each slot owns a + * session timeline plus the protocol live state bound to that timeline. + * `active` points at the slot whose job the scheduler is stepping right + * now; per-request code reaches its session and live state through it, + * which keeps the single-worker ownership model explicit. */ + server_slot *slots; + int n_slots; + server_slot *active; + uint64_t sched_seq; + double sched_prefill_slice_sec; + int sched_decode_tokens; int default_tokens; kv_disk_cache kv; tool_memory tool_mem; - live_tool_state responses_live; - live_tool_state anthropic_live; - visible_live_state thinking_live; bool disable_exact_dsml_tool_replay; bool enable_cors; pthread_mutex_t tool_mu; @@ -7977,18 +8015,18 @@ static void visible_live_free(visible_live_state *st) { static void thinking_live_clear(server *s) { if (!s) return; pthread_mutex_lock(&s->tool_mu); - visible_live_clear_locked(&s->thinking_live); + visible_live_clear_locked(&s->active->thinking_live); pthread_mutex_unlock(&s->tool_mu); } static void thinking_live_remember(server *s, const char *visible_text) { if (!s || !visible_text || !visible_text[0]) return; pthread_mutex_lock(&s->tool_mu); - visible_live_clear_locked(&s->thinking_live); - s->thinking_live.visible_text = xstrdup(visible_text); - s->thinking_live.visible_len = strlen(visible_text); - s->thinking_live.live_tokens = ds4_session_pos(s->session); - s->thinking_live.valid = true; + visible_live_clear_locked(&s->active->thinking_live); + s->active->thinking_live.visible_text = xstrdup(visible_text); + s->active->thinking_live.visible_len = strlen(visible_text); + s->active->thinking_live.live_tokens = ds4_session_pos(s->active->session); + s->active->thinking_live.valid = true; pthread_mutex_unlock(&s->tool_mu); } @@ -7996,50 +8034,55 @@ static void responses_live_remember(server *s, const char *visible_text, const tool_calls *calls) { if (!s || !visible_text || !visible_text[0]) return; pthread_mutex_lock(&s->tool_mu); - live_tool_state_clear_locked(&s->responses_live); - s->responses_live.visible_text = xstrdup(visible_text); - s->responses_live.visible_len = strlen(visible_text); + live_tool_state_clear_locked(&s->active->responses_live); + s->active->responses_live.visible_text = xstrdup(visible_text); + s->active->responses_live.visible_len = strlen(visible_text); if (calls) { for (int i = 0; i < calls->len; i++) { - id_list_push_unique(&s->responses_live.call_ids, calls->v[i].id); + id_list_push_unique(&s->active->responses_live.call_ids, calls->v[i].id); } } - s->responses_live.live_tokens = ds4_session_pos(s->session); - s->responses_live.valid = true; + s->active->responses_live.live_tokens = ds4_session_pos(s->active->session); + s->active->responses_live.valid = true; pthread_mutex_unlock(&s->tool_mu); } static void anthropic_live_remember(server *s, const tool_calls *calls) { if (!s || !calls || calls->len == 0) return; pthread_mutex_lock(&s->tool_mu); - live_tool_state_clear_locked(&s->anthropic_live); + live_tool_state_clear_locked(&s->active->anthropic_live); for (int i = 0; i < calls->len; i++) { - id_list_push_unique(&s->anthropic_live.call_ids, calls->v[i].id); + id_list_push_unique(&s->active->anthropic_live.call_ids, calls->v[i].id); } - s->anthropic_live.live_tokens = ds4_session_pos(s->session); - s->anthropic_live.valid = s->anthropic_live.call_ids.len > 0; + s->active->anthropic_live.live_tokens = ds4_session_pos(s->active->session); + s->active->anthropic_live.valid = s->active->anthropic_live.call_ids.len > 0; pthread_mutex_unlock(&s->tool_mu); } static void responses_live_clear(server *s) { if (!s) return; pthread_mutex_lock(&s->tool_mu); - live_tool_state_clear_locked(&s->responses_live); + live_tool_state_clear_locked(&s->active->responses_live); pthread_mutex_unlock(&s->tool_mu); } static void anthropic_live_clear(server *s) { if (!s) return; pthread_mutex_lock(&s->tool_mu); - live_tool_state_clear_locked(&s->anthropic_live); + live_tool_state_clear_locked(&s->active->anthropic_live); pthread_mutex_unlock(&s->tool_mu); } static bool responses_live_has_call_id(server *s, const char *id) { if (!s || !id || !id[0]) return false; + /* Called from the request parser on a client thread: the job has not been + * bound to a slot yet, so ask every slot's live state. */ pthread_mutex_lock(&s->tool_mu); - bool found = s->responses_live.valid && - id_list_contains(&s->responses_live.call_ids, id); + bool found = false; + for (int i = 0; !found && i < s->n_slots; i++) { + found = s->slots[i].responses_live.valid && + id_list_contains(&s->slots[i].responses_live.call_ids, id); + } pthread_mutex_unlock(&s->tool_mu); return found; } @@ -8047,8 +8090,11 @@ static bool responses_live_has_call_id(server *s, const char *id) { static bool anthropic_live_has_call_id(server *s, const char *id) { if (!s || !id || !id[0]) return false; pthread_mutex_lock(&s->tool_mu); - bool found = s->anthropic_live.valid && - id_list_contains(&s->anthropic_live.call_ids, id); + bool found = false; + for (int i = 0; !found && i < s->n_slots; i++) { + found = s->slots[i].anthropic_live.valid && + id_list_contains(&s->slots[i].anthropic_live.call_ids, id); + } pthread_mutex_unlock(&s->tool_mu); return found; } @@ -8057,11 +8103,11 @@ static bool responses_live_matches_request(server *s, const stop_list *ids, int live_tokens) { if (!s || !ids || ids->len == 0) return false; pthread_mutex_lock(&s->tool_mu); - bool ok = s->responses_live.valid && - s->responses_live.live_tokens == live_tokens && - s->responses_live.call_ids.len == ids->len; + bool ok = s->active->responses_live.valid && + s->active->responses_live.live_tokens == live_tokens && + s->active->responses_live.call_ids.len == ids->len; for (int i = 0; ok && i < ids->len; i++) { - ok = id_list_contains(&s->responses_live.call_ids, ids->v[i]); + ok = id_list_contains(&s->active->responses_live.call_ids, ids->v[i]); } pthread_mutex_unlock(&s->tool_mu); return ok; @@ -8071,11 +8117,11 @@ static bool anthropic_live_matches_request(server *s, const stop_list *ids, int live_tokens) { if (!s || !ids || ids->len == 0) return false; pthread_mutex_lock(&s->tool_mu); - bool ok = s->anthropic_live.valid && - s->anthropic_live.live_tokens == live_tokens && - s->anthropic_live.call_ids.len == ids->len; + bool ok = s->active->anthropic_live.valid && + s->active->anthropic_live.live_tokens == live_tokens && + s->active->anthropic_live.call_ids.len == ids->len; for (int i = 0; ok && i < ids->len; i++) { - ok = id_list_contains(&s->anthropic_live.call_ids, ids->v[i]); + ok = id_list_contains(&s->active->anthropic_live.call_ids, ids->v[i]); } pthread_mutex_unlock(&s->tool_mu); return ok; @@ -8710,7 +8756,7 @@ static bool kv_cache_store_live_prefix_text(server *s, const ds4_tokens *tokens, const char *cache_text_key) { char err[160] = {0}; ds4_kvstore_trailer_hooks hooks = kv_cache_tool_map_hooks(s, NULL); - return ds4_kvstore_store_live_prefix_text(&s->kv, s->engine, s->session, + return ds4_kvstore_store_live_prefix_text(&s->kv, s->engine, s->active->session, tokens, store_len, reason, cache_text_override, cache_text_ext, @@ -8725,27 +8771,27 @@ static bool kv_cache_store_live_prefix(server *s, const ds4_tokens *tokens, } static void kv_cache_store_current(server *s, const char *reason) { - const ds4_tokens *tokens = ds4_session_tokens(s->session); + const ds4_tokens *tokens = ds4_session_tokens(s->active->session); if (!tokens) return; char *visible_text = NULL; uint8_t visible_ext = 0; const char *visible_key = NULL; pthread_mutex_lock(&s->tool_mu); - if (s->responses_live.valid && - s->responses_live.live_tokens == tokens->len && - s->responses_live.visible_text && - s->responses_live.visible_text[0]) + if (s->active->responses_live.valid && + s->active->responses_live.live_tokens == tokens->len && + s->active->responses_live.visible_text && + s->active->responses_live.visible_text[0]) { - visible_text = xstrdup(s->responses_live.visible_text); + visible_text = xstrdup(s->active->responses_live.visible_text); visible_ext = KV_EXT_RESPONSES_VISIBLE; visible_key = "responses-visible"; - } else if (s->thinking_live.valid && - s->thinking_live.live_tokens == tokens->len && - s->thinking_live.visible_text && - s->thinking_live.visible_text[0]) + } else if (s->active->thinking_live.valid && + s->active->thinking_live.live_tokens == tokens->len && + s->active->thinking_live.visible_text && + s->active->thinking_live.visible_text[0]) { - visible_text = xstrdup(s->thinking_live.visible_text); + visible_text = xstrdup(s->active->thinking_live.visible_text); visible_ext = KV_EXT_THINKING_VISIBLE; visible_key = "thinking-visible"; } @@ -8791,12 +8837,12 @@ static void kv_cache_discard_failed_disk_entry(server *s, const char *path) { path, strerror(errno)); } s->kv.continued_last_store_tokens = 0; - ds4_session_invalidate(s->session); + ds4_session_invalidate(s->active->session); } static void kv_cache_maybe_store_continued(server *s) { kv_disk_cache *kc = &s->kv; - const ds4_tokens *tokens = ds4_session_tokens(s->session); + const ds4_tokens *tokens = ds4_session_tokens(s->active->session); if (!tokens) return; const int target = kv_cache_continued_store_target(kc, tokens->len); if (target == 0) return; @@ -8821,7 +8867,7 @@ static int kv_cache_try_load_text(server *s, const char *prompt_text, if (loaded_ext_flags_out) *loaded_ext_flags_out = 0; ds4_kvstore_load_result lr = {0}; ds4_kvstore_trailer_hooks hooks = kv_cache_tool_map_hooks(s, NULL); - int loaded = ds4_kvstore_try_load_text(&s->kv, s->engine, s->session, + int loaded = ds4_kvstore_try_load_text(&s->kv, s->engine, s->active->session, prompt_text, effective_prompt, &lr, &hooks, responses_protocol); if (loaded > 0) { @@ -8846,7 +8892,7 @@ static int kv_cache_try_load(server *s, const request *req, static int live_text_prefix_prompt(server *s, const request *req, ds4_tokens *effective_prompt) { if (!s || !req || !req->prompt_text || !effective_prompt) return 0; - const ds4_tokens *live_tokens = ds4_session_tokens(s->session); + const ds4_tokens *live_tokens = ds4_session_tokens(s->active->session); if (!live_tokens || live_tokens->len <= 0) return 0; size_t live_text_len = 0; @@ -8886,7 +8932,7 @@ static int responses_live_continuation_prompt(server *s, const request *req, if (!responses_live_matches_request(s, &req->responses_live_call_ids, live_pos)) return 0; - const ds4_tokens *live_tokens = ds4_session_tokens(s->session); + const ds4_tokens *live_tokens = ds4_session_tokens(s->active->session); if (!live_tokens || live_tokens->len != live_pos) return 0; build_prompt_from_exact_prefix_and_text_suffix( @@ -8912,7 +8958,7 @@ static int anthropic_live_continuation_prompt(server *s, const request *req, if (!anthropic_live_matches_request(s, &req->anthropic_live_call_ids, live_pos)) return 0; - const ds4_tokens *live_tokens = ds4_session_tokens(s->session); + const ds4_tokens *live_tokens = ds4_session_tokens(s->active->session); if (!live_tokens || live_tokens->len != live_pos) return 0; build_prompt_from_exact_prefix_and_text_suffix( @@ -8944,18 +8990,18 @@ static int responses_live_visible_prefix_prompt(server *s, const request *req, const size_t prompt_len = strlen(req->prompt_text); size_t visible_len = 0; pthread_mutex_lock(&s->tool_mu); - bool ok = s->responses_live.valid && - s->responses_live.live_tokens == live_pos && - s->responses_live.visible_text && - s->responses_live.visible_len < prompt_len && + bool ok = s->active->responses_live.valid && + s->active->responses_live.live_tokens == live_pos && + s->active->responses_live.visible_text && + s->active->responses_live.visible_len < prompt_len && byte_prefix_match(req->prompt_text, prompt_len, - s->responses_live.visible_text, - s->responses_live.visible_len); - if (ok) visible_len = s->responses_live.visible_len; + s->active->responses_live.visible_text, + s->active->responses_live.visible_len); + if (ok) visible_len = s->active->responses_live.visible_len; pthread_mutex_unlock(&s->tool_mu); if (!ok) return 0; - const ds4_tokens *live_tokens = ds4_session_tokens(s->session); + const ds4_tokens *live_tokens = ds4_session_tokens(s->active->session); if (!live_tokens || live_tokens->len != live_pos) return 0; build_prompt_from_exact_prefix_and_text_suffix( @@ -8986,18 +9032,18 @@ static int thinking_live_visible_prefix_prompt(server *s, const request *req, const size_t prompt_len = strlen(req->prompt_text); size_t visible_len = 0; pthread_mutex_lock(&s->tool_mu); - bool ok = s->thinking_live.valid && - s->thinking_live.live_tokens == live_pos && - s->thinking_live.visible_text && - s->thinking_live.visible_len < prompt_len && + bool ok = s->active->thinking_live.valid && + s->active->thinking_live.live_tokens == live_pos && + s->active->thinking_live.visible_text && + s->active->thinking_live.visible_len < prompt_len && byte_prefix_match(req->prompt_text, prompt_len, - s->thinking_live.visible_text, - s->thinking_live.visible_len); - if (ok) visible_len = s->thinking_live.visible_len; + s->active->thinking_live.visible_text, + s->active->thinking_live.visible_len); + if (ok) visible_len = s->active->thinking_live.visible_len; pthread_mutex_unlock(&s->tool_mu); if (!ok) return 0; - const ds4_tokens *live_tokens = ds4_session_tokens(s->session); + const ds4_tokens *live_tokens = ds4_session_tokens(s->active->session); if (!live_tokens || live_tokens->len != live_pos) return 0; build_prompt_from_exact_prefix_and_text_suffix( @@ -9468,7 +9514,7 @@ static int chat_think_tool_recovery(server *s, ds4_tokens toks = {0}; ds4_tokenize_rendered_chat(s->engine, inject, &toks); - const int room = ds4_session_ctx(s->session) - ds4_session_pos(s->session); + const int room = ds4_session_ctx(s->active->session) - ds4_session_pos(s->active->session); if (toks.len <= 0 || toks.len >= room || *completion + toks.len >= max_tokens) { @@ -9481,7 +9527,7 @@ static int chat_think_tool_recovery(server *s, } for (int i = 0; i < toks.len; i++) { - if (ds4_session_eval(s->session, toks.v[i], err, errlen) != 0) { + if (ds4_session_eval(s->active->session, toks.v[i], err, errlen) != 0) { ds4_tokens_free(&toks); return -1; } @@ -9554,7 +9600,7 @@ static bool append_rendered_suffix_to_live_session(server *s, const char *suffix char *err, size_t errlen) { if (tokens_appended) *tokens_appended = 0; if (!s || !suffix || !suffix[0]) return true; - const ds4_tokens *live = ds4_session_tokens(s->session); + const ds4_tokens *live = ds4_session_tokens(s->active->session); if (!live) { if (err && errlen) snprintf(err, errlen, "live session is unavailable"); return false; @@ -9562,10 +9608,10 @@ static bool append_rendered_suffix_to_live_session(server *s, const char *suffix ds4_tokens target = {0}; build_prompt_from_exact_prefix_and_text_suffix(s->engine, live, suffix, &target); - const int before = ds4_session_pos(s->session); - bool ok = ds4_session_sync(s->session, &target, err, errlen) == 0; + const int before = ds4_session_pos(s->active->session); + bool ok = ds4_session_sync(s->active->session, &target, err, errlen) == 0; if (ok && tokens_appended) { - int delta = ds4_session_pos(s->session) - before; + int delta = ds4_session_pos(s->active->session) - before; *tokens_appended = delta > 0 ? delta : 0; } ds4_tokens_free(&target); @@ -9806,10 +9852,10 @@ static void remember_thinking_checkpoint(server *s, const job *j, const char *ct thinking_live_remember(s, visible); server_log(DS4_LOG_KVCACHE, "ds4-server: thinking live checkpoint remembered ctx=%s live=%d visible=%zu", - ctx, ds4_session_pos(s->session), strlen(visible)); + ctx, ds4_session_pos(s->active->session), strlen(visible)); trace_event(s, trace_id, "thinking live checkpoint remembered: live=%d visible=%zu", - ds4_session_pos(s->session), strlen(visible)); + ds4_session_pos(s->active->session), strlen(visible)); free(visible); } @@ -9831,12 +9877,12 @@ static void canonicalize_tool_checkpoint(server *s, const job *j, const char *ct ds4_tokens canonical = {0}; ds4_tokenize_rendered_chat(s->engine, rendered.ptr ? rendered.ptr : "", &canonical); - const int live_len = ds4_session_pos(s->session); - const int common = ds4_session_common_prefix(s->session, &canonical); + const int live_len = ds4_session_pos(s->active->session); + const int common = ds4_session_common_prefix(s->active->session, &canonical); if (common == live_len && canonical.len == live_len) goto done; size_t live_text_len = 0; - char *live_text = render_tokens_text(s->engine, ds4_session_tokens(s->session), &live_text_len); + char *live_text = render_tokens_text(s->engine, ds4_session_tokens(s->active->session), &live_text_len); if (live_text_len == rendered.len && (live_text_len == 0 || memcmp(live_text, rendered.ptr, live_text_len) == 0)) { @@ -9857,7 +9903,7 @@ static void canonicalize_tool_checkpoint(server *s, const job *j, const char *ct char err[160] = {0}; ds4_session_rewrite_result rr = - ds4_session_rewrite_from_common(s->session, &canonical, common, + ds4_session_rewrite_from_common(s->active->session, &canonical, common, err, sizeof(err)); if (rr == DS4_SESSION_REWRITE_OK) { server_log(DS4_LOG_KVCACHE, @@ -9875,7 +9921,7 @@ static void canonicalize_tool_checkpoint(server *s, const job *j, const char *ct ds4_tokens effective = {0}; int loaded = kv_cache_try_load_text(s, rendered.ptr ? rendered.ptr : "", &effective, &path, NULL, false); - if (loaded == 0) ds4_session_invalidate(s->session); + if (loaded == 0) ds4_session_invalidate(s->active->session); char sync_err[160] = {0}; const ds4_tokens *sync_prompt = loaded > 0 ? &effective : &canonical; @@ -9921,11 +9967,11 @@ static void canonicalize_tool_checkpoint(server *s, const job *j, const char *ct .headers_sent = true, }; snprintf(rebuild_progress.ctx, sizeof(rebuild_progress.ctx), "%s", rebuild_ctx); - ds4_session_set_progress(s->session, server_progress_cb, &rebuild_progress); - ds4_session_set_display_progress(s->session, server_progress_cb, &rebuild_progress); - if (ds4_session_sync(s->session, sync_prompt, sync_err, sizeof(sync_err)) == 0) { - ds4_session_set_progress(s->session, NULL, NULL); - ds4_session_set_display_progress(s->session, NULL, NULL); + ds4_session_set_progress(s->active->session, server_progress_cb, &rebuild_progress); + ds4_session_set_display_progress(s->active->session, server_progress_cb, &rebuild_progress); + if (ds4_session_sync(s->active->session, sync_prompt, sync_err, sizeof(sync_err)) == 0) { + ds4_session_set_progress(s->active->session, NULL, NULL); + ds4_session_set_display_progress(s->active->session, NULL, NULL); const double rebuild_sec = now_sec() - rebuild_t0; if (loaded > 0) { server_log(DS4_LOG_KVCACHE, @@ -9943,8 +9989,8 @@ static void canonicalize_tool_checkpoint(server *s, const job *j, const char *ct common, live_len, canonical.len, err); } } else { - ds4_session_set_progress(s->session, NULL, NULL); - ds4_session_set_display_progress(s->session, NULL, NULL); + ds4_session_set_progress(s->active->session, NULL, NULL); + ds4_session_set_display_progress(s->active->session, NULL, NULL); server_log(DS4_LOG_KVCACHE, "ds4-server: tool checkpoint rebuild failed ctx=%s request_ctx=%s source=%s cached=%d replay=%d target=%d error=\"%s\"", rebuild_ctx, ctx, source, loaded, replay_tokens, @@ -10005,7 +10051,7 @@ typedef enum { GEN_STEP_REDECODE, /* tool-error recovery wants another decode round */ } gen_step; -typedef struct { +struct gen_state { char err[160]; int old_pos; int common; @@ -10032,6 +10078,7 @@ typedef struct { char req_flags[64]; int cold_store_len; int suppressed_continued_last; + int prefill_stage; /* 0 = cold checkpoint sync pending, 1 = main sync */ char id[96]; bool structured_stream; anthropic_stream anthropic_live; @@ -10064,16 +10111,16 @@ typedef struct { size_t think_recovery_scan_from; bool think_tool_recovery_enabled; dsml_decode_tracker dsml_tracker; -} gen_state; +}; /* Resolve continuations and caches, account usage, register the prefill * progress callback and plan the cold disk checkpoint. Returns false when * the request was already answered (continuation state unavailable). */ static bool job_begin(server *s, job *j, gen_state *gs) { gs->err[0] = '\0'; - gs->old_pos = ds4_session_pos(s->session); - gs->common = ds4_session_common_prefix(s->session, &j->req.prompt); - trace_cache_capture(&gs->cache_diag, ds4_session_tokens(s->session), + gs->old_pos = ds4_session_pos(s->active->session); + gs->common = ds4_session_common_prefix(s->active->session, &j->req.prompt); + trace_cache_capture(&gs->cache_diag, ds4_session_tokens(s->active->session), &j->req.prompt, gs->old_pos, gs->common); gs->prompt_for_sync = &j->req.prompt; gs->responses_protocol = j->req.api == API_RESPONSES; @@ -10266,8 +10313,8 @@ static bool job_begin(server *s, job *j, gen_state *gs) { gs->ctx_span, gs->req_flags[0] ? " " : "", gs->req_flags); - ds4_session_set_progress(s->session, server_progress_cb, &gs->progress); - ds4_session_set_display_progress(s->session, server_progress_cb, &gs->progress); + ds4_session_set_progress(s->active->session, server_progress_cb, &gs->progress); + ds4_session_set_display_progress(s->active->session, server_progress_cb, &gs->progress); gs->cold_store_len = 0; if (gs->cached == 0 && @@ -10296,53 +10343,79 @@ static bool job_begin(server *s, job *j, gen_state *gs) { return true; } +enum { + JOB_PREFILL_DONE = 0, + JOB_PREFILL_YIELD, /* interrupted at a chunk boundary; call again */ + JOB_PREFILL_FAILED, +}; + +/* True when the running sync was interrupted because the scheduler asked this + * slot to yield (as opposed to a real failure). */ +static bool slot_take_yield(server *s) { + server_slot *sl = s->active; + if (!sl || !sl->yield_requested) return false; + sl->yield_requested = false; + return true; +} + /* Run the cold-prefix sync (when a cold disk checkpoint is planned) and the - * main prompt sync. Returns false when prefill failed and the failure - * response was sent. */ -static bool job_prefill(server *s, job *j, gen_state *gs) { - if (s->kv.enabled && - gs->cold_store_len >= s->kv.opt.min_tokens && - gs->cold_store_len < gs->prompt_for_sync->len) - { - ds4_tokens prefix = {0}; - tokens_copy_prefix(&prefix, gs->prompt_for_sync, gs->cold_store_len); - if (ds4_session_sync(s->session, &prefix, gs->err, sizeof(gs->err)) != 0) { + * main prompt sync. Both syncs resume from the session checkpoint, so a + * JOB_PREFILL_YIELD result just means "call me again on this job's next + * turn". On failure the response has already been sent. */ +static int job_prefill(server *s, job *j, gen_state *gs) { + if (gs->prefill_stage == 0) { + if (s->kv.enabled && + gs->cold_store_len >= s->kv.opt.min_tokens && + gs->cold_store_len < gs->prompt_for_sync->len) + { + ds4_tokens prefix = {0}; + tokens_copy_prefix(&prefix, gs->prompt_for_sync, gs->cold_store_len); + const int rc = ds4_session_sync(s->active->session, &prefix, gs->err, sizeof(gs->err)); ds4_tokens_free(&prefix); - ds4_tokens_free(&gs->effective_prompt); - ds4_session_set_progress(s->session, NULL, NULL); - ds4_session_set_display_progress(s->session, NULL, NULL); - kv_cache_restore_suppressed_continued(&s->kv, gs->suppressed_continued_last, - gs->cold_store_len); - kv_cache_discard_failed_disk_entry(s, gs->disk_cache_path); - free(gs->disk_cache_path); - trace_event(s, gs->trace_id, "prefill failed: %s", gs->err); - send_prefill_failure_response(s, j, &gs->progress, gs->ctx_span, gs->req_flags, gs->err); - return false; - } - if (kv_cache_store_live_prefix(s, gs->prompt_for_sync, gs->cold_store_len, "cold")) { - kv_cache_note_store(&s->kv, gs->cold_store_len); - gs->suppressed_continued_last = -1; - } else { - kv_cache_restore_suppressed_continued(&s->kv, gs->suppressed_continued_last, - gs->cold_store_len); - gs->suppressed_continued_last = -1; + if (rc == DS4_SESSION_SYNC_INTERRUPTED && slot_take_yield(s)) { + return JOB_PREFILL_YIELD; + } + if (rc != 0) { + ds4_tokens_free(&gs->effective_prompt); + ds4_session_set_progress(s->active->session, NULL, NULL); + ds4_session_set_display_progress(s->active->session, NULL, NULL); + kv_cache_restore_suppressed_continued(&s->kv, gs->suppressed_continued_last, + gs->cold_store_len); + kv_cache_discard_failed_disk_entry(s, gs->disk_cache_path); + free(gs->disk_cache_path); + trace_event(s, gs->trace_id, "prefill failed: %s", gs->err); + send_prefill_failure_response(s, j, &gs->progress, gs->ctx_span, gs->req_flags, gs->err); + return JOB_PREFILL_FAILED; + } + if (kv_cache_store_live_prefix(s, gs->prompt_for_sync, gs->cold_store_len, "cold")) { + kv_cache_note_store(&s->kv, gs->cold_store_len); + gs->suppressed_continued_last = -1; + } else { + kv_cache_restore_suppressed_continued(&s->kv, gs->suppressed_continued_last, + gs->cold_store_len); + gs->suppressed_continued_last = -1; + } } - ds4_tokens_free(&prefix); + gs->prefill_stage = 1; } - if (ds4_session_sync(s->session, gs->prompt_for_sync, gs->err, sizeof(gs->err)) != 0) { + const int rc = ds4_session_sync(s->active->session, gs->prompt_for_sync, gs->err, sizeof(gs->err)); + if (rc == DS4_SESSION_SYNC_INTERRUPTED && slot_take_yield(s)) { + return JOB_PREFILL_YIELD; + } + if (rc != 0) { ds4_tokens_free(&gs->effective_prompt); - ds4_session_set_progress(s->session, NULL, NULL); - ds4_session_set_display_progress(s->session, NULL, NULL); + ds4_session_set_progress(s->active->session, NULL, NULL); + ds4_session_set_display_progress(s->active->session, NULL, NULL); kv_cache_restore_suppressed_continued(&s->kv, gs->suppressed_continued_last, gs->cold_store_len); kv_cache_discard_failed_disk_entry(s, gs->disk_cache_path); free(gs->disk_cache_path); trace_event(s, gs->trace_id, "prefill failed: %s", gs->err); send_prefill_failure_response(s, j, &gs->progress, gs->ctx_span, gs->req_flags, gs->err); - return false; + return JOB_PREFILL_FAILED; } - return true; + return JOB_PREFILL_DONE; } /* Post-prefill bookkeeping: clear stale live bindings, store checkpoints, @@ -10355,8 +10428,8 @@ static bool job_start_decode(server *s, job *j, gen_state *gs) { if (!gs->responses_live_continuation) responses_live_clear(s); if (!gs->anthropic_live_continuation) anthropic_live_clear(s); if (!gs->thinking_live_continuation) thinking_live_clear(s); - ds4_session_set_progress(s->session, NULL, NULL); - ds4_session_set_display_progress(s->session, NULL, NULL); + ds4_session_set_progress(s->active->session, NULL, NULL); + ds4_session_set_display_progress(s->active->session, NULL, NULL); kv_cache_maybe_store_continued(s); server_log(DS4_LOG_PREFILL, "ds4-server: %s ctx=%s%s%s prompt done %.3fs", @@ -10451,7 +10524,7 @@ static void job_decode_round_init(server *s, job *j, gen_state *gs) { gs->finish = "length"; gs->completion = 0; gs->max_tokens = j->req.max_tokens; - int room = ds4_session_ctx(s->session) - ds4_session_pos(s->session); + int room = ds4_session_ctx(s->active->session) - ds4_session_pos(s->active->session); gs->saw_tool_start = false; gs->saw_tool_end = false; gs->saw_orphan_tool_end = false; @@ -10479,7 +10552,7 @@ static void job_decode_round_init(server *s, job *j, gen_state *gs) { * Returns true while more tokens should be generated in this round. */ static bool job_decode_step(server *s, job *j, gen_state *gs) { if (g_stop_requested || gs->completion >= gs->max_tokens || - ds4_session_pos(s->session) >= ds4_session_ctx(s->session)) { + ds4_session_pos(s->active->session) >= ds4_session_ctx(s->active->session)) { return false; } dsml_decode_state dsml_state = j->req.kind == REQ_CHAT && j->req.has_tools ? @@ -10501,7 +10574,7 @@ static bool job_decode_step(server *s, job *j, gen_state *gs) { if (in_tool_call && !dsml_decode_state_uses_payload_sampling(dsml_state)) { temperature = 0.0f; } - int token = ds4_session_sample(s->session, temperature, top_k, top_p, min_p, &gs->rng); + int token = ds4_session_sample(s->active->session, temperature, top_k, top_p, min_p, &gs->rng); if (token == ds4_token_eos(s->engine)) { gs->finish = "stop"; return false; @@ -10513,7 +10586,7 @@ static bool job_decode_step(server *s, job *j, gen_state *gs) { ds4_engine_mtp_draft_tokens(s->engine) > 1 && getenv("DS4_MTP_SPEC_DISABLE") == NULL) { - ntok = ds4_session_eval_speculative_argmax(s->session, + ntok = ds4_session_eval_speculative_argmax(s->active->session, token, gs->max_tokens - gs->completion, ds4_token_eos(s->engine), @@ -10526,7 +10599,7 @@ static bool job_decode_step(server *s, job *j, gen_state *gs) { return false; } } else { - if (ds4_session_eval(s->session, token, gs->err, sizeof(gs->err)) != 0) { + if (ds4_session_eval(s->active->session, token, gs->err, sizeof(gs->err)) != 0) { gs->finish = "error"; return false; } @@ -10710,7 +10783,7 @@ static bool job_decode_step(server *s, job *j, gen_state *gs) { gs->finish = "stop"; gs->text.len = stop_pos; gs->text.ptr[gs->text.len] = '\0'; - ds4_session_invalidate(s->session); + ds4_session_invalidate(s->active->session); stop_decode = true; break; } @@ -11135,18 +11208,139 @@ static gen_step job_finish(server *s, job *j, gen_state *gs) { return GEN_STEP_DONE; } -static void generate_job(server *s, job *j) { - gen_state gs; - memset(&gs, 0, sizeof(gs)); - if (!job_begin(s, j, &gs)) return; - if (!job_prefill(s, j, &gs)) return; - if (!job_start_decode(s, j, &gs)) return; - for (;;) { - job_decode_round_init(s, j, &gs); - while (job_decode_step(s, j, &gs)) { +/* Cooperative prefill preemption. Registered once per session; consulted by + * ds4_session_sync() at chunk boundaries. Yield only when the slice expired + * and some other job actually needs the worker. Reads of other slots' job + * pointers and of the queue head are benign races: the worst case is one + * extra prefill chunk before yielding or one spurious yield. */ +static bool slot_prefill_cancel_cb(void *ud) { + server_slot *sl = ud; + server *s = sl->srv; + if (sl->yield_requested) return true; + if (s->n_slots < 2) return false; + bool contended = s->head != NULL; + for (int i = 0; !contended && i < s->n_slots; i++) { + if (&s->slots[i] != sl && s->slots[i].job) contended = true; + } + if (!contended) return false; + if (now_sec() < sl->prefill_deadline) return false; + sl->yield_requested = true; + return true; +} + +static void slot_complete(server *s, server_slot *sl) { + (void)s; + job *j = sl->job; + sl->job = NULL; + sl->phase = SLOT_IDLE; + pthread_mutex_lock(&j->mu); + j->done = true; + pthread_cond_signal(&j->cv); + pthread_mutex_unlock(&j->mu); +} + +/* Run one scheduling turn for a slot: a whole phase transition, one prefill + * slice, or up to sched_decode_tokens decode steps. */ +static void slot_step(server *s, server_slot *sl) { + sl->last_step_seq = ++s->sched_seq; + s->active = sl; + s->kv.continued_last_store_tokens = sl->kv_continued_last_store_tokens; + job *j = sl->job; + gen_state *gs = sl->gs; + switch (sl->phase) { + case SLOT_IDLE: + break; + case SLOT_BEGIN: + if (!job_begin(s, j, gs)) { + slot_complete(s, sl); + break; + } + sl->phase = SLOT_PREFILL; + break; + case SLOT_PREFILL: { + sl->yield_requested = false; + sl->prefill_deadline = now_sec() + s->sched_prefill_slice_sec; + const int rc = job_prefill(s, j, gs); + if (rc == JOB_PREFILL_FAILED) { + slot_complete(s, sl); + break; } - if (job_finish(s, j, &gs) != GEN_STEP_REDECODE) break; + if (rc == JOB_PREFILL_YIELD) break; + sl->phase = SLOT_START_DECODE; + break; + } + case SLOT_START_DECODE: + if (!job_start_decode(s, j, gs)) { + slot_complete(s, sl); + break; + } + job_decode_round_init(s, j, gs); + sl->phase = SLOT_DECODE; + break; + case SLOT_DECODE: { + int budget = s->sched_decode_tokens > 0 ? s->sched_decode_tokens : 1; + bool more = true; + while (budget-- > 0 && (more = job_decode_step(s, j, gs))) { + } + if (!more) sl->phase = SLOT_FINISH; + break; + } + case SLOT_FINISH: + if (job_finish(s, j, gs) == GEN_STEP_REDECODE) { + job_decode_round_init(s, j, gs); + sl->phase = SLOT_DECODE; + } else { + slot_complete(s, sl); + } + break; } + sl->kv_continued_last_store_tokens = s->kv.continued_last_store_tokens; + s->active = NULL; +} + +static bool sched_any_runnable(server *s) { + for (int i = 0; i < s->n_slots; i++) { + if (s->slots[i].job) return true; + } + return false; +} + +/* Round-robin over slots with work: pick the one that ran longest ago. */ +static server_slot *sched_next_slot(server *s) { + server_slot *best = NULL; + for (int i = 0; i < s->n_slots; i++) { + server_slot *sl = &s->slots[i]; + if (!sl->job) continue; + if (!best || sl->last_step_seq < best->last_step_seq) best = sl; + } + return best; +} + +/* Session affinity for new jobs: the idle slot whose live timeline shares the + * longest token prefix with the prompt keeps caches warm (the shared system + * prompt, live tool continuations). Ties go to the least recently used. */ +static server_slot *sched_pick_idle_slot(server *s, job *j) { + server_slot *best = NULL; + int best_common = -1; + for (int i = 0; i < s->n_slots; i++) { + server_slot *sl = &s->slots[i]; + if (sl->job) continue; + const int common = ds4_session_common_prefix(sl->session, &j->req.prompt); + if (!best || common > best_common || + (common == best_common && sl->last_step_seq < best->last_step_seq)) { + best = sl; + best_common = common; + } + } + return best; +} + +static void slot_attach(server *s, server_slot *sl, job *j) { + sl->job = j; + sl->phase = SLOT_BEGIN; + memset(sl->gs, 0, sizeof(*sl->gs)); + sl->yield_requested = false; + sl->last_step_seq = ++s->sched_seq; } @@ -11163,31 +11357,30 @@ static bool enqueue(server *s, job *j) { return true; } -static job *dequeue(server *s) { - pthread_mutex_lock(&s->mu); - while (!s->head && !s->stopping) pthread_cond_wait(&s->cv, &s->mu); - if (!s->head) { - pthread_mutex_unlock(&s->mu); - return NULL; - } - job *j = s->head; - s->head = j->next; - if (!s->head) s->tail = NULL; - pthread_mutex_unlock(&s->mu); - j->next = NULL; - return j; -} - static void *worker_main(void *arg) { server *s = arg; for (;;) { - job *j = dequeue(s); - if (!j) break; - generate_job(s, j); - pthread_mutex_lock(&j->mu); - j->done = true; - pthread_cond_signal(&j->cv); - pthread_mutex_unlock(&j->mu); + pthread_mutex_lock(&s->mu); + while (!s->stopping && !s->head && !sched_any_runnable(s)) { + pthread_cond_wait(&s->cv, &s->mu); + } + if (s->stopping && !s->head && !sched_any_runnable(s)) { + pthread_mutex_unlock(&s->mu); + break; + } + /* Hand queued jobs to idle slots (affinity first, then LRU). */ + while (s->head) { + server_slot *sl = sched_pick_idle_slot(s, s->head); + if (!sl) break; + job *j = s->head; + s->head = j->next; + if (!s->head) s->tail = NULL; + j->next = NULL; + slot_attach(s, sl, j); + } + pthread_mutex_unlock(&s->mu); + server_slot *sl = sched_next_slot(s); + if (sl) slot_step(s, sl); } return NULL; } @@ -11324,7 +11517,7 @@ static void append_model_json(buf *b, const server *s, const char *id) { append_model_json_values(b, id, ds4_engine_model_name(s->engine), - ds4_session_ctx(s->session), + ds4_session_ctx(s->slots[0].session), s->default_tokens); } @@ -11395,7 +11588,7 @@ static void *client_main(void *arg) { request req; char err[160]; bool ok = false; - const int ctx_size = ds4_session_ctx(s->session); + const int ctx_size = ds4_session_ctx(s->slots[0].session); if (!strcmp(hr.method, "POST") && !strcmp(hr.path, "/v1/messages")) { ok = parse_anthropic_request(s->engine, s, hr.body, s->default_tokens, ctx_size, &req, err, sizeof(err)); @@ -11516,6 +11709,7 @@ typedef struct { bool disable_exact_dsml_tool_replay; int tool_memory_max_ids; bool enable_cors; + int parallel; } server_config; static int parse_int_arg(const char *s, const char *opt) { @@ -11580,15 +11774,21 @@ static void server_close_resources(server *s) { } kv_cache_close(&s->kv); tool_memory_free(&s->tool_mem); - live_tool_state_free(&s->responses_live); - live_tool_state_free(&s->anthropic_live); - visible_live_free(&s->thinking_live); + for (int i = 0; i < s->n_slots; i++) { + live_tool_state_free(&s->slots[i].responses_live); + live_tool_state_free(&s->slots[i].anthropic_live); + visible_live_free(&s->slots[i].thinking_live); + if (s->slots[i].session) ds4_session_free(s->slots[i].session); + free(s->slots[i].gs); + } + free(s->slots); + s->slots = NULL; + s->n_slots = 0; pthread_mutex_destroy(&s->tool_mu); pthread_mutex_destroy(&s->trace_mu); pthread_cond_destroy(&s->clients_cv); pthread_cond_destroy(&s->cv); pthread_mutex_destroy(&s->mu); - ds4_session_free(s->session); ds4_engine_close(s->engine); memset(s, 0, sizeof(*s)); } @@ -11676,6 +11876,9 @@ static server_config parse_options(int argc, char **argv) { c.engine.mtp_margin = parse_float_arg(need_arg(&i, argc, argv, arg), arg, 0.0f, 1000.0f); } else if (!strcmp(arg, "-c") || !strcmp(arg, "--ctx")) { c.ctx_size = parse_int_arg(need_arg(&i, argc, argv, arg), arg); + } else if (!strcmp(arg, "--parallel")) { + c.parallel = parse_int_arg(need_arg(&i, argc, argv, arg), arg); + if (c.parallel < 1) c.parallel = 1; } else if (!strcmp(arg, "-n") || !strcmp(arg, "--tokens")) { c.default_tokens = parse_int_arg(need_arg(&i, argc, argv, arg), arg); } else if (!strcmp(arg, "-t") || !strcmp(arg, "--threads")) { @@ -11838,10 +12041,12 @@ int main(int argc, char **argv) { return rc; } - ds4_session *session = NULL; - if (ds4_session_create(&session, engine, cfg.ctx_size) != 0) { - server_log(DS4_LOG_DEFAULT, "ds4-server: failed to create %s session", - ds4_backend_name(cfg.engine.backend)); + const int n_slots = cfg.parallel > 0 ? cfg.parallel : 1; + if (n_slots > 1 && cfg.engine.ssd_streaming) { + server_log(DS4_LOG_DEFAULT, + "ds4-server: --parallel %d requires resident experts; " + "it cannot be combined with --ssd-streaming", + n_slots); ds4_engine_close(engine); return 1; } @@ -11849,7 +12054,45 @@ int main(int argc, char **argv) { server s; memset(&s, 0, sizeof(s)); s.engine = engine; - s.session = session; + s.slots = xmalloc((size_t)n_slots * sizeof(*s.slots)); + memset(s.slots, 0, (size_t)n_slots * sizeof(*s.slots)); + s.n_slots = n_slots; + { + const char *ms = getenv("DS4_SERVER_SCHED_PREFILL_MS"); + const char *dt = getenv("DS4_SERVER_SCHED_DECODE_TOKENS"); + s.sched_prefill_slice_sec = ms ? atof(ms) / 1000.0 : 1.5; + s.sched_decode_tokens = dt ? atoi(dt) : 6; + if (s.sched_prefill_slice_sec <= 0) s.sched_prefill_slice_sec = 1.5; + if (s.sched_decode_tokens < 1) s.sched_decode_tokens = 1; + } + for (int i = 0; i < n_slots; i++) { + server_slot *sl = &s.slots[i]; + sl->srv = &s; + sl->id = i; + sl->gs = xmalloc(sizeof(*sl->gs)); + memset(sl->gs, 0, sizeof(*sl->gs)); + if (ds4_session_create(&sl->session, engine, cfg.ctx_size) != 0) { + server_log(DS4_LOG_DEFAULT, + "ds4-server: failed to create %s session %d/%d", + ds4_backend_name(cfg.engine.backend), i + 1, n_slots); + free(sl->gs); + while (i-- > 0) { + ds4_session_free(s.slots[i].session); + free(s.slots[i].gs); + } + free(s.slots); + ds4_engine_close(engine); + return 1; + } + ds4_session_set_cancel(sl->session, slot_prefill_cancel_cb, sl); + } + if (n_slots > 1) { + server_log(DS4_LOG_DEFAULT, + "ds4-server: %d parallel sessions, ctx %d each " + "(prefill slice %.0f ms, decode burst %d tokens)", + n_slots, cfg.ctx_size, + s.sched_prefill_slice_sec * 1000.0, s.sched_decode_tokens); + } s.default_tokens = cfg.default_tokens; s.disable_exact_dsml_tool_replay = cfg.disable_exact_dsml_tool_replay; s.tool_mem.max_entries = cfg.tool_memory_max_ids; @@ -11943,13 +12186,19 @@ int main(int argc, char **argv) { while (s.clients > 0) pthread_cond_wait(&s.clients_cv, &s.mu); pthread_mutex_unlock(&s.mu); - const ds4_tokens *tokens = ds4_session_tokens(s.session); - if (s.kv.enabled && tokens && tokens->len >= s.kv.opt.min_tokens) { - server_log(DS4_LOG_KVCACHE, - "ds4-server: persisting current KV cache before shutdown tokens=%d", - tokens->len); - kv_cache_store_current(&s, "shutdown"); + for (int i = 0; i < s.n_slots; i++) { + s.active = &s.slots[i]; + s.kv.continued_last_store_tokens = + s.slots[i].kv_continued_last_store_tokens; + const ds4_tokens *tokens = ds4_session_tokens(s.active->session); + if (s.kv.enabled && tokens && tokens->len >= s.kv.opt.min_tokens) { + server_log(DS4_LOG_KVCACHE, + "ds4-server: persisting slot %d KV cache before shutdown tokens=%d", + i, tokens->len); + kv_cache_store_current(&s, "shutdown"); + } } + s.active = NULL; server_close_resources(&s); return 0; }