diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 73d11ac27..a137d032d 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -6829,9 +6829,11 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, #if defined(GGML_USE_HIP) // Drop last graph's mmb activation cache: keyed by raw data pointers the pool recycles. + // Empty containers are left alone: graphs without the qwen4exp profile never fill them, + // and another backend may be computing on another thread. ggml_cuda_mmb_begin_graph(); - g_hc_marked_xn.clear(); - g_hc_q8_xn.clear(); + if (!g_hc_marked_xn.empty()) g_hc_marked_xn.clear(); + if (!g_hc_q8_xn.empty()) g_hc_q8_xn.clear(); // xn is a view of HC_COMBINE_NORM, so reads() keys on (view_src, view_offs); a plain pointer test misses consumers. if (ggml_cuda_qwen4exp_enabled() && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmb.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmb.cu index 9110d4a41..0c40c2ed8 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmb.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmb.cu @@ -786,8 +786,23 @@ void ggml_cuda_mmb_mark_f16_only(const ggml_tensor * t) { if (mmb_arch()) g_mmb_ bool ggml_cuda_mmb_is_f16_only(const ggml_tensor * t) { return g_mmb_f16_only.count(t) > 0; } void ggml_cuda_mmb_mark_bf16_only(const ggml_tensor * t) { if (mmb_arch()) g_mmb_bf16_only.insert(t); } bool ggml_cuda_mmb_is_bf16_only(const ggml_tensor * t) { return g_mmb_bf16_only.count(t) > 0; } -void ggml_cuda_mmb_begin_graph() { for (auto & e : g_mmb_cache) delete e.buf; g_mmb_cache.clear(); for (auto & e : g_mmb_f16_cache) delete e.buf; g_mmb_f16_cache.clear(); for (auto & e : g_mmb_slots) { e.root = nullptr; e.data = nullptr; e.n = 0; } g_mmb_bf16_only.clear(); g_mmb_f16_only.clear(); g_mmb_q8_root = nullptr; g_mmb_q8_data = nullptr; } +// True when no graph left MMB state behind. Only graphs computed under the qwen4exp profile create any, so a graph +// without it (another model's backend, e.g. the PFlash drafter scoring on its own thread beside a serving target) +// then only reads these process-wide containers instead of clearing them under the other thread. +static bool mmb_graph_state_empty() { + if (!g_mmb_cache.empty() || !g_mmb_f16_cache.empty() || !g_mmb_bf16_only.empty() || !g_mmb_f16_only.empty()) return false; + if (g_mmb_q8_root || g_mmb_q8_data) return false; + for (const auto & e : g_mmb_slots) if (e.root || e.data || e.n) return false; + return true; +} +static bool mmb_buffers_empty() { + if (g_mmb_q8_buf || !g_mmb_shadow.empty()) return false; + for (const auto & e : g_mmb_slots) if (e.buf) return false; + return true; +} +void ggml_cuda_mmb_begin_graph() { if (!mmb_enabled() && mmb_graph_state_empty()) return; for (auto & e : g_mmb_cache) delete e.buf; g_mmb_cache.clear(); for (auto & e : g_mmb_f16_cache) delete e.buf; g_mmb_f16_cache.clear(); for (auto & e : g_mmb_slots) { e.root = nullptr; e.data = nullptr; e.n = 0; } g_mmb_bf16_only.clear(); g_mmb_f16_only.clear(); g_mmb_q8_root = nullptr; g_mmb_q8_data = nullptr; } void ggml_cuda_mmb_release_all() { + if (mmb_graph_state_empty() && mmb_buffers_empty()) return; ggml_cuda_mmb_begin_graph(); delete g_mmb_q8_buf; g_mmb_q8_buf = nullptr; g_mmb_q8_cap = 0; for (int i = 0; i < 4; ++i) { if (g_mmb_slots[i].buf) delete g_mmb_slots[i].buf; g_mmb_slots[i].buf = nullptr; g_mmb_slot_cap[i] = 0; } diff --git a/server/src/common/backend_args.h b/server/src/common/backend_args.h index 0334901f8..7b593eed9 100644 --- a/server/src/common/backend_args.h +++ b/server/src/common/backend_args.h @@ -105,6 +105,9 @@ struct BackendArgs { // Total paged K/V pool in tokens shared by all slots (--kv-pool-tokens; // model-page-rounded). 0 = derive capacity from available device memory. long long kv_pool_tokens = 0; + // Device memory an automatically sized pool leaves free for a model that + // runs beside concurrent decode (the PFlash drafter). Bytes; 0 = none. + long long kv_pool_reserve_bytes = 0; int kq_stride_pad = 32; int draft_block_size = 0; // 0 = drafter metadata int draft_swa_window = 0; diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 0f015b27c..98bc05a59 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -209,6 +209,7 @@ std::unique_ptr construct_backend( cfg.paged_attention = cache.paged_attention; cfg.max_concurrency = execution.max_concurrency; cfg.kv_pool_tokens = cache.kv_pool_tokens; + cfg.kv_pool_reserve_bytes = cache.kv_pool_reserve_bytes; cfg.kq_stride_pad = cache.kq_stride_pad; cfg.draft_block_size = speculation.draft_block_size; cfg.draft_swa_window = cache.draft_swa_window; diff --git a/server/src/common/backend_factory.h b/server/src/common/backend_factory.h index ff08ddb70..776ed682a 100644 --- a/server/src/common/backend_factory.h +++ b/server/src/common/backend_factory.h @@ -53,6 +53,7 @@ class BackendPlan final { int fa_window = 0; bool paged_attention = false; long long kv_pool_tokens = 0; + long long kv_pool_reserve_bytes = 0; ggml_type cache_type_k = GGML_TYPE_COUNT; ggml_type cache_type_v = GGML_TYPE_COUNT; int kq_stride_pad = 32; diff --git a/server/src/common/backend_plan.cpp b/server/src/common/backend_plan.cpp index d5b8bd85a..8bad935cd 100644 --- a/server/src/common/backend_plan.cpp +++ b/server/src/common/backend_plan.cpp @@ -144,6 +144,7 @@ BackendPreparation BackendPlanBuilder::resolve( plan.cache_.fa_window = args.fa_window; plan.cache_.paged_attention = args.paged_attention; plan.cache_.kv_pool_tokens = args.kv_pool_tokens; + plan.cache_.kv_pool_reserve_bytes = args.kv_pool_reserve_bytes; plan.cache_.cache_type_k = args.cache_type_k; plan.cache_.cache_type_v = args.cache_type_v; plan.cache_.kq_stride_pad = args.kq_stride_pad; diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index 901f17ce7..fda354dbe 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -254,9 +254,14 @@ std::string check_feature_compatibility( if (args.fa_window != 0) { return "--paged-attention requires full attention (--fa-window 0)"; } - if (admission.pflash_enabled) { - return "--paged-attention cannot be combined with PFlash prefill " - "compression"; + // Only concurrent dense Qwen3.5 serving compresses beside paged + // decode: its drafter scores on a thread of its own while the slots + // keep decoding, in memory the KV pool leaves free for it. + if (admission.pflash_enabled && + (args.max_concurrency <= 1 || arch != "qwen35")) { + return "--paged-attention combines with PFlash prefill compression " + "only for concurrent dense Qwen3.5 serving " + "(--max-concurrency > 1)"; } if (admission.fixed_kvflash_requested()) { return "--paged-attention cannot be combined with KVFlash"; diff --git a/server/src/errors.cpp b/server/src/errors.cpp index dca47dfc6..22cb1b80a 100644 --- a/server/src/errors.cpp +++ b/server/src/errors.cpp @@ -1,27 +1,27 @@ -// Thread-safe last-error string used by loaders and graph builders. -// Consumed by tests and the test_dflash driver via luce_last_error(). +// Last-error string used by loaders and graph builders, one per thread: a +// caller reads the error its own failing call set, and the pointer stays +// valid until that thread sets another. A concurrent server scores PFlash +// prompts on one thread while another serves, and neither may read, or +// free, the other's. Consumed by tests and the test_dflash driver via +// luce_last_error(). #include "luce.h" #include "internal.h" -#include #include namespace luce::common { namespace { -std::mutex g_err_mu; -std::string g_last_error; +thread_local std::string g_last_error; } void set_last_error(std::string msg) { - std::lock_guard lk(g_err_mu); g_last_error = std::move(msg); } } // namespace luce::common extern "C" const char * luce_last_error(void) { - std::lock_guard lk(luce::common::g_err_mu); return luce::common::g_last_error.c_str(); } diff --git a/server/src/kv_quant.cpp b/server/src/kv_quant.cpp index d16ae5efa..3342225fb 100644 --- a/server/src/kv_quant.cpp +++ b/server/src/kv_quant.cpp @@ -171,7 +171,8 @@ void validate_kv_pair_or_abort(ggml_type k, ggml_type v, const char * who) { } void resolve_kv_types(ggml_type & k_out, ggml_type & v_out, - ggml_type k_override, ggml_type v_override) { + ggml_type k_override, ggml_type v_override, + bool tq3_shorthand) { ggml_type k = GGML_TYPE_Q4_0; ggml_type v = GGML_TYPE_Q4_0; @@ -182,7 +183,7 @@ void resolve_kv_types(ggml_type & k_out, ggml_type & v_out, if (const char * s = std::getenv("LUCE_KV_Q4")) { if (std::atoi(s) != 0) { k = GGML_TYPE_Q4_0; v = GGML_TYPE_Q4_0; } } - if (const char * s = std::getenv("LUCE_KV_TQ3")) { + if (const char * s = tq3_shorthand ? std::getenv("LUCE_KV_TQ3") : nullptr) { if (std::atoi(s) != 0) { k = GGML_TYPE_TQ3_0; v = GGML_TYPE_TQ3_0; } } diff --git a/server/src/kv_quant.h b/server/src/kv_quant.h index 453d871c9..e20958c93 100644 --- a/server/src/kv_quant.h +++ b/server/src/kv_quant.h @@ -27,13 +27,15 @@ void validate_kv_pair_or_abort(ggml_type k, ggml_type v, const char * who); // Precedence (high -> low): // 0. Explicit k_override / v_override (COUNT means unspecified) // 1. LUCE_KV_K= / LUCE_KV_V= (independent override) -// 2. LUCE_KV_F16 / _KV_Q4 / _KV_TQ3 (legacy shorthand, K==V) +// 2. LUCE_KV_F16 / _KV_Q4 / _KV_TQ3 (legacy shorthand, K==V; +// tq3_shorthand=false ignores LUCE_KV_TQ3) // 3. Default: GGML_TYPE_Q4_0 for both (with FWHT K-rotation) // On invalid input or unsupported (K,V) pair, prints an explanatory message // and calls std::abort(). Returns the resolved pair via out params. void resolve_kv_types(ggml_type & k_out, ggml_type & v_out, ggml_type k_override = GGML_TYPE_COUNT, - ggml_type v_override = GGML_TYPE_COUNT); + ggml_type v_override = GGML_TYPE_COUNT, + bool tq3_shorthand = true); // KV reservation bytes per token for a hybrid attention/SSM model. Single source // of truth for both qwen35 (dense) and qwen35moe expert-placement budgeting. diff --git a/server/src/pflash/qwen35_drafter.cpp b/server/src/pflash/qwen35_drafter.cpp index 29cc60e31..5dd861edb 100644 --- a/server/src/pflash/qwen35_drafter.cpp +++ b/server/src/pflash/qwen35_drafter.cpp @@ -19,6 +19,7 @@ #include "common/gguf_inspect.h" #include "anchor_params.h" #include "internal.h" +#include "kv_quant.h" #include "ggml.h" #include "ggml-alloc.h" @@ -52,36 +53,17 @@ static size_t causal_mask_bytes(int kv_start, int n_tokens) { return (size_t)align_up_i(kv_start + n_tokens, 256) * align_up_i(n_tokens, 32) * sizeof(uint16_t); } -// create_target_cache honours LUCE_KV_TQ3; the drafter cache never wants -// the TurboQuant rotation, so force it off while the cache is created. -struct ScopedKvTq3Off { - ScopedKvTq3Off() { -#if defined(_WIN32) - char * raw = nullptr; - size_t len = 0; - _dupenv_s(&raw, &len, "LUCE_KV_TQ3"); - had_ = raw != nullptr; - old_ = had_ ? raw : ""; - free(raw); - _putenv_s("LUCE_KV_TQ3", "0"); -#else - const char * raw = std::getenv("LUCE_KV_TQ3"); - had_ = raw != nullptr; - old_ = had_ ? raw : ""; - setenv("LUCE_KV_TQ3", "0", 1); -#endif - } - ~ScopedKvTq3Off() { -#if defined(_WIN32) - // _putenv_s with empty value removes the variable on MSVCRT. - _putenv_s("LUCE_KV_TQ3", had_ ? old_.c_str() : ""); -#else - if (had_) setenv("LUCE_KV_TQ3", old_.c_str(), 1); - else unsetenv("LUCE_KV_TQ3"); -#endif - } - bool had_ = false; - std::string old_; +// The drafter cache's KV types: the target's resolution without the +// LUCE_KV_TQ3 shorthand, since the drafter never wants the TurboQuant +// rotation. Resolved without touching the environment, which other threads +// read while a concurrent server scores a prompt. +struct DrafterKvTypes { + ggml_type k = GGML_TYPE_COUNT; + ggml_type v = GGML_TYPE_COUNT; + DrafterKvTypes() { + luce::resolve_kv_types(k, v, GGML_TYPE_COUNT, GGML_TYPE_COUNT, + /*tq3_shorthand=*/false); + } }; } // namespace @@ -144,11 +126,11 @@ std::vector qwen35_score_and_compress( std::vector running_max((size_t)n_lookahead * S, -INFINITY); TargetCache cache; - { - ScopedKvTq3Off tq3_off; - if (!create_target_cache(w, S, 0, w.backend, cache, true)) { - return {}; - } + const DrafterKvTypes kv; + if (!create_target_cache(w, S, 0, w.backend, cache, true, 0, + /*paged_attention=*/false, 1, + /*concurrent_tree=*/false, kv.k, kv.v)) { + return {}; } ggml_init_params act_ip{}; @@ -621,13 +603,14 @@ int scoring_session_limit() { bool allocate_scoring_session(TargetWeights & w, int capacity, Qwen35ScoringSession & session) { - { - ScopedKvTq3Off tq3_off; - if (!create_target_cache_partial(w, capacity, 0, w.backend, session.cache, - /*prefill_only=*/true, 0, kQwen35HeadBlock, - /*allocate_target_feat=*/false)) { - return false; - } + const DrafterKvTypes kv; + if (!create_target_cache_partial(w, capacity, 0, w.backend, session.cache, + /*prefill_only=*/true, 0, kQwen35HeadBlock, + /*allocate_target_feat=*/false, 0, + /*f32_ssm_intermediates=*/false, + /*paged_attention=*/false, 1, + /*concurrent_tree=*/false, kv.k, kv.v)) { + return false; } if (!ensure_ssm_snapshot(session.cache, w.backend)) { free_qwen35_scoring_session(session); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index d303cb14f..5fdea8a0c 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -509,6 +509,16 @@ bool Qwen35Backend::init() { // pool's index space) as the write target of dead decode-batch rows. int64_t pool_tokens = 0; if (n_slots > 1) { + PagedKvAutoBudget budget; + // TODO: Size tensor-parallel pools from each device's free memory + // and account for fixed cache tensors mirrored by the meta backend. + budget.free_bytes = (int64_t)gpu_free; + budget.bytes_per_token = kvf_budget.bytes_per_token; + budget.reserve_bytes = + kvf_budget.reserve_bytes + cfg_.kv_pool_reserve_bytes; + budget.fixed_cache_bytes = concurrent_fixed_cache_bytes( + w_, cfg_.device.max_ctx, n_slots, budget.bytes_per_token, + scratch_tokens, fixed_chain.enabled); if (cfg_.kv_pool_tokens > 0) { const int64_t max_pool_tokens = ((int64_t)INT32_MAX - scratch_tokens) / @@ -516,16 +526,25 @@ bool Qwen35Backend::init() { pool_tokens = (int64_t)paged_token_capacity( (int)std::min( cfg_.kv_pool_tokens, max_pool_tokens)); + // Memory set aside for a model beside the slots (the PFlash + // drafter) bounds an explicit pool too. + if (cfg_.kv_pool_reserve_bytes > 0 && budget.bytes_per_token > 0) { + const int64_t fit = (std::max)(0, + budget.free_bytes - budget.fixed_cache_bytes - + budget.reserve_bytes) / budget.bytes_per_token; + if (pool_tokens > fit) { + set_last_error( + "--kv-pool-tokens " + std::to_string(pool_tokens) + + " leaves less than the " + + std::to_string(cfg_.kv_pool_reserve_bytes >> 20) + + " MiB reserved for the PFlash drafter; at most " + + std::to_string(fit / PAGED_BLOCK_SIZE * PAGED_BLOCK_SIZE) + + " tokens fit"); + std::fprintf(stderr, "[parallel] %s\n", luce_last_error()); + return false; + } + } } else { - PagedKvAutoBudget budget; - // TODO: Size tensor-parallel pools from each device's free memory - // and account for fixed cache tensors mirrored by the meta backend. - budget.free_bytes = (int64_t)gpu_free; - budget.bytes_per_token = kvf_budget.bytes_per_token; - budget.reserve_bytes = kvf_budget.reserve_bytes; - budget.fixed_cache_bytes = concurrent_fixed_cache_bytes( - w_, cfg_.device.max_ctx, n_slots, budget.bytes_per_token, - scratch_tokens, fixed_chain.enabled); pool_tokens = paged_kv_auto_pool_tokens( cfg_.device.max_ctx, n_slots, budget); const int64_t one_context = @@ -1439,9 +1458,14 @@ std::vector Qwen35Backend::compress_batch( // Synchronize all backends to flush any outstanding async CUDA work // before loading the drafter. Without this, pending operations on - // target/draft streams can corrupt the drafter's allocations. - ggml_backend_synchronize(target_backend_); - if (draft_backend_) ggml_backend_synchronize(draft_backend_); + // target/draft streams can corrupt the drafter's allocations. Only + // parking frees their memory: a drafter beside resident models + // allocates its own, and concurrent serving scores on a thread of its + // own while the scheduler keeps the target busy. + if (should_park) { + ggml_backend_synchronize(target_backend_); + if (draft_backend_) ggml_backend_synchronize(draft_backend_); + } // Load drafter with its OWN backend (not target_backend_). // Matches test_dflash.cpp: separate backend supports multi-GPU diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index 702e6448d..f364b4f60 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -75,6 +75,9 @@ struct Qwen35Config { // (--kv-pool-tokens; rounded up to whole 16-token blocks). // 0 derives the shared physical capacity from available device memory. int64_t kv_pool_tokens = 0; + // Device memory the automatically sized pool leaves free for the PFlash + // drafter, which scores beside concurrent decode (bytes). + int64_t kv_pool_reserve_bytes = 0; // Draft int draft_block_size = 0; // 0 = use drafter metadata diff --git a/server/src/server/admission.h b/server/src/server/admission.h index 796a8edf0..f4749ecc3 100644 --- a/server/src/server/admission.h +++ b/server/src/server/admission.h @@ -6,6 +6,21 @@ #include +// Device memory concurrent serving leaves free beside the paged KV pool for the +// PFlash drafter, which scores on its own thread while the slots decode: its +// weights plus a workspace that grows linearly with the prompt it scores +// (activations, its own cache, the largest ubatch's attention). The +// per-token figure is the Qwen3.5-0.8B drafter's measured peak, 53 KiB per +// token from 33K to 206K tokens, rounded up; its fixed part (0.6 GiB) is +// under the weights. +inline constexpr long long kPFlashScoringBytesPerToken = 56LL * 1024; + +inline long long pflash_concurrent_reserve_bytes(long long drafter_weight_bytes, int max_prompt_tokens) +{ + return (std::max)(0LL, drafter_weight_bytes) + + (long long)(std::max)(0, max_prompt_tokens) * kPFlashScoringBytesPerToken; +} + // The reply a request can get: what it asked for, trimmed to the context its // prompt leaves (0 when the prompt fills it). Clients ask for their usual cap // whatever the prompt, Claude Code 32000 tokens, so a long session would diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index c181b2fd2..3ff542b0b 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -4524,7 +4524,7 @@ json build_openai_completion_response( // the OpenAI shape (ignored by clients that don't recognize it). // spec_decode_ran disambiguates a real zero-acceptance speculative run // from an autoregressive fallback. See docs/specs/thinking-budget.md §6.3. - const int prompt_tokens = (int) req.prompt_tokens.size(); + const int prompt_tokens = usage_prompt_tokens(req); const json usage = { {"prompt_tokens", prompt_tokens}, {"completion_tokens", counts.total}, @@ -4604,7 +4604,7 @@ json build_anthropic_response( } const json usage = { - {"input_tokens", (int) req.prompt_tokens.size()}, + {"input_tokens", usage_prompt_tokens(req)}, {"output_tokens", counts.total}, {"timings", build_timings_json(timings, counts.total)}, {"accept_rate", result.accept_rate}, @@ -4630,7 +4630,7 @@ json build_responses_api_response( const json output = responses_output_items( req.response_id + "_msg", emitter.accumulated_text(), emitter.tool_calls()); - const int prompt_tokens = (int) req.prompt_tokens.size(); + const int prompt_tokens = usage_prompt_tokens(req); const json usage = { {"input_tokens", prompt_tokens}, {"output_tokens", counts.total}, @@ -4943,7 +4943,12 @@ std::string HttpServer::apply_pflash_compression( const ParsedRequest & req, PreparedPrompt & prepared) { const bool selection_environment = luce::pflash::has_pflash_selection_environment(); - auto [full_slot, full_len] = prefix_cache_.lookup_full(req.prompt_tokens); + // Concurrent serving measured the cache on the scheduler thread, which + // owns the prefix cache, before handing the request over. + const bool cache_measured = req.pflash_cached_prefix.has_value(); + auto [full_slot, full_len] = cache_measured + ? std::pair{-1, 0} + : prefix_cache_.lookup_full(req.prompt_tokens); if (http_detail::pflash_full_cache_restore_allowed( selection_environment) && full_slot >= 0) { std::fprintf(stderr, @@ -6117,8 +6122,10 @@ std::string HttpServer::apply_pflash_compression( existing.view_gen_begin, input_tokens, true, 0}; } } else { - const auto restorable = restorable_prefix_tokens( - req.prompt_tokens, req.disk_cache_policy); + const auto restorable = cache_measured + ? *req.pflash_cached_prefix + : restorable_prefix_tokens( + req.prompt_tokens, req.disk_cache_policy); const int cached_target = restorable.tokens; const int cached_drafter = prompt_tokens > 0 ? (std::min)(input_tokens, (int) std::floor( diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 71c695663..d12fa1fa7 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -46,6 +46,7 @@ #include #include #include +#include #include #include #if !defined(_WIN32) @@ -846,8 +847,21 @@ struct ParsedRequest { // The prompt ends with tool results: an agent's step, which clients only // append to. bool ends_with_tool_result = false; + // Concurrent serving: the restorable prefix of prompt_tokens (target + // tokens) the scheduler measured before it handed the request to the + // PFlash thread, which never reads the prefix cache. Empty = not measured. + std::optional pflash_cached_prefix; + // The client prompt's size when prompt_tokens holds the compressed + // prompt the concurrent scheduler serves; -1: prompt_tokens is the + // client's. usage.prompt_tokens reports it (usage_prompt_tokens). + int client_prompt_tokens = -1; }; +inline int usage_prompt_tokens(const ParsedRequest & req) { + return req.client_prompt_tokens >= 0 ? req.client_prompt_tokens + : (int) req.prompt_tokens.size(); +} + // Resident budget the PrefixCache enforces: the concurrent limit in paged // mode, else --prefix-cache-max-mib, with auto sized from the backend's // snapshot at full context. @@ -1010,6 +1024,14 @@ class HttpServer { // Prompt preparation keeps the FlowKV and whole-prompt PFlash policies // out of the decode path while preserving their shared precedence rules. PreparedPrompt prepare_prompt(const ParsedRequest & req); + // Concurrent dense Qwen3.5 serving compresses on the scheduler's PFlash + // thread (scheduler.cpp), for prompts up to --max-ctx; other engines keep + // their paths. + bool pflash_concurrent_serving() { + return backend_.seq_engine() != nullptr && config_.arch == "qwen35" && + config_.pflash_mode != ServerConfig::PflashMode::OFF && + drafter_tokenizer_ != nullptr; + } void apply_flowkv_compression(const ParsedRequest & req, PreparedPrompt & prepared); std::string apply_pflash_compression(const ParsedRequest & req, @@ -1428,6 +1450,13 @@ struct ServerJob { // First concurrent-scheduler attempt; retained across busy deferrals so // server-side prefill/elapsed telemetry does not erase queueing delay. std::chrono::steady_clock::time_point parallel_started_at{}; + // Concurrent PFlash (scheduler.cpp): the job came back from the + // compression thread. When it was compressed, req.prompt_tokens holds the + // served prompt and req.client_prompt_tokens the client prompt's size. + bool pflash_prepared = false; + nlohmann::json pflash_stats; // usage.timings.pflash + int pflash_error_status = 0; // prompt preparation failed + std::string pflash_error; }; // ─── Parse session_id from a chat-completion JSON body ────────────────── diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 96d8da137..8bb1e19fb 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -16,11 +16,14 @@ #include #include +#include #include #include +#include #include #include #include +#include #include #include @@ -81,6 +84,15 @@ enum class AdmissionDisposition { Retired, // job finished or failed during admission; no slot was taken }; +// Long prompts handed to the PFlash compression thread: the scheduler +// pushes, the compression thread pops. +struct PflashStageQueue { + std::mutex mu; + std::condition_variable cv; + std::deque jobs; + bool stop = false; +}; + } // namespace void HttpServer::scheduler_loop(SeqEngine & engine) { @@ -336,10 +348,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { /*prefilled_tokens=*/prompt_tokens - s.cached_prefix_tokens, /*effective_prompt_tokens=*/prompt_tokens, }; + gen_timings.pflash = s.job->pflash_stats; if (!s.error) { PerfRecord perf; perf.prompt_tokens = (int)req.prompt_tokens.size(); + perf.pflash = req.client_prompt_tokens >= 0; perf.completion_tokens = s.completion_tokens; const int computed_prompt_tokens = prompt_tokens - s.cached_prefix_tokens; @@ -460,7 +474,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { if (n_gen_cap < 1) { // Degenerate ask: reply with an empty completion, no slot needed. SseEmitter emitter(req.format, req.response_id, req.model, - (int)req.prompt_tokens.size(), req.tools, + usage_prompt_tokens(req), req.tools, &tool_memory_, req.stop_sequences, req.started_in_thinking); GenTimings t{ @@ -756,7 +770,9 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } if (ar.status == SeqEngine::AdmitResult::Status::busy) { prepared_capture.cancel(); - if (job->report_admission) { + // A routed prompt the drafter already scored waits for this + // model rather than being scored again for another. + if (job->report_admission && !job->pflash_prepared) { job->admission = RoutingAdmission::busy; finish_job(job); return AdmissionDisposition::Retired; @@ -787,7 +803,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // can still try another model; no tokenizer/model identity is on wire. auto emitter = std::make_unique( req.format, req.response_id, req.model, - (int)req.prompt_tokens.size(), req.tools, &tool_memory_, + usage_prompt_tokens(req), req.tools, &tool_memory_, req.stop_sequences, req.started_in_thinking); if (req.stream) { bool ok = send_sse_headers(job); @@ -865,6 +881,113 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // has them until it ends. std::chrono::steady_clock::time_point decode_only_until{}; + // PFlash. A prompt over the compression threshold leaves the line for + // the compression thread: live slots keep decoding while the drafter + // scores it, and it comes back to the head of the queue with the prompt + // to serve. That thread owns what prompt preparation touches (the + // drafter, the chat-view store); the prefix cache stays with this + // thread, which measures the prompt's restorable prefix before the + // hand-off. + const bool pflash_concurrent = pflash_concurrent_serving(); + PflashStageQueue pflash_stage; + auto pflash_wants = [&](const ServerJob * job) { + if (!pflash_concurrent || job->pflash_prepared || job->req.images) { + return false; + } + return config_.pflash_mode == ServerConfig::PflashMode::ALWAYS || + (int)job->req.prompt_tokens.size() >= config_.pflash_threshold; + }; + auto pflash_prepare = [this](ServerJob * job) { + if (job->client_disconnected.load(std::memory_order_acquire)) return; + const auto started = std::chrono::steady_clock::now(); + const int raw_tokens = (int)job->req.prompt_tokens.size(); + PreparedPrompt prepared = prepare_prompt(job->req); + if (prepared.error_status != 0) { + job->pflash_error_status = prepared.error_status; + job->pflash_error = std::move(prepared.error); + return; + } + job->pflash_stats = std::move(prepared.pflash_stats); + if (prepared.compressed) { + job->req.client_prompt_tokens = raw_tokens; + job->req.prompt_tokens = std::move(prepared.tokens); + // A view's generation prompt is where its next turn branches off. + if (job->req.pin_end_token <= 0 && prepared.snapshot_cut > 0) { + job->req.pin_end_token = prepared.snapshot_cut; + } + // A compressed prompt is not raw history: never persist it. + job->req.disk_cache_policy.mode = DiskPrefixCacheMode::Off; + } + std::fprintf(stderr, + "[pflash-concurrent] %s: %d -> %zu tokens in %.2fs\n", + job->req.response_id.c_str(), raw_tokens, + job->req.prompt_tokens.size(), + std::chrono::duration( + std::chrono::steady_clock::now() - started).count()); + }; + // Prepared jobs return to the head of the queue, behind the prepared + // ones already waiting there, so they keep their order. + auto requeue_prepared = [this](ServerJob * job) { + std::lock_guard lk(queue_mu_); + ServerJob ** link = &queue_head_; + while (*link && (*link)->pflash_prepared) link = &(*link)->next; + job->next = *link; + *link = job; + if (!job->next) queue_tail_ = job; + queue_cv_.notify_one(); + }; + std::thread pflash_thread; + if (pflash_concurrent) { + pflash_thread = std::thread([&pflash_stage, pflash_prepare, + requeue_prepared] { + for (;;) { + ServerJob * job = nullptr; + { + std::unique_lock lk(pflash_stage.mu); + pflash_stage.cv.wait(lk, [&] { + return pflash_stage.stop || !pflash_stage.jobs.empty(); + }); + if (pflash_stage.stop) return; + job = pflash_stage.jobs.front(); + pflash_stage.jobs.pop_front(); + } + pflash_prepare(job); + job->pflash_prepared = true; + requeue_prepared(job); + } + }); + } + auto stop_pflash_thread = [&]() { + if (!pflash_thread.joinable()) return; + { + std::lock_guard lk(pflash_stage.mu); + pflash_stage.stop = true; + } + pflash_stage.cv.notify_all(); + pflash_thread.join(); + }; + struct StopOnExit { + decltype(stop_pflash_thread) & stop; + ~StopOnExit() { stop(); } + } pflash_thread_guard{stop_pflash_thread}; + // No response bytes are committed before admission, so a failed + // preparation answers like a failed admission. + auto fail_preparation = [&](ServerJob * job) { + const ParsedRequest & req = job->req; + std::fprintf(stderr, "[server] request failed: %s\n", + job->pflash_error.c_str()); + const ResponseError error = job->pflash_error_status == 400 + ? ResponseError::invalid_request("invalid_request", job->pflash_error) + : job->pflash_error_status == 503 + ? ResponseError::unavailable("unavailable", job->pflash_error) + : ResponseError::internal("server_error", job->pflash_error); + const json body = build_error_response( + req.format, error, req.response_id); + send_response(job->fd, response_error_http_status(error), + "application/json", body.dump() + "\n"); + finish_job(job); + }; + while (true) { reap_readers(); // Phase 1 — Admission: deferred job first (FIFO), then the queue. @@ -912,6 +1035,24 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { finish_job(job); continue; } + if (pflash_wants(job)) { + if (job->parallel_started_at == + std::chrono::steady_clock::time_point{}) { + job->parallel_started_at = std::chrono::steady_clock::now(); + } + job->req.pflash_cached_prefix = restorable_prefix_tokens( + job->req.prompt_tokens, job->req.disk_cache_policy); + { + std::lock_guard lk(pflash_stage.mu); + pflash_stage.jobs.push_back(job); + } + pflash_stage.cv.notify_one(); + continue; + } + if (job->pflash_error_status != 0) { + fail_preparation(job); + continue; + } const AdmissionDisposition outcome = admit_job(job); if (outcome == AdmissionDisposition::Loading) { deferred = job; @@ -1364,6 +1505,21 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } } + // A prompt still being compressed returns to the queue below; the ones + // not started yet are answered here. + stop_pflash_thread(); + for (ServerJob * waiting : pflash_stage.jobs) { + const ParsedRequest & req = waiting->req; + const ResponseError error = ResponseError::unavailable( + "server_shutting_down", "server shutting down"); + const json body = build_error_response( + req.format, error, req.response_id); + send_response(waiting->fd, response_error_http_status(error), + "application/json", body.dump() + "\n"); + finish_job(waiting); + } + pflash_stage.jobs.clear(); + drop_disk_load(); for (auto & reader : finished_readers) reader.reader.join(); finished_readers.clear(); diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 239ec2594..1366b4351 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -12,6 +12,7 @@ // [--max-tokens 4096] [--target-device auto:0] #include "http_server.h" +#include "admission.h" #include "chat_template.h" #include "launch_profiles.h" #include "model_card.h" @@ -1389,6 +1390,52 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod std::fprintf(stderr, "[server] invalid PFlash configuration: %s\n", error.c_str()); return 1; } + // Concurrent dense Qwen3.5 serving compresses on a thread of its own + // while the slots decode: the target stays resident, and the paged + // pool leaves the drafter the memory to score a prompt that fills the + // context (longer ones are refused, see validate_request_context). + // That memory was measured for the trained head's scorer with the + // drafter released after each prompt. + if (bargs.max_concurrency > 1 && + inspect_gguf_model_info(bargs.model_path.c_str()).arch == "qwen35") { + const char * head = std::getenv("PFLASH_SCORING_HEAD_GGUF"); + const char * legacy = std::getenv("PFLASH_QWEN35_LEGACY_SCORER"); + luce::pflash::PFlashSelectionConfig selection; + std::string selection_error; + if (!head || !*head || (legacy && std::strcmp(legacy, "1") == 0) || + !luce::pflash::resolve_pflash_selection(1 << 16, 32, selection, selection_error) || + selection.mode == luce::pflash::PFlashSelectionMode::Legacy || + selection.scorer != luce::pflash::PFlashScorer::Head) { + std::fprintf(stderr, + "[server] PFlash with --max-concurrency > 1 needs the release profile's " + "strict selection with the trained head's scorer (set " + "PFLASH_SCORING_HEAD_GGUF; no legacy or split scorer)\n"); + return 1; + } + if (sconfig.draft_residency == DraftResidencyPolicy::Persistent) { + std::fprintf(stderr, + "[server] PFlash with --max-concurrency > 1 needs the drafter released after " + "each prompt: --draft-residency persistent keeps its scoring sessions on the " + "device, outside the memory reserved for it\n"); + return 1; + } + if (!sconfig.pflash_skip_park) { + sconfig.pflash_skip_park = true; + std::fprintf(stderr, + "[server] PFlash with --max-concurrency %d: the drafter scores beside the " + "resident target (--prefill-skip-park)\n", bargs.max_concurrency); + } + std::error_code size_error; + const auto drafter_bytes = + std::filesystem::file_size(sconfig.pflash_drafter_path, size_error); + bargs.kv_pool_reserve_bytes = pflash_concurrent_reserve_bytes( + size_error ? 0 : (long long)drafter_bytes, bargs.device.max_ctx); + std::fprintf(stderr, + "[server] PFlash with --max-concurrency %d: %.2f GiB of device memory " + "reserved for the drafter, for prompts up to %d tokens\n", + bargs.max_concurrency, bargs.kv_pool_reserve_bytes / 1073741824.0, + bargs.device.max_ctx); + } } auto & spark_autotune = model.spark_autotune; auto & spark_slots = model.spark_slots; @@ -1611,6 +1658,13 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod sconfig.pflash_drafter_gpu = pflash_placement.drafter_gpu; sconfig.pflash_remote_drafter = pflash_placement.remote_drafter; sconfig.pflash_remote = pflash_placement.remote; + if (sconfig.pflash_remote_drafter && backend_execution.max_concurrency > 1 && + arch == "qwen35") { + std::fprintf(stderr, + "[server] PFlash with --max-concurrency > 1 needs an in-process drafter " + "(the IPC drafter is single-request only)\n"); + return 2; + } // ── Apply environment defaults ───────────────────────────────────── // --freq / --collect-routing: enable routing stats via env vars that diff --git a/server/test/unit/test_admission.cpp b/server/test/unit/test_admission.cpp index a73a19f6e..0bad6f892 100644 --- a/server/test/unit/test_admission.cpp +++ b/server/test/unit/test_admission.cpp @@ -1,4 +1,4 @@ -// Unit tests for should_reject_oversized and fit_reply_to_context — pure, GPU-free. +// Unit tests for the admission gate, reply fitting and the concurrent PFlash reserve — pure, GPU-free. #include "CppUnitTestFramework.hpp" #include "server/admission.h" @@ -81,3 +81,18 @@ TEST_CASE(AdmissionFixture, test_oversized_with_compression_respects_max_ctx) { CHECK(should_reject_oversized(204800, 1, 204800, false)); CHECK(!should_reject_oversized(1000, 1, 139264, false)); } + +// The drafter's reserve covers its measured peak (Qwen3.5-0.8B BF16 drafter, R9700, peak minus the floor before the +// request) at every prompt length that was measured. +TEST_CASE(AdmissionFixture, test_pflash_concurrent_reserve_covers_measured_drafter_peaks) { + const long long drafter_bytes = 1516744736; // Qwen3.5-0.8B-BF16.gguf + const long long gib = 1LL << 30; + const struct { int tokens; double peak_gib; } measured[] = { + {33513, 2.26}, {57775, 3.43}, {83928, 4.83}, {136684, 7.38}, {206205, 11.04}}; + for (const auto & m : measured) { + CHECK(pflash_concurrent_reserve_bytes(drafter_bytes, m.tokens) >= (long long)(m.peak_gib * gib)); + } + CHECK(pflash_concurrent_reserve_bytes(drafter_bytes, 139264) == + drafter_bytes + 139264LL * kPFlashScoringBytesPerToken); + CHECK(pflash_concurrent_reserve_bytes(-1, -1) == 0); +} diff --git a/server/test/unit/test_feature_gate.cpp b/server/test/unit/test_feature_gate.cpp index 41a7f576e..0d820c945 100644 --- a/server/test/unit/test_feature_gate.cpp +++ b/server/test/unit/test_feature_gate.cpp @@ -603,6 +603,24 @@ void test_feature_gate_parallel_and_kv_pool_rules() { parallel.max_concurrency = 65; CHECK(!gate_result(parallel, "qwen35", PlacementBackend::Cuda).empty()); + // PFlash runs beside paged decode only in concurrent dense Qwen3.5 + // serving; single-request paged serving still rejects it. + BackendAdmissionContext pflash; + pflash.pflash_enabled = true; + pflash.pflash_drafter_configured = true; + for (const auto backend : {PlacementBackend::Cuda, PlacementBackend::Hip}) { + BackendArgs concurrent = paged; + concurrent.max_concurrency = 4; + CHECK(gate_result(concurrent, "qwen35", backend, pflash).empty()); + concurrent.max_concurrency = 1; + CHECK(gate_result(concurrent, "qwen35", backend, pflash) + .find("--max-concurrency > 1") != std::string::npos); + } + BackendArgs ds4_pflash = gate_args_hip_deepseek4(); + ds4_pflash.paged_attention = true; + ds4_pflash.max_concurrency = 2; + CHECK(!gate_result(ds4_pflash, "deepseek4", PlacementBackend::Hip, pflash).empty()); + BackendArgs ds4 = gate_args_hip_deepseek4(); ds4.paged_attention = true; ds4.max_concurrency = DEEPSEEK4_MAX_PAGED_SEQUENCES; diff --git a/server/test/unit/test_kv_quant.cpp b/server/test/unit/test_kv_quant.cpp index ca2d01f72..7e0a40d42 100644 --- a/server/test/unit/test_kv_quant.cpp +++ b/server/test/unit/test_kv_quant.cpp @@ -191,6 +191,15 @@ static void t4_tq3_shorthand() { ggml_type k, v; luce::resolve_kv_types(k, v); assert(k == GGML_TYPE_TQ3_0 && v == GGML_TYPE_TQ3_0); + // The PFlash drafter's cache resolves without the shorthand; the other + // settings still apply. + luce::resolve_kv_types(k, v, GGML_TYPE_COUNT, GGML_TYPE_COUNT, + /*tq3_shorthand=*/false); + assert(k == GGML_TYPE_Q4_0 && v == GGML_TYPE_Q4_0); + setenv("LUCE_KV_F16", "1", 1); + luce::resolve_kv_types(k, v, GGML_TYPE_COUNT, GGML_TYPE_COUNT, + /*tq3_shorthand=*/false); + assert(k == GGML_TYPE_F16 && v == GGML_TYPE_F16); clear_kv_env(); std::puts("T4 PASS"); diff --git a/server/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index ead7521d9..8a5d1e6a1 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -19485,3 +19485,234 @@ TEST_CASE(ServerUnitFixture, test_unified_distinct_tokenizers_keep_lines_whole) TEST_ASSERT_MSG(content_broke.empty(), content_broke + "\n" + content); TEST_ASSERT(content.find("item 33 value kiwikiwikiwikiwi\n") != std::string::npos); } + +// PFlash on the concurrent scheduler. The engine counts the steps that ran +// while the drafter scored. +class PflashStageEngine final : public SeqEngine { +public: + explicit PflashStageEngine(const std::atomic & scoring) : scoring_(scoring) {} + int slot_count() const override { return 2; } + int max_context() const override { return 4096; } + bool token_is_eos(int32_t) const override { return false; } + StepPlanLimits step_plan_limits(int) const override { + StepPlanLimits limits; + limits.max_prefill_sequences = 2; + limits.max_prefill_tokens_per_sequence = 16; + limits.max_prefill_tokens_total = 32; + limits.prefill_allocation_quantum = 16; + return limits; + } + AdmitResult admit(uint64_t, const std::vector & prompt, const SamplerCfg &) override { + AdmitResult result; + const int slot = !slots_[0].active ? 0 : !slots_[1].active ? 1 : -1; + if (slot < 0) { + result.status = AdmitResult::Status::busy; + return result; + } + slots_[(size_t)slot] = {true, (int)prompt.size()}; + { + std::lock_guard lock(mu); + admitted_sizes.push_back((int)prompt.size()); + } + result.status = AdmitResult::Status::admitted; + result.slot = slot; + return result; + } + StepResult step(const StepPlan & plan) override { + if (scoring_.load()) steps_beside_drafter.fetch_add(1); + StepResult result; + for (const StepInput & input : plan.decode) result.decode.push_back({input.slot, 7, false, {}}); + for (const PrefillSlice & slice : plan.prefills) { + int & left = slots_[(size_t)slice.slot].remaining; + left -= std::min(left, slice.max_tokens); + PrefillOutput out; + out.slot = slice.slot; + out.status = left == 0 ? PrefillOutput::Status::completed + : PrefillOutput::Status::advanced; + if (left == 0) out.token = 7; + result.prefills.push_back(out); + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + return result; + } + void retire(int slot) override { slots_[(size_t)slot] = {}; } + + std::mutex mu; + std::vector admitted_sizes; + std::atomic steps_beside_drafter{0}; + +private: + struct Slot { bool active = false; int remaining = 0; }; + const std::atomic & scoring_; + std::array slots_{}; +}; + +struct PflashStageBackend : MockPflashBudgetBackend { + std::atomic scoring{false}; + PflashStageEngine engine{scoring}; + SeqEngine * seq_engine() override { return &engine; } + CompressResult compress(const CompressRequest & request) override { + scoring = true; + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + CompressResult result = MockPflashBudgetBackend::compress(request); + scoring = false; + return result; + } +}; + +struct PflashStageRun { + bool done = false; + int raw_tokens = 0; + int compress_calls = 0; + int steps_beside_drafter = 0; + std::vector admitted_sizes; + std::vector usage_prompt_tokens; // per job, from its JSON reply + std::vector admissions; +}; + +// Long prompts (PFlash) and short chats (served as they are) on a 2-slot +// engine; a chat starts first and the long prompts arrive while it decodes. +// Routed long prompts report their admission to a model router. +static PflashStageRun run_pflash_stage_case(int chats, int long_prompts, + bool routed = false) { + luce_test::ScopedEnvVar mode{"PFLASH_SELECT_MODE", "budget_only"}; + luce_test::ScopedEnvVar chunk{"PFLASH_SELECT_CHUNK_SIZE", "4"}; + const auto prompt = pflash_long_system_prompt_case(); + const std::string path = + write_pflash_bpe_tokenizer_fixture(prompt.vocab, prompt.rendered); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::vector ids = tokenizer.encode(prompt.rendered); + + auto backend_owner = std::make_unique(); + PflashStageBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.arch = "qwen35"; + config.pflash_mode = ServerConfig::PflashMode::AUTO; + config.pflash_threshold = 64; + config.pflash_keep_ratio = 0.3f; + config.max_ctx = 4096; + config.prefix_cache_cap = 0; + config.prefill_cache_cap = 0; + config.admission_coalesce_ms = 0; + PflashStageRun run; + run.raw_tokens = (int)ids.size(); + { + HttpServer server(engine, tokenizer, config); + server.set_drafter_tokenizer(&tokenizer); + const int n = chats + long_prompts; + std::vector fds((size_t)n * 2, -1); + std::vector> jobs; + for (int i = 0; i < n; ++i) { + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, &fds[(size_t)i * 2]) == 0); + auto job = std::make_unique(); + job->fd = fds[(size_t)i * 2]; + job->req.format = ApiFormat::OPENAI_CHAT; + job->req.stream = false; + job->req.model = "scheduler-test"; + job->req.response_id = "job-" + std::to_string(i); + if (i < chats) { + job->req.prompt_tokens.assign(ids.begin(), ids.begin() + 4); + job->req.max_output = 60; + } else { + job->req.messages = prompt.messages; + job->req.prompt_tokens = ids; + job->req.max_output = 1; + job->report_admission = routed; + } + jobs.push_back(std::move(job)); + } + for (int i = 0; i < chats; ++i) SchedulerTestHarness::enqueue(server, jobs[(size_t)i].get()); + std::thread scheduler([&] { SchedulerTestHarness::run(server, backend.engine); }); + if (chats > 0) std::this_thread::sleep_for(std::chrono::milliseconds(20)); + for (int i = chats; i < n; ++i) SchedulerTestHarness::enqueue(server, jobs[(size_t)i].get()); + run.done = true; + for (auto & job : jobs) { + std::unique_lock lock(job->mu); + run.done &= job->cv.wait_for(lock, std::chrono::seconds(10), [&] { return job->done; }); + run.admissions.push_back(job->admission); + } + SchedulerTestHarness::stop(server); + scheduler.join(); + for (int i = 0; i < n; ++i) { + std::string reply; + char buf[4096]; + ssize_t got; + while ((got = recv(fds[(size_t)i * 2 + 1], buf, sizeof(buf), MSG_DONTWAIT)) > 0) { + reply.append(buf, (size_t)got); + } + const size_t body = reply.find("\r\n\r\n"); + TEST_ASSERT(body != std::string::npos); + const auto j = nlohmann::json::parse(reply.substr(body + 4)); + run.usage_prompt_tokens.push_back(j["usage"]["prompt_tokens"].get()); + } + for (const int fd : fds) close(fd); + } + run.compress_calls = backend.compress_calls; + run.steps_beside_drafter = backend.engine.steps_beside_drafter.load(); + run.admitted_sizes = backend.engine.admitted_sizes; + unlink(path.c_str()); + return run; +} + +// A long prompt on an idle server is scored, then served compressed; its +// usage counts the prompt the client sent. +TEST_CASE(ServerUnitFixture, test_scheduler_pflash_serves_a_long_prompt_on_an_idle_server) { + const auto run = run_pflash_stage_case(/*chats=*/0, /*long_prompts=*/1); + TEST_ASSERT(run.done); + TEST_ASSERT(run.compress_calls == 1); + TEST_ASSERT(run.admitted_sizes.size() == 1); + TEST_ASSERT(run.admitted_sizes[0] < run.raw_tokens); + TEST_ASSERT(run.usage_prompt_tokens.size() == 1); + TEST_ASSERT(run.usage_prompt_tokens[0] == run.raw_tokens); +} + +// The chat already decoding keeps its slot stepping while the drafter +// scores the long prompt beside it. +TEST_CASE(ServerUnitFixture, test_scheduler_pflash_decodes_while_the_drafter_scores) { + const auto run = run_pflash_stage_case(/*chats=*/1, /*long_prompts=*/1); + TEST_ASSERT(run.done); + TEST_ASSERT(run.compress_calls == 1); + TEST_ASSERT(run.steps_beside_drafter > 0); + TEST_ASSERT(run.admitted_sizes.size() == 2); + TEST_ASSERT(run.admitted_sizes[1] < run.raw_tokens); +} + +// Two long prompts that arrive together are both scored and served +// compressed. +TEST_CASE(ServerUnitFixture, test_scheduler_pflash_serves_two_long_prompts) { + const auto run = run_pflash_stage_case(/*chats=*/0, /*long_prompts=*/2); + TEST_ASSERT(run.done); + TEST_ASSERT(run.compress_calls == 2); + TEST_ASSERT(run.admitted_sizes.size() == 2); + TEST_ASSERT(run.admitted_sizes[0] < run.raw_tokens); + TEST_ASSERT(run.admitted_sizes[1] < run.raw_tokens); +} + +// A routed long prompt is scored too, and served by the model that scored it. +TEST_CASE(ServerUnitFixture, test_scheduler_pflash_compresses_a_routed_prompt) { + const auto run = run_pflash_stage_case(/*chats=*/0, /*long_prompts=*/1, + /*routed=*/true); + TEST_ASSERT(run.done); + TEST_ASSERT(run.compress_calls == 1); + TEST_ASSERT(run.admitted_sizes.size() == 1); + TEST_ASSERT(run.admitted_sizes[0] < run.raw_tokens); + TEST_ASSERT(run.admissions.size() == 1); + TEST_ASSERT(run.admissions[0] == RoutingAdmission::handled); + TEST_ASSERT(run.usage_prompt_tokens[0] == run.raw_tokens); +} + +// Each thread reads the error its own failing call set: the PFlash thread +// scores beside the scheduler, and neither may see, or free, the other's. +TEST_CASE(ServerUnitFixture, test_last_error_is_per_thread) { + luce::common::set_last_error("scheduler error"); + std::string seen_by_other; + std::thread other([&] { + luce::common::set_last_error("drafter error"); + seen_by_other = luce_last_error(); + }); + other.join(); + TEST_ASSERT(seen_by_other == "drafter error"); + TEST_ASSERT(std::string(luce_last_error()) == "scheduler error"); +}