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/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 34a9d5084..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, @@ -9988,184 +10034,261 @@ 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; + +struct gen_state { 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; + int prefill_stage; /* 0 = cold checkpoint sync pending, 1 = main sync */ + 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; +}; + +/* 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->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; + 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 +10300,265 @@ 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->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 && 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; +} - if (s->kv.enabled && - cold_store_len >= s->kv.opt.min_tokens && - cold_store_len < 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) { +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. 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(&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; + 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; + } } - 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; - } else { - kv_cache_restore_suppressed_continued(&s->kv, suppressed_continued_last, - cold_store_len); - 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); - 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; + gs->prefill_stage = 1; + } + + 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; } - free(disk_cache_path); + 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; + } + return JOB_PREFILL_DONE; +} + +/* 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); - ds4_session_set_progress(s->session, NULL, NULL); - ds4_session_set_display_progress(s->session, NULL, NULL); + 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->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", 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; - 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 = + 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->active->session) - ds4_session_pos(s->active->session); + 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->active->session) >= ds4_session_ctx(s->active->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 +10574,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->active->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]; @@ -10419,100 +10586,100 @@ static void generate_job(server *s, job *j) { 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, - 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->active->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 +10687,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 +10709,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'; - ds4_session_invalidate(s->session); + gs->finish = "stop"; + gs->text.len = stop_pos; + gs->text.ptr[gs->text.len] = '\0'; + ds4_session_invalidate(s->active->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 +10815,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 +10824,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 +10965,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 +10974,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 +11013,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 +11063,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 +11080,270 @@ 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; +} + +/* 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 (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; +} + + static bool enqueue(server *s, job *j) { pthread_mutex_lock(&s->mu); if (s->stopping) { @@ -11049,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; } @@ -11210,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); } @@ -11281,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)); @@ -11402,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) { @@ -11466,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)); } @@ -11562,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")) { @@ -11724,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; } @@ -11735,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; @@ -11829,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; } 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; +}