diff --git a/ds4_server.c b/ds4_server.c index f4f563ecb..5fa14e000 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -8163,6 +8163,8 @@ static void server_log(ds4_log_type type, const char *fmt, ...) { typedef struct job job; typedef struct server_slot server_slot; +static bool server_prefill_enter(server *s, server_slot *slot); +static void server_prefill_leave(server *s); typedef ds4_kvstore_entry kv_entry; typedef ds4_kvstore_options kv_cache_options; @@ -8280,6 +8282,11 @@ struct server { pthread_cond_t model_cv; bool model_busy; bool model_stopping; + /* A bounded prefill quantum must hand the next model turn to an already + * active generator. Without an explicit handoff, the prefill worker can + * repeatedly reacquire model_mu before decode workers enqueue their next + * token, turning a 128-token quantum into an unbounded decode stall. */ + bool decode_handoff; int decode_pending; int active_generations; int last_prefill_slot; @@ -9348,6 +9355,19 @@ static void kv_cache_store_current(server *s, server_slot *slot, } } +/* Replacing a resident session with a disk checkpoint first persists the + * current frontier. Treat that eviction save as model work so a long prefill + * cannot repeatedly reacquire inference_mu ahead of another resident slot. + * The raw store helper remains recursive because prefill progress callbacks + * also use it while their slot already owns the model turn. */ +static bool kv_cache_store_current_scheduled(server *s, server_slot *slot, + const char *reason) { + if (!server_prefill_enter(s, slot)) return false; + kv_cache_store_current(s, slot, reason); + server_prefill_leave(s); + return true; +} + #ifdef DS4_SERVER_TEST static int kv_cache_suppress_continued_store(kv_disk_cache *kc, int tokens) { return ds4_kvstore_suppress_continued_store(kc, tokens); @@ -9442,13 +9462,13 @@ static int kv_cache_try_load_text(server *s, server_slot *slot, 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); - pthread_mutex_lock(&s->inference_mu); + if (!server_prefill_enter(s, slot)) return 0; pthread_mutex_lock(&s->kv_mu); int loaded = ds4_kvstore_try_load_text(&s->kv, s->engine, slot->session, prompt_text, effective_prompt, &lr, &hooks, responses_protocol); pthread_mutex_unlock(&s->kv_mu); - pthread_mutex_unlock(&s->inference_mu); + server_prefill_leave(s); if (loaded > 0) { if (loaded_path_out && lr.path) *loaded_path_out = xstrdup(lr.path); if (loaded_ext_flags_out) *loaded_ext_flags_out = lr.ext_flags; @@ -10245,7 +10265,7 @@ static bool server_prefill_enter(server *s, server_slot *slot) { slot->prefill_waiting = true; pthread_cond_broadcast(&s->model_cv); while (!g_stop_requested && - (s->model_busy || s->decode_pending > 0 || + (s->model_busy || s->decode_handoff || s->decode_pending > 0 || server_next_prefill_slot_locked(s) != slot->id)) { pthread_cond_wait(&s->model_cv, &s->model_mu); } @@ -10268,7 +10288,20 @@ static void server_prefill_leave(server *s) { if (!s->batched_mode) return; pthread_mutex_lock(&s->model_mu); s->model_busy = false; + if (s->active_generations > 0 && !s->model_stopping && + !g_stop_requested) { + s->decode_handoff = true; + } pthread_cond_broadcast(&s->model_cv); + while (s->decode_handoff && s->active_generations > 0 && + !s->model_stopping && !g_stop_requested) { + pthread_cond_wait(&s->model_cv, &s->model_mu); + } + if (s->active_generations == 0 || s->model_stopping || + g_stop_requested) { + s->decode_handoff = false; + pthread_cond_broadcast(&s->model_cv); + } pthread_mutex_unlock(&s->model_mu); } @@ -10309,10 +10342,8 @@ static int server_session_sync(server *s, server_slot *slot, return rc; } - pthread_mutex_lock(&s->inference_mu); int live = ds4_session_pos(slot->session); int common = ds4_session_common_prefix(slot->session, prompt); - pthread_mutex_unlock(&s->inference_mu); int done = common == live && prompt->len >= live ? live : 0; bool called = false; @@ -10795,6 +10826,11 @@ static void server_generation_enter(server *s) { if (!s || !s->batched_mode) return; pthread_mutex_lock(&s->model_mu); s->active_generations++; + /* Reserve the next model turn immediately. A different prefill may have + * acquired the executor after this request's final prefill quantum but + * before it registered as a generator; waiting until that prefill leaves + * can otherwise repeat the first-token starvation race. */ + s->decode_handoff = true; pthread_cond_broadcast(&s->model_cv); pthread_mutex_unlock(&s->model_mu); } @@ -10803,6 +10839,7 @@ static void server_generation_leave(server *s) { if (!s || !s->batched_mode) return; pthread_mutex_lock(&s->model_mu); if (s->active_generations > 0) s->active_generations--; + if (s->active_generations == 0) s->decode_handoff = false; pthread_cond_broadcast(&s->model_cv); pthread_mutex_unlock(&s->model_mu); } @@ -10926,6 +10963,7 @@ static void *decode_worker_main(void *arg) { pthread_mutex_lock(&s->model_mu); s->model_busy = false; + s->decode_handoff = false; for (int i = 0; i < count; i++) { server_slot *slot = members[i]; slot->decode_in_flight = false; @@ -11075,7 +11113,10 @@ static void generate_job(server *s, server_slot *slot, job *j) { /* 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, slot, "evict"); + if (!kv_cache_store_current_scheduled(s, slot, "evict")) { + ds4_tokens_free(&effective_prompt); + return; + } } if (cached == 0) { disk_cached = kv_cache_try_load(s, slot, &j->req, &effective_prompt, @@ -13181,6 +13222,101 @@ static void test_batched_prefill_round_robin(void) { TEST_ASSERT(server_next_prefill_slot_locked(&s) == -1); } +typedef struct { + server *srv; + pthread_mutex_t mu; + pthread_cond_t cv; + bool entered; + bool done; +} test_prefill_handoff_state; + +static void *test_prefill_handoff_worker(void *arg) { + test_prefill_handoff_state *state = arg; + pthread_mutex_lock(&state->srv->inference_mu); + pthread_mutex_lock(&state->mu); + state->entered = true; + pthread_cond_broadcast(&state->cv); + pthread_mutex_unlock(&state->mu); + + server_prefill_leave(state->srv); + + pthread_mutex_lock(&state->mu); + state->done = true; + pthread_cond_broadcast(&state->cv); + pthread_mutex_unlock(&state->mu); + return NULL; +} + +static void test_batched_prefill_hands_off_to_active_decode(void) { + server s = {0}; + s.batched_mode = true; + s.model_busy = true; + pthread_mutex_init(&s.inference_mu, NULL); + pthread_mutex_init(&s.model_mu, NULL); + pthread_cond_init(&s.model_cv, NULL); + + server_generation_enter(&s); + TEST_ASSERT(s.active_generations == 1); + TEST_ASSERT(s.decode_handoff); + + test_prefill_handoff_state state = {.srv = &s}; + pthread_mutex_init(&state.mu, NULL); + pthread_cond_init(&state.cv, NULL); + + /* Hold model_mu until the worker has entered server_prefill_leave(), then + * observe the handoff while that worker waits for a decode completion. */ + pthread_mutex_lock(&s.model_mu); + pthread_t thread; + int create_rc = pthread_create(&thread, NULL, + test_prefill_handoff_worker, &state); + TEST_ASSERT(create_rc == 0); + if (create_rc != 0) { + pthread_mutex_unlock(&s.model_mu); + pthread_cond_destroy(&state.cv); + pthread_mutex_destroy(&state.mu); + pthread_cond_destroy(&s.model_cv); + pthread_mutex_destroy(&s.model_mu); + pthread_mutex_destroy(&s.inference_mu); + return; + } + pthread_mutex_lock(&state.mu); + while (!state.entered) pthread_cond_wait(&state.cv, &state.mu); + pthread_mutex_unlock(&state.mu); + pthread_mutex_unlock(&s.model_mu); + + pthread_mutex_lock(&s.model_mu); + struct timespec deadline; + clock_gettime(CLOCK_REALTIME, &deadline); + timespec_add_us(&deadline, 1000000); + while (s.model_busy || !s.decode_handoff) { + int rc = pthread_cond_timedwait(&s.model_cv, &s.model_mu, &deadline); + if (rc == ETIMEDOUT) break; + } + TEST_ASSERT(!s.model_busy); + TEST_ASSERT(s.decode_handoff); + pthread_mutex_lock(&state.mu); + TEST_ASSERT(!state.done); + pthread_mutex_unlock(&state.mu); + + /* This is the state transition performed after a decode batch. */ + s.decode_handoff = false; + pthread_cond_broadcast(&s.model_cv); + pthread_mutex_unlock(&s.model_mu); + pthread_join(thread, NULL); + + pthread_mutex_lock(&state.mu); + TEST_ASSERT(state.done); + pthread_mutex_unlock(&state.mu); + server_generation_leave(&s); + TEST_ASSERT(s.active_generations == 0); + TEST_ASSERT(!s.decode_handoff); + pthread_cond_destroy(&state.cv); + pthread_mutex_destroy(&state.mu); + pthread_cond_destroy(&s.model_cv); + pthread_mutex_destroy(&s.model_mu); + pthread_mutex_destroy(&s.inference_mu); +} + static void test_batched_live_continuation_slot_binding(void) { server s = {0}; server_slot slots[3] = {0}; @@ -17413,6 +17549,7 @@ static void test_thinking_canonical_non_thinking_mode_noop(void) { static void ds4_server_unit_tests_run(void) { test_batched_prefill_round_robin(); + test_batched_prefill_hands_off_to_active_decode(); test_batched_live_continuation_slot_binding(); test_request_defaults_use_min_p_filtering(); test_reasoning_effort_mapping();