From 1507f877dcd3b5dd79805d98e226613a480224dd Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 10 Oct 2026 06:16:04 +0000 Subject: [PATCH 1/2] feat(server): serve several models from one process and swap them on a shared GPU Load every named model block and pin named requests to its capacity, while unnamed and auto requests retain primary-first balancing. Report OpenAI chat logprobs from raw logits, including tokens substituted by the thinking budget. Load tied-embedding Qwen3.5 models without a separate output head and support the Qwen3.5-9B DFlash drafter. Expose prefill residual streams through POST /v1/hidden_states, with an optional prefix checkpoint for a subsequent chat request. Run eviction and reinstatement on each model's serial worker, retaining host prefix snapshots and rebuilding the layout fixed at load. Swap residency drains overlapping models and admits requests only to the named resident. Keep selected Qwen weights warm, retry cold if necessary, restore displaced models on failure and enforce a shared 30-second failure cooldown. Load and evict initially inactive blocks before loading the resident set, so the models never need to fit together during startup. DeepSeek4 has partial and cold eviction with frozen expert ownership. A partial eviction releases the primary device and keeps the secondary expert tier, streamed cache and host snapshots. Control operations carry an eviction level; the manager evicts a model partially when its secondary devices do not overlap the target, keeps those devices as retained, and cold-evicts a retained or warm-held model whose devices the target needs. One model block may set --expert-device and --peer-access, a later block no longer turns the peer-access opt-in off, and a DeepSeek4 block that would inherit the expert device from the environment is refused. Residency names auto:N devices by the compiled backend and requires --draft-device for a draft, so aliases and unplaced drafts cannot hide a shared GPU. An automatic request that exceeds every resident model's context is rejected instead of waiting for a swap it cannot trigger. On integrated GPUs the DeepSeek4 memory queries are clamped to system RAM and available host memory, so the streamed cache no longer sizes itself from the 1.25 TiB the driver reports without a carve and exhausts host memory. On an R9700 (PCIe Gen5 x16) with a Strix Halo, Qwen 3.8 27B + DFlash2 and DeepSeek V4.1 Flash (ds41-lucebox placement) share the R9700 in one process. DeepSeek V4.1 evicts partially in 2-58 ms and reinstates in 9-14 s without reloading its 90 GiB Strix Halo tier; the 27B reinstates in 7.6-8.4 s and decodes at the same speed afterwards. Ten partial cycles returned the R9700 to the same free memory with identical restored snapshots, and greedy outputs matched across swaps. Earlier on an R9700 over PCIe Gen3 x4, the 27B and Qwen3.5-9B reinstated cold in 5.1 s and 2.1-2.35 s and warm in 5-9 ms, and a 52.7k-token session returned in 2.5 s instead of about 80 s. Co-Authored-By: Claude Opus 5.5 --- docs/design/model-swap.md | 196 ++++ server/CMakeLists.txt | 9 +- server/docs/API.md | 130 ++- server/docs/MODEL_LOAD_BALANCING.md | 189 +++- server/docs/RECOMMENDED_SETUPS.md | 14 + server/include/luce.h | 3 + server/scripts/check_hidden_states.py | 173 ++++ server/src/common/backend_factory.cpp | 2 +- server/src/common/concurrency/seq_engine.h | 36 + server/src/common/dflash2_batch.cpp | 84 +- server/src/common/dflash2_head.cpp | 4 + server/src/common/dflash2_head.h | 14 +- server/src/common/dflash_capture.cpp | 15 + server/src/common/dflash_capture.h | 14 + server/src/common/dspark_head.cpp | 26 +- server/src/common/dspark_head.h | 5 + server/src/common/generation_types.h | 4 + .../src/common/geometric_draft_topk_cuda.cu | 4 + server/src/common/geometric_draft_topk_cuda.h | 4 + server/src/common/geometric_sampler_cuda.cu | 4 + server/src/common/geometric_sampler_cuda.h | 5 + server/src/common/hidden_states.h | 85 ++ server/src/common/model_backend.h | 83 ++ server/src/common/moe_hybrid_expert_cache.cpp | 57 +- server/src/common/moe_hybrid_expert_cache.h | 24 +- server/src/common/moe_hybrid_storage.cpp | 202 ++++- server/src/common/moe_hybrid_storage.h | 71 ++ server/src/common/moe_hybrid_stream.cpp | 37 +- server/src/common/moe_hybrid_stream.h | 7 +- server/src/common/sampler.h | 7 + server/src/common/step_graph.h | 5 + server/src/common/token_logprobs.cpp | 55 ++ server/src/common/token_logprobs.h | 33 + server/src/deepseek4/deepseek4_backend.cpp | 606 +++++++++++-- server/src/deepseek4/deepseek4_backend.h | 105 ++- server/src/deepseek4/deepseek4_graph.cpp | 27 +- server/src/deepseek4/deepseek4_hc_cuda.cu | 38 +- server/src/deepseek4/deepseek4_hc_cuda.h | 5 + server/src/deepseek4/deepseek4_internal.h | 8 + server/src/deepseek4/deepseek4_loader.cpp | 1 + server/src/internal.h | 25 +- server/src/qwen3/qwen3_drafter.cpp | 6 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 149 ++- .../qwen35/concurrency/qwen35_seq_engine.h | 31 +- server/src/qwen35/gguf_target_loader.cpp | 47 +- server/src/qwen35/graph_builders.cpp | 25 +- server/src/qwen35/graph_builders.h | 7 +- server/src/qwen35/qwen35_backend.cpp | 536 +++++++++-- server/src/qwen35/qwen35_backend.h | 52 ++ server/src/qwen35/qwen35_target_graph.cpp | 74 +- server/src/qwen35moe/qwen35moe_backend.h | 4 + server/src/server/http_server.cpp | 845 +++++++++++++++++- server/src/server/http_server.h | 187 +++- server/src/server/scheduler.cpp | 113 ++- server/src/server/server_main.cpp | 271 +++++- server/test/test_model_names.cmake | 23 + .../test/unit/test_dflash_capture_layers.cpp | 42 + server/test/unit/test_hidden_states.cpp | 216 +++++ server/test/unit/test_model_routing.cpp | 226 ++++- server/test/unit/test_model_swap_gpu.cpp | 463 ++++++++++ server/test/unit/test_moe_hybrid_storage.cpp | 207 +++++ server/test/unit/test_seq_batch_plan.cpp | 6 +- server/test/unit/test_server_unit.cpp | 139 +++ server/test/unit/test_token_logprobs.cpp | 129 +++ 64 files changed, 5837 insertions(+), 377 deletions(-) create mode 100644 docs/design/model-swap.md create mode 100755 server/scripts/check_hidden_states.py create mode 100644 server/src/common/hidden_states.h create mode 100644 server/src/common/token_logprobs.cpp create mode 100644 server/src/common/token_logprobs.h create mode 100644 server/test/unit/test_dflash_capture_layers.cpp create mode 100644 server/test/unit/test_hidden_states.cpp create mode 100644 server/test/unit/test_model_swap_gpu.cpp create mode 100644 server/test/unit/test_token_logprobs.cpp diff --git a/docs/design/model-swap.md b/docs/design/model-swap.md new file mode 100644 index 0000000000..d0e6762f2f --- /dev/null +++ b/docs/design/model-swap.md @@ -0,0 +1,196 @@ +# Swapping models on a shared GPU in one luce_server process + +Status: implemented for qwen35 (cold and warm eviction) and DeepSeek V4 +(partial and cold eviction), both driven by swap residency. User-facing flags are documented in +[MODEL_LOAD_BALANCING.md](../../server/docs/MODEL_LOAD_BALANCING.md#swap-residency). + +## Goal + +One process serves several model blocks (`--model-routing name`) that share +GPUs but do not fit on them together. Only models whose device sets do not +overlap are resident at a time. A request for an evicted model swaps it in: + +1. drain the conflicting resident models; +2. release their device memory while keeping their host-RAM prefix snapshots; +3. bring the requested model back with the device layout it had at load. + +Compared with restarting a process per model, an in-process swap keeps every +model's prefix snapshots, so a long session that comes back after a swap +restores its prefix instead of prefilling it again. That is most of the value; +the weight upload itself can only approach the PCIe floor. + +Out of scope: choosing which model serves a request (the client names it), +co-residence of two models' working sets on one GPU, and GPU-to-GPU +migration. + +## Measured constraints (R9700, PCIe Gen3 x4) + +- About 3.5 GB/s each way, pinned or pageable alike. The weight upload floor + is about 4.1 s for Qwen 3.8 27B + DFlash2 and 1.8 s for Qwen3.5-9B + DFlash. +- A process restart takes 6.2 s and 2.9 s for the two, and loses every prefix + snapshot. +- Snapshot restore at 52k tokens takes 2.2 s (27B) and 1.0 s (9B); a cold + prefill at 52k takes 77 s and 29.5 s. + +Measured with this change: cold reinstatement 5.1 s (27B) and 2.1-2.35 s (9B), +warm reinstatement 5-9 ms, a 52.7k-token session back in 2.5 s instead of about +80 s, greedy output identical after a snapshot restore, and 26.9 GB peak VRAM +with both models' weights kept plus a 52k-token prefill at 64k context. + +## Backend contract + +`ModelBackend` gains, separate from `park()` (which a backend uses for its own +transient memory juggling): + +- `supports_eviction(reason)`: whether this configuration can be evicted. +- `evict(level)`: releases device state on the model's worker while it is + idle. `rejected` means nothing was torn down and the model is unchanged. + - `cold`: every device allocation (weights, draft, KV/recurrent caches, + paged pools, graphs, scratch, backend contexts and their pools). + - `warm`: keeps the weights on the device and releases the rest, so a + reinstatement only rebuilds caches. + - `partial`: releases the primary device only and keeps a secondary + device's expert tier (DeepSeek V4 with in-process experts). +- `reinstate(error)`: rebuilds exactly the layout frozen at init (context, + KV types, paged-pool size, draft configuration), never re-deriving it from + the free memory of the moment, so snapshots taken before the eviction + restore afterwards. On failure the model stays evicted with its partial + state released. +- `is_evicted()`, `weights_resident()`, `weight_device_bytes()` for the + manager and `/status/json`. + +While evicted, generation, hidden-state readouts and snapshot saves fail with +an explicit error, and size and memory queries answer without touching the +freed tensors. + +### qwen35 + +Supported for the single-device serial worker with the draft on the target +device; `supports_eviction()` refuses `--max-concurrency` > 1, paged +attention, KVFlash, a remote draft, tensor parallelism, vision and the other +Qwen3.5 variants. + +`init()` is split so the device-allocation steps run both at first init and in +`reinstate()`. Eviction materializes any deferred live snapshot to host, then +releases consumers before what they reference: the DFlash adapter, drafter, +step graphs, DFlash2 selector graphs and the worker's thread-local device +scratch (DFlash2, SpecLA, GPU sampler and draft top-k, BSA persistent +buffers), the draft KV and feature mirror, the weights, the target cache, and +finally the backend contexts: on the HIP VMM pool only `ggml_backend_free` +returns pool memory. KV caches are zeroed on the device at reinstatement +instead of uploading host zeros. + +Warm eviction keeps the weight buffers (they belong to the device's buffer +type, not to the backend context) and points the rebuilt model at them, which +makes a reinstatement a cache rebuild of a few milliseconds. + +### DeepSeek V4 + +Supported for the classic serial worker with in-process experts on two GPUs +of one runtime (the `ds41-lucebox` profile); `supports_eviction()` names what +else refuses (paged serving, vision, PFlash, remote experts, spare primary +cache slots, mixed-qtype experts). + +| State | Primary device | Secondary expert tier | Host snapshots and model state | +| --- | --- | --- | --- | +| resident | present | present, running | present | +| partially evicted | released | present, suspended | present | +| fully evicted | released | released | present | + +- Settings the backend used to reread from the environment (expert + parallelism, draft placement, fused modes) are resolved once at init into + `DeepSeek4ResolvedSettings`, so a reinstated model comes back with the + devices and modes it was loaded with. +- The expert placement is frozen at init (`FrozenPlacement`): ordered hot and + cold expert ids, physical and decode owner maps, streamed count, spill + sizes and the selection bias after routing adjustments. A reinstatement + never places experts again, since a placement recomputed from current free + memory would move rows under the retained secondary stack and cache, and + verifies the rebuilt storage against the record. +- Router-bias and protected-expert inputs are parsed once and applied once per + set of weights. +- Partial eviction order: finish mailbox-dependent launches while the + resolver and loaders still run; suspend the streamed expert cache (no new + jobs, issued copies finish, slots and LRU kept); synchronize every context + and stream; destroy graph consumers before tensors (fused graphs, storage + graph caches and prefill arenas, worker-local runtime, DSpark head scratch, + direct HC scratch, stream engine); release the primary payloads and backend + context. +- Partial reinstatement recreates the primary backend, reloads the dense + tensors, rebuilds the primary expert stacks from the recorded shapes and + ordered ids, reapplies the routing adjustments, recreates the cache with the + frozen geometry and resumes the expert cache without resizing it. +- Primitives used: `MoeHybridStorage::release_primary_experts()` / + `rebuild_primary_experts()` / `rebind_cold_backend()` / + `release_prefill_allocators()`, expert cache `suspend()` / `resume()` / + `release_graphs()`, and per-device release of the DSpark head and HC + scratch. + +## Control operations on the worker + +The worker's job queue carries control operations (evict warm or cold, +reinstate) with a completion the manager waits on. The serial worker runs +them between jobs, when nothing else touches the GPU, and publishes a fresh +memory report afterwards. Models served by the concurrent scheduler +(`--max-concurrency` > 1) refuse control operations for now, so swap +residency needs the serial worker. + +## Swap manager + +The multi-model front keeps, per model block, the device set (target, layer +split GPUs, draft device, expert device), a residency state +(`resident | evicted | failed`), waiter counts and swap statistics. + +- Startup: cheap checks first (explicit target devices, and draft devices + for blocks with a draft), then every later + block whose command-line device set overlaps a model that stays resident + loads and is evicted on its own, before any resident block loads. The + eviction is partial when the block has an expert device that no resident + block and no other retained tier uses, otherwise cold. The first block + stays resident, as does any block on devices of its own. Keep-weights + applies to runtime swaps. + Overlapping resolved device sets of resident models are rejected. +- A pinned request for an evicted model registers as a waiter and drives the + swap: one swap at a time; conflicting residents first serve the requests + that were waiting when the swap was requested, then close admission and + drain; each is evicted (partial when its secondary devices do not overlap + the target's, warm when listed in `--swap-keep-weights` and both models are + single-device, otherwise cold); the target is reinstated. +- A partially evicted model keeps its secondary devices as + `retained_devices`. A model that needs one of them, or a warm-held model + whose devices a multi-device target needs whole, is evicted cold before + the target reinstates. +- If a reinstatement fails while other models keep weights warm on the same + devices, those are evicted cold and the reinstatement retried once. +- A failed eviction or reinstatement reinstates what this swap evicted, marks + the target failed (no retry for 30 s) and answers its waiters 503. +- Unnamed and `auto` requests only use resident models and never swap; when + no model is resident and no swap runs they get 503 instead of waiting. +- A requester whose model became resident admits at once. +- Weight residency and kept-weight bytes for `/status/json` are read on the + worker when a transition finishes and published under the manager's lock, + so no HTTP thread queries a backend that may be mid-swap. + +## Remaining work + +- Launch profiles that set environment variables are refused in a + multi-model server, so a mixed server sets the `ds41-lucebox` environment + and flags itself, and Qwen runs under it (including + `ROCBLAS_USE_HIPBLASLT=0`). Qwen performance under that environment is + not qualified beyond the 27B decode check below. +- Control operations through the concurrent scheduler, so models served with + `--max-concurrency` > 1 can swap. + +Measured on an R9700 (PCIe Gen5 x16) with a Strix Halo (no carve, 125 GiB) +serving Qwen 3.8 27B + DFlash2 and DeepSeek V4.1 Flash (`ds41-lucebox` +placement, Q2K-Q4K DSpark drafter): DeepSeek V4.1 starts partially evicted +with 9761 experts (90 GiB) held on the Strix Halo. Its partial eviction takes +2-58 ms and its partial reinstatement 9-14 s (dense weights, primary experts +and drafter reread from SSD); the Strix Halo tier is never reloaded. The 27B +reinstates cold in 7.6-8.4 s and decodes at the same speed after a swap. +Across ten partial cycles the R9700 returned to the same free memory, and +greedy outputs and restored snapshots matched the first run. + +GPU graphs: multi-model processes disable graph capture because another +worker's blocking calls break a capture in flight. With swap residency and one +resident model per device this hazard may disappear; left for later. diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 99550d2d63..28402c146f 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -663,6 +663,7 @@ add_library(luce_common STATIC src/qwen35/layer_split_daemon_loop.cpp src/qwen35/qwen35_daemon.cpp src/common/sampler.cpp + src/common/token_logprobs.cpp src/common/backend_precision.cpp src/common/daemon_loop.cpp src/common/gguf_inspect.cpp @@ -2329,6 +2330,7 @@ if(LUCE_TESTS) test/unit/test_paged_kv_offload.cpp test/unit/test_anchor_params.cpp test/unit/test_derived_scalars.cpp + test/unit/test_dflash_capture_layers.cpp test/unit/test_adaptive_spec_width.cpp test/unit/test_immutable_graph_input_pool.cpp test/unit/test_mmq_tile_selection.cpp @@ -2343,6 +2345,8 @@ if(LUCE_TESTS) test/unit/test_bandit_integration.cpp test/unit/test_admission.cpp test/unit/test_restore_delta.cpp + test/unit/test_token_logprobs.cpp + test/unit/test_hidden_states.cpp test/unit/test_chain_rollback_policy.cpp test/unit/test_ddtree_tau.cpp test/unit/test_anchor_transitive.cpp @@ -2455,11 +2459,14 @@ if(LUCE_TESTS) test/smoke/smoke_load_target.cpp test/smoke/smoke_load_target_laguna.cpp test/unit/test_paged_attention.cpp - test/unit/test_decode_kv_offload_gpu.cpp) + test/unit/test_decode_kv_offload_gpu.cpp + test/unit/test_model_swap_gpu.cpp) target_include_directories(test_model_smoke PRIVATE ${LUCE_SRC_INCLUDE_DIRS} ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src) + target_compile_definitions(test_model_smoke PRIVATE + LUCE_TEST_SHARE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/share") if(LUCE_GPU_BACKEND STREQUAL "hip") target_compile_definitions(test_model_smoke PRIVATE LUCE_BACKEND_HIP=1 diff --git a/server/docs/API.md b/server/docs/API.md index 8fb6ff3d58..98da8b4306 100644 --- a/server/docs/API.md +++ b/server/docs/API.md @@ -15,6 +15,7 @@ backend (qwen35, qwen3, gemma4, laguna). | POST | `/v1/chat/completions` | OpenAI Chat Completions | ✅ | | POST | `/v1/messages` | Anthropic Messages | ✅ | | POST | `/v1/responses` | OpenAI Responses API | ✅ | +| POST | `/v1/hidden_states` | Prefill residual-stream activations | ✅ qwen35 | --- @@ -78,8 +79,8 @@ Successful requests expose the measured result under `usage.timings`: | `chat_template_kwargs` | object | — | Direct template control (`{"enable_thinking":true}`). Jinja templates (e.g. qwen4exp) also accept `preserve_thinking` (bool): whether earlier assistant turns replay their recorded `reasoning_content` inside `...` (default: template's own default, typically true) | ✅ | | `stop` | string/array | — | Stop sequences | ✅ | | `n` | int | — | Number of completions | ❌ TODO | -| `logprobs` | bool | — | Return log probabilities | ❌ TODO | -| `top_logprobs` | int | — | Number of top logprobs per token | ❌ TODO | +| `logprobs` | bool | `false` | Return log probabilities (non-streaming, qwen35 targets; see [Logprobs](#logprobs)) | ✅ | +| `top_logprobs` | int | 0 | Number of top logprobs per token (0–20; ignored without `logprobs: true`) | ✅ | | `response_format` | object | — | JSON mode / structured output | ❌ TODO | | `tool_choice` | string/object | — | Tool choice / force tool usage | ✅ | | `logit_bias` | object | — | Per-token logit adjustments | ❌ TODO | @@ -101,7 +102,38 @@ Successful requests expose the measured result under `usage.timings`: | `usage.prompt_tokens` | ✅ | | `usage.completion_tokens` | ✅ | | `usage.total_tokens` | ✅ | -| `choices[].logprobs` | ❌ TODO | +| `choices[].logprobs` | ✅ when `logprobs: true` | + +### Logprobs + +`logprobs: true` returns, for every generated token, its log-probability and the `top_logprobs` most likely alternatives in the OpenAI shape. Values come from the raw logits row the token was chosen from: temperature, top-k/top-p and penalties change which token is picked, not the reported distribution. A token that is only part of a UTF-8 character has its exact bytes in `bytes` and U+FFFD in `token`. + +```bash +curl -s http://localhost:8080/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{"messages":[{"role":"user","content":"Rate this task: rename a variable"}], + "max_tokens":1,"temperature":0,"logprobs":true,"top_logprobs":3}' +``` + +```json +{ + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "easy"}, + "finish_reason": "length", + "logprobs": {"content": [{ + "token": "easy", "logprob": -0.0312, "bytes": [101, 97, 115, 121], + "top_logprobs": [ + {"token": "easy", "logprob": -0.0312, "bytes": [101, 97, 115, 121]}, + {"token": "medium", "logprob": -3.52, "bytes": [109, 101, 100, 105, 117, 109]}, + {"token": "hard", "logprob": -6.87, "bytes": [104, 97, 114, 100]} + ] + }]} + }] +} +``` + +Only `logprobs: true` on a non-streaming `/v1/chat/completions` request served by the single-device qwen35 backend returns logprobs. Elsewhere the field is ignored and the response carries none, as before logprobs existed: `stream: true`, other endpoints (including the legacy completions integer `logprobs`), a `top_logprobs` without `logprobs: true`, and other backends (layer-split, qwen35moe, DeepSeek4, Laguna, Gemma4 and Qwen3 targets). With `logprobs: true`, a `top_logprobs` that is not an integer in 0–20 returns HTTP 400. A logprobs request decodes without speculation, and on the single-request path it bypasses the prefix caches; with `--max-concurrency` it still uses prefix checkpoints, which always leave the last prompt token to prefill. --- @@ -182,6 +214,96 @@ to render in its native format. --- +## POST `/v1/hidden_states` (prefill activations) + +Runs the prompt through **prefill only** on the named model and returns the +residual stream (pre-final-norm, f32) after chosen blocks, at the last prompt +token and optionally mean-pooled over every prompt token. Nothing is decoded +or sampled. Use it to read a model's internal representation of a prompt, +for example as features for a separate classifier. + +```bash +curl -s http://127.0.0.1:8080/v1/hidden_states -H 'Content-Type: application/json' -d '{ + "model": "qwen3.5-2b", + "messages": [{"role": "user", "content": "Prove that sqrt(2) is irrational."}], + "chat_template_kwargs": {"enable_thinking": false}, + "layers": [11, 17, -1], + "pooling": "both", + "generation_prompt": "bare" +}' +``` + +```json +{ + "object": "hidden_states", "model": "qwen3.5-2b", + "n_layers": 24, "n_embd": 2048, "prompt_tokens": 19, + "pooling": "both", "generation_prompt": "bare", "last_token": "\n", + "layers": { + "11": {"last": [0.013, -0.402, ...], "mean": [0.021, -0.117, ...]}, + "17": {"last": [...], "mean": [...]}, + "-1": {"last": [...], "mean": [...]} + }, + "timings": {"prefill_ms": 31.4} +} +``` + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `model` | string | — | Model block name (`--model-name`). Required on a multi-model listener; an unknown name is 404 even with `--unknown-model primary`. | +| `messages` | array | — | OpenAI chat messages, rendered with the model's chat template like `/v1/chat/completions` (tools, reasoning and `chat_template_kwargs` apply the same way). | +| `layers` | int[] | — | 0-based block indices; entry `L` is the residual stream **after** block `L`. Negative counts from the end (`-1` = last block). HF `output_hidden_states=True` has `hidden_states[L+1]` for block `L`. | +| `pooling` | string | `"last"` | `"last"`, `"mean"` or `"both"`. `last` is always returned; `mean` (over all prompt tokens) only for `mean`/`both`. | +| `generation_prompt` | string | `"template"` | `"template"` keeps whatever the template appends after the assistant header. `"bare"` ends the prompt exactly at `<\|im_start\|>assistant\n`, dropping e.g. Qwen3.5's `\n\n\n\n` (thinking off) or `\n` (thinking on), so the features do not depend on the thinking mode of a later request. | +| `cache_prefix` | bool | `false` | Also store this prefill in the prefix cache, so a following `/v1/chat/completions` on the same model whose prompt extends this one (typically the `"bare"` prompt plus the template's think block) restores it instead of prefilling it again. See below. | + +Response `layers` is keyed by each layer exactly as requested. `last_token` is +the decoded text of the final prompt token (where `last` was read). +`timings.prefill_ms` is wall time from admission to the end of prefill. +With `cache_prefix: true` the response also carries +`"cache_prefix": {"cached_tokens": N}`: how many leading tokens of this prompt +the prefix cache holds once the request is done (`prompt_tokens` when the +checkpoint was stored or already present, 0 when it could not be stored). + +Behavior: + +- Always a full prefill from position 0: prefix-cache snapshots are never + restored, and PFlash compression never applies. By default nothing is + captured either. +- `cache_prefix: true` (opt-in) captures one prefix-cache checkpoint at the + **end** of the prompt, through the same machinery chat requests use: an + engine prefix checkpoint (paged K/V copy plus DeltaNet/conv recurrent + state) under `--max-concurrency` > 1, or an end-of-prefill inline snapshot + on the single-request worker. The cut is the prompt end, where the last + prefill chunk stops anyway, so the readout is computed exactly as without + the flag. A later chat request whose tokens extend this prompt finds the + entry by exact token prefix and reports the hit in `timings` + (`cache_hit`, `cached_prefix_tokens`). Storing is best effort: it is + skipped when the prefix cache is disabled (`--prefix-cache-slots 0`, + single-sequence `--paged-attention`), when another capture is in flight, + or when the entry would not fit the resident budget. Entries age out like + any other prefix-cache entry. The restored chat prefills its suffix in a + different chunk split than a cold prefill would, so greedy output matches + a no-cache run only up to the same floating-point tolerance as any other + prefix-cache hit. +- Runs through the model's own queue and capacity like a completion. With + `--max-concurrency` > 1 the request takes one engine slot and its prompt + chunks are batched with other slots' prefill/decode; the readout is a + per-slot row range of the shared step graph. Single-request models run it + on the worker between generations. +- Cost: one `[n_embd]` row copy per requested layer per prefill chunk, plus a + transpose + row sum when `mean` is requested. The LM head still runs for + the last row only. + +Errors (400): missing/empty/non-integer or repeated `layers`, a layer outside +the model, unknown `pooling`/`generation_prompt`, empty `messages`, image +input, `generation_prompt: "bare"` with a template that does not end in an +open `<|im_start|>assistant` turn, an unnamed model on a multi-model listener, +and backends without support. Only the single-device dense qwen35 backend +(no tensor parallelism, no KVFlash) supports it; layer-split, qwen35moe and +other architectures return 400. Unknown model under name routing: 404. + +--- + ## Sampling Chain The sampler applies penalties and sampling in this order: @@ -264,7 +386,7 @@ support causes errors or silent feature degradation. | Feature | Notes | |---------|-------| -| **`logprobs` / `top_logprobs`** | Token probabilities in response. Debugging/analysis only. | +| **Streaming `logprobs`** | Non-streaming chat completions return logprobs; SSE chunks do not carry them yet. | | **`n` (multiple completions)** | Generate N choices per request. No known agent uses this. | | **`logit_bias`** | Per-token logit adjustments. | | **`user`** | End-user identifier (tracking only). | diff --git a/server/docs/MODEL_LOAD_BALANCING.md b/server/docs/MODEL_LOAD_BALANCING.md index b7fa7e0196..318913c6b4 100644 --- a/server/docs/MODEL_LOAD_BALANCING.md +++ b/server/docs/MODEL_LOAD_BALANCING.md @@ -7,9 +7,11 @@ them. If all eligible models are busy, requests enter a bounded waiting queue. Model placement, primary GPU selection and enabling load balancing are separate settings. For example, Qwen can stay on R9700 and DeepSeek4 on Strix Halo while -either GPU is selected as primary. Every generation request follows this -capacity-based policy, regardless of the request's `model` field. The response -identifies the model that actually generated the answer. +either GPU is selected as primary. By default every generation request follows +this capacity-based policy, regardless of the request's `model` field; +`--model-routing name` instead pins a request that names a model to that model +(see [Name routing](#name-routing)). The response identifies the model that +actually generated the answer. Each model keeps its own tokenizer, defaults, backend and scheduler. Clients use the existing Chat Completions, Messages and Responses endpoints. @@ -64,11 +66,187 @@ limit or a promise that every sequence will fit in the KV pool. {"model":"auto","messages":[{"role":"user","content":"Write a Python parser."}],"max_tokens":512,"stream":true} ``` +## Name routing + +`--model-routing name` hosts several different models behind one listener and +lets the client choose among them, like separate servers would, but in one +process. It loads every model block (it implies `--load-balancing`, with the +same per-block rules below) and applies one rule per request: + +| Request `model` | Generation | `count_tokens` | `/v1/hidden_states` | +| --- | --- | --- | --- | +| a block's `--model-name` | that block only; waits for its capacity | that block | that block only; waits for its capacity | +| omitted, empty or `auto` | balanced as below (primary first) | 400 | 400 | +| anything else | 404 listing the loaded names | 404 | 404, even with `--unknown-model primary` | + +`/v1/hidden_states` pins by name under either `--model-routing` mode: +activations only mean something for the model that produced them. + +A pinned request never spills onto another model: when its block is full it +waits in the shared routing queue (`--routing-queue-limit`) until that block +has capacity, the client disconnects or the server stops. A request that can +never fit that block's context or KV pool gets 400 without waiting. +`--unknown-model primary` serves unrecognized names on the primary instead of +answering 404 (the default is `reject`). The option is accepted only with +`--model-routing name`. `/status/json` and `/props` report +`routing: "by-name"`; the per-model `models[].{id, capacity, in_flight}` entries +are unchanged, so a client can read each model's load by its name. + +Example: three models in one process, two of them sharing the R9700: + +```bash +luce_server --model-routing name ~/models/Qwen3.8-27B-UD-IQ4_XS.gguf \ + --host 127.0.0.1 --port 8080 --routing-queue-limit 64 --model-name qwen3.8-27b \ + --target-device hip:0 \ + --draft ~/models/qwen38-dflash2-q8_0.gguf --draft-device hip:0 --draft-block-size 16 \ + --cache-type-k q8_0 --cache-type-v q8_0 \ + --max-concurrency 4 --kv-pool-tokens 131072 --max-ctx 65536 --max-tokens 32768 \ + --model ~/models/Qwen3.5-0.8B-Q8_0.gguf --model-name qwen3.5-0.8b \ + --target-device hip:0 \ + --max-concurrency 4 --kv-pool-tokens 131072 --max-ctx 32768 \ + --model ~/models/Qwen3.5-2B-Q8_0.gguf --model-name qwen3.5-2b \ + --target-device hip:1 \ + --max-concurrency 4 --kv-pool-tokens 262144 --max-ctx 122880 +``` + +Two blocks share the R9700 here. Give every block on a shared GPU an explicit +`--kv-pool-tokens`: without it Qwen sizes its pool from the memory free when +that block loads, starving the blocks loaded after it. + +Per-block options include placement (`--target-device`, `--draft`, +`--draft-device`, `--draft-block-size`), capacity (`--max-concurrency`, +`--kv-pool-tokens`, `--max-ctx`), output defaults (`--max-tokens`, +`--default-max-tokens` and the thinking budgets), Qwen's `--cache-type-k/v`, +prefix-cache sizes and `--admission-coalesce-ms`. Process-wide: the listener +(`--host`, `--port`, `--no-cors`, `--routing-queue-limit`, first block only), +`--model-routing`, `--unknown-model`, environment variables, and GPU graph +capture (next section). Each batched block also keeps its own copied prefix +checkpoints, bounded by 4 GiB of host RAM per block. + +Request features are checked by the block that serves the request, after +routing: for example `logprobs` is accepted when the named block's backend +reports per-token log-probabilities, independently of the primary. + +## Swap residency + +`--swap-residency` (with `--model-routing name`) serves models that do not fit +on their GPUs together. Models whose device sets overlap are never resident at +the same time; a request that names an evicted model swaps it in: + +1. the resident models that share a device with it stop admitting and drain + (requests that were already waiting for them are served first); +2. each of them is evicted on its worker: every device allocation it owns is + released (weights, draft, KV and recurrent caches, graphs, scratch, backend + context and pools), while its host-RAM prefix snapshots stay; +3. the requested model is reinstated with exactly the device layout it had at + load (same context, KV types and pool sizes), so its prefix snapshots + restore as before. + +The first model block starts resident. A later block whose devices overlap a +resident one loads on its own and is evicted cold before any resident block +loads, so the models never need to fit together; a block on devices of its +own stays resident. Keep-weights applies to runtime swaps. Every block needs an explicit `--target-device`, and `--draft-device` when it has a draft. Requests that omit +the model or ask for `auto` only use resident models and never trigger a swap; +when no model is resident and no swap is running they get 503. A failed +reinstatement reinstates the models evicted for it, answers its waiters 503 +and is not retried for 30 s. Only one swap runs at a time. + +`--swap-keep-weights ` evicts the listed models warm: their weights +stay in device memory and only the KV cache, graphs and backend context are +released, so swapping them back in takes milliseconds instead of a weight +upload. If the model being swapped in then does not fit, the kept weights are +dropped and the reinstatement retried. Warm eviction applies only between +single-device models. List the main model, not occasional ones: kept weights +shrink the room every other model has. + +```bash +luce_server --model-routing name --swap-residency --swap-keep-weights qwen3.8-27b \ + ~/models/Qwen3.8-27B-UD-IQ4_XS.gguf --model-name qwen3.8-27b \ + --target-device hip:0 \ + --draft ~/models/qwen38-dflash2-q8_0.gguf --draft-device hip:0 \ + --max-ctx 65536 \ + --model ~/models/Qwen3.5-9B-Q4_K_M.gguf --model-name qwen3.5-9b \ + --target-device hip:0 \ + --draft ~/models/qwen35-9b-dflash-q8_0.gguf --draft-device hip:0 \ + --max-ctx 65536 +``` + +`Qwen3.5-9B-Q4_K_M.gguf` is from `unsloth/Qwen3.5-9B-GGUF`; the 9B drafter is +`z-lab/Qwen3.5-9B-DFlash` converted to a Q8_0 GGUF. + +Measured on one R9700 (PCIe Gen3 x4) with these two models: a cold +reinstatement takes 5.1 s for the 27B and 2.1-2.35 s for the 9B, a warm one +5-9 ms. A 52.7k-token session that comes back after a swap restores its prefix +snapshot in 2.5 s instead of about 80 s of cold prefill, with greedy output +identical to a run without the swap. + +`/status/json` reports each model's `residency` (`resident`, `evicted` or +`failed`), `devices`, `weights_resident`, `weight_device_bytes`, `draining`, +`waiting_for_residency`, `swaps_in`, `last_evict_ms`, `last_reinstate_ms` and +`last_residency_error`, plus process-wide `swap_residency` and `swap_active`. + +Eviction is implemented for the single-device qwen35 backend on the serial +worker (no `--max-concurrency` > 1, paged attention, KVFlash, tensor +parallelism or vision, and the draft on the target device); a model whose +backend cannot be evicted makes `--swap-residency` refuse to start. + +DeepSeek V4 with `--expert-device` is evicted partially: only its primary +device is released, and its secondary expert tier, streamed expert cache and +host snapshots stay, so swapping it back rereads the primary weights only. +`/status/json` lists the kept devices as `retained_devices`; another model +that needs one of them evicts DeepSeek V4 cold first. One model block may set +`--expert-device` and `--peer-access` (both change process-wide state); a +DeepSeek V4 block without `--expert-device` is refused while `LUCE_DS4_MOE_TP` +is set, since it would inherit that expert device. Launch +profiles that set environment variables are refused with several models, so +set that environment for the process. For Qwen 3.8 27B on the R9700 beside +DeepSeek V4.1 with its expert tier on Strix Halo: + +```bash +# ds41-lucebox environment, process-wide (see launch_profiles.h) +export LUCE_EXPERT_BUDGET_MB=10500 LUCE_DS4_FUSED_VERIFY=1 ROCBLAS_USE_HIPBLASLT=0 # ... +luce_server --model-routing name --swap-residency \ + ~/models/Qwen3.8-27B-UD-IQ4_XS.gguf --model-name q27 --target-device hip:0 \ + --draft ~/models/Qwen3.8-27B-DFlash2-Q8_0.gguf --draft-device hip:0 --max-ctx 65536 \ + --model ~/models/ds41-lucebox-final.gguf --model-name ds41 \ + --draft ~/models/DeepSeek-V4.1-Flash-DSpark-draft-Q2K-Q4K.gguf --draft-device hip:0 \ + --target-device hip:0 --expert-device hip:1 --peer-access --max-ctx 131072 \ + --chunk 4096 --ds4-prefill dense \ + --ds4-expert-placement share/deepseek41/placement_lucebox.json \ + --ds4-router-bias share/deepseek41/router_bias_lucebox_40x384_f32.bin \ + --ds4-protected-experts share/deepseek41/massive_experts.json +``` + +DeepSeek V4.1 starts partially evicted with its Strix Halo tier loaded (run +it with unlimited locked memory). Swapping it in takes 9-14 s, swapping the +27B back 7.6-8.4 s, and evicting either takes milliseconds. + +## GPU graphs with several models + +A process that loads more than one model disables HIP/CUDA graph capture +(`GGML_CUDA_DISABLE_GRAPHS=1`) unless `LUCE_MULTI_MODEL_GRAPHS=1` is set. Each +model captures from its own worker thread with relaxed capture mode, and a +blocking runtime call from another worker (the `cudaStreamPerThread` +copy-and-synchronize in ggml's tensor get/set, prefix-checkpoint copies) +invalidates a capture in flight; the capturing worker then aborts. PR #770 +moved the helper copies and replay-log commits off the legacy stream, but still +measured about two failed runs in three with the two models on different GPUs +(R9700 and Strix Halo). So separate devices or separate graph caches do not +make capture safe, and a per-model or per-device relaxation (possible +mechanically through `GGML_CUDA_DISABLE_GRAPHS_DEVICES` or the per-thread +override) would still expose the capturing model to the other workers' blocking +calls. A safe relaxation needs the capture window to exclude other workers' +blocking calls, which does not exist yet. On the balanced two-27B aggregate +eager launches cost nothing measurable; small, launch-bound models such as a +0.8B/2B lose more, which is unmeasured here. `LUCE_MULTI_MODEL_GRAPHS=1` +remains an at-your-own-risk experiment. + ## How requests are balanced - Every generation request follows operator-configured priority. An omitted, empty, `auto`, explicit model name, or other client alias has the same - load-balancing behavior. There is no model pinning. With balancing enabled, requests try the + load-balancing behavior. There is no model pinning unless + `--model-routing name` is set. With balancing enabled, requests try the primary and then the remaining models; without it, only the primary serves. This is preference with capacity fallback, not round-robin. - The listener reserves an available model slot. Its scheduler then calls @@ -111,7 +289,8 @@ model identity, regardless of the client alias. With balancing enabled, model names must be unique, nonempty and cannot be `auto`. With balancing enabled, startup rejects model-block options that mutate shared process policy, including KVFlash, Spark, PFlash, -remote drafts and target splitting. Environment variables remain process-wide; +remote drafts and target splitting; `--expert-device` and `--peer-access` are +accepted on one model block. Environment variables remain process-wide; this does not establish arbitrary model/environment isolation. Qwen's explicit KV types are stored per backend rather than in shared environment state. With balancing disabled, the selected primary retains single-model CLI diff --git a/server/docs/RECOMMENDED_SETUPS.md b/server/docs/RECOMMENDED_SETUPS.md index cff233dfe5..b9e0067f1a 100644 --- a/server/docs/RECOMMENDED_SETUPS.md +++ b/server/docs/RECOMMENDED_SETUPS.md @@ -19,6 +19,20 @@ For long contexts on the 32 GB R9700, `--cache-type-k f16 --cache-type-v f16` de If a machine has both Strix Halo and an R9700, set `HIP_VISIBLE_DEVICES=` before using an R9700-only profile. The selected card is then `hip:0` inside the process. +## Small Qwen 3.5 models on Strix Halo + +Qwen 3.5 0.8B and 2B chat GGUFs, and classifiers fine-tuned from them, load through the qwen35 backend. They tie the output projection to the token embeddings (the GGUF has no `output.weight`), so the loader uploads `token_embd.weight` as the LM head. Run them without a draft; `--max-concurrency` serves several requests at once from a paged KV pool. On a machine that also has an R9700, check the Strix Halo index with `luce_server --list-devices`. + +```bash +luce_server /path/to/Qwen3.5-0.8B-Q8_0.gguf \ + --target-device hip:1 \ + --max-concurrency 4 \ + --max-ctx 8192 \ + --port 8294 --model-name qwen3.5-0.8b +``` + +To read a classifier's first-token distribution, ask for `max_tokens: 1`, `temperature: 0` and `logprobs: true` with `top_logprobs`; see [Logprobs](API.md#logprobs). + ## DeepSeek V4 on Strix Halo `--profile ds4-strix` is the qualified launch: it sets `--max-ctx 131072 --chunk 8192 --ds4-fused-decode --ds4-fused-verify-f16-kv --ds4-expert-top-k 6 --ds4-prefill sparse`, and the `gfx1151` device profile installs every kernel and policy default at start (fused five-row verifier, verify width from the DSpark confidence head, sparse prefill kernels), so there is nothing to tune. Any flag you pass replaces the profile's value. Use the adaptive ROCmFPX artifact with all six routed experts on the 128 GB part. Measured this way: 42 tok/s decode and 320 tok/s prefill at 8K, 36 tok/s at 123K, 39 tok/s on code and math, 25 tok/s on prose ([PR #729](https://github.com/Luce-Org/lucebox/pull/729); details in the [DeepSeek V4 guide](DS4.md#experimental-amd-q5-verifier)). diff --git a/server/include/luce.h b/server/include/luce.h index 3d62b57861..7e616ffa55 100644 --- a/server/include/luce.h +++ b/server/include/luce.h @@ -35,6 +35,9 @@ extern "C" { #define LUCE_DRAFT_LAYERS 5 #define LUCE_DRAFT_BLOCK_SIZE 16 #define LUCE_DRAFT_N_TARGET_LAYERS 5 // fc projects 5*hidden -> hidden +// Capacity of the target-side capture-layer table. Drafters may capture more +// layers than the 5-layer default (z-lab Qwen3.5-9B DFlash captures 8). +#define LUCE_DRAFT_MAX_TARGET_LAYERS 16 #define LUCE_DRAFT_MASK_TOKEN_ID 248070 // target_layer_ids = {1, 16, 31, 46, 61} (0-indexed into target layers) diff --git a/server/scripts/check_hidden_states.py b/server/scripts/check_hidden_states.py new file mode 100755 index 0000000000..0798de02cb --- /dev/null +++ b/server/scripts/check_hidden_states.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""Compare luce_server's /v1/hidden_states against Hugging Face transformers. + +The server reports the residual stream after block L (before the final norm) +at the last prompt token and, with pooling=both, its mean over the prompt. +This script renders the same chat with the HF tokenizer, runs the HF model, +reads each requested decoder block's output through a forward hook (HF's +output_hidden_states puts block L at hidden_states[L+1], but its last entry +already has the final norm applied, so hooks are used for every layer), and +prints the cosine similarity per layer. A correct server gives > 0.99 (the +GGUF is quantized, so values are not bit-identical). + +Needs the base HF weights of the served GGUF, for example: + hf download Qwen/Qwen3.5-2B --local-dir ~/models/hf/Qwen3.5-2B + +Example: + python3 server/scripts/check_hidden_states.py \ + --url http://127.0.0.1:8080 --model qwen3.5-2b \ + --hf ~/models/hf/Qwen3.5-2B --layers 0 11 -1 --bare +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +import urllib.request + +ASSISTANT_HEADER = "<|im_start|>assistant\n" + +DEFAULT_PROMPTS = [ + "Prove that the square root of 2 is irrational.", + "Write a haiku about autumn.", + "def fib(n):\n # complete this function and explain its complexity", +] + + +def post(url: str, body: dict) -> dict: + req = urllib.request.Request( + url.rstrip("/") + "/v1/hidden_states", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(req, timeout=600) as resp: + return json.loads(resp.read()) + + +def cosine(a, b) -> float: + dot = sum(x * y for x, y in zip(a, b)) + na = math.sqrt(sum(x * x for x in a)) + nb = math.sqrt(sum(y * y for y in b)) + return dot / (na * nb) if na and nb else float("nan") + + +def find_decoder_layers(model, n_hint: int | None): + import torch + + best = None + for name, module in model.named_modules(): + if isinstance(module, torch.nn.ModuleList) and name.endswith("layers"): + if "vision" in name or "visual" in name: + continue + if best is None or len(module) > len(best[1]): + best = (name, module) + if best is None: + sys.exit("could not find the decoder layer list in the HF model") + if n_hint is not None and len(best[1]) != n_hint: + print(f"warning: HF has {len(best[1])} layers, server reports {n_hint}") + return best[1] + + +def load_model(path: str, dtype): + import transformers + + last: Exception | None = None + + for cls_name in ("AutoModelForCausalLM", "AutoModelForImageTextToText", "AutoModel"): + cls = getattr(transformers, cls_name, None) + if cls is None: + continue + try: + return cls.from_pretrained(path, torch_dtype=dtype) + except (ValueError, KeyError, OSError) as exc: + last = exc + sys.exit(f"could not load {path}: {last}") + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--url", default="http://127.0.0.1:8080") + ap.add_argument("--model", default="qwen35-2b", help="server model block name") + ap.add_argument("--hf", required=True, help="HF checkpoint directory or repo id") + ap.add_argument("--layers", type=int, nargs="+", default=[0, 11, -1]) + ap.add_argument("--prompt", action="append", help="user message (repeatable)") + ap.add_argument("--bare", action="store_true", help='generation_prompt="bare"') + ap.add_argument("--thinking", action="store_true", help="enable_thinking=true") + ap.add_argument("--threshold", type=float, default=0.99) + ap.add_argument("--device", default="cpu") + ap.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"]) + args = ap.parse_args() + + import torch + from transformers import AutoTokenizer + + dtype = getattr(torch, args.dtype) + hf_path = os.path.expanduser(args.hf) + tok = AutoTokenizer.from_pretrained(hf_path) + model = load_model(hf_path, dtype).to(args.device).eval() + + worst = 1.0 + for prompt in args.prompt or DEFAULT_PROMPTS: + messages = [{"role": "user", "content": prompt}] + resp = post(args.url, { + "model": args.model, + "messages": messages, + "chat_template_kwargs": {"enable_thinking": args.thinking}, + "layers": args.layers, + "pooling": "both", + "generation_prompt": "bare" if args.bare else "template", + }) + n_layers = resp["n_layers"] + + text = tok.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True, + enable_thinking=args.thinking) + if args.bare: + text = text[: text.rindex(ASSISTANT_HEADER) + len(ASSISTANT_HEADER)] + ids = tok(text, return_tensors="pt", add_special_tokens=False).input_ids + if ids.shape[1] != resp["prompt_tokens"]: + print(f"warning: HF prompt has {ids.shape[1]} tokens, server " + f"{resp['prompt_tokens']}; templates differ, compare with care") + hf_last_tok = tok.decode(ids[0, -1:]) + if hf_last_tok != resp["last_token"]: + print(f"warning: last token HF={hf_last_tok!r} server={resp['last_token']!r}") + + blocks = find_decoder_layers(model, n_layers) + captured: dict[int, torch.Tensor] = {} + hooks = [] + for requested in args.layers: + block = requested % len(blocks) + + def hook(_mod, _inp, out, block=block): + h = out[0] if isinstance(out, (tuple, list)) else out + captured[block] = h.detach().float()[0] + + hooks.append(blocks[block].register_forward_hook(hook)) + with torch.no_grad(): + model(input_ids=ids.to(args.device)) + for h in hooks: + h.remove() + + print(f"\nprompt={prompt[:48]!r} tokens={resp['prompt_tokens']} " + f"last_token={resp['last_token']!r} prefill_ms={resp['timings']['prefill_ms']:.1f}") + for requested in args.layers: + block = requested % len(blocks) + ref = captured[block] + srv = resp["layers"][str(requested)] + c_last = cosine(srv["last"], ref[-1].tolist()) + c_mean = cosine(srv["mean"], ref.mean(dim=0).tolist()) + # A zero or non-finite vector has no cosine (nan); min() would skip it. + worst = min(worst, *(c if math.isfinite(c) else -math.inf for c in (c_last, c_mean))) + print(f" layer {requested:>4} (block {block:>2}): " + f"cos(last)={c_last:.5f} cos(mean)={c_mean:.5f}") + + ok = worst > args.threshold + print(f"\nworst cosine {worst:.5f} -> {'PASS' if ok else 'FAIL'} (threshold {args.threshold})") + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 694d33cff5..fc6b2e02ee 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -226,7 +226,7 @@ std::unique_ptr construct_backend( auto backend = std::make_unique(std::move(cfg)); if (!backend->init()) { - std::fprintf(stderr, "[backend_factory] Qwen35Backend init failed\n"); + std::fprintf(stderr, "[backend_factory] Qwen35Backend init failed: %s\n", luce_last_error()); return nullptr; } return backend; diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index f58de81d4b..547347dcf4 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -51,11 +51,14 @@ #include #include #include +#include #include #include +#include "common/hidden_states.h" #include "common/image_prompt.h" #include "common/sampler.h" +#include "common/token_logprobs.h" #include "prefix_store.h" namespace luce::common { @@ -205,6 +208,25 @@ class SeqEngine { return result; } + // Hidden-states requests (/v1/hidden_states): prefill only. The slot + // never decodes: its completed PrefillOutput carries hidden_states and + // the scheduler retires it. The prompt is always prefilled from position + // zero: admission never restores a checkpoint. A valid `capture` + // (cache_prefix; only from an engine reporting supports_prefix_store()) + // is armed like admit_with_prefix's and reported in + // AdmitResult::prefix_store.capture when accepted. + virtual bool supports_hidden_states() const { return false; } + virtual AdmitResult admit_hidden_states(uint64_t request_id, + const std::vector & prompt, + const HiddenStatesSpec & spec, + PrefixCaptureTicket capture) { + (void) request_id; (void) prompt; (void) spec; (void) capture; + AdmitResult result; + result.status = AdmitResult::Status::failed; + result.error = "this engine does not report hidden states"; + return result; + } + // Optional prefix-checkpoint admission. Unsupported engines remain on // cold admission and never receive a plan from the scheduler. virtual bool supports_prefix_store() const { return false; } @@ -222,6 +244,12 @@ class SeqEngine { // Release engine-owned payload without touching server policy metadata. virtual void discard_prefix_store(PrefixStoreRef) {} + // Re-read the last step's logits for a scheduler-substituted token. + // Called on the worker before the next step or retirement of this slot. + virtual std::optional token_logprobs(int, int32_t) { + return std::nullopt; + } + struct StepInput { int slot = -1; int32_t token = -1; // token to commit at this slot's next position @@ -239,6 +267,9 @@ class SeqEngine { // per-request error instead of silently truncating generation. std::string error; std::vector committed_tokens; + // Log-probabilities of `token` when the slot's sampler asked for them + // (such slots never speculate, so committed_tokens stays empty). + std::optional logprobs; }; struct PrefillOutput { @@ -253,6 +284,11 @@ class SeqEngine { // Present only for completed: the request's first sampled token, // pending until the scheduler feeds it into the next decode step. int32_t token = -1; + // Log-probabilities of `token`, as for DecodeOutput. + std::optional logprobs; + // Present only for completed hidden-states slots (see + // admit_hidden_states); `token` is then meaningless. + std::optional hidden_states; // Present only for failed. std::string error; // A capture ending on this successfully-computed prefill boundary. diff --git a/server/src/common/dflash2_batch.cpp b/server/src/common/dflash2_batch.cpp index 516ec47744..6fe4c53cc9 100644 --- a/server/src/common/dflash2_batch.cpp +++ b/server/src/common/dflash2_batch.cpp @@ -230,44 +230,55 @@ bool dflash2_select_chains_batched( const DraftSelectorWeights & selector = dw.selector; const int n_lanes = static_cast(hidden_by_lane.size()); const int n_cand = q_len - 1; - const int K = selector.top_k; - const int rank = selector.rank; + // A DFlash 1 drafter has no selector: its chain is the per-position + // argmax of the shared lm_head, i.e. the top-1 candidate. + const bool greedy = !selector.enabled; + const int K = greedy ? 1 : selector.top_k; + const int rank = greedy ? 1 : selector.rank; const int hdim = dw.n_embd; - if (!selector.enabled || !selector.hproj || !selector.pred_cb || + if (greedy) { + if (!backend || !lm_head || n_lanes <= 0 || + static_cast(last_tokens.size()) != n_lanes || + n_cand <= 0 || hdim <= 0) { + return false; + } + } else if (!selector.hproj || !selector.pred_cb || !selector.succ_cb || !backend || !lm_head || n_lanes <= 0 || static_cast(last_tokens.size()) != n_lanes || n_cand <= 0 || K <= 0 || rank <= 0 || hdim <= 0) { return false; } - DFlash2SelectorLayout selector_layout; - selector_layout.rank = rank; - selector_layout.top_k = K; - selector_layout.hproj_rank = selector.hproj->ne[1]; - selector_layout.pred_rank = selector.pred_cb->ne[0]; - selector_layout.pred_vocab = selector.pred_cb->ne[1]; - selector_layout.succ_rank = selector.succ_cb->ne[0]; - selector_layout.succ_vocab = selector.succ_cb->ne[1]; - selector_layout.target_output_vocab = lm_head->ne[1]; - std::string selector_error; - if (!validate_dflash2_selector_layout( - selector_layout, selector_error)) { - std::fprintf(stderr, "dflash2_select_chains_batched: %s\n", - selector_error.c_str()); - return false; - } for (const float * hidden : hidden_by_lane) { if (!hidden) return false; } - for (size_t lane = 0; lane < last_tokens.size(); ++lane) { - const int32_t token = last_tokens[lane]; - if (token < 0 || token >= selector_layout.pred_vocab) { - std::fprintf(stderr, - "dflash2_select_chains_batched: lane %zu seed token " - "%d is outside codebook vocab %lld\n", - lane, token, - (long long)selector_layout.pred_vocab); + if (!greedy) { + DFlash2SelectorLayout selector_layout; + selector_layout.rank = rank; + selector_layout.top_k = K; + selector_layout.hproj_rank = selector.hproj->ne[1]; + selector_layout.pred_rank = selector.pred_cb->ne[0]; + selector_layout.pred_vocab = selector.pred_cb->ne[1]; + selector_layout.succ_rank = selector.succ_cb->ne[0]; + selector_layout.succ_vocab = selector.succ_cb->ne[1]; + selector_layout.target_output_vocab = lm_head->ne[1]; + std::string selector_error; + if (!validate_dflash2_selector_layout( + selector_layout, selector_error)) { + std::fprintf(stderr, "dflash2_select_chains_batched: %s\n", + selector_error.c_str()); return false; } + for (size_t lane = 0; lane < last_tokens.size(); ++lane) { + const int32_t token = last_tokens[lane]; + if (token < 0 || token >= selector_layout.pred_vocab) { + std::fprintf(stderr, + "dflash2_select_chains_batched: lane %zu seed token " + "%d is outside codebook vocab %lld\n", + lane, token, + (long long)selector_layout.pred_vocab); + return false; + } + } } const int n_positions = n_lanes * n_cand; @@ -324,6 +335,20 @@ bool dflash2_select_chains_batched( candidate_log_probs.data(), candidate_ids.data(), 1.0f); } + if (greedy) { + draft_tokens.assign( + (size_t) n_lanes, + std::vector((size_t) q_len)); + for (int lane = 0; lane < n_lanes; ++lane) { + draft_tokens[(size_t) lane][0] = last_tokens[(size_t) lane]; + for (int depth = 0; depth < n_cand; ++depth) { + draft_tokens[(size_t) lane][(size_t) depth + 1] = + candidate_ids[(size_t) (lane * n_cand + depth)]; + } + } + return true; + } + BatchedSelectorGraph & graph = batched_selector_graph(); if (!ensure_selector_graph( graph, dw, backend, n_lanes, n_cand, K)) { @@ -409,4 +434,9 @@ bool dflash2_select_chains_batched( return true; } +void dflash2_batch_release_thread_graphs() { + free_projection_graph(projection_graph()); + free_selector_graph(batched_selector_graph()); +} + } // namespace luce::common diff --git a/server/src/common/dflash2_head.cpp b/server/src/common/dflash2_head.cpp index 2250155afc..e3a2498b2c 100644 --- a/server/src/common/dflash2_head.cpp +++ b/server/src/common/dflash2_head.cpp @@ -69,6 +69,10 @@ uint64_t dflash2_selector_graph_generation() { return s_selector_generation.load(std::memory_order_relaxed); } +void dflash2_selector_release_thread_graph() { + selector_graph_free(selector_graph()); +} + bool dflash2_score_candidates(const DraftWeights & dw, ggml_backend_t backend, DFlashTarget & target, diff --git a/server/src/common/dflash2_head.h b/server/src/common/dflash2_head.h index c1976a7a06..7286b61dcb 100644 --- a/server/src/common/dflash2_head.h +++ b/server/src/common/dflash2_head.h @@ -16,6 +16,17 @@ namespace luce::common { void dflash2_selector_graph_invalidate(); uint64_t dflash2_selector_graph_generation(); +// Free the calling thread's cached selector, projection and batched-selector +// graphs together with their device buffers. Invalidation only frees lazily on +// the next call; model eviction needs the memory back now, so it calls this on +// the worker thread before releasing the model's backends. +void dflash2_selector_release_thread_graph(); +void dflash2_batch_release_thread_graphs(); +inline void dflash2_release_thread_graphs() { + dflash2_selector_release_thread_graph(); + dflash2_batch_release_thread_graphs(); +} + // DFlash 2 candidate selector for greedy chain drafting. // // For every drafted block position the target lm_head logits are reduced to @@ -37,7 +48,8 @@ bool dflash2_select_chain(const DraftWeights & dw, // Same selector, batched over host-resident drafter hidden blocks and using a // local target lm_head tensor. The expensive lm_head projection covers every // (lane, depth) in one graph, GPU top-K is invoked once, and selector -// projections/readback are shared across the cohort. +// projections/readback are shared across the cohort. A drafter without a +// selector (DFlash 1) gets the per-position lm_head argmax chain instead. bool dflash2_select_chains_batched( const DraftWeights & dw, ggml_backend_t backend, diff --git a/server/src/common/dflash_capture.cpp b/server/src/common/dflash_capture.cpp index 59b2b8f29a..ad4ebe654a 100644 --- a/server/src/common/dflash_capture.cpp +++ b/server/src/common/dflash_capture.cpp @@ -12,4 +12,19 @@ int target_capture_index(const int * capture_layer_ids, return -1; } +bool adopt_drafter_capture_layers(const std::vector & drafter_ids, + int target_n_layer, + int max_slots, + int * capture_layer_ids, + int & n_capture_layers) { + const int n = (int)drafter_ids.size(); + if (n <= 0 || n > max_slots || !capture_layer_ids) return false; + for (int id : drafter_ids) { + if (id < 0 || id >= target_n_layer) return false; + } + for (int k = 0; k < n; k++) capture_layer_ids[k] = drafter_ids[(size_t)k]; + n_capture_layers = n; + return true; +} + } // namespace luce::common diff --git a/server/src/common/dflash_capture.h b/server/src/common/dflash_capture.h index 553a01f3d7..424e8995bb 100644 --- a/server/src/common/dflash_capture.h +++ b/server/src/common/dflash_capture.h @@ -7,6 +7,8 @@ #pragma once +#include + namespace luce::common { // Linear search for layer_idx in capture_layer_ids[0..n_capture_layers). @@ -15,4 +17,16 @@ int target_capture_index(const int * capture_layer_ids, int n_capture_layers, int layer_idx); +// Adopt a drafter's explicit capture layers (GGUF dflash.target_layer_ids) +// into a target capture table of `max_slots` entries. The count follows the +// drafter (its fc width), not the 5-layer default: capturing different layers +// than the drafter was trained on silently destroys acceptance. Returns false +// and leaves the table untouched when the list is empty, longer than the +// table, or names a layer outside [0, target_n_layer). +bool adopt_drafter_capture_layers(const std::vector & drafter_ids, + int target_n_layer, + int max_slots, + int * capture_layer_ids, + int & n_capture_layers); + } // namespace luce::common diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp index 6bb68b164c..06d6174dc0 100644 --- a/server/src/common/dspark_head.cpp +++ b/server/src/common/dspark_head.cpp @@ -13,6 +13,10 @@ namespace luce::common { namespace { +// Worker-local graph arenas of the head paths below; released by +// dspark_head_release_thread_scratch(). +thread_local ggml_gallocr_t g_dspark_step_alloc = nullptr; + bool dspark_step(const DraftWeights & dw, ggml_backend_t backend, int32_t prev_token, @@ -76,7 +80,7 @@ bool dspark_step(const DraftWeights & dw, ggml_build_forward_expand(gf, conf); } - static thread_local ggml_gallocr_t galloc = nullptr; + ggml_gallocr_t & galloc = g_dspark_step_alloc; if (!galloc) { galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); } @@ -246,6 +250,9 @@ struct MarkovChainGraphCache { } }; +thread_local MarkovChainGraphCache g_dspark_chain_cache; +thread_local ggml_gallocr_t g_dspark_topk_alloc = nullptr; + bool dspark_chain_graph_cache_disabled() { static const bool disabled = [] { const char * value = std::getenv("LUCE_DSPARK_NO_CHAIN_GRAPH_CACHE"); @@ -408,7 +415,7 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, const int hdim = dw.n_embd; const int n_cand = q_len - 1; - static thread_local MarkovChainGraphCache cache; + MarkovChainGraphCache & cache = g_dspark_chain_cache; const bool want_confidence = confidence_out != nullptr; const bool want_logit_margin = logit_margin_out != nullptr; if (confidence_out) confidence_out->clear(); @@ -532,7 +539,7 @@ bool dspark_markov_project_topk(const DraftWeights & dw, return false; } - static thread_local ggml_gallocr_t galloc_topk = nullptr; + ggml_gallocr_t & galloc_topk = g_dspark_topk_alloc; if (!galloc_topk) { galloc_topk = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); } @@ -575,4 +582,17 @@ bool dspark_markov_project_topk(const DraftWeights & dw, return true; } +void dspark_head_release_thread_scratch() { + for (ggml_gallocr_t * alloc : {&g_dspark_step_alloc, &g_dspark_topk_alloc}) { + if (*alloc) { + ggml_gallocr_free(*alloc); + *alloc = nullptr; + } + } + MarkovChainGraphCache & cache = g_dspark_chain_cache; + cache.invalidate(); + if (cache.allocator) ggml_gallocr_free(cache.allocator); + cache = MarkovChainGraphCache{}; +} + } // namespace luce::common diff --git a/server/src/common/dspark_head.h b/server/src/common/dspark_head.h index c407800dea..573833c0c3 100644 --- a/server/src/common/dspark_head.h +++ b/server/src/common/dspark_head.h @@ -17,6 +17,11 @@ uint64_t dspark_drafter_generation(); // Number of fused chain graphs built so far (diagnostic; lets tests assert a // cache hit or miss without reaching into the cache). uint64_t dspark_chain_graph_build_count(); +// Frees the calling thread's head graph arenas (step, fused chain and top-k +// paths), which hold device memory on the backend they last ran on. Call on +// the worker that ran the head, after its device work is synchronized, when +// that device's memory must be returned (eviction). +void dspark_head_release_thread_scratch(); bool dspark_markov_correct_greedy_chain(const DraftWeights & dw, ggml_backend_t backend, diff --git a/server/src/common/generation_types.h b/server/src/common/generation_types.h index a526b79c4c..ecd0b866a4 100644 --- a/server/src/common/generation_types.h +++ b/server/src/common/generation_types.h @@ -11,6 +11,7 @@ #include "image_prompt.h" #include "sampler.h" #include "thinking_budget.h" +#include "token_logprobs.h" namespace luce::common { @@ -99,6 +100,9 @@ struct GenerateResult { // A producer must explicitly call succeed() before returning success. std::optional error = GenerateError{}; std::vector tokens; + // One entry per token when sampler.wants_logprobs() and the backend + // supports it; empty otherwise. + std::vector logprobs; double prefill_s = 0.0; double decode_s = 0.0; // Prompt tokens confirmed by the backend's physical snapshot restore. diff --git a/server/src/common/geometric_draft_topk_cuda.cu b/server/src/common/geometric_draft_topk_cuda.cu index 0e0f896297..722087d205 100644 --- a/server/src/common/geometric_draft_topk_cuda.cu +++ b/server/src/common/geometric_draft_topk_cuda.cu @@ -429,4 +429,8 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, return ok; } +void geometric_draft_topk_release_thread_scratch() { + scratch.reset(); +} + } // namespace luce::common diff --git a/server/src/common/geometric_draft_topk_cuda.h b/server/src/common/geometric_draft_topk_cuda.h index d3ea8819a6..2902251b91 100644 --- a/server/src/common/geometric_draft_topk_cuda.h +++ b/server/src/common/geometric_draft_topk_cuda.h @@ -38,4 +38,8 @@ bool geometric_extract_draft_topk_cuda(const void * d_logits, int32_t * out_token_ids, float temperature); +// Free the calling thread's device scratch; see +// geometric_sampler_release_thread_scratch(). +void geometric_draft_topk_release_thread_scratch(); + } // namespace luce::common diff --git a/server/src/common/geometric_sampler_cuda.cu b/server/src/common/geometric_sampler_cuda.cu index 9d7b36446c..9f376ec757 100644 --- a/server/src/common/geometric_sampler_cuda.cu +++ b/server/src/common/geometric_sampler_cuda.cu @@ -571,4 +571,8 @@ bool geometric_compute_probs_cuda(const float * logits, return ok; } +void geometric_sampler_release_thread_scratch() { + scratch.reset(); +} + } // namespace luce::common diff --git a/server/src/common/geometric_sampler_cuda.h b/server/src/common/geometric_sampler_cuda.h index 667d60a901..5f5c2ef62e 100644 --- a/server/src/common/geometric_sampler_cuda.h +++ b/server/src/common/geometric_sampler_cuda.h @@ -77,6 +77,11 @@ bool geometric_compute_probs_cuda(const float * logits, // gate the GPU path at runtime. bool gpu_sampler_enabled(); +// Free the calling thread's device scratch (raw device allocations reused +// across calls). Model eviction calls it on the worker before the device is +// handed to another model; the next sample reallocates it. +void geometric_sampler_release_thread_scratch(); + // Whether geometric_sample_logits_cuda can (and should) handle this config // entirely on the GPU. false for top_k>0 or top_p in (0,1) — not because // geometric_sample_logits_cuda can't be asked to try (it always returns -1 diff --git a/server/src/common/hidden_states.h b/server/src/common/hidden_states.h new file mode 100644 index 0000000000..ffd89ec598 --- /dev/null +++ b/server/src/common/hidden_states.h @@ -0,0 +1,85 @@ +// Residual-stream activations read out of a prefill (/v1/hidden_states). +// +// A hidden-states request runs the prompt through prefill only — no decode, +// no sampling — and reports, for each requested block L, the residual stream +// after block L (before the final norm) at the last prompt token and, +// optionally, averaged over every prompt token. Prefill is chunked, so the +// backend hands each chunk's rows to HiddenStatesAccumulator, which keeps the +// running per-layer sum and the latest last row. + +#pragma once + +#include +#include + +namespace luce::common { + +struct HiddenStatesSpec { + // Block indices in [0, n_layers), in response order. Entries may repeat + // when two requested spellings (e.g. 23 and -1) name the same block. + std::vector layers; + // Also report the mean over all prompt tokens. + bool mean = false; +}; + +struct HiddenStates { + // One n_embd vector per HiddenStatesSpec::layers entry, in the same order. + std::vector> last; + // Parallel to `last`; empty unless HiddenStatesSpec::mean. + std::vector> mean; + // Prompt tokens the activations were read over. + int n_tokens = 0; +}; + +class HiddenStatesAccumulator { +public: + HiddenStatesAccumulator() = default; + HiddenStatesAccumulator(HiddenStatesSpec spec, int n_embd) + : spec_(std::move(spec)), n_embd_(n_embd), + last_(spec_.layers.size(), std::vector((size_t)n_embd, 0.0f)), + sum_(spec_.mean ? spec_.layers.size() : 0, + std::vector((size_t)n_embd, 0.0)) {} + + const HiddenStatesSpec & spec() const { return spec_; } + int n_embd() const { return n_embd_; } + + // Record layer entry `k` of one prefill chunk: its last row and, when the + // spec asks for the mean, the sum of its rows. Call add_rows() once per + // chunk after every entry is recorded. + void record(size_t k, const float * last_row, const float * row_sum) { + last_[k].assign(last_row, last_row + n_embd_); + if (!spec_.mean || !row_sum) return; + std::vector & sum = sum_[k]; + for (int i = 0; i < n_embd_; ++i) sum[(size_t)i] += row_sum[i]; + } + void add_rows(int n_rows) { rows_ += n_rows; } + int rows() const { return rows_; } + + HiddenStates finish() const { + HiddenStates out; + out.n_tokens = rows_; + out.last = last_; + if (spec_.mean) { + out.mean.reserve(sum_.size()); + for (const std::vector & sum : sum_) { + std::vector mean(sum.size(), 0.0f); + if (rows_ > 0) { + for (size_t i = 0; i < sum.size(); ++i) { + mean[i] = (float)(sum[i] / (double)rows_); + } + } + out.mean.push_back(std::move(mean)); + } + } + return out; + } + +private: + HiddenStatesSpec spec_; + int n_embd_ = 0; + int rows_ = 0; + std::vector> last_; + std::vector> sum_; +}; + +} // namespace luce::common diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 9bcc9df271..1819d61dc6 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -23,6 +23,7 @@ #include "ggml.h" #include "ggml-backend.h" #include "generation_types.h" +#include "hidden_states.h" #include "sampler.h" #include "image_prompt.h" #include "concurrency/seq_engine.h" @@ -169,6 +170,55 @@ struct ModelBackend { virtual bool unpark(ParkTarget target) = 0; virtual bool is_target_parked() const = 0; + // ── Eviction (runtime model swapping) ──────────────────────────── + // evict() releases every device allocation the model owns (weights, + // drafts, KV/recurrent caches, graphs, scratch and its backend contexts) + // so another model can use the device. Host-side prefix snapshots stay + // valid; reinstate() rebuilds the device layout that existed at init + // (same context, KV types and sizes), so they restore afterwards. + // Both run on the model's worker thread while it is idle. Unlike park(), + // which a backend uses for its own transient memory juggling, eviction + // releases the working state; the level controls which weights remain. + // cold: release everything. warm: keep the weights in device memory and + // release the rest (KV/recurrent cache, graphs, scratch, backend context + // with its pools), so a later reinstate only rebuilds the cache. + // partial: release the primary device only, keeping a secondary device's + // expert tier (DeepSeek V4 with in-process experts; others reject it). + enum class EvictLevel { cold, warm, partial }; + enum class EvictStatus { + ok, // device state released (or already evicted) + rejected, // refused before any teardown; the model is unchanged + }; + struct EvictResult { + EvictStatus status = EvictStatus::rejected; + std::string error; + // Deferred snapshots that could not be copied to host and were freed; + // their prefix-cache entries no longer restore. + int snapshots_lost = 0; + }; + // Whether this configuration supports eviction; `reason` explains a no. + virtual bool supports_eviction(std::string & reason) const { + reason = "this backend does not support eviction"; + return false; + } + virtual EvictResult evict(EvictLevel level = EvictLevel::cold) { + (void) level; + EvictResult r; + supports_eviction(r.error); + return r; + } + // On failure the model is left evicted with its partial device state + // released, so a retry or another model can use the device. + virtual bool reinstate(std::string & error) { + supports_eviction(error); + return false; + } + virtual bool is_evicted() const { return false; } + // Evicted warm: the weights still occupy device memory. + virtual bool weights_resident() const { return !is_evicted(); } + // Device bytes of the weights (target + draft) while they are resident. + virtual size_t weight_device_bytes() const { return 0; } + // ── Generation ─────────────────────────────────────────────────── // Run a full prefill + decode cycle. Backend owns the strategy // (autoregressive, speculative, DDTree, …). @@ -194,6 +244,39 @@ struct ModelBackend { virtual GenerateResult generate_impl(const GenerateRequest & req, const DaemonIO & io) = 0; + // Whether generate() (with force_ar_decode) and seq_engine() fill + // per-token log-probabilities when sampler.wants_logprobs(). The server + // ignores logprobs for backends that do not support them. + virtual bool supports_logprobs() const { return false; } + + // ── Hidden states (/v1/hidden_states) ──────────────────────────── + // Whether this backend can report the residual stream after chosen + // blocks for a prompt: through compute_hidden_states() when it serves + // one request at a time, and through SeqEngine::admit_hidden_states() + // when seq_engine() is non-null. The server rejects the request + // otherwise. + virtual bool supports_hidden_states() const { return false; } + struct HiddenStateShape { + int n_layers = 0; // blocks a request may index + int n_embd = 0; // floats per reported vector + }; + virtual HiddenStateShape hidden_state_shape() const { return {}; } + // Prefill `prompt` from position zero -- never from a prefix snapshot -- + // and read `spec`'s layers. No token is decoded. When `snap_slot` >= 0 + // the state after the whole prompt is also saved into that snapshot + // slot (an end-of-prefill save, which leaves the prefill itself + // unchanged). Serves the single-request worker only. + virtual bool compute_hidden_states(const std::vector & prompt, + const HiddenStatesSpec & spec, + int snap_slot, + const DaemonIO & io, + HiddenStates & out, + std::string & error) { + (void) prompt; (void) spec; (void) snap_slot; (void) io; (void) out; + error = "this model backend does not report hidden states"; + return false; + } + // ── Concurrent serving ─────────────────────────────────────────── // Backends that can hold several live sequences at once and execute a // batched decode over paged KV expose them as decode slots through a diff --git a/server/src/common/moe_hybrid_expert_cache.cpp b/server/src/common/moe_hybrid_expert_cache.cpp index a011e9fbba..e6c0760c76 100644 --- a/server/src/common/moe_hybrid_expert_cache.cpp +++ b/server/src/common/moe_hybrid_expert_cache.cpp @@ -314,6 +314,49 @@ void MoeStreamedExpertCache::destroy() { tick_ = 0; stats_ = {}; stopping_ = false; + suspended_ = false; + loads_in_flight_ = 0; +} + +bool MoeStreamedExpertCache::suspend(std::string * err) { + if (!ready()) return true; + const auto refuse = [&](const char * why) { + if (err) *err = why; + return false; + }; + // Launches and calls wait on the loaders: they must end before the + // loaders stop. + if (in_call_.load(std::memory_order_acquire)) return refuse("an expert-cache call is running"); + if (mailbox_.busy()) return refuse("a mailbox launch has not ended"); + std::unique_lock lk(mu_); + if (!staged_.empty() || !acquired_.empty()) return refuse("expert-cache slots are still pinned"); + for (const Slot & s : slots_) { + if (s.pins > 0) return refuse("expert-cache slots are still pinned"); + } + suspended_ = true; + // The loads already taken finish (load_slot synchronizes its stream), so + // no copy into the pool is left in flight once this returns. + cv_.wait(lk, [&] { return loads_in_flight_ == 0; }); + return true; +} + +void MoeStreamedExpertCache::resume() { + { + std::lock_guard lk(mu_); + if (!suspended_) return; + suspended_ = false; + } + cv_.notify_all(); +} + +void MoeStreamedExpertCache::release_graphs() { + for (auto & kv : graphs_) kv.second.free(); + graphs_.clear(); +} + +bool MoeStreamedExpertCache::suspended() const { + std::lock_guard lk(mu_); + return suspended_; } MoeStreamedExpertCache::Stats MoeStreamedExpertCache::stats() const { @@ -334,7 +377,8 @@ void MoeStreamedExpertCache::loader_main(Loader * self) { std::unique_lock lk(mu_); while (true) { cv_.wait(lk, [&] { - return stopping_ || !jobs_.empty() || (warm_next_ < warm_.size() && !warm_blocked_); + return stopping_ || + (!suspended_ && (!jobs_.empty() || (warm_next_ < warm_.size() && !warm_blocked_))); }); if (stopping_) break; // Demand and prefetch loads first; the warm start only fills idle time. @@ -360,10 +404,12 @@ void MoeStreamedExpertCache::loader_main(Loader * self) { } // Read under the lock: stage() may retouch a loading slot meanwhile. const bool bulk = slots_[(size_t) slot].bulk; + ++loads_in_flight_; lk.unlock(); uint64_t read_us = 0, upload_us = 0; const bool ok = load_slot(*self, slot, bulk, &read_us, &upload_us); lk.lock(); + --loads_in_flight_; slots_[(size_t) slot].state = SlotState::Ready; if (warm) { --warm_loading_; @@ -1031,7 +1077,7 @@ void MoeStreamedMailbox::destroy() { step_ = nullptr; channels_.clear(); jobs_.clear(); - pending_ = launch_done_ = stopping_ = false; + pending_ = launch_done_ = open_ = stopping_ = false; error_.clear(); cache_ = nullptr; } @@ -1050,14 +1096,21 @@ void MoeStreamedMailbox::begin(const std::vector & jobs, int32_t invalid_ro invalid_route_ = invalid_route; error_.clear(); launch_done_ = false; + open_ = true; pending_ = true; cv_.notify_all(); } +bool MoeStreamedMailbox::busy() const { + std::lock_guard lk(mu_); + return open_; +} + bool MoeStreamedMailbox::end(std::string * err) { std::unique_lock lk(mu_); launch_done_ = true; cv_.wait(lk, [&] { return !pending_; }); + open_ = false; const bool ok = error_.empty(); if (!ok && err) *err = error_; lk.unlock(); diff --git a/server/src/common/moe_hybrid_expert_cache.h b/server/src/common/moe_hybrid_expert_cache.h index d841a28a4b..9b558a43a9 100644 --- a/server/src/common/moe_hybrid_expert_cache.h +++ b/server/src/common/moe_hybrid_expert_cache.h @@ -116,6 +116,8 @@ class MoeStreamedMailbox { // releases the launch's slots. False with *err when a layer could not be // answered exactly (a load failed or its experts exceed the pool). bool end(std::string * err); + // A launch was begun and not ended yet. + bool busy() const; private: void resolver_main(); @@ -129,13 +131,14 @@ class MoeStreamedMailbox { std::vector channels_; std::vector slot_of_; - std::mutex mu_; + mutable std::mutex mu_; std::condition_variable cv_; std::thread thread_; std::vector jobs_; int32_t invalid_route_ = -1; bool pending_ = false; // a launch is being answered bool launch_done_ = false; // end() was called for it + bool open_ = false; // begin() was called, end() not yet bool stopping_ = false; std::string error_; }; @@ -192,6 +195,23 @@ class MoeStreamedExpertCache { // into empty or bulk slots). Loads only move bytes: outputs never change. void set_bulk(bool bulk); + // Pauses the background work, for a caller that is about to release + // device state the cache's users depend on (DS4 partial eviction): the + // loaders take no new job or warm load, and the loads already running + // finish their copies. The slots keep their experts, queued loads stay + // queued (their slots stay Loading) and the warm start keeps its place, so + // resume() continues exactly where suspend() stopped. Fails, changing + // nothing, while an eval/acquire runs, slots are pinned or a mailbox + // launch has not ended: finish those first, since they wait on the + // loaders. Idempotent; a cache that is not ready suspends trivially. + bool suspend(std::string * err = nullptr); + void resume(); + bool suspended() const; + // Drops the cached compute graphs (rebuilt on demand). They are built + // over the caller's layer descriptors, so a caller that reloads the + // tensors behind those drops them first. Only while no call runs. + void release_graphs(); + // Adds the weighted output of the streamed routes to out // ([n_embd, n_tokens], host). `selected` / `weights` are [n_used, // n_tokens]; routes the storage does not stream are skipped. @@ -339,6 +359,8 @@ class MoeStreamedExpertCache { std::chrono::steady_clock::time_point warm_t0_; uint64_t tick_ = 0; bool stopping_ = false; + bool suspended_ = false; // loaders take no work (suspend()) + int loads_in_flight_ = 0; // loads a loader has taken and not finished std::vector> predicted_; // per layer, until its eval std::vector staged_; // slots stage() pinned for staged_layer_ std::vector acquired_; // slots acquire() pinned diff --git a/server/src/common/moe_hybrid_storage.cpp b/server/src/common/moe_hybrid_storage.cpp index d46eac8cd0..bc2c08abe0 100644 --- a/server/src/common/moe_hybrid_storage.cpp +++ b/server/src/common/moe_hybrid_storage.cpp @@ -203,22 +203,196 @@ void MoeHybridStorage::release_graph_caches() { } } +void MoeHybridStorage::release_prefill_allocators() { + for (ggml_gallocr_t * alloc : {&prefill_route_alloc, &prefill_hot_alloc, &prefill_cold_alloc}) { + if (*alloc) { + ggml_gallocr_free(*alloc); + *alloc = nullptr; + } + } +} + +static bool is_mix_expert_type(const ggml_tensor * t) { + return t && (t->type == GGML_TYPE_Q3_1_ROCMFP3_MIX || t->type == GGML_TYPE_Q2_1_ROCMFP2_MIX); +} + +static void free_hot_stack(MoeHybridLayerStorage & layer) { + if (layer.hot_buf) { + ggml_backend_buffer_free(layer.hot_buf); + layer.hot_buf = nullptr; + } + if (layer.hot_ctx) { + ggml_free(layer.hot_ctx); + layer.hot_ctx = nullptr; + } + layer.gate_hot = nullptr; + layer.up_hot = nullptr; + layer.down_hot = nullptr; + layer.gate_up_hot = nullptr; +} + +bool MoeHybridStorage::release_primary_experts(std::string * err) { + if (primary_released_) return true; + const auto refuse = [&](const char * why) { + if (err) *err = why; + return false; + }; + if (!has_mmap() || layer_regions.size() != layers.size()) { + return refuse("primary experts can only be rebuilt from a retained file mapping"); + } + for (const MoeHybridLayerStorage & layer : layers) { + if (layer.cache_slots != 0) return refuse("primary spare cache slots cannot be rebuilt"); + for (const ggml_tensor * t : {layer.gate_hot, layer.up_hot, layer.down_hot, layer.gate_up_hot}) { + if (is_mix_expert_type(t)) return refuse("mixed-qtype primary experts cannot be rebuilt"); + } + } + // Graphs and arenas sized over (or pointing at) the hot tensors go first. + release_prefill_allocators(); + release_graph_caches(); + for (MoeHybridLayerStorage & layer : layers) { + ggml_tensor * const hot[4] = {layer.gate_hot, layer.up_hot, layer.down_hot, layer.gate_up_hot}; + for (int i = 0; i < 4; ++i) { + MoeHybridLayerStorage::StackShape & shape = layer.released_hot_shapes[i]; + shape = {}; + if (!hot[i]) continue; + shape.type = (int) hot[i]->type; + for (int d = 0; d < 4; ++d) shape.ne[d] = hot[i]->ne[d]; + } + free_hot_stack(layer); + // Not resident until the rebuild has uploaded every stack. + std::fill(layer.expert_vram_mask.begin(), layer.expert_vram_mask.end(), 0); + } + primary_released_ = true; + return true; +} + +bool MoeHybridStorage::rebuild_primary_experts(ggml_backend_t gpu_backend, const std::string & path, + std::string * err) { + if (!primary_released_) return true; + const auto fail = [&](const std::string & why) { + for (MoeHybridLayerStorage & layer : layers) free_hot_stack(layer); + if (err) *err = why; + return false; + }; + if (!gpu_backend) return fail("no primary backend to rebuild the expert stacks on"); + std::vector spans; + uint64_t bytes = 0; + for (size_t il = 0; il < layers.size(); ++il) { + MoeHybridLayerStorage & layer = layers[il]; + const LayerExpertRegions & regions = layer_regions[il]; + bool any = false; + for (const auto & shape : layer.released_hot_shapes) any = any || shape.type >= 0; + if (!any) continue; + // The maps the retained cold stack and decode tables index with must + // still describe exactly this ordered hot set. + for (size_t i = 0; i < layer.hot_expert_ids.size(); ++i) { + const int32_t expert = layer.hot_expert_ids[i]; + if (expert < 0 || (size_t) expert >= layer.hot_local_by_global.size() || + layer.hot_local_by_global[(size_t) expert] != (int32_t) i) { + return fail("layer " + std::to_string(il) + ": hot expert map no longer matches its ordered ids"); + } + } + ggml_init_params ip{}; + ip.mem_size = 24 * ggml_tensor_overhead(); + ip.no_alloc = true; + layer.hot_ctx = ggml_init(ip); + if (!layer.hot_ctx) return fail("failed to init hot_ctx"); + ggml_tensor ** const hot[4] = {&layer.gate_hot, &layer.up_hot, &layer.down_hot, &layer.gate_up_hot}; + for (int i = 0; i < 4; ++i) { + const auto & shape = layer.released_hot_shapes[i]; + if (shape.type < 0) continue; + if (shape.ne[2] != (int64_t) layer.hot_expert_ids.size()) { + return fail("layer " + std::to_string(il) + ": recorded hot stack does not hold its experts"); + } + *hot[i] = ggml_new_tensor(layer.hot_ctx, (ggml_type) shape.type, 4, shape.ne); + } + layer.hot_buf = ggml_backend_alloc_ctx_tensors(layer.hot_ctx, gpu_backend); + if (!layer.hot_buf) { + return fail("failed to allocate hot expert GPU buffer (layer " + std::to_string(il) + ")"); + } + const auto add = [&](ggml_tensor * stack, const ExpertFileRegion & region, size_t expert_bytes) { + if (!stack) return true; + for (size_t i = 0; i < layer.hot_expert_ids.size(); ++i) { + const size_t rel = expert_bytes * (size_t) layer.hot_expert_ids[i]; + if (expert_bytes == 0 || rel + expert_bytes > region.size) return false; + const uint64_t file_offset = (uint64_t) region.offset + rel; + if (!spans.empty() && spans.back().tensor == stack && + spans.back().tensor_offset + spans.back().size == expert_bytes * i && + spans.back().file_offset + spans.back().size == file_offset) { + spans.back().size += expert_bytes; + } else { + spans.push_back({stack, expert_bytes * i, file_offset, expert_bytes}); + } + bytes += expert_bytes; + } + return true; + }; + if (!add(layer.gate_hot, regions.gate_exps, layer.gate_expert_bytes) || + !add(layer.up_hot, regions.up_exps, layer.up_expert_bytes) || + !add(layer.down_hot, regions.down_exps, layer.down_expert_bytes) || + !add(layer.gate_up_hot, regions.gate_up_exps, layer.gate_up_expert_bytes)) { + return fail("layer " + std::to_string(il) + ": hot expert outside its file region"); + } + } + if (path.empty()) { + const auto * file = static_cast(mmap_data); + for (const TensorFileSpan & span : spans) { + if (span.file_offset + span.size > mmap_size) return fail("hot expert outside the mapping"); + ggml_backend_tensor_set(span.tensor, file + span.file_offset, span.tensor_offset, span.size); + } + } else { + std::string load_err; + if (!load_tensor_spans(path, mmap_data, mmap_size, spans, &load_err)) { + return fail("reading the primary experts: " + load_err); + } + } + for (MoeHybridLayerStorage & layer : layers) { + std::fill(layer.expert_vram_mask.begin(), layer.expert_vram_mask.end(), 0); + for (int32_t expert : layer.hot_expert_ids) layer.set_expert_hot(expert); + } + primary_upload_bytes += bytes; + primary_released_ = false; + return true; +} + +MoeHybridOwnershipRecord MoeHybridOwnershipRecord::capture(const MoeHybridStorage & storage) { + MoeHybridOwnershipRecord record; + record.layers.reserve(storage.layers.size()); + for (const MoeHybridLayerStorage & layer : storage.layers) { + record.layers.push_back({layer.hot_expert_ids, layer.cold_expert_ids, + layer.hot_local_by_global, layer.cold_local_by_global, + layer.decode_hot_local_by_global, layer.decode_cold_local_by_global, + layer.n_streamed}); + } + return record; +} + +bool MoeHybridOwnershipRecord::matches(const MoeHybridStorage & storage, std::string * err) const { + if (storage.layers.size() != layers.size()) { + if (err) *err = "expert ownership has a different layer count"; + return false; + } + const MoeHybridOwnershipRecord now = capture(storage); + for (size_t il = 0; il < layers.size(); ++il) { + if (!(now.layers[il] == layers[il])) { + if (err) *err = "expert ownership of layer " + std::to_string(il) + " changed"; + return false; + } + } + return true; +} + +void MoeHybridStorage::rebind_cold_backend(ggml_backend_t backend) { + if (cold_backend_kind != MoeHybridColdBackend::Gpu) return; + cold_backend = backend; + for (MoeHybridLayerStorage & layer : layers) layer.cold_backend = backend; +} + MoeHybridStorage::~MoeHybridStorage() { // Registry entries point into the owner buffers, so remove them before the // buffers can be released or their addresses reused. unregister_mix_tensors(); - if (prefill_route_alloc) { - ggml_gallocr_free(prefill_route_alloc); - prefill_route_alloc = nullptr; - } - if (prefill_hot_alloc) { - ggml_gallocr_free(prefill_hot_alloc); - prefill_hot_alloc = nullptr; - } - if (prefill_cold_alloc) { - ggml_gallocr_free(prefill_cold_alloc); - prefill_cold_alloc = nullptr; - } + release_prefill_allocators(); release_graph_caches(); for (auto & layer : layers) { if (layer.hot_buf) { @@ -738,6 +912,8 @@ bool build_moe_hybrid_storage_from_file( return false; } + out.primary_upload_bytes += (uint64_t) hot_count * + (dst.gate_expert_bytes + dst.up_expert_bytes + dst.down_expert_bytes + dst.gate_up_expert_bytes); if (hot_count > 0 && dst.fused_gate_up) { if (!copy_experts(dst.gate_up_hot, fd.gate_up_exps, dst.hot_expert_ids, dst.gate_up_expert_bytes)) return false; if (!copy_experts(dst.down_hot, fd.down_exps, dst.hot_expert_ids, dst.down_expert_bytes)) return false; @@ -777,6 +953,8 @@ bool build_moe_hybrid_storage_from_file( return false; } + out.secondary_upload_bytes += (uint64_t) cold_count * + (dst.gate_expert_bytes + dst.up_expert_bytes + dst.down_expert_bytes + dst.gate_up_expert_bytes); if (dst.fused_gate_up) { if (!copy_experts(dst.gate_up_cold, fd.gate_up_exps, dst.cold_expert_ids, dst.gate_up_expert_bytes)) return false; if (!copy_experts(dst.down_cold, fd.down_exps, dst.cold_expert_ids, dst.down_expert_bytes)) return false; diff --git a/server/src/common/moe_hybrid_storage.h b/server/src/common/moe_hybrid_storage.h index 8053113783..cedf02cc29 100644 --- a/server/src/common/moe_hybrid_storage.h +++ b/server/src/common/moe_hybrid_storage.h @@ -84,6 +84,15 @@ struct MoeHybridLayerStorage { ggml_tensor * down_hot = nullptr; ggml_tensor * gate_up_hot = nullptr; + // Shape of each hot tensor (gate, up, down, gate_up), recorded by + // MoeHybridStorage::release_primary_experts() so the stack is rebuilt + // with the same rows; type -1 for an absent tensor. + struct StackShape { + int type = -1; + int64_t ne[4] = {0, 0, 0, 0}; + }; + StackShape released_hot_shapes[4]; + ggml_context * cold_ctx = nullptr; ggml_backend_buffer_t cold_buf = nullptr; ggml_tensor * gate_cold = nullptr; @@ -264,9 +273,42 @@ struct MoeHybridStorage { // Per-layer file region metadata for streaming (populated when mmap is active). std::vector layer_regions; +private: + bool primary_released_ = false; +public: + + // Expert bytes uploaded into the primary (hot) and secondary (cold) + // stacks over this storage's life: the build, then every + // rebuild_primary_experts(). A partial eviction keeps the secondary count. + uint64_t primary_upload_bytes = 0; + uint64_t secondary_upload_bytes = 0; + // Remove decode-table registrations while their owner tensors are alive. // Safe to call repeatedly, including during failed partial registration. void unregister_mix_tensors(); + + // Primary-device release for a partial eviction. Frees the hot stacks + // (and the graph caches and prefill arenas built over them) while the + // expert maps, the cold stacks, the file mapping and the regions stay, so + // this object keeps its address for the streamed cache that points at it. + // Refused (false, nothing freed) for storage it cannot rebuild exactly: + // spare cache slots (a mutable primary population), mixed-qtype experts + // (their decode tables come from a sidecar) or no retained mapping. + bool release_primary_experts(std::string * err = nullptr); + bool primary_experts_released() const { return primary_released_; } + // Rebuilds the released hot stacks on `gpu_backend`: the recorded shapes, + // the same ordered hot_expert_ids (so every local index and map stays + // valid), read from `path` (the mapped file) or, with an empty path, + // copied out of the mapping. On failure the partial stacks are freed + // again and the storage stays released. + bool rebuild_primary_experts(ggml_backend_t gpu_backend, const std::string & path, + std::string * err = nullptr); + // The secondary owner is computed on a caller-owned backend; a caller + // that recreates that backend (its buffers belong to the device, not the + // backend) points the storage and its layers at the new one. + void rebind_cold_backend(ggml_backend_t backend); + // Frees the three prefill arenas (route, hot, cold). + void release_prefill_allocators(); bool matches(const MoeHybridConfig & cfg) const; bool empty() const; bool has_mmap() const { return mmap_data != nullptr && mmap_size > 0; } @@ -276,6 +318,35 @@ struct MoeHybridStorage { void release_graph_caches(); }; +// The expert ownership a storage resolved: per layer the ordered ids of both +// stacks (their order is each stack's row order), the physical and decode +// owner maps and the streamed count. Captured once a model is placed, it +// proves that a reinstated storage computes on exactly the same rows. +struct MoeHybridOwnershipRecord { + struct Layer { + std::vector hot_expert_ids; + std::vector cold_expert_ids; + std::vector hot_local_by_global; + std::vector cold_local_by_global; + std::vector decode_hot_local_by_global; + std::vector decode_cold_local_by_global; + int n_streamed = 0; + bool operator==(const Layer & o) const { + return hot_expert_ids == o.hot_expert_ids && cold_expert_ids == o.cold_expert_ids && + hot_local_by_global == o.hot_local_by_global && + cold_local_by_global == o.cold_local_by_global && + decode_hot_local_by_global == o.decode_hot_local_by_global && + decode_cold_local_by_global == o.decode_cold_local_by_global && + n_streamed == o.n_streamed; + } + }; + std::vector layers; + + static MoeHybridOwnershipRecord capture(const MoeHybridStorage & storage); + // False with the first differing layer in *err. + bool matches(const MoeHybridStorage & storage, std::string * err = nullptr) const; +}; + // Expert tensor file data for split loading (one entry per expert tensor). struct ExpertTensorFileData { const uint8_t * data = nullptr; diff --git a/server/src/common/moe_hybrid_stream.cpp b/server/src/common/moe_hybrid_stream.cpp index bb92c955ac..7041cb2e97 100644 --- a/server/src/common/moe_hybrid_stream.cpp +++ b/server/src/common/moe_hybrid_stream.cpp @@ -24,14 +24,14 @@ MoeHybridStreamEngine::~MoeHybridStreamEngine() { MoeHybridStreamEngine::MoeHybridStreamEngine(MoeHybridStreamEngine && o) noexcept : pinned_buf_(o.pinned_buf_), pinned_size_(o.pinned_size_), gpu_scratch_(o.gpu_scratch_), scratch_size_(o.scratch_size_), - backend_(o.backend_), + scratch_device_(o.scratch_device_), backend_(o.backend_), scratch_gate_(o.scratch_gate_), scratch_up_(o.scratch_up_), scratch_down_(o.scratch_down_), last_gate_bytes_(o.last_gate_bytes_), last_up_bytes_(o.last_up_bytes_), last_down_bytes_(o.last_down_bytes_), stats_(o.stats_) { o.stats_ = {}; o.pinned_buf_ = nullptr; o.pinned_size_ = 0; - o.gpu_scratch_ = nullptr; o.scratch_size_ = 0; + o.gpu_scratch_ = nullptr; o.scratch_size_ = 0; o.scratch_device_ = -1; o.backend_ = nullptr; o.scratch_gate_ = nullptr; o.scratch_up_ = nullptr; o.scratch_down_ = nullptr; o.last_gate_bytes_ = 0; o.last_up_bytes_ = 0; o.last_down_bytes_ = 0; @@ -42,6 +42,7 @@ MoeHybridStreamEngine & MoeHybridStreamEngine::operator=(MoeHybridStreamEngine & destroy(); pinned_buf_ = o.pinned_buf_; pinned_size_ = o.pinned_size_; gpu_scratch_ = o.gpu_scratch_; scratch_size_ = o.scratch_size_; + scratch_device_ = o.scratch_device_; backend_ = o.backend_; scratch_gate_ = o.scratch_gate_; scratch_up_ = o.scratch_up_; scratch_down_ = o.scratch_down_; @@ -49,7 +50,7 @@ MoeHybridStreamEngine & MoeHybridStreamEngine::operator=(MoeHybridStreamEngine & last_down_bytes_ = o.last_down_bytes_; stats_ = o.stats_; o.stats_ = {}; o.pinned_buf_ = nullptr; o.pinned_size_ = 0; - o.gpu_scratch_ = nullptr; o.scratch_size_ = 0; + o.gpu_scratch_ = nullptr; o.scratch_size_ = 0; o.scratch_device_ = -1; o.backend_ = nullptr; o.scratch_gate_ = nullptr; o.scratch_up_ = nullptr; o.scratch_down_ = nullptr; o.last_gate_bytes_ = 0; o.last_up_bytes_ = 0; o.last_down_bytes_ = 0; @@ -67,8 +68,28 @@ MoeHybridStreamEngine & MoeHybridStreamEngine::operator=(MoeHybridStreamEngine & static constexpr size_t kRegionPad = 4096; static constexpr size_t kRegions = 3; // gate, up, down (or gate_up, down) +namespace { + +// Makes `device` current for this scope (no-op for -1), restoring the +// caller's device afterwards: a raw allocation belongs to the current device. +struct ScopedDevice { + int previous = -1; + explicit ScopedDevice(int device) { + if (device < 0 || cudaGetDevice(&previous) != cudaSuccess) { + previous = -1; + return; + } + if (previous == device || cudaSetDevice(device) != cudaSuccess) previous = -1; + } + ~ScopedDevice() { + if (previous >= 0) (void) cudaSetDevice(previous); + } +}; + +} // namespace + bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, size_t max_expert_bytes, - std::string * err) { + std::string * err, int device) { destroy(); if (!gpu_backend || max_expert_bytes == 0) { if (err) *err = "invalid arguments to stream engine init"; @@ -87,7 +108,10 @@ bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, size_t max_expert_b pinned_size_ = alloc_bytes; // Allocate GPU scratch buffer - cuda_err = cudaMalloc(&gpu_scratch_, alloc_bytes); + { + const ScopedDevice scoped(device); + cuda_err = cudaMalloc(&gpu_scratch_, alloc_bytes); + } if (cuda_err != cudaSuccess) { if (err) *err = std::string("cudaMalloc scratch failed: ") + cudaGetErrorString(cuda_err); cudaFreeHost(pinned_buf_); @@ -96,6 +120,7 @@ bool MoeHybridStreamEngine::init(ggml_backend_t gpu_backend, size_t max_expert_b return false; } scratch_size_ = alloc_bytes; + scratch_device_ = device; backend_ = gpu_backend; return true; } @@ -106,9 +131,11 @@ bool MoeHybridStreamEngine::is_ready() const { void MoeHybridStreamEngine::destroy() { if (gpu_scratch_) { + const ScopedDevice scoped(scratch_device_); cudaFree(gpu_scratch_); gpu_scratch_ = nullptr; } + scratch_device_ = -1; if (pinned_buf_) { cudaFreeHost(pinned_buf_); pinned_buf_ = nullptr; diff --git a/server/src/common/moe_hybrid_stream.h b/server/src/common/moe_hybrid_stream.h index 4d1cb9638d..a1eab84ddb 100644 --- a/server/src/common/moe_hybrid_stream.h +++ b/server/src/common/moe_hybrid_stream.h @@ -38,7 +38,11 @@ class MoeHybridStreamEngine { // Initialize the engine with a maximum expert size (bytes for one expert's // gate+up+down tensors). Allocates pinned host buffer and GPU scratch. - bool init(ggml_backend_t gpu_backend, size_t max_expert_bytes, std::string * err = nullptr); + // The scratch is a raw allocation on `device` (the GPU of gpu_backend), + // selected for the allocation and again for the free; -1 keeps the + // calling thread's current device, which then must be gpu_backend's. + bool init(ggml_backend_t gpu_backend, size_t max_expert_bytes, std::string * err = nullptr, + int device = -1); bool is_ready() const; void destroy(); @@ -94,6 +98,7 @@ class MoeHybridStreamEngine { void * gpu_scratch_ = nullptr; // GPU device memory for one expert size_t scratch_size_ = 0; + int scratch_device_ = -1; // device gpu_scratch_ was allocated on ggml_backend_t backend_ = nullptr; // Offsets into scratch for last-streamed expert diff --git a/server/src/common/sampler.h b/server/src/common/sampler.h index 05c963f1dd..ae57124418 100644 --- a/server/src/common/sampler.h +++ b/server/src/common/sampler.h @@ -34,6 +34,13 @@ struct SamplerCfg { float freq_pen = 0.0f; float pres_pen = 0.0f; + // Per-token log-probability reporting (OpenAI `logprobs`), carried here + // because the sampler is what reaches every decode path. -1 = off; else + // the number of top alternatives per position (0..kMaxTopLogprobs). + // It does not change which token is chosen. + int logprobs_top_n = -1; + bool wants_logprobs() const { return logprobs_top_n >= 0; } + // True when any logit modifier is active (penalties or stochastic sampling). // Backends should use CPU sample_logits() path when this returns true. bool needs_logit_processing() const { diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index d617833ff3..666a4caa94 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -108,6 +108,9 @@ struct StepGraph { std::vector delta_captures; ggml_tensor * tree_features = nullptr; std::vector moe_selected; + // Residual-stream readouts (qwen35 build_target_step hidden_captures). + std::vector hidden_last; + std::vector hidden_sum; }; // Reset the per-call graph state (ctx + graph + tensor handles) but KEEP the @@ -153,6 +156,8 @@ inline void step_graph_free(StepGraph & sg) { sg.delta_captures.clear(); sg.tree_features = nullptr; sg.moe_selected.clear(); + sg.hidden_last.clear(); + sg.hidden_sum.clear(); sg.paged_tree_key.reset(); } diff --git a/server/src/common/token_logprobs.cpp b/server/src/common/token_logprobs.cpp new file mode 100644 index 0000000000..9ac31fd745 --- /dev/null +++ b/server/src/common/token_logprobs.cpp @@ -0,0 +1,55 @@ +#include "token_logprobs.h" + +#include +#include +#include + +namespace luce::common { + +TokenLogprobs compute_token_logprobs(const float * logits, int vocab, + int32_t chosen, int top_n) { + TokenLogprobs out; + out.chosen.token = chosen; + if (!logits || vocab <= 0) { + out.chosen.logprob = -std::numeric_limits::infinity(); + return out; + } + + float max_logit = -std::numeric_limits::infinity(); + for (int v = 0; v < vocab; ++v) { + if (logits[v] > max_logit) max_logit = logits[v]; + } + // Accumulate in double: a 250k-entry vocab loses precision in float. + double sum = 0.0; + for (int v = 0; v < vocab; ++v) { + sum += std::exp((double)logits[v] - (double)max_logit); + } + const double log_z = (double)max_logit + std::log(sum); + auto logprob_of = [&](int v) { + return (float)((double)logits[v] - log_z); + }; + + out.chosen.logprob = chosen >= 0 && chosen < vocab + ? logprob_of(chosen) + : -std::numeric_limits::infinity(); + + const int n = std::min({std::max(top_n, 0), kMaxTopLogprobs, vocab}); + if (n == 0) return out; + // Insertion into a short sorted list: O(vocab * n) worst case, no + // vocab-sized allocation. Ascending scan plus upper_bound keeps the lower + // id first among equal logits. + std::vector best; + best.reserve((size_t)n + 1); + auto higher = [&](float x, int v) { return x > logits[v]; }; + for (int v = 0; v < vocab; ++v) { + const float x = logits[v]; + if ((int)best.size() == n && !(x > logits[best.back()])) continue; + best.insert(std::upper_bound(best.begin(), best.end(), x, higher), v); + if ((int)best.size() > n) best.pop_back(); + } + out.top.reserve(best.size()); + for (int v : best) out.top.push_back({(int32_t)v, logprob_of(v)}); + return out; +} + +} // namespace luce::common diff --git a/server/src/common/token_logprobs.h b/server/src/common/token_logprobs.h new file mode 100644 index 0000000000..6fe0ecfecd --- /dev/null +++ b/server/src/common/token_logprobs.h @@ -0,0 +1,33 @@ +// Per-token log-probabilities reported to API clients (OpenAI `logprobs`). +// +// Backends compute them from the raw logits row a token was chosen from, so +// the values describe the model's distribution before temperature, top-k/p +// and penalties. The HTTP layer turns token ids into text. + +#pragma once + +#include +#include + +namespace luce::common { + +// OpenAI caps top_logprobs at 20. +inline constexpr int kMaxTopLogprobs = 20; + +struct TokenLogprob { + int32_t token = -1; + float logprob = 0.0f; +}; + +struct TokenLogprobs { + TokenLogprob chosen; + // Most likely tokens, descending; equal logits keep the lower id first. + std::vector top; +}; + +// Log-softmax of `logits[0..vocab)`, returning `chosen`'s log-probability +// and the `top_n` most likely tokens (clamped to [0, kMaxTopLogprobs]). +TokenLogprobs compute_token_logprobs(const float * logits, int vocab, + int32_t chosen, int top_n); + +} // namespace luce::common diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index ac4bf5ea64..2f63126f26 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -16,10 +16,12 @@ #include "common/io_utils.h" #include "common/peer_access.h" #include "common/platform_env.h" +#include "common/platform_io.h" #include "common/sampler.h" #include "common/dspark_head.h" #include "common/moe_hybrid_expert_cache.h" #include "common/moe_hybrid_routing_stats.h" +#include "deepseek4_hc_cuda.h" #include #include @@ -43,6 +45,9 @@ #include #include #include +#if defined(__linux__) +#include +#endif #include #include #include @@ -626,17 +631,6 @@ static void configure_gfx1201_hybrid_sub_batch_default(int gpu) { #endif } -struct Ds4MoeTpConfig { - bool requested = false; - bool in_process = false; - bool backend_valid = true; - PlacementBackend secondary_backend = PlacementBackend::Auto; - int secondary_gpu = 0; - bool all_on_secondary = false; - bool concentrate_secondary = false; - bool profile_hot_on_secondary = false; -}; - static Ds4MoeTpConfig ds4_moe_tp_config(int local_gpu) { Ds4MoeTpConfig result; result.requested = env_flag_enabled("LUCE_DS4_MOE_TP"); @@ -693,6 +687,32 @@ static bool ds4_draft_backend(PlacementBackend & out) { out != PlacementBackend::Auto; } +static DeepSeek4ResolvedSettings resolve_deepseek4_settings(const DeepSeek4BackendConfig & cfg) { + DeepSeek4ResolvedSettings out; + out.tp = ds4_moe_tp_config(cfg.device.gpu); + out.force_full_load = env_flag_enabled("LUCE_DS4_FORCE_FULL_LOAD"); + out.dense_tp_mask = env_flag_enabled("LUCE_DS4_DENSE_TP_MASK"); + out.fused_decode_env = env_flag_enabled("LUCE_DS4_FUSED_DECODE"); + out.decode_all_cold = env_flag_enabled("LUCE_DS4_DECODE_ALL_COLD"); + out.spec_requested = !cfg.draft_path.empty() || env_flag_enabled("LUCE_DS4_SPEC"); + if (const char * path = std::getenv("LUCE_DS4_DRAFT")) out.spec_env_draft_path = path; + // server_main makes an explicit --draft-device concrete (auto:N included), + // so an auto backend here means "not placed": environment, then target. + out.draft_gpu = cfg.device.gpu; + if (cfg.draft_device.backend != PlacementBackend::Auto) { + out.draft_backend = cfg.draft_device.backend; + out.draft_gpu = cfg.draft_device.gpu; + } else { + if (const char * gpu = std::getenv("LUCE_DS4_DRAFT_GPU")) { + out.draft_gpu = std::max(0, std::atoi(gpu)); + } + out.draft_backend_valid = ds4_draft_backend(out.draft_backend); + } + out.draft_separate_stream = env_flag_enabled("LUCE_DS4_DRAFT_SEPARATE_STREAM"); + out.draft_low_priority = env_flag_enabled("LUCE_DS4_DRAFT_LOW_PRIORITY"); + return out; +} + static double gib(uint64_t bytes) { return (double) bytes / 1024.0 / 1024.0 / 1024.0; } @@ -1685,7 +1705,7 @@ bool DeepSeek4Backend::validate_model_features() const { // 16,384 compressed rows) does not match the reference yet. unsupported = "a CUDA build (V4.1 is qualified on ROCm)"; #endif - if (cfg_.fused_decode || env_flag_enabled("LUCE_DS4_FUSED_DECODE")) { + if (cfg_.fused_decode || settings_.fused_decode_env) { unsupported = "fused decode"; } else if (cfg_.fused_verify_f16_kv) { unsupported = "the fused verifier's F16 K/V"; @@ -1750,8 +1770,19 @@ bool DeepSeek4Backend::init_engram() { // --ds4-protected-experts ({"": [expert ids], ...}) into the weights. bool DeepSeek4Backend::load_routing_adjustments() { const size_t n = (size_t) w_.n_layer * (size_t) w_.n_expert; - w_.router_bias_delta.clear(); - w_.protected_experts.clear(); + if (routing_inputs_.parsed) { + // Parsed once: a reload uses the inputs the model was loaded with, + // whatever the files hold now. + if (routing_inputs_.n_layer != w_.n_layer || routing_inputs_.n_expert != w_.n_expert) { + std::fprintf(stderr, "[deepseek4] routing adjustments were parsed for [%d][%d], the model is [%d][%d]\n", + routing_inputs_.n_layer, routing_inputs_.n_expert, w_.n_layer, w_.n_expert); + return false; + } + w_.router_bias_delta = routing_inputs_.router_bias_delta; + w_.protected_experts = routing_inputs_.protected_experts; + return true; + } + routing_inputs_ = RoutingInputs{}; if (!cfg_.router_bias_path.empty()) { std::ifstream f(cfg_.router_bias_path, std::ios::binary | std::ios::ate); if (!f || (size_t) f.tellg() != n * sizeof(float)) { @@ -1759,11 +1790,11 @@ bool DeepSeek4Backend::load_routing_adjustments() { cfg_.router_bias_path.c_str(), n, w_.n_layer, w_.n_expert); return false; } - w_.router_bias_delta.resize(n); + routing_inputs_.router_bias_delta.resize(n); f.seekg(0); - f.read(reinterpret_cast(w_.router_bias_delta.data()), (std::streamsize) (n * sizeof(float))); + f.read(reinterpret_cast(routing_inputs_.router_bias_delta.data()), (std::streamsize) (n * sizeof(float))); if (!f) return false; - for (float v : w_.router_bias_delta) { + for (float v : routing_inputs_.router_bias_delta) { if (!std::isfinite(v)) { std::fprintf(stderr, "[deepseek4] router bias %s holds a non-finite value\n", cfg_.router_bias_path.c_str()); @@ -1786,7 +1817,7 @@ bool DeepSeek4Backend::load_routing_adjustments() { cfg_.protected_experts_path.c_str()); return false; } - w_.protected_experts.assign(n, 0); + routing_inputs_.protected_experts.assign(n, 0); for (auto it = j.begin(); it != j.end(); ++it) { int layer = -1; const std::string & key = it.key(); @@ -1802,10 +1833,15 @@ bool DeepSeek4Backend::load_routing_adjustments() { std::fprintf(stderr, "[deepseek4] protected experts: bad expert in layer %d\n", layer); return false; } - w_.protected_experts[(size_t) layer * (size_t) w_.n_expert + (size_t) expert] = 1; + routing_inputs_.protected_experts[(size_t) layer * (size_t) w_.n_expert + (size_t) expert] = 1; } } } + routing_inputs_.n_layer = w_.n_layer; + routing_inputs_.n_expert = w_.n_expert; + routing_inputs_.parsed = true; + w_.router_bias_delta = routing_inputs_.router_bias_delta; + w_.protected_experts = routing_inputs_.protected_experts; return true; } @@ -1848,6 +1884,12 @@ bool DeepSeek4Backend::upload_protected_routing() { // Adds the router bias delta to every layer's selection bias once, so routing // pays nothing per token. Mixing weights are unaffected (they never read it). bool DeepSeek4Backend::apply_routing_adjustments() { + // The delta goes onto the biases exactly as loaded from the file; adding + // it to already adjusted ones would shift routing on every reload. + if (w_.router_bias_applied) { + std::fprintf(stderr, "[deepseek4] routing adjustments are already applied to these weights\n"); + return false; + } if (!load_routing_adjustments()) return false; if (!w_.router_bias_delta.empty()) { std::vector bias((size_t) w_.n_expert); @@ -1864,6 +1906,7 @@ bool DeepSeek4Backend::apply_routing_adjustments() { ggml_backend_tensor_set(t, bias.data(), 0, sizeof(float) * bias.size()); } } + w_.router_bias_applied = true; if (!upload_protected_routing()) return false; // Host routing reads a selection bias per layer and token; keep a copy. w_.selection_bias_host.clear(); @@ -1937,6 +1980,8 @@ bool DeepSeek4Backend::apply_expert_ownership(bool secondary_owner, int secondar const uint64_t keep = ds4_host_reserve_bytes(w_, cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192); const uint64_t host_want = avail > keep ? avail - keep : 0; host_bytes = ggml_backend_cuda_set_host_spill(secondary_gpu, carve_reserve, host_want); + frozen_.spill_carve_reserve = carve_reserve; + frozen_.spill_host_want = host_want; if (host_bytes + (1ULL << 30) < host_want) { std::fprintf(stderr, "[deepseek4] device %d: only %.2f of %.2f GiB of host memory can be locked " "(RLIMIT_MEMLOCK); the rest of these experts stream from the drive. Run as a " @@ -2084,12 +2129,12 @@ bool DeepSeek4Backend::load_model() { // backends, so forcing a full load would disable the requested placement // before the TP runtime can initialize. init() has already rejected paged // deployments outside the qualified R9700 + Strix Halo topology. - const bool force_full = env_flag_enabled("LUCE_DS4_FORCE_FULL_LOAD"); - const bool heterogeneous_tp = env_flag_enabled("LUCE_DS4_MOE_TP"); + const bool force_full = settings_.force_full_load; + const bool heterogeneous_tp = settings_.tp.requested; if (!cfg_.mmproj_path.empty()) { // Images run on one HIP GPU holding the whole model, or on two HIP GPUs // that split the experts in process. Both need batched sparse prefill. - const auto tp = ds4_moe_tp_config(cfg_.device.gpu); + const Ds4MoeTpConfig & tp = settings_.tp; const bool two_gpu_ok = tp.in_process && tp.backend_valid && tp.secondary_backend == PlacementBackend::Hip && tp.secondary_gpu != cfg_.device.gpu && !tp.all_on_secondary && !force_full; @@ -2098,7 +2143,7 @@ bool DeepSeek4Backend::load_model() { if (target_backend != PlacementBackend::Hip || cfg_.device.is_layer_split() || (cfg_.prefill_mode != PrefillAttentionMode::Sparse && !cfg_.paged_attention) || (tp.requested && !two_gpu_ok) || - env_flag_enabled("LUCE_DS4_DENSE_TP_MASK")) { + settings_.dense_tp_mask) { std::fprintf(stderr, "[deepseek4] --mmproj requires a HIP target with --ds4-prefill sparse, " "on one GPU or with in-process expert owners on two distinct GPUs\n"); return false; @@ -2211,25 +2256,14 @@ bool DeepSeek4Backend::load_spec_drafter() { } ggml_backend_t draft_backend = backend_; - int draft_gpu = cfg_.device.gpu; - const bool separate_draft_stream = - env_flag_enabled("LUCE_DS4_DRAFT_SEPARATE_STREAM"); - PlacementBackend draft_kind = PlacementBackend::Auto; - // server_main makes an explicit --draft-device concrete (auto:N included), - // so an auto backend here means "not placed": environment, then target. - if (cfg_.draft_device.backend != PlacementBackend::Auto) { - draft_kind = cfg_.draft_device.backend; - draft_gpu = cfg_.draft_device.gpu; - } else { - if (const char * gpu = std::getenv("LUCE_DS4_DRAFT_GPU")) { - draft_gpu = std::max(0, std::atoi(gpu)); - } - if (!ds4_draft_backend(draft_kind)) { - std::fprintf(stderr, - "[deepseek4] invalid LUCE_DS4_DRAFT_BACKEND; " - "expected cuda or hip\n"); - return false; - } + const int draft_gpu = settings_.draft_gpu; + const bool separate_draft_stream = settings_.draft_separate_stream; + const PlacementBackend draft_kind = settings_.draft_backend; + if (!settings_.draft_backend_valid) { + std::fprintf(stderr, + "[deepseek4] invalid LUCE_DS4_DRAFT_BACKEND; " + "expected cuda or hip\n"); + return false; } const PlacementBackend target_kind = placement_backend_of(backend_); if (draft_kind != target_kind || draft_gpu != cfg_.device.gpu || @@ -2245,8 +2279,7 @@ bool DeepSeek4Backend::load_spec_drafter() { return false; } draft_backend = spec_backend_; - const bool low_priority = separate_draft_stream && - env_flag_enabled("LUCE_DS4_DRAFT_LOW_PRIORITY"); + const bool low_priority = separate_draft_stream && settings_.draft_low_priority; const bool priority_configured = low_priority && backend_pair_capabilities(backend_, spec_backend_).same_runtime && ggml_backend_cuda_set_low_priority_stream(spec_backend_); @@ -2549,12 +2582,13 @@ bool DeepSeek4Backend::supports_batched_spec_feature_capture( bool DeepSeek4Backend::init() { // --draft selects DSpark directly; LUCE_DS4_SPEC + LUCE_DS4_DRAFT remain // the environment spelling of the same request. - spec_requested_ = !cfg_.draft_path.empty() || env_flag_enabled("LUCE_DS4_SPEC"); + settings_ = resolve_deepseek4_settings(cfg_); + const bool spec_requested = settings_.spec_requested; if (cfg_.paged_attention) { const PlacementBackend target_backend = cfg_.device.backend == PlacementBackend::Auto ? compiled_placement_backend() : cfg_.device.backend; - const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); + const Ds4MoeTpConfig & tp = settings_.tp; if (target_backend != PlacementBackend::Hip) { std::fprintf(stderr, "[deepseek4] paged concurrency requires a local HIP target\n"); @@ -2585,17 +2619,17 @@ bool DeepSeek4Backend::init() { // Install the gfx1151 DSpark verifier profile first: the MMVQ crossover // below reads LUCE_DS4_Q5_VERIFY. - if (!configure_gfx1151_dspark_verifier_defaults(cfg_.device.gpu, spec_requested_)) { + if (!configure_gfx1151_dspark_verifier_defaults(cfg_.device.gpu, spec_requested)) { return false; } // The shared MMVQ/MMQ crossover defaults to q=3 for NVIDIA. On gfx1151, // DSpark q=4 is faster through MMVQ. Keep AR and other devices unchanged, // and preserve LUCE_MMVQ_MAX_NCOLS as an explicit override. - if (!configure_dspark_mmvq_defaults(cfg_.device.gpu, spec_requested_)) { + if (!configure_dspark_mmvq_defaults(cfg_.device.gpu, spec_requested)) { return false; } configure_gfx1151_paged_mmvq_default(cfg_.device.gpu, cfg_.paged_attention, - spec_requested_); + spec_requested); if (!configure_gfx1151_sparse_prefill_kernel_defaults( cfg_.device.gpu, cfg_.prefill_mode)) { return false; @@ -2609,8 +2643,8 @@ bool DeepSeek4Backend::init() { (cfg_.prefill_mode != PrefillAttentionMode::Exact && cfg_.prefill_mode != PrefillAttentionMode::Sparse) || cfg_.fused_decode || cfg_.fused_verify_f16_kv || - env_flag_enabled("LUCE_DS4_FUSED_DECODE") || - spec_requested_)) { + settings_.fused_decode_env || + spec_requested)) { std::fprintf(stderr, "[deepseek4] paged serving requires 1..%d local slots, exact " "or sparse prefill, and autoregressive non-fused decode\n", @@ -2706,7 +2740,7 @@ bool DeepSeek4Backend::init() { cache_.prefill_mode = cfg_.prefill_mode; } - if (env_flag_enabled("LUCE_DS4_MOE_TP") && !init_moe_tensor_parallel()) { + if (settings_.tp.requested && !init_moe_tensor_parallel()) { return false; } if (cfg_.paged_attention && expert_runtime_.compute) { @@ -2780,10 +2814,9 @@ bool DeepSeek4Backend::init() { "DSpark drafter\n", cfg_.draft_path.c_str()); return false; } - } else if (!cfg_.paged_attention && spec_requested_) { - const char * dp = std::getenv("LUCE_DS4_DRAFT"); - if (dp && *dp) { - spec_draft_path_ = dp; + } else if (!cfg_.paged_attention && spec_requested) { + if (!settings_.spec_env_draft_path.empty()) { + spec_draft_path_ = settings_.spec_env_draft_path; if (!load_spec_drafter()) { std::fprintf(stderr, "[deepseek4] DSpark drafter load failed; " @@ -2797,6 +2830,14 @@ bool DeepSeek4Backend::init() { if (!size_hybrid_prefill_chunk()) return false; if (!setup_draft_swap()) return false; image_capable_ = vision_ != nullptr; + if (moe_hybrid_) { + // What a reinstatement must reproduce exactly (see FrozenPlacement). + frozen_.placement = moe_placement_; + frozen_.decode_placement = moe_decode_placement_; + frozen_.ownership = MoeHybridOwnershipRecord::capture(*moe_hybrid_); + frozen_.selection_bias_host = w_.selection_bias_host; + frozen_.valid = true; + } return true; } @@ -2809,6 +2850,26 @@ static size_t ds4_device_headroom_bytes(int device) { return integrated ? ((size_t) 4 << 30) : ((size_t) 3 << 29); } +void ds4_device_memory(int device, size_t * free_b, size_t * total_b) { + ggml_backend_cuda_get_device_memory(device, free_b, total_b); +#if defined(__linux__) + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, device) != cudaSuccess || !prop.integrated) return; + const uint64_t ram = (uint64_t) sysconf(_SC_PHYS_PAGES) * (uint64_t) sysconf(_SC_PAGE_SIZE); + if (ram == 0 || *total_b <= ram) return; + *total_b = (size_t) ram; + *free_b = std::min(*free_b, (size_t) host_available_bytes()); +#endif +} + +// Whether `device` draws on host memory rather than a carve of its own. +static bool ds4_device_shares_host_memory(int device) { + size_t free_b = 0, total_b = 0, raw_free = 0, raw_total = 0; + ds4_device_memory(device, &free_b, &total_b); + ggml_backend_cuda_get_device_memory(device, &raw_free, &raw_total); + return total_b != raw_total; +} + // Streamed expert cache, sized from what is left after // the model, the owner stacks, the caches and the drafter are in place. bool DeepSeek4Backend::init_streamed_expert_tier() { @@ -2825,7 +2886,11 @@ bool DeepSeek4Backend::init_streamed_expert_tier() { cache_opts.direct_all = ggml_backend_cuda_host_spill_bytes(cache_opts.device) > 0; { size_t free_b = 0, total_b = 0; - ggml_backend_cuda_get_device_memory(cache_opts.device, &free_b, &total_b); + ds4_device_memory(cache_opts.device, &free_b, &total_b); + // Out of host memory, the cache also leaves the host its reserve. + if (ds4_device_shares_host_memory(cache_opts.device)) { + cache_opts.reserve_bytes += ds4_host_reserve_bytes(w_, cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192); + } cache_opts.pool_bytes = free_b > cache_opts.reserve_bytes ? free_b - cache_opts.reserve_bytes : 0; if (cache_opts.pool_bytes > 0 && init_deepseek4_streamed_expert_cache(w_, *hybrid, cache_opts, expert_cache_, &err)) { @@ -2865,7 +2930,7 @@ bool DeepSeek4Backend::log_device_memory(const char * when) const { bool ok = true; for (int device : devices) { size_t free_b = 0, total_b = 0; - ggml_backend_cuda_get_device_memory(device, &free_b, &total_b); + ds4_device_memory(device, &free_b, &total_b); const size_t headroom = ds4_device_headroom_bytes(device); const size_t need = when ? headroom / 2 : headroom; std::fprintf(stderr, "[deepseek4] device %d memory%s%s: %.2f of %.2f GiB used, %.2f GiB free " @@ -2967,7 +3032,7 @@ bool DeepSeek4Backend::size_hybrid_prefill_chunk() { const HybridPrefillScratch per_token = hybrid_prefill_scratch_per_token(w_, max_ctx, chunk); const auto fit = [this](int device, size_t bytes, size_t fixed) { size_t free_b = 0, total_b = 0; - ggml_backend_cuda_get_device_memory(device, &free_b, &total_b); + ds4_device_memory(device, &free_b, &total_b); if (device == cfg_.device.gpu) free_b += prefill_sizing_extra_free_; return hybrid_prefill_fit_tokens(free_b, ds4_device_headroom_bytes(device) / 2 + fixed, bytes); }; @@ -2975,7 +3040,7 @@ bool DeepSeek4Backend::size_hybrid_prefill_chunk() { // A context the smallest chunk cannot prefill at its end is refused // at load rather than failing an allocation in the middle of a prompt. size_t free_b = 0, total_b = 0; - ggml_backend_cuda_get_device_memory(cfg_.device.gpu, &free_b, &total_b); + ds4_device_memory(cfg_.device.gpu, &free_b, &total_b); const size_t keep = ds4_device_headroom_bytes(cfg_.device.gpu) / 2 + per_token.target_fixed; const size_t smallest = 64 * per_token.target; if (free_b < keep + smallest) { @@ -2989,7 +3054,7 @@ bool DeepSeek4Backend::size_hybrid_prefill_chunk() { int tokens = fit(cfg_.device.gpu, per_token.target, per_token.target_fixed); if (stream_cache_device_ >= 0 && stream_cache_device_ != cfg_.device.gpu) { size_t free_b = 0, total_b = 0; - ggml_backend_cuda_get_device_memory(stream_cache_device_, &free_b, &total_b); + ds4_device_memory(stream_cache_device_, &free_b, &total_b); const size_t keep = ds4_device_headroom_bytes(stream_cache_device_) / 2; if (free_b < keep + 64 * per_token.second) { std::fprintf(stderr, @@ -3026,7 +3091,7 @@ bool DeepSeek4Backend::init_moe_tensor_parallel() { return false; } - const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); + const Ds4MoeTpConfig & tp = settings_.tp; if (tp.in_process) { if (!expert_backend_ || !moe_hybrid_->materialized_cold_experts || moe_hybrid_->cold_backend != expert_backend_) { @@ -3111,7 +3176,7 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & kv_bytes += paged_plan.total_persistent_bytes; } Ds4HybridBudgetInfo budget; - const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); + const Ds4MoeTpConfig & tp = settings_.tp; if (!compute_ds4_hybrid_budget_info( w, backend_, cfg_.device.gpu, kv_bytes, tp.all_on_secondary, vision_ != nullptr, budget, err)) { @@ -3344,7 +3409,13 @@ bool DeepSeek4Backend::init_hybrid_model() { std::string err; const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; - if (!compute_uniform_hybrid_placement( + // A reinstatement places the experts exactly as init() did; the free + // memory and usage profile of the moment would place them differently. + const bool reuse_placement = frozen_.valid; + if (reuse_placement) { + moe_placement_ = frozen_.placement; + moe_decode_placement_ = frozen_.decode_placement; + } else if (!compute_uniform_hybrid_placement( w_, max_ctx, moe_placement_, &moe_decode_placement_, &err)) { std::fprintf(stderr, "[deepseek4] failed to compute hybrid placement: %s\n", err.c_str()); return false; @@ -3365,7 +3436,7 @@ bool DeepSeek4Backend::init_hybrid_model() { return true; } - const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); + const Ds4MoeTpConfig & tp = settings_.tp; const bool inprocess_tp = tp.requested && tp.in_process; const PlacementBackend local_kind = cfg_.device.backend == PlacementBackend::Auto @@ -3479,9 +3550,19 @@ bool DeepSeek4Backend::init_hybrid_model() { hybrid_cfg.cold_expert_backend = MoeHybridColdBackend::Gpu; } const int host_spill_gpu = same_runtime_tp ? tp.secondary_gpu : -1; - if (!apply_expert_ownership(inprocess_tp, host_spill_gpu, hybrid_cfg)) { - if (host_spill_gpu >= 0) ggml_backend_cuda_set_host_spill(host_spill_gpu, 0, 0); - return fail_hybrid_init(); + if (reuse_placement) { + if (inprocess_tp) hybrid_cfg.cold_expert_ids = frozen_.cold_expert_ids; + if (host_spill_gpu >= 0 && frozen_.spill_host_want > 0) { + ggml_backend_cuda_set_host_spill(host_spill_gpu, frozen_.spill_carve_reserve, + frozen_.spill_host_want); + } + } else { + frozen_.spill_carve_reserve = frozen_.spill_host_want = 0; + if (!apply_expert_ownership(inprocess_tp, host_spill_gpu, hybrid_cfg)) { + if (host_spill_gpu >= 0) ggml_backend_cuda_set_host_spill(host_spill_gpu, 0, 0); + return fail_hybrid_init(); + } + frozen_.cold_expert_ids = hybrid_cfg.cold_expert_ids; } if (vision_) { #if defined(LUCE_BACKEND_HIP) || defined(GGML_USE_HIP) @@ -3608,7 +3689,7 @@ bool DeepSeek4Backend::init_hybrid_model() { } } - if (env_flag_enabled("LUCE_DS4_DECODE_ALL_COLD")) { + if (settings_.decode_all_cold) { for (int il = 0; il < w_.n_layer; ++il) { MoeHybridLayerStorage & layer = hybrid->layers[(size_t) il]; if (layer.cold_local_by_global.size() != @@ -3642,7 +3723,8 @@ bool DeepSeek4Backend::init_hybrid_model() { std::fprintf(stderr, "[deepseek4] failed to compute streaming expert size\n"); return fail_hybrid_init(); } - if (!stream_engine_.init(backend_, max_expert_bytes, &err)) { + frozen_.stream_expert_bytes = max_expert_bytes; + if (!stream_engine_.init(backend_, max_expert_bytes, &err, cfg_.device.gpu)) { std::fprintf(stderr, "[deepseek4] failed to init cold-expert stream engine: %s\n", err.c_str()); return fail_hybrid_init(); @@ -3695,6 +3777,7 @@ void DeepSeek4Backend::print_ready_banner() const { } bool DeepSeek4Backend::park(ParkTarget target) { + if (is_evicted()) return true; // nothing of it is on the primary const bool want_draft = park_target_includes_draft_model(target); const bool want_target_model = park_target_includes_target_model(target); @@ -3747,6 +3830,10 @@ bool DeepSeek4Backend::park(ParkTarget target) { } bool DeepSeek4Backend::unpark(ParkTarget target) { + if (is_evicted()) { + std::fprintf(stderr, "[deepseek4] unpark: model is evicted; reinstate it first\n"); + return false; + } const bool want_draft = park_target_includes_draft_model(target); const bool want_target_model = park_target_includes_target_model(target); @@ -3787,7 +3874,7 @@ bool DeepSeek4Backend::unpark(ParkTarget target) { return false; } - if (env_flag_enabled("LUCE_DS4_MOE_TP") && + if (settings_.tp.requested && !init_moe_tensor_parallel()) { free_deepseek4_cache(cache_); release_vision(); @@ -3840,6 +3927,338 @@ bool DeepSeek4Backend::unpark(ParkTarget target) { return true; } +// ── Eviction (runtime model swapping) ─────────────────────────────────── + +bool DeepSeek4Backend::supports_eviction(std::string & reason) const { + // An evicted model was supported when it was evicted; its configuration + // has not changed since. + if (residency_ != Residency::resident) { + reason.clear(); + return true; + } + const PlacementBackend local_kind = cfg_.device.backend == PlacementBackend::Auto + ? compiled_placement_backend() : cfg_.device.backend; + const Ds4MoeTpConfig & tp = settings_.tp; + // First milestone: one classic worker, fixed in-process expert owners on + // two GPUs of one runtime, DSpark (if any) on the target. Every other mode + // owns device state this teardown does not cover. + if (cfg_.paged_attention || seq_engine_ || cfg_.max_concurrency > 1) { + reason = "eviction supports the classic serial worker only (no paged attention or --max-concurrency > 1)"; + } else if (!cfg_.mmproj_path.empty() || vision_) { + reason = "eviction does not support vision (--mmproj)"; + } else if (pflash_drafter_loaded_) { + reason = "eviction does not support a loaded PFlash drafter"; + } else if (cfg_.device.is_layer_split() || cfg_.device.is_tensor_parallel() || w_.dense_split_buf) { + reason = "eviction does not support layer or tensor split"; + } else if (expert_runtime_.compute) { + reason = "eviction does not support out-of-process (remote) experts"; + } else if (!tp.requested || !tp.in_process || !tp.backend_valid || + tp.secondary_backend != local_kind || tp.secondary_gpu == cfg_.device.gpu) { + reason = "eviction requires in-process experts on a second GPU of the same runtime (--expert-device)"; + } else if (!moe_hybrid_ || !frozen_.valid || !moe_hybrid_->materialized_cold_experts) { + reason = "eviction requires the hybrid expert placement with a materialized secondary stack"; + } else if (!moe_hybrid_->has_mmap()) { + reason = "eviction requires the expert file mapping to rebuild the primary experts"; + } else if (spec_backend_) { + reason = "eviction requires DSpark on the target device and stream"; + } else if (!snap_backend_) { + reason = "eviction requires host-memory snapshots"; + } else { + for (const MoeHybridLayerStorage & layer : moe_hybrid_->layers) { + if (layer.cache_slots != 0) { + reason = "eviction does not support primary spare cache slots"; + return false; + } + for (const ggml_tensor * t : {layer.gate_hot, layer.up_hot, layer.down_hot, layer.gate_up_hot}) { + if (t && (t->type == GGML_TYPE_Q3_1_ROCMFP3_MIX || t->type == GGML_TYPE_Q2_1_ROCMFP2_MIX)) { + reason = "eviction does not support mixed-qtype experts"; + return false; + } + } + } + reason.clear(); + return true; + } + return false; +} + +// Releases everything this model holds on its primary device. Runs on the +// worker with no request in flight and the streamed expert cache suspended; +// also the cleanup of a failed reinstatement, so every step tolerates state +// that was never built. The secondary stack, the streamed cache and the +// locked host tier stay. +void DeepSeek4Backend::release_primary_state() { + // Work queued for this model on any context finishes first. + for (ggml_backend_t b : {backend_, expert_backend_, spec_backend_}) { + if (b) ggml_backend_synchronize(b); + } + maybe_save_routing_stats(); + last_logits_.clear(); + last_logits_pos_ = -1; + spec_feat_window_.clear(); + cache_has_images_ = false; + // Consumers before what they reference. The cache owns the fused + // decode/verify graphs (their native graph keys are retired with them), + // the layer runtime and its arenas. + free_deepseek4_cache(cache_); + if (moe_hybrid_) { + moe_hybrid_->release_graph_caches(); + moe_hybrid_->release_prefill_allocators(); + } + // Built over the layer descriptors of the weights about to go. + expert_cache_.release_graphs(); + stream_engine_.destroy(); + // Draft weights, cloned heads and this worker's draft graph caches. + release_spec_drafter(/*mark_parked=*/false); + dspark_head_release_thread_scratch(); + if (moe_hybrid_) { + std::string err; + if (!moe_hybrid_->release_primary_experts(&err)) { + std::fprintf(stderr, "[evict] primary experts kept: %s\n", err.c_str()); + } + } + // Also this worker's runtime graphs keyed by the weights' context. + free_deepseek4_weights(w_); + deepseek4_cuda_hc_release_device(cfg_.device.gpu); + // Contexts last; they own the device pools, streams and BLAS handles. + // The expert backend's buffers belong to its device, not to the context: + // freeing the context drops whatever pools it grew while joining with the + // primary, and reinstate() recreates it over the same buffers. + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + if (moe_hybrid_) moe_hybrid_->rebind_cold_backend(nullptr); + } + if (backend_) { + ggml_backend_free(backend_); + backend_ = nullptr; + } +} + +// Releases the retained tier: the streamed cache, the secondary stack with +// its locked host memory, the file mapping and the secondary scratch. +void DeepSeek4Backend::release_secondary_state() { + if (expert_backend_) ggml_backend_synchronize(expert_backend_); + expert_cache_.destroy(); // the cache points at the storage: first + expert_runtime_.reset(); + moe_hybrid_.reset(); + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + } + if (settings_.tp.secondary_gpu != cfg_.device.gpu) { + deepseek4_cuda_hc_release_device(settings_.tp.secondary_gpu); + } + moe_placement_ = {}; + moe_decode_placement_ = {}; + stream_cache_device_ = -1; +} + +ModelBackend::EvictResult DeepSeek4Backend::evict(EvictLevel level) { + EvictResult result; + if (level == EvictLevel::warm) { + result.error = "DeepSeek V4 has no warm eviction; use partial or cold"; + return result; + } + const bool full = level == EvictLevel::cold; + if (residency_ == Residency::fully_evicted || + (residency_ == Residency::partially_evicted && !full)) { + result.status = EvictStatus::ok; + return result; + } + if (residency_ == Residency::partially_evicted) { + release_secondary_state(); + residency_ = Residency::fully_evicted; + std::printf("[evict] deepseek4 retained expert tier released\n"); + std::fflush(stdout); + result.status = EvictStatus::ok; + return result; + } + if (!supports_eviction(result.error)) return result; + if (!full && stream_cache_device_ == cfg_.device.gpu) { + result.error = "partial eviction requires the streamed expert cache on the secondary device"; + return result; + } + // Launches that wait on the expert loaders end with their request, so the + // loaders can stop now without stranding one. Refused, nothing changed. + std::string busy; + if (!expert_cache_.suspend(&busy)) { + result.error = "the streamed expert cache is busy: " + busy; + return result; + } + const auto t0 = Clock::now(); + reload_drafter_on_reinstate_ = spec_drafter_ != nullptr; + release_primary_state(); + const double primary_ms = elapsed_s(t0) * 1000.0; + if (full) release_secondary_state(); + residency_ = full ? Residency::fully_evicted : Residency::partially_evicted; + result.status = EvictStatus::ok; + std::printf("[evict] deepseek4 %s: primary released in %.0f ms%s\n", + full ? "cold" : "partial", primary_ms, + full ? ", expert tier released" : "; secondary experts and streamed cache kept"); + std::fflush(stdout); + return result; +} + +bool DeepSeek4Backend::verify_frozen_state(std::string & error) const { + std::string err; + if (!moe_hybrid_ || !frozen_.ownership.matches(*moe_hybrid_, &err)) { + error = "expert placement differs from the one the model was loaded with: " + err; + return false; + } + if (w_.selection_bias_host != frozen_.selection_bias_host) { + error = "routing biases differ from the ones the model was loaded with"; + return false; + } + return true; +} + +// Partial reinstatement: the primary comes back around the retained tier, +// with the placement init() resolved, and the cache resumes last. +bool DeepSeek4Backend::reinstate_primary(std::string & error) { + const auto fail = [&](const std::string & what) { + error = what; + // The new primary state goes; the retained tier stays suspended. + release_primary_state(); + return false; + }; + auto t = Clock::now(); + double stage_ms[6] = {}; // backend, dense, experts, routing, cache, draft + const auto lap = [&](int stage) { + const auto now = Clock::now(); + stage_ms[stage] = std::chrono::duration(now - t).count(); + t = now; + }; + backend_ = ggml_backend_cuda_init(cfg_.device.gpu); + if (!backend_) return fail("target backend init failed"); + std::string err; + const Ds4MoeTpConfig & tp = settings_.tp; + expert_backend_ = init_placement_backend(tp.secondary_backend, tp.secondary_gpu, &err); + if (!expert_backend_) return fail("expert backend init failed: " + err); + moe_hybrid_->rebind_cold_backend(expert_backend_); + lap(0); + // The dense tensors only; the routed experts are rebuilt below. + TargetLoadPlan plan; + plan.skip_expert_tensors = true; + if (!load_deepseek4_gguf_partial(cfg_.model_path, backend_, plan, w_)) { + return fail(std::string("dense load: ") + luce_last_error()); + } + // What load_model() and init_hybrid_model() set for this placement. + w_.moe_hybrid = true; + w_.mixed_mmq_policy = moe_hybrid_->mixed_mmq_policy; + w_.routed_expert_top_k = cfg_.expert_top_k; + w_.fused_decode = false; + w_.fused_verify_f16_kv = false; + lap(1); + if (!moe_hybrid_->rebuild_primary_experts(backend_, cfg_.model_path, &err)) { + return fail("primary experts: " + err); + } + lap(2); + if (!apply_routing_adjustments()) return fail("routing adjustments"); + // As init() and unpark(): the reload re-read the Engram metadata. + if (!init_engram()) return fail("Engram tables"); + lap(3); + const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; + if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { + return fail("cache allocation (ctx=" + std::to_string(max_ctx) + ")"); + } + cache_.prefill_mode = cfg_.prefill_mode; + if (frozen_.stream_expert_bytes > 0 && + !stream_engine_.init(backend_, frozen_.stream_expert_bytes, &err, cfg_.device.gpu)) { + return fail("cold-expert stream engine: " + err); + } + if (!init_moe_tensor_parallel()) return fail("expert parallelism"); + lap(4); + if (reload_drafter_on_reinstate_ && !load_spec_drafter()) return fail("DSpark drafter load"); + lap(5); + if (!verify_frozen_state(err)) return fail(err); + expert_cache_.resume(); + std::printf("[reinstate] deepseek4 primary restored: backend %.0f ms, dense %.0f ms, " + "experts %.0f ms, routing %.0f ms, cache %.0f ms, draft %.0f ms\n", + stage_ms[0], stage_ms[1], stage_ms[2], stage_ms[3], stage_ms[4], stage_ms[5]); + std::fflush(stdout); + log_device_memory("after reinstate"); + return true; +} + +// Full reinstatement: the load path of init() with the frozen placement +// (init_hybrid_model reuses it), then a fresh streamed cache. +bool DeepSeek4Backend::reinstate_full(std::string & error) { + const auto fail = [&](const std::string & what) { + error = what; + release_primary_state(); + release_secondary_state(); + return false; + }; + const auto t0 = Clock::now(); + backend_ = ggml_backend_cuda_init(cfg_.device.gpu); + if (!backend_) return fail("target backend init failed"); + if (!load_model()) return fail("model load"); + if (!moe_hybrid_) return fail("the reload did not produce the hybrid expert placement"); + if (!apply_routing_adjustments()) return fail("routing adjustments"); + if (!init_engram()) return fail("Engram tables"); + const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; + if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { + return fail("cache allocation (ctx=" + std::to_string(max_ctx) + ")"); + } + cache_.prefill_mode = cfg_.prefill_mode; + if (settings_.tp.requested && !init_moe_tensor_parallel()) return fail("expert parallelism"); + if (reload_drafter_on_reinstate_ && !load_spec_drafter()) return fail("DSpark drafter load"); + if (!init_streamed_expert_tier()) return fail("streamed expert cache"); + std::string err; + if (!verify_frozen_state(err)) return fail(err); + std::printf("[reinstate] deepseek4 fully restored in %.0f ms\n", elapsed_s(t0) * 1000.0); + std::fflush(stdout); + log_device_memory("after reinstate"); + return true; +} + +bool DeepSeek4Backend::reinstate(std::string & error) { + if (residency_ == Residency::resident) return true; + const bool partial = residency_ == Residency::partially_evicted; + if (!(partial ? reinstate_primary(error) : reinstate_full(error))) { + std::fprintf(stderr, "[reinstate] deepseek4: %s\n", error.c_str()); + return false; + } + residency_ = Residency::resident; + return true; +} + +size_t DeepSeek4Backend::weight_device_bytes() const { + const auto size = [](ggml_backend_buffer_t b) { return b ? ggml_backend_buffer_get_size(b) : (size_t) 0; }; + size_t bytes = 0; + if (residency_ == Residency::resident) { + bytes += size(w_.buf) + size(w_.dense_split_buf) + size(w_.routing_buf); + if (spec_drafter_) bytes += size(spec_drafter_->core.buf) + size(spec_drafter_->head_buf); + } + if (residency_ != Residency::fully_evicted && moe_hybrid_) { + for (const MoeHybridLayerStorage & layer : moe_hybrid_->layers) { + if (residency_ == Residency::resident) bytes += size(layer.hot_buf); + if (moe_hybrid_->cold_backend_kind == MoeHybridColdBackend::Gpu) bytes += size(layer.cold_buf); + } + } + return bytes; +} + +DeepSeek4Backend::ExpertTierReport DeepSeek4Backend::expert_tier_report() const { + ExpertTierReport report; + report.partially_evicted = residency_ == Residency::partially_evicted; + if (moe_hybrid_) { + report.primary_upload_bytes = moe_hybrid_->primary_upload_bytes; + report.secondary_upload_bytes = moe_hybrid_->secondary_upload_bytes; + report.storage = moe_hybrid_.get(); + for (const MoeHybridLayerStorage & layer : moe_hybrid_->layers) { + if (layer.cold_buf) { + report.secondary_base = ggml_backend_buffer_get_base(layer.cold_buf); + break; + } + } + } + report.stream_cache_slots = expert_cache_.ready() ? expert_cache_.n_slots() : 0; + report.stream_cache_suspended = expert_cache_.ready() && expert_cache_.suspended(); + return report; +} + int deepseek4_hybrid_prefill_chunk_tokens( int requested_chunk, int context_end, @@ -4612,6 +5031,10 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, GenerateResult DeepSeek4Backend::generate_from_state( const GenerateRequest & req, const DaemonIO & io, int kv_offset) { GenerateResult result; + if (is_evicted()) { + result.fail(GenerateErrorCode::ModelParked, "model is evicted"); + return result; + } DaemonIO out_io = io.with_token_callback(req.on_token); auto t0 = Clock::now(); sampler_ = req.sampler; @@ -4907,7 +5330,7 @@ GenerateResult DeepSeek4Backend::generate_from_state( // ── Snapshots ─────────────────────────────────────────────────────────── bool DeepSeek4Backend::snapshot_save(int slot) { - if (cache_has_images_) return false; + if (cache_has_images_ || is_evicted()) return false; if (slot < 0 || slot >= PREFIX_SLOTS || !snap_backend_ || cache_.cur_pos <= 0 || last_logits_pos_ != cache_.cur_pos || w_.n_vocab <= 0 || @@ -5045,7 +5468,27 @@ void add_deepseek4_cache_memory(ModelBackend::MemoryReport::Cache & out, ModelBackend::MemoryReport DeepSeek4Backend::memory_report() const { MemoryReport report; const bool paged = paged_cache_.buf != nullptr; - if (!paged && !cache_.buf) return report; + const auto add_snapshots = [&] { + for (int slot = 0; slot < PREFIX_SLOTS; ++slot) { + const auto & snap = snapshots_[slot]; + if (!snap.ctx || !snap.buf) continue; + const auto & aux = snapshot_aux_[slot]; + const uint64_t host_copies = + (aux.last_logits.capacity() + aux.spec_feat_window.capacity()) * + sizeof(float); + report.snapshots.push_back({slot, snap.cur_pos, + ggml_backend_buffer_get_size(snap.buf), + ggml_backend_buffer_is_host(snap.buf), + host_copies}); + } + }; + if (!paged && !cache_.buf) { + // Evicted (or not initialized): no device cache, but the host + // snapshots are retained and still count. + add_snapshots(); + report.available = !report.snapshots.empty(); + return report; + } report.available = true; report.cache.host = ggml_backend_buffer_is_host(paged ? paged_cache_.buf : cache_.buf); if (paged) { @@ -5071,18 +5514,7 @@ ModelBackend::MemoryReport DeepSeek4Backend::memory_report() const { // Trimming a vector keeps its allocation, so count capacity, not size. report.cache.host_state_bytes = (last_logits_.capacity() + spec_feat_window_.capacity()) * sizeof(float); - for (int slot = 0; slot < PREFIX_SLOTS; ++slot) { - const auto & snap = snapshots_[slot]; - if (!snap.ctx || !snap.buf) continue; - const auto & aux = snapshot_aux_[slot]; - const uint64_t host_copies = - (aux.last_logits.capacity() + aux.spec_feat_window.capacity()) * - sizeof(float); - report.snapshots.push_back({slot, snap.cur_pos, - ggml_backend_buffer_get_size(snap.buf), - ggml_backend_buffer_is_host(snap.buf), - host_copies}); - } + add_snapshots(); return report; } @@ -5185,6 +5617,10 @@ bool DeepSeek4Backend::snapshot_restore(int slot) { GenerateResult DeepSeek4Backend::restore_and_generate_impl( int slot, const GenerateRequest & req, const DaemonIO & io) { GenerateResult result; + if (is_evicted()) { + result.fail(GenerateErrorCode::ModelParked, "model is evicted"); + return result; + } if (req.images) { result.fail(GenerateErrorCode::PrefillFailed, "image requests cannot restore token-only snapshots"); return result; diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index e042ae786b..6f2742c9b5 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -73,6 +73,42 @@ bool deepseek4_mix_mmq_prefill_default( ggml_mixed_mmq_policy deepseek4_mix_mmq_prefill_policy( PrefillAttentionMode mode, const char * gcn_arch, const char * explicit_value); +// MoE expert parallelism requested through LUCE_DS4_MOE_TP* (server_main's +// --expert-device writes them). +struct Ds4MoeTpConfig { + bool requested = false; + bool in_process = false; + bool backend_valid = true; + PlacementBackend secondary_backend = PlacementBackend::Auto; + int secondary_gpu = 0; + bool all_on_secondary = false; + bool concentrate_secondary = false; + bool profile_hot_on_secondary = false; +}; + +// What a DeepSeek4Backend reads from the process environment, resolved once +// when init() starts. Everything after that, a reinstatement included, uses +// these values: the environment can change under a long-lived server (the +// launch profile and --expert-device write it), and a model that comes back +// must come back with the devices and modes it was loaded with. +struct DeepSeek4ResolvedSettings { + Ds4MoeTpConfig tp; + bool force_full_load = false; // LUCE_DS4_FORCE_FULL_LOAD + bool dense_tp_mask = false; // LUCE_DS4_DENSE_TP_MASK + bool fused_decode_env = false; // LUCE_DS4_FUSED_DECODE + bool decode_all_cold = false; // LUCE_DS4_DECODE_ALL_COLD + // DSpark: --draft, or LUCE_DS4_SPEC with LUCE_DS4_DRAFT. + bool spec_requested = false; + std::string spec_env_draft_path; + // DSpark placement: an explicit --draft-device, else LUCE_DS4_DRAFT_GPU / + // LUCE_DS4_DRAFT_BACKEND, else the target. + bool draft_backend_valid = true; + PlacementBackend draft_backend = PlacementBackend::Auto; + int draft_gpu = 0; + bool draft_separate_stream = false; // LUCE_DS4_DRAFT_SEPARATE_STREAM + bool draft_low_priority = false; // LUCE_DS4_DRAFT_LOW_PRIORITY +}; + class DeepSeek4Backend : public ModelBackend { public: explicit DeepSeek4Backend(DeepSeek4BackendConfig cfg); @@ -98,6 +134,38 @@ class DeepSeek4Backend : public ModelBackend { bool unpark(ParkTarget target) override; bool is_target_parked() const override { return parked_; } + // Eviction (runtime model swapping). Supported for the classic serial + // worker with in-process experts on two GPUs of one runtime (the + // ds41-lucebox profile); supports_eviction() names what else refuses. + // partial: release the primary device (dense weights, primary expert + // stacks, KV cache, draft, graphs, scratch, backend context); + // keep the secondary expert stack, the streamed expert cache + // (suspended) and the locked host tier. + // cold: release both devices. + // warm: refused (DS4 keeps no weights warm). + // Host prefix snapshots survive both and restore after reinstate(). + // partial -> cold releases the retained tier; cold -> partial and a + // repeated level do nothing. + bool supports_eviction(std::string & reason) const override; + EvictResult evict(EvictLevel level = EvictLevel::cold) override; + bool reinstate(std::string & error) override; + bool is_evicted() const override { return residency_ != Residency::resident; } + // Partially evicted, the secondary expert stack is still on its device. + bool weights_resident() const override { return residency_ != Residency::fully_evicted; } + size_t weight_device_bytes() const override; + + // Where the expert tiers stand, for logs and the eviction tests. + struct ExpertTierReport { + bool partially_evicted = false; + uint64_t primary_upload_bytes = 0; // into the primary stacks, since load + uint64_t secondary_upload_bytes = 0; // into the secondary stack, since load + const void * storage = nullptr; // identity of the expert storage + const void * secondary_base = nullptr; // first secondary stack buffer + int stream_cache_slots = 0; + bool stream_cache_suspended = false; + }; + ExpertTierReport expert_tier_report() const; + GenerateResult generate_impl(const GenerateRequest & req, const DaemonIO & io) override; @@ -135,6 +203,7 @@ class DeepSeek4Backend : public ModelBackend { private: DeepSeek4BackendConfig cfg_; + DeepSeek4ResolvedSettings settings_; ggml_backend_t backend_ = nullptr; ggml_backend_t snap_backend_ = nullptr; ggml_backend_t expert_backend_ = nullptr; @@ -184,7 +253,6 @@ class DeepSeek4Backend : public ModelBackend { // DSpark speculative decode (opt-in: --draft , or // LUCE_DS4_SPEC=1 + LUCE_DS4_DRAFT=). - bool spec_requested_ = false; bool spec_enabled_ = false; bool spec_drafter_parked_ = false; // LUCE_DS4_DRAFT_SWAP: a pinned host mirror of the drafter's core @@ -341,6 +409,31 @@ class DeepSeek4Backend : public ModelBackend { const DaemonIO & io, ThinkingBudget & budget); + enum class Residency { resident, partially_evicted, fully_evicted }; + Residency residency_ = Residency::resident; + // DSpark was loaded when the model was evicted: reinstate() reloads it. + bool reload_drafter_on_reinstate_ = false; + // The placement init() resolved, kept so a reinstatement never places the + // experts again (a placement recomputed from the free memory of the + // moment would move rows under the retained secondary stack and cache). + struct FrozenPlacement { + bool valid = false; // captured at the end of init() + MoeHybridPlacement placement; + MoeHybridPlacement decode_placement; + std::vector> cold_expert_ids; // explicit secondary set + uint64_t spill_carve_reserve = 0; // secondary host spill, when set + uint64_t spill_host_want = 0; + size_t stream_expert_bytes = 0; // fallback stream engine scratch + MoeHybridOwnershipRecord ownership; + std::vector selection_bias_host; // after the routing adjustments + }; + FrozenPlacement frozen_; + // Eviction steps; each runs on the worker and tolerates partly built state. + void release_primary_state(); + void release_secondary_state(); + bool reinstate_primary(std::string & error); + bool reinstate_full(std::string & error); + bool verify_frozen_state(std::string & error) const; bool load_model(); bool init_hybrid_model(); bool init_streamed_expert_tier(); @@ -351,6 +444,16 @@ class DeepSeek4Backend : public ModelBackend { bool validate_prefill_mode() const; bool validate_model_features() const; bool init_engram(); + // --ds4-router-bias / --ds4-protected-experts as parsed from their files + // the first time (load_routing_adjustments), [n_layer * n_expert] each. + struct RoutingInputs { + bool parsed = false; + int n_layer = 0; + int n_expert = 0; + std::vector router_bias_delta; + std::vector protected_experts; + }; + RoutingInputs routing_inputs_; bool load_routing_adjustments(); bool apply_routing_adjustments(); bool upload_protected_routing(); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 84674dd35f..343af81d77 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -6656,7 +6656,7 @@ static size_t ds4_decode_attn_cache_budget(DeepSeek4LayerRangeCache & rc) { if (rc.device >= 0 && rc.backend && ggml_backend_is_cuda(rc.backend)) { size_t free_bytes = 0; size_t total_bytes = 0; - ggml_backend_cuda_get_device_memory(rc.device, &free_bytes, &total_bytes); + ds4_device_memory(rc.device, &free_bytes, &total_bytes); (void) total_bytes; budget = free_bytes / 4; } @@ -11121,7 +11121,7 @@ bool deepseek4_step_layer_range( if (required_bytes > attn_bytes_before) { size_t free_bytes = 0; size_t total_bytes = 0; - ggml_backend_cuda_get_device_memory( + ds4_device_memory( device, &free_bytes, &total_bytes); (void) total_bytes; // The route and hot-owner arenas belong to the @@ -11178,7 +11178,7 @@ bool deepseek4_step_layer_range( // arena from growing by only a few dozen MiB. ds4_free_route_and_hot_arenas(moe_hybrid); size_t free_after_evict = 0; - ggml_backend_cuda_get_device_memory( + ds4_device_memory( device, &free_after_evict, &total_bytes); std::fprintf(stderr, "[deepseek4] evicted fused decode graphs " @@ -11955,7 +11955,7 @@ bool deepseek4_prefill_layer_major( constexpr size_t kResidualHeadroom = (size_t) 1536 << 20; const size_t bytes = (size_t) n_tokens * hc_dim * sizeof(float); size_t free_b = 0, total_b = 0; - ggml_backend_cuda_get_device_memory(device, &free_b, &total_b); + ds4_device_memory(device, &free_b, &total_b); if (free_b > bytes + kResidualHeadroom) { ggml_init_params p{}; p.mem_size = 2 * ggml_tensor_overhead(); @@ -12258,7 +12258,7 @@ void deepseek4_release_prefill_scratch( size_t free_before = 0; size_t total_bytes = 0; if (runtime && runtime->device >= 0) { - ggml_backend_cuda_get_device_memory( + ds4_device_memory( runtime->device, &free_before, &total_bytes); } @@ -12278,20 +12278,7 @@ void deepseek4_release_prefill_scratch( ds4_layer_major_meta_arena.shrink_to_fit(); ds4_layer_major_meta_owner = nullptr; } - if (moe_hybrid) { - if (moe_hybrid->prefill_route_alloc) { - ggml_gallocr_free(moe_hybrid->prefill_route_alloc); - moe_hybrid->prefill_route_alloc = nullptr; - } - if (moe_hybrid->prefill_hot_alloc) { - ggml_gallocr_free(moe_hybrid->prefill_hot_alloc); - moe_hybrid->prefill_hot_alloc = nullptr; - } - if (moe_hybrid->prefill_cold_alloc) { - ggml_gallocr_free(moe_hybrid->prefill_cold_alloc); - moe_hybrid->prefill_cold_alloc = nullptr; - } - } + if (moe_hybrid) moe_hybrid->release_prefill_allocators(); static const bool report_release = [] { const char * value = std::getenv("LUCE_DS4_TIMING"); @@ -12300,7 +12287,7 @@ void deepseek4_release_prefill_scratch( }(); if (report_release && runtime && runtime->device >= 0) { size_t free_after = 0; - ggml_backend_cuda_get_device_memory( + ds4_device_memory( runtime->device, &free_after, &total_bytes); std::fprintf(stderr, "[deepseek4] released prefill scratch: " diff --git a/server/src/deepseek4/deepseek4_hc_cuda.cu b/server/src/deepseek4/deepseek4_hc_cuda.cu index 13a24113db..941d1e3865 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.cu +++ b/server/src/deepseek4/deepseek4_hc_cuda.cu @@ -36,18 +36,24 @@ struct HcCudaScratch { explicit HcCudaScratch(int device = -1) : owner_device(device) {} - ~HcCudaScratch() { + ~HcCudaScratch() { release(); } + + // Frees every allocation on the owning device; the next ensure() + // allocates again. Restores the calling thread's device. + void release() { + int previous = -1; + const bool restore = owner_device >= 0 && cudaGetDevice(&previous) == cudaSuccess && + previous != owner_device; if (owner_device >= 0) { (void) cudaSetDevice(owner_device); } - if (d_state) cudaFree(d_state); - if (d_sums) cudaFree(d_sums); - if (d_mix) cudaFree(d_mix); - if (d_scale) cudaFree(d_scale); - if (d_base) cudaFree(d_base); - if (d_working) cudaFree(d_working); - if (d_post) cudaFree(d_post); - if (d_comb) cudaFree(d_comb); + for (float ** p : {&d_state, &d_sums, &d_mix, &d_scale, &d_base, &d_working, &d_post, &d_comb}) { + if (*p) cudaFree(*p); + *p = nullptr; + } + state_cap = working_cap = 0; + batch_cap = 0; + if (restore) (void) cudaSetDevice(previous); } bool ensure(size_t hc_dim, size_t n_embd, int n_tokens = 1) { @@ -404,6 +410,20 @@ bool hc_pre_device_locked(HcCudaScratch & scratch, } // namespace +void deepseek4_cuda_hc_release_device(int device) { + HcScratchSlot * slot = nullptr; + { + std::lock_guard lock(g_scratch_slots_mutex); + if (device < 0 || (size_t) device >= g_scratch_slots.size()) return; + slot = g_scratch_slots[(size_t) device].get(); + } + if (!slot) return; + // The slot itself stays (other threads may hold its address); a helper + // running on it finishes first. + std::lock_guard lock(slot->mutex); + slot->scratch.release(); +} + bool deepseek4_cuda_hc_set_device(int device) { if (device < 0) { return false; diff --git a/server/src/deepseek4/deepseek4_hc_cuda.h b/server/src/deepseek4/deepseek4_hc_cuda.h index b10555f956..46fedbe7bc 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.h +++ b/server/src/deepseek4/deepseek4_hc_cuda.h @@ -8,6 +8,11 @@ namespace luce::common { // before invoking the direct HC helpers. bool deepseek4_cuda_hc_set_device(int device); +// Frees the process-wide direct-helper scratch of `device` (a later helper +// call on it allocates again). For eviction: the scratch is owned by no model +// object, so a model releasing a device must drop it explicitly. +void deepseek4_cuda_hc_release_device(int device); + bool deepseek4_cuda_hc_pre_mix(const float * hc_state_host, const void * fn_device, int n_embd, diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index f2bd8fccfe..1eacb2ba5b 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -374,6 +374,9 @@ struct DeepSeek4Weights { // bias applied), [n_layer * n_expert], taken once after the adjustments. // Empty when a layer has no F32 bias; host routing then reads the device. std::vector selection_bias_host; + // The router bias delta was added to these ffn_exp_probs_b (a fresh load + // starts false; apply_routing_adjustments refuses a second application). + bool router_bias_applied = false; bool fused_decode = false; bool fused_verify_f16_kv = false; }; @@ -1023,4 +1026,9 @@ struct DeepSeek4Snapshot { }; +// Free and total memory of a DS4 device. An integrated GPU without a carve +// reports its GTT window, which the kernel's TTM limit can set far past the +// machine's RAM; it then has only what the host has available. +void ds4_device_memory(int device, size_t * free_b, size_t * total_b); + } // namespace luce::common diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index 009ed16265..1695ea2802 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -2895,6 +2895,7 @@ void free_deepseek4_weights(DeepSeek4Weights & w) { if (w.routing_ctx) { ggml_free(w.routing_ctx); w.routing_ctx = nullptr; } w.layers.clear(); w.selection_bias_host.clear(); + w.router_bias_applied = false; w.embedder.tok_embd_owned.clear(); w.embedder.tok_embd_bytes = nullptr; w.moe_hybrid = false; diff --git a/server/src/internal.h b/server/src/internal.h index dfd52d8759..f4e57c0ca7 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -255,7 +255,9 @@ struct TargetWeights { // Computed from n_layer at load time: step = (n_layer - 2) / (N - 1), // ids[k] = 1 + k * step. E.g. 27B→{1,16,31,46,61}, 9B→{1,8,15,22,29}. int n_capture_layers = LUCE_DRAFT_N_TARGET_LAYERS; - int capture_layer_ids[LUCE_DRAFT_N_TARGET_LAYERS] = {1, 16, 31, 46, 61}; + // Drafters with explicit target_layer_ids may raise the count up to + // LUCE_DRAFT_MAX_TARGET_LAYERS (see resolve_drafter_capture_layers). + int capture_layer_ids[LUCE_DRAFT_MAX_TARGET_LAYERS] = {1, 16, 31, 46, 61}; }; // Check if a token is an end-of-sequence marker for the given target weights. @@ -813,6 +815,8 @@ bool migrate_prefill_cache(const TargetWeights & w, // the next state-resident verify (§5.2). Commits the bank the just-run verify // wrote (1 - specla_pending_bank) into durable SSM and conv state, so it must // be called before the host-side bank rotation. accepted_idx is in path order. +// Free the calling thread's SpecLA commit-graph allocator (model eviction). +void specla_release_thread_scratch(); bool specla_commit_accepted(TargetCache & cache, ggml_backend_t backend, const int32_t * accepted_idx, @@ -876,6 +880,17 @@ struct QwenPrefillSegment { int seq_slot = 0; }; +// Residual-stream readout for /v1/hidden_states: the output of block `layer` +// over token rows [row_begin, row_begin + n_rows) of this graph (one prompt +// chunk). The graph exposes that range's last row and, when `sum`, the sum of +// its rows, both f32 [n_embd] (QwenGraphOutputs::hidden_last/hidden_sum). +struct QwenHiddenCapture { + int layer = -1; + int row_begin = 0; + int n_rows = 0; + bool sum = false; +}; + struct QwenGraphInputs { ggml_tensor * inp_embed; // [hidden, n_tokens, 1] f32 — pre-embedded by the caller ggml_tensor * positions; // [4 * n_tokens] i32 (M-RoPE needs 4 per token) @@ -978,6 +993,10 @@ struct QwenGraphInputs { int specla_n_waves = 0; int specla_n_boundaries = 0; int specla_max_parallel_chains = 0; + + // Optional residual-stream readouts (see QwenHiddenCapture). + const QwenHiddenCapture * hidden_captures = nullptr; + int n_hidden_captures = 0; }; struct QwenGraphOutputs { @@ -994,6 +1013,10 @@ struct QwenGraphOutputs { // One entry per target layer. Populated only when capture_moe_router is // true; qwen35 dense layers and non-MoE models leave entries null. std::vector moe_selected; + // Parallel to QwenGraphInputs::hidden_captures: f32 [n_embd] outputs; + // hidden_sum entries are null where the capture asked for no sum. + std::vector hidden_last; + std::vector hidden_sum; }; struct QwenLayerPrefnOutputs { diff --git a/server/src/qwen3/qwen3_drafter.cpp b/server/src/qwen3/qwen3_drafter.cpp index 27979e9396..4900fa52b7 100644 --- a/server/src/qwen3/qwen3_drafter.cpp +++ b/server/src/qwen3/qwen3_drafter.cpp @@ -197,7 +197,11 @@ bool load_drafter(const std::string & gguf_path, int /*gpu_layers*/, if (arch == DrafterArch::Qwen35_0p8b) { auto * st = new Qwen35DrafterState(); - if (!load_target_gguf(gguf_path, out.backend, st->weights)) { + // Scoring runs the layers only; skip output_norm + LM head (the + // 0.8B ties it to token_embd, which would otherwise be uploaded). + TargetLoadPlan plan; + plan.load_output = false; + if (!load_target_gguf_partial(gguf_path, out.backend, plan, st->weights)) { delete st; return false; } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index c26af182e1..91aa5e390c 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -61,6 +61,7 @@ Qwen35SeqEngine::Qwen35SeqEngine( scratch_row_(scratch_row), fixed_chain_(fixed_chain) { const int n_slots = slots_.slot_count(); + slot_logits_row_.assign(static_cast(n_slots), -1); slot_draft_kv_.resize(static_cast(n_slots)); seq_lens_.assign(static_cast(n_slots), 0); reserve_growth_.assign(static_cast(n_slots), 0); @@ -279,6 +280,7 @@ bool Qwen35SeqEngine::chain_spec_input_capable( } const Qwen35Slot & slot = slots_.slot(input.slot); return slot.decoding() && !slot.sampler.needs_logit_processing() && + !slot.sampler.wants_logprobs() && slot.cur_pos >= 1 && slot.cur_pos + (width > 0 ? width : fixed_chain_.width) <= slots_.max_context(); @@ -340,8 +342,12 @@ Qwen35SeqEngine::prepare_chain_drafts( std::vector lanes; lanes.reserve(inputs.size()); const int hidden = b_.w_.n_embd; + // The drafter GGUF's MASK id wins over the target default, as in the + // single-stream path (the z-lab Qwen3.5-9B drafter trains on 248077). + const int32_t mask_id = b_.dw_.mask_token_id >= 0 + ? b_.dw_.mask_token_id : b_.w_.mask_token_id; std::vector noise( - static_cast(draft_width), b_.w_.mask_token_id); + static_cast(draft_width), mask_id); std::vector noise_embed( static_cast(hidden) * draft_width); @@ -377,7 +383,7 @@ Qwen35SeqEngine::prepare_chain_drafts( } noise[0] = input.token; std::fill( - noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + noise.begin() + 1, noise.end(), mask_id); if (!b_.w_.embedder.embed( noise.data(), draft_width, noise_embed.data())) { reset_lanes(); @@ -415,7 +421,7 @@ Qwen35SeqEngine::prepare_chain_drafts( dummies.push_back(std::move(dummy)); } noise[0] = lanes.front().root; - std::fill(noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + std::fill(noise.begin() + 1, noise.end(), mask_id); if (!b_.w_.embedder.embed( noise.data(), draft_width, noise_embed.data())) { reset_lanes(); @@ -491,6 +497,44 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( return result; } +bool Qwen35SeqEngine::supports_hidden_states() const { + return b_.supports_hidden_states(); +} + +SeqEngine::AdmitResult Qwen35SeqEngine::admit_hidden_states( + uint64_t request_id, + const std::vector & prompt, + const HiddenStatesSpec & spec, + PrefixCaptureTicket capture) { + AdmitResult refused; + refused.status = AdmitResult::Status::failed; + if (!supports_hidden_states()) { + refused.error = "this model does not report hidden states"; + return refused; + } + for (int layer : spec.layers) { + if (layer < 0 || layer >= b_.w_.n_layer) { + refused.error = "hidden-state layer out of range"; + return refused; + } + } + // Greedy, logprob-free sampler: the prefill's committing row is sampled + // like any other, and the token is discarded. + AdmitResult result = admit(request_id, prompt, SamplerCfg{}); + if (result.status != AdmitResult::Status::admitted) return result; + if (slot_hidden_.size() < static_cast(slots_.slot_count())) { + slot_hidden_.resize(static_cast(slots_.slot_count())); + } + slot_hidden_[static_cast(result.slot)].emplace(spec, b_.w_.n_embd); + // cache_prefix: store the prefill as a checkpoint, never restore one. + // The scheduler places the capture at the prompt end, which is already + // where the last chunk stops, so the readout's chunking is unchanged. + if (arm_capture(result.slot, capture, /*restored_tokens=*/0)) { + result.prefix_store.capture = capture; + } + return result; +} + bool Qwen35SeqEngine::supports_images() const { return b_.supports_images(); } @@ -686,9 +730,11 @@ PrefixStoreEvent Qwen35SeqEngine::capture_prefix( int32_t Qwen35SeqEngine::sample_graph_row( int slot, int logits_row, const int32_t * cached_argmax, - std::vector * logits_scratch) { + std::vector * logits_scratch, + std::optional * logprobs_out) { const TargetWeights & w = b_.w_; const int vocab = w.n_vocab; + slot_logits_row_[static_cast(slot)] = logits_row; Qwen35Slot & seq = slots_.slot(slot); int32_t token = -1; if (seq.sampler.needs_logit_processing()) { @@ -712,9 +758,36 @@ int32_t Qwen35SeqEngine::sample_graph_row( (size_t)logits_row * sizeof(int32_t), sizeof(int32_t)); ggml_backend_synchronize(b_.target_backend_); } - return b_.apply_min_tokens_floor( + token = b_.apply_min_tokens_floor( token, seq.generated_tokens(), (size_t)logits_row * (size_t)vocab * sizeof(float)); + if (logprobs_out && token >= 0 && seq.sampler.wants_logprobs()) { + std::vector local_logits; + std::vector & logits = logits_scratch + ? *logits_scratch + : local_logits; + logits.resize((size_t)vocab); + ggml_backend_tensor_get( + b_.sg_.logits, logits.data(), + (size_t)logits_row * (size_t)vocab * sizeof(float), + sizeof(float) * (size_t)vocab); + *logprobs_out = compute_token_logprobs( + logits.data(), vocab, token, seq.sampler.logprobs_top_n); + } + return token; +} + +std::optional Qwen35SeqEngine::token_logprobs(int slot, int32_t token) { + if (slot < 0 || slot >= slot_count() || !slots_.is_active(slot) || + slot_logits_row_[static_cast(slot)] < 0) return std::nullopt; + const SamplerCfg & sampler = slots_.slot(slot).sampler; + if (!sampler.wants_logprobs()) return std::nullopt; + const int vocab = b_.w_.n_vocab; + logits_buf_.resize(static_cast(vocab)); + ggml_backend_tensor_get(b_.sg_.logits, logits_buf_.data(), + static_cast(slot_logits_row_[static_cast(slot)]) * vocab * sizeof(float), + static_cast(vocab) * sizeof(float)); + return compute_token_logprobs(logits_buf_.data(), vocab, token, sampler.logprobs_top_n); } bool Qwen35SeqEngine::upload_block_table_delta( @@ -963,6 +1036,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( int position = -1; int64_t physical_row = -1; int32_t pending = -1; + std::optional logprobs; }; const int full_width = fixed_chain_.width; @@ -1443,7 +1517,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( ArLane & lane = ar_lanes[static_cast(lane_index)]; lane.pending = sample_graph_row( lane.slot, lane_index, - &posterior[static_cast(lane_index)], &logits_buf_); + &posterior[static_cast(lane_index)], &logits_buf_, + &lane.logprobs); if (lane.pending < 0) { result.error = "compact AR sampling failed"; return result; @@ -1478,8 +1553,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step_chain_spec( proposal.tokens.begin() + 1, proposal.tokens.begin() + proposal.accepted); } else { - output.token = - ar_lanes[static_cast(ar_for_input[i])].pending; + ArLane & lane = ar_lanes[static_cast(ar_for_input[i])]; + output.token = lane.pending; + output.logprobs = std::move(lane.logprobs); } result.decode.push_back(std::move(output)); } @@ -1679,6 +1755,27 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { : std::max(1, n_commits)) : 0; + // Residual-stream readouts for hidden-states slots in this step, with + // the (prefill index, layer entry) each one feeds. + std::vector hidden_captures; + std::vector> hidden_owners; + { + int row = 0; + for (size_t i = 0; i < prefills.size(); ++i) { + const int slot = plan.prefills[i].slot; + if (hidden_slot(slot)) { + const HiddenStatesSpec & spec = + slot_hidden_[static_cast(slot)]->spec(); + for (size_t k = 0; k < spec.layers.size(); ++k) { + hidden_captures.push_back( + {spec.layers[k], row, prefills[i].chunk, spec.mean}); + hidden_owners.emplace_back(i, k); + } + } + row += prefills[i].chunk; + } + } + bool built = false; if (with_prefill) { built = build_target_step( @@ -1697,7 +1794,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { /*paged_max_kv_len=*/max_kv_len, /*n_prefill_tokens=*/n_prefill, segments.data(), (int)segments.size(), gather_rows, - /*compact_slots=*/with_decode); + /*compact_slots=*/with_decode, + hidden_captures.data(), (int)hidden_captures.size()); } else { built = build_target_step( sg, w, b_.cache_, b_.target_backend_, @@ -1899,13 +1997,36 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { ggml_backend_synchronize(b_.target_backend_); } + if (!hidden_captures.empty()) { + std::vector last((size_t)hidden), sum((size_t)hidden); + for (size_t c = 0; c < hidden_captures.size(); ++c) { + const auto [i, k] = hidden_owners[c]; + ggml_backend_tensor_get(sg.hidden_last[c], last.data(), 0, + sizeof(float) * last.size()); + const bool has_sum = sg.hidden_sum[c] != nullptr; + if (has_sum) { + ggml_backend_tensor_get(sg.hidden_sum[c], sum.data(), 0, + sizeof(float) * sum.size()); + } + slot_hidden_[static_cast(plan.prefills[i].slot)]->record( + k, last.data(), has_sum ? sum.data() : nullptr); + } + for (size_t i = 0; i < prefills.size(); ++i) { + const int slot = plan.prefills[i].slot; + if (hidden_slot(slot)) { + slot_hidden_[static_cast(slot)]->add_rows(prefills[i].chunk); + } + } + } + for (size_t oi = 0; oi < inputs.size(); ++oi) { DecodeOutput & out = decode_outputs[oi]; if (out.failed) continue; slots_.commit_step(out.slot); const int row = decode_row0 + output_rows_[oi]; out.token = sample_graph_row( - out.slot, row, &argmax_buf_[(size_t)row], &logits_buf_); + out.slot, row, &argmax_buf_[(size_t)row], &logits_buf_, + &out.logprobs); } int commit_row = 0; @@ -1923,9 +2044,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { out.status = PrefillOutput::Status::completed; out.token = sample_graph_row( slot, commit_row, &argmax_buf_[(size_t)commit_row], - &logits_buf_); + &logits_buf_, &out.logprobs); ++commit_row; slots_.commit_prefill(slot); + if (hidden_slot(slot)) { + out.hidden_states = + slot_hidden_[static_cast(slot)]->finish(); + } } prefill_outputs.push_back(std::move(out)); } @@ -1978,6 +2103,8 @@ bool Qwen35SeqEngine::evict_kv(int slot, int32_t pending_token, void Qwen35SeqEngine::retire(int slot) { offload_.discard(slot); clear_slot_images(slot); + clear_slot_hidden(slot); + slot_logits_row_[static_cast(slot)] = -1; if (!slots_.is_active(slot)) return; reset_slot_draft_kv(slot); slots_.retire(slot); diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 420f61631f..83b167af76 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -88,6 +88,7 @@ class Qwen35SeqEngine final : public SeqEngine { const PrefixStorePlan & plan) override; StepResult step(const StepPlan & plan) override; + std::optional token_logprobs(int slot, int32_t token) override; // A step that carries a prompt chunk decodes without the drafter, one // token per request. Beside a prompt of more than 8192 new tokens (a // document, not an agent's next step), the decoding requests get half @@ -132,6 +133,17 @@ class Qwen35SeqEngine final : public SeqEngine { bool token_is_eos(int32_t token) const override; + bool supports_hidden_states() const override; + AdmitResult admit_hidden_states(uint64_t request_id, + const std::vector & prompt, + const HiddenStatesSpec & spec, + PrefixCaptureTicket capture) override; + // A hidden-states slot accumulates over one pass of its prompt; a + // re-prefill would count rows twice. + bool kv_recomputable(int slot) const override { + return !hidden_slot(slot); + } + private: struct PrefillStage { @@ -155,6 +167,19 @@ class Qwen35SeqEngine final : public SeqEngine { int rope_delta = 0; }; std::vector slot_images_; + + // Hidden-states slots (admit_hidden_states): the running readout, fed + // by every prefill chunk of the slot and emitted on its commit. + std::vector> slot_hidden_; + bool hidden_slot(int slot) const { + return slot >= 0 && slot < static_cast(slot_hidden_.size()) && + slot_hidden_[static_cast(slot)].has_value(); + } + void clear_slot_hidden(int slot) { + if (slot >= 0 && slot < static_cast(slot_hidden_.size())) { + slot_hidden_[static_cast(slot)].reset(); + } + } int rope_delta(int slot) const { return slot >= 0 && slot < static_cast(slot_images_.size()) ? slot_images_[static_cast(slot)].rope_delta : 0; @@ -186,9 +211,12 @@ class Qwen35SeqEngine final : public SeqEngine { const char * client_message); PrefillStage stage_prefill_chunk(int slot, int max_tokens, std::vector & outputs); + // `logprobs_out`, when non-null, receives the chosen token's + // log-probabilities if the slot's sampler asks for them. int32_t sample_graph_row(int slot, int logits_row, const int32_t * cached_argmax = nullptr, - std::vector * logits_scratch = nullptr); + std::vector * logits_scratch = nullptr, + std::optional * logprobs_out = nullptr); std::vector select_chain_lanes( const StepPlan & plan) const; // width: the block the round drafts (0 = the configured block). @@ -291,6 +319,7 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector rows_buf_; std::vector argmax_buf_; std::vector logits_buf_; + std::vector slot_logits_row_; }; } // namespace luce::common diff --git a/server/src/qwen35/gguf_target_loader.cpp b/server/src/qwen35/gguf_target_loader.cpp index 3d35711012..16eb6c8065 100644 --- a/server/src/qwen35/gguf_target_loader.cpp +++ b/server/src/qwen35/gguf_target_loader.cpp @@ -10,6 +10,8 @@ // token_embd.weight [hidden, vocab] // output_norm.weight [hidden] F32 // output.weight [hidden, vocab] Q6_K (lm_head) +// absent when embeddings are tied (0.8B, +// 2B): token_embd.weight is the LM head. // // Per layer blk. (full-attention layers, i.e. i % 4 == 3): // attn_norm.weight [hidden] F32 @@ -49,6 +51,9 @@ #include "common/gguf_inspect.h" #include "common/layer_split_utils.h" #include "common/gguf_mmap.h" + +#include +#include #include "common/tensor_file_reader.h" #include "common/gguf_bounds.h" @@ -171,12 +176,16 @@ static bool is_expert_tensor_name(const char * name) { (len == 16 && std::strncmp(base, "ffn_gate_up_exps", 16) == 0); } +// token_embd.weight normally stays host-only for the CPU embedder. With tied +// embeddings (no output.weight, e.g. Qwen3.5-0.8B/2B) it doubles as the LM +// head, so the stage that owns the output also uploads it. static bool should_load_target_tensor(const char * name, int layer_begin, int layer_end, bool load_output, + bool tied_output, bool skip_expert_tensors = false) { - if (std::strcmp(name, "token_embd.weight") == 0) return false; + if (std::strcmp(name, "token_embd.weight") == 0) return tied_output && load_output; if (std::strcmp(name, "output_norm.weight") == 0 || std::strcmp(name, "output.weight") == 0) { return load_output; @@ -334,6 +343,13 @@ bool load_target_gguf_partial(const std::string & path, ggml_backend_t backend, const TargetLoadPlan & plan_in, TargetWeights & out) { + // LUCE_LOAD_PROFILE=1 prints the time of each load stage (swap profiling). + static const bool profile = std::getenv("LUCE_LOAD_PROFILE") != nullptr; + using clock = std::chrono::steady_clock; + const auto t_begin = clock::now(); + const auto ms_since = [](clock::time_point t) { + return std::chrono::duration(clock::now() - t).count(); + }; // ── 1. Parse metadata + create a ggml_context holding tensor descriptors ─ ggml_context * meta_ctx = nullptr; @@ -577,11 +593,17 @@ bool load_target_gguf_partial(const std::string & path, out.tok_embd = g("token_embd.weight"); out.out_norm = g("output_norm.weight"); out.output = g("output.weight"); - if (!out.tok_embd || !out.out_norm || !out.output) { - set_last_error("missing top-level tensors (token_embd/output_norm/output)"); + if (!out.tok_embd || !out.out_norm) { + set_last_error("missing top-level tensors (token_embd/output_norm)"); gguf_free(gctx); return false; } + const bool tied_output = out.output == nullptr; + if (tied_output) { + out.output = out.tok_embd; + std::printf("[loader] no output.weight: LM head tied to token_embd.weight (%s)\n", + ggml_type_name(out.tok_embd->type)); + } out.n_vocab = (int)out.tok_embd->ne[1]; for (int il = 0; il < (int)n_layer; il++) { @@ -692,7 +714,7 @@ bool load_target_gguf_partial(const std::string & path, for (int64_t tid = 0; tid < n_tensors; tid++) { const char * tname = gguf_get_tensor_name(gctx, tid); ggml_tensor * t = ggml_get_tensor(meta_ctx, tname); - if (!t || !should_load_target_tensor(tname, plan.layer_begin, plan.layer_end, plan.load_output, plan.skip_expert_tensors)) { + if (!t || !should_load_target_tensor(tname, plan.layer_begin, plan.layer_end, plan.load_output, tied_output, plan.skip_expert_tensors)) { continue; } TargetTensorAlloc a; @@ -964,10 +986,12 @@ bool load_target_gguf_partial(const std::string & path, return true; } + const double ms_meta = ms_since(t_begin); // ── 4. mmap the file and copy tensor bytes to CUDA ──────────────── // // SKIP uploading token_embd.weight — it stays on CPU for embedding - // lookup (CUDA get_rows doesn't support k-quants). Its bytes are copied + // lookup (CUDA get_rows doesn't support k-quants) — unless it is the tied + // LM head, in which case the output stage gets a device copy too. Its bytes are copied // into owned host memory below (step 5), so the mmap is released when this // local goes out of scope. GgufMmap mm; @@ -1002,16 +1026,17 @@ bool load_target_gguf_partial(const std::string & path, tok_embd_off = off; tok_embd_sz = sz; tok_embd_type = gguf_get_tensor_type(gctx, tid); - continue; + if (!tied_output) continue; } ggml_tensor * t = ggml_get_tensor(meta_ctx, tname); if (!t) continue; - if (!should_load_target_tensor(tname, plan.layer_begin, plan.layer_end, plan.load_output, plan.skip_expert_tensors)) { + if (!should_load_target_tensor(tname, plan.layer_begin, plan.layer_end, plan.load_output, tied_output, plan.skip_expert_tensors)) { continue; } spans.push_back({t, 0, off, sz}); total += sz; } + const auto t_upload = clock::now(); { std::string read_err; if (!load_tensor_spans(path, mm_addr, mm_len, spans, &read_err)) { @@ -1026,6 +1051,8 @@ bool load_target_gguf_partial(const std::string & path, } } + const double ms_upload = ms_since(t_upload); + const auto t_post = clock::now(); // ── 4b. Read NVFP4 per-tensor weight scales (optional; 1.0 for non-NVFP4). // // Scale tensors are F32 shape [1] — a single float per matmul weight. @@ -1132,6 +1159,12 @@ bool load_target_gguf_partial(const std::string & path, out.embedder.n_vocab = out.n_vocab; out.embedder.row_bytes = tok_embd_sz / (size_t)out.n_vocab; + if (profile) { + std::fprintf(stderr, "[load-profile] %s: metadata+alloc %.0f ms, upload %.0f ms " + "(%.2f GiB, %zu spans, %.2f GB/s), post %.0f ms, total %.0f ms\n", + path.c_str(), ms_meta, ms_upload, total / 1073741824.0, spans.size(), + ms_upload > 0 ? total / ms_upload / 1e6 : 0.0, ms_since(t_post), ms_since(t_begin)); + } // Stash the total for callers that want to print it char summary[192]; std::snprintf(summary, sizeof(summary), diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index c986c79275..12cff3a8a1 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -383,8 +383,20 @@ bool build_target_step( const QwenPrefillSegment * prefill_segments, int n_prefill_segments, int n_logits_rows, - bool compact_slots) { + bool compact_slots, + const QwenHiddenCapture * hidden_captures, + int n_hidden_captures) { step_graph_free(sg); + if (n_hidden_captures < 0 || (n_hidden_captures > 0 && !hidden_captures)) { + return false; + } + for (int k = 0; k < n_hidden_captures; ++k) { + const QwenHiddenCapture & hc = hidden_captures[k]; + if (hc.layer < 0 || hc.layer >= w.n_layer || hc.n_rows <= 0 || + hc.row_begin < 0 || hc.row_begin + hc.n_rows > n_tokens) { + return false; + } + } // Compact n_seqs is a decode graph bucket width, not the physical // slot count. active_slot_ids maps live rows to cache columns and uses -1 @@ -493,6 +505,13 @@ bool build_target_step( for (int i = 0; i < n_prefill_segments; ++i) { shape.push_back(prefill_segments[i].n_tokens); } + // Hidden-state readouts add graph nodes; keep those shapes apart. + for (int k = 0; k < n_hidden_captures; ++k) { + const QwenHiddenCapture & hc = hidden_captures[k]; + shape.push_back(hc.layer); + shape.push_back(hc.row_begin); + shape.push_back(hc.n_rows * 2 + (hc.sum ? 1 : 0)); + } uint64_t hash = 1469598103934665603ull; for (int value : shape) { hash ^= (uint32_t)value; @@ -675,12 +694,16 @@ bool build_target_step( gi.specla_n_waves = hld_schedule.n_waves; gi.specla_n_boundaries = hld_schedule.n_boundaries; gi.specla_max_parallel_chains = hld_schedule.max_parallel_chains; + gi.hidden_captures = hidden_captures; + gi.n_hidden_captures = n_hidden_captures; QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); if (!go.logits) return false; sg.logits = go.logits; sg.delta_captures = std::move(go.delta_captures); sg.moe_selected = std::move(go.moe_selected); + sg.hidden_last = std::move(go.hidden_last); + sg.hidden_sum = std::move(go.hidden_sum); ggml_set_output(sg.logits); sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index 02dbff0a10..83b501727c 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -168,6 +168,9 @@ bool build_hybrid_full_layer_step( // input mapping every token to its slot-local feature-ring destination. This // keeps accepted-path replay graph-stable and leaves legacy offset capture // unchanged for callers that do not use paged serving. +// `hidden_captures` — optional residual-stream readouts over row ranges of +// this graph (see QwenHiddenCapture); results land in sg.hidden_last / +// sg.hidden_sum, parallel to the array. bool build_target_step( StepGraph & sg, const TargetWeights & w, @@ -192,7 +195,9 @@ bool build_target_step( const QwenPrefillSegment * prefill_segments = nullptr, int n_prefill_segments = 0, int n_logits_rows = 0, - bool compact_slots = false); + bool compact_slots = false, + const QwenHiddenCapture * hidden_captures = nullptr, + int n_hidden_captures = 0); // Full target forward: DDTree tree-verify mode. bool build_target_step_tree( diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 51f3b90d17..fb6709531f 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -17,6 +17,7 @@ #include "prefill_helpers.h" #include "common/sampler.h" #ifdef LUCE_HAVE_GPU_SAMPLER +#include "common/geometric_draft_topk_cuda.h" #include "common/geometric_sampler_cuda.h" #include #endif @@ -36,6 +37,7 @@ #include "flashprefill.h" #include +#include #include #include #include @@ -238,21 +240,23 @@ static bool qwen35_empty_visible_output(const std::vector & tokens, static void apply_drafter_capture_layer_ids(const DraftWeights & dw, TargetWeights & w) { if (dw.capture_layer_ids.empty()) return; const int n = (int)dw.capture_layer_ids.size(); - bool ok = (n == w.n_capture_layers); - for (int k = 0; ok && k < n; k++) - ok = dw.capture_layer_ids[k] >= 0 && dw.capture_layer_ids[k] < w.n_layer; - if (!ok) { + std::vector before(w.capture_layer_ids, + w.capture_layer_ids + w.n_capture_layers); + // The drafter's fc consumes n_target_layers*n_embd features, so its id + // list must match that width; the table then takes the drafter's count + // (e.g. 8 for the z-lab Qwen3.5-9B DFlash) before any feature buffer is + // sized from w.n_capture_layers. + if (n != dw.n_target_layers || + !luce::common::adopt_drafter_capture_layers( + dw.capture_layer_ids, w.n_layer, LUCE_DRAFT_MAX_TARGET_LAYERS, + w.capture_layer_ids, w.n_capture_layers)) { std::fprintf(stderr, - "[draft] drafter target_layer_ids invalid (n=%d, slots=%d); " - "keeping derived capture layers\n", n, w.n_capture_layers); + "[draft] drafter target_layer_ids invalid (n=%d, fc layers=%d, " + "max=%d); keeping derived capture layers\n", + n, dw.n_target_layers, LUCE_DRAFT_MAX_TARGET_LAYERS); return; } - bool changed = false; - for (int k = 0; k < n; k++) { - changed |= w.capture_layer_ids[k] != dw.capture_layer_ids[k]; - w.capture_layer_ids[k] = dw.capture_layer_ids[k]; - } - if (changed) { + if (before != dw.capture_layer_ids) { std::printf("[draft] target capture layers from drafter GGUF:"); for (int k = 0; k < n; k++) std::printf(" %d", w.capture_layer_ids[k]); std::printf("\n"); @@ -357,59 +361,7 @@ bool Qwen35Backend::init() { std::printf("[draft] remote ipc ready gpu=%d cap=%d\n", cfg_.draft_gpu, cap); } else if (cfg_.draft_path) { - const std::string & dp = *cfg_.draft_path; - bool draft_ok = (dp.size() >= 5 && dp.substr(dp.size() - 5) == ".gguf") - ? load_draft_gguf(*cfg_.draft_path, draft_backend_, dw_, &w_) - : load_draft_safetensors(*cfg_.draft_path, draft_backend_, dw_, &w_); - if (!draft_ok) { - std::fprintf(stderr, "draft load: %s\n", luce_last_error()); - return false; - } - std::printf("[draft] loaded\n"); - apply_drafter_capture_layer_ids(dw_, w_); - - if (cfg_.draft_swa_window > 0) { - const DraftSwaOverrideResult swa = - apply_draft_swa_window_override(dw_, cfg_.draft_swa_window); - std::printf("[draft] SWA layers: %d/%d (window=%d)\n", - swa.swa_layers, swa.total_layers, swa.effective_window); - } - - // Legacy 8-layer drafter YaRN from config flags. Applied here AND in - // unpark() so the rotary encoding cannot flip mid-process (it used - // to switch from plain RoPE to YaRN on the first park/unpark). - if (dw_.rope_ext_factor == 0.0f && dw_.n_layer == 8 && dw_.n_embd == 2048) { - float yf = cfg_.draft_yarn_factor > 1.0f ? cfg_.draft_yarn_factor : 64.0f; - dw_.rope_freq_scale = 1.0f / yf; - dw_.rope_ext_factor = 1.0f; dw_.rope_attn_factor = 1.0f; - dw_.rope_beta_fast = cfg_.draft_yarn_beta_fast; - dw_.rope_beta_slow = cfg_.draft_yarn_beta_slow; - dw_.rope_n_ctx_orig = cfg_.draft_yarn_orig_ctx; - } - - // The checkpoint metadata is the drafter's published proposal - // horizon. Greedy chain verification keeps output byte-identical to - // plain decode at any width, so widening only risks acceptance - // depth; it is allowed up to 2x the horizon (measured on Qwen3.8 - // DFlash2: byte-identical completions across widths 8-24, commits - // grow through 16, step time cliffs past 16 with no commit gain). - if (cfg_.draft_block_size != 0) { - const int checkpoint_block_size = dw_.block_size; - if (!draft_block_size_override_supported( - cfg_.draft_block_size, checkpoint_block_size)) { - std::fprintf(stderr, - "[draft] --draft-block-size must be in [2, %d] for this " - "drafter (up to 2x the checkpoint horizon); got %d\n", - 2 * checkpoint_block_size, cfg_.draft_block_size); - return false; - } - std::printf("[draft] block size override: %d -> %d%s\n", - checkpoint_block_size, cfg_.draft_block_size, - cfg_.draft_block_size > checkpoint_block_size - ? " (beyond the published horizon; exact greedy verify)" - : ""); - dw_.block_size = cfg_.draft_block_size; - } + if (!load_local_draft()) return false; } // Create KV cache @@ -484,7 +436,7 @@ bool Qwen35Backend::init() { !tensor_parallel && !split_gpus_ && target_backend_ == draft_backend_ && cfg_.device.gpu == cfg_.draft_gpu && - dw_.selector.enabled && dw_.block_size > 1 && + dw_.block_size > 1 && dw_.block_size <= 16 && w_.output; if (fixed_chain.enabled) { fixed_chain.width = dw_.block_size; @@ -492,9 +444,8 @@ bool Qwen35Backend::init() { } if (n_slots > 1 && cfg_.draft_path && !fixed_chain.enabled) { set_last_error( - "concurrent paged DFlash2 requires a selector-enabled local " - "same-device draft with block size in [2, 16] and a target " - "lm_head"); + "concurrent paged DFlash requires a local same-device draft " + "with block size in [2, 16] and a target lm_head"); return false; } const int64_t scratch_tokens = PAGED_BLOCK_SIZE + @@ -563,6 +514,8 @@ bool Qwen35Backend::init() { std::fprintf(stderr, "cache: %s\n", luce_last_error()); return false; } + layout_max_verify_tokens_ = max_verify_tokens; + layout_ctx_alloc_ = ctx_alloc; if (cfg_.paged_attention) { const auto supported_type = [](ggml_type type) { return type == GGML_TYPE_F16 || @@ -656,10 +609,74 @@ bool Qwen35Backend::init() { std::fflush(stdout); } + init_feature_mirror(fixed_chain.enabled); + + return true; +} + +bool Qwen35Backend::load_local_draft() { + const std::string & dp = *cfg_.draft_path; + bool draft_ok = (dp.size() >= 5 && dp.substr(dp.size() - 5) == ".gguf") + ? load_draft_gguf(*cfg_.draft_path, draft_backend_, dw_, &w_) + : load_draft_safetensors(*cfg_.draft_path, draft_backend_, dw_, &w_); + if (!draft_ok) { + std::fprintf(stderr, "draft load: %s\n", luce_last_error()); + return false; + } + std::printf("[draft] loaded\n"); + apply_drafter_capture_layer_ids(dw_, w_); + + if (cfg_.draft_swa_window > 0) { + const DraftSwaOverrideResult swa = + apply_draft_swa_window_override(dw_, cfg_.draft_swa_window); + std::printf("[draft] SWA layers: %d/%d (window=%d)\n", + swa.swa_layers, swa.total_layers, swa.effective_window); + } + + // Legacy 8-layer drafter YaRN from config flags. Applied here AND in + // unpark() so the rotary encoding cannot flip mid-process (it used + // to switch from plain RoPE to YaRN on the first park/unpark). + if (dw_.rope_ext_factor == 0.0f && dw_.n_layer == 8 && dw_.n_embd == 2048) { + float yf = cfg_.draft_yarn_factor > 1.0f ? cfg_.draft_yarn_factor : 64.0f; + dw_.rope_freq_scale = 1.0f / yf; + dw_.rope_ext_factor = 1.0f; dw_.rope_attn_factor = 1.0f; + dw_.rope_beta_fast = cfg_.draft_yarn_beta_fast; + dw_.rope_beta_slow = cfg_.draft_yarn_beta_slow; + dw_.rope_n_ctx_orig = cfg_.draft_yarn_orig_ctx; + } + + // The checkpoint metadata is the drafter's published proposal + // horizon. Greedy chain verification keeps output byte-identical to + // plain decode at any width, so widening only risks acceptance + // depth; it is allowed up to 2x the horizon (measured on Qwen3.8 + // DFlash2: byte-identical completions across widths 8-24, commits + // grow through 16, step time cliffs past 16 with no commit gain). + if (cfg_.draft_block_size != 0) { + const int checkpoint_block_size = dw_.block_size; + if (!draft_block_size_override_supported( + cfg_.draft_block_size, checkpoint_block_size)) { + std::fprintf(stderr, + "[draft] --draft-block-size must be in [2, %d] for this " + "drafter (up to 2x the checkpoint horizon); got %d\n", + 2 * checkpoint_block_size, cfg_.draft_block_size); + return false; + } + std::printf("[draft] block size override: %d -> %d%s\n", + checkpoint_block_size, cfg_.draft_block_size, + cfg_.draft_block_size > checkpoint_block_size + ? " (beyond the published horizon; exact greedy verify)" + : ""); + dw_.block_size = cfg_.draft_block_size; + } + return true; +} + +void Qwen35Backend::init_feature_mirror(bool fixed_chain_enabled) { + const bool use_remote_draft = cfg_.remote_draft.enabled(); // Init feature mirror when draft model is available (needed for spec decode). // On single-GPU, this is an F32 conversion buffer; on split-GPU, a cross-device mirror. if (cfg_.draft_path && !use_remote_draft && - !fixed_chain.enabled) { + !fixed_chain_enabled) { const int mirror_cap = std::min({cfg_.draft_ctx_max, cfg_.device.max_ctx, cache_.target_feat_cap > 0 ? cache_.target_feat_cap : cfg_.device.max_ctx}); if (!draft_feature_mirror_init(feature_mirror_, draft_backend_, @@ -669,8 +686,6 @@ bool Qwen35Backend::init() { std::fprintf(stderr, "warning: feature mirror init failed, spec decode will use AR fallback\n"); } } - - return true; } bool Qwen35Backend::load_target_model(ggml_backend_t backend, TargetWeights & out) { @@ -863,6 +878,15 @@ int32_t Qwen35Backend::apply_min_tokens_floor(int32_t tok, int generated, return alt; } +void Qwen35Backend::record_ar_logprobs(size_t logits_offset, int32_t token) { + const int vocab = w_.n_vocab; + std::vector row((size_t)vocab); + ggml_backend_tensor_get(sg_.logits, row.data(), logits_offset, + sizeof(float) * (size_t)vocab); + ar_logprobs_.push_back(compute_token_logprobs( + row.data(), vocab, token, sampler_.logprobs_top_n)); +} + SeqEngine * Qwen35Backend::seq_engine() { return seq_engine_.get(); } @@ -878,6 +902,7 @@ void Qwen35Backend::print_ready_banner() const { // ── Park / unpark ─────────────────────────────────────────────────────── bool Qwen35Backend::park(ParkTarget target) { + if (evicted_) return true; // nothing is resident materialize_live_snapshot(); live_snapshot_slot_ = -1; const bool want_draft_model = park_target_includes_draft_model(target); @@ -909,6 +934,10 @@ bool Qwen35Backend::park(ParkTarget target) { } bool Qwen35Backend::unpark(ParkTarget target) { + if (evicted_) { + std::fprintf(stderr, "[unpark] model is evicted; reinstate it first\n"); + return false; + } materialize_live_snapshot(); live_snapshot_slot_ = -1; const bool want_target_model = park_target_includes_target_model(target); @@ -925,6 +954,10 @@ bool Qwen35Backend::unpark(ParkTarget target) { free_target_weights(w_); return false; } + // The reload re-derived evenly spaced capture layers; a resident + // drafter still needs the layers it was trained on. + if (!draft_parked_ && !use_remote_draft && cfg_.draft_path) + apply_drafter_capture_layer_ids(dw_, w_); kvflash_drafter_failed_ = false; // fresh VRAM: allow a retry target_parked_ = false; std::printf("[unpark] target restored\n"); std::fflush(stdout); @@ -991,6 +1024,7 @@ bool Qwen35Backend::unpark(ParkTarget target) { // ── Snapshots ─────────────────────────────────────────────────────────── bool Qwen35Backend::snapshot_save(int slot) { + if (evicted_) return false; if (cfg_.paged_attention) { static bool warned = false; if (!warned) { @@ -1036,6 +1070,7 @@ int Qwen35Backend::snapshot_granularity() const { } bool Qwen35Backend::snapshot_save_deferred(int slot) { + if (evicted_) return false; if (generating_ || cfg_.paged_attention || kvflash_active() || slot < 0 || slot >= PREFIX_SLOTS || cache_.cur_pos <= 0) { return snapshot_save(slot); @@ -1090,7 +1125,194 @@ int Qwen35Backend::snapshot_cur_pos(int slot) const { return prefix_snapshots_[slot].cur_pos; } +// ── Eviction (runtime model swapping) ─────────────────────────────────── + +bool Qwen35Backend::supports_eviction(std::string & reason) const { + // First milestone: one classic worker, dense KV, local same-device draft. + // Every other mode owns device state this teardown does not cover yet. + if (typeid(*this) != typeid(Qwen35Backend)) { + reason = "eviction is not implemented for this Qwen3.5 variant"; + } else if (concurrent_slots() > 1 || seq_engine_) { + reason = "eviction does not support --max-concurrency > 1 yet"; + } else if (cfg_.paged_attention) { + reason = "eviction does not support paged attention yet"; + } else if (kvflash_active()) { + reason = "eviction does not support KVFlash"; + } else if (cfg_.remote_draft.enabled()) { + reason = "eviction does not support a remote draft"; + } else if (cfg_.device.is_tensor_parallel() || tensor_parallel_) { + reason = "eviction does not support tensor parallelism"; + } else if (split_gpus_) { + reason = "eviction requires the draft on the target device"; + } else if (!cfg_.mmproj_path.empty()) { + reason = "eviction does not support vision (--mmproj)"; + } else if (snap_backend_ == target_backend_) { + // Unified memory keeps snapshots in the compute backend we free. + reason = "eviction requires host-memory snapshots (discrete GPU)"; + } else { + reason.clear(); + return true; + } + return false; +} + +void Qwen35Backend::release_device_state(bool keep_weights) { + if (target_backend_) ggml_backend_synchronize(target_backend_); + // Consumers before what they reference: the DFlash adapter captured the + // backend handle, graphs reference weights and cache tensors. + dflash_target_.reset(); + free_drafter(); + step_graph_destroy(sg_); + step_graph_destroy(draft_sg_); + step_graph_destroy(proj_sg_); + dflash2_selector_graph_invalidate(); + // Thread-local device scratch of this (worker) thread. + dflash2_release_thread_graphs(); + specla_release_thread_scratch(); +#ifdef LUCE_HAVE_GPU_SAMPLER + geometric_sampler_release_thread_scratch(); + geometric_draft_topk_release_thread_scratch(); +#endif +#ifdef LUCE_HAVE_BSA + flashprefill::dflash_bsa_free_persistent(); +#endif + draft_kv_free(draft_kv_); + draft_feature_mirror_free(feature_mirror_); + if (!keep_weights) { + // A failed reload can own partial allocations while still parked. + free_draft_weights(dw_); + free_target_weights(w_); + } + vision_.reset(); + free_target_cache(cache_); + // The backend context owns the device pools (the VMM pool can only be + // returned by destroying it), streams and BLAS handles. The draft shares + // it in the supported configuration; free each context once. + if (draft_backend_ && draft_backend_ != target_backend_) { + ggml_backend_free(draft_backend_); + } + draft_backend_ = nullptr; + // Weight buffers belong to the device's buffer type, not to the backend + // context, so a warm eviction keeps them across this free. + if (target_backend_) ggml_backend_free(target_backend_); + target_backend_ = nullptr; + if (!keep_weights) { + target_parked_ = true; + draft_parked_ = true; + } +} + +size_t Qwen35Backend::weight_device_bytes() const { + if (!weights_resident()) return 0; + size_t bytes = 0; + if (!target_parked_ && w_.buf) bytes += ggml_backend_buffer_get_size(w_.buf); + if (!target_parked_ && w_.gate_buf) bytes += ggml_backend_buffer_get_size(w_.gate_buf); + if (!draft_parked_ && dw_.buf) bytes += ggml_backend_buffer_get_size(dw_.buf); + return bytes; +} + +ModelBackend::EvictResult Qwen35Backend::evict(EvictLevel level) { + EvictResult result; + if (level == EvictLevel::partial) { + result.error = "partial eviction keeps a secondary expert tier; this model has none"; + return result; + } + const bool warm = level == EvictLevel::warm; + if (evicted_) { + // Warm to cold drops the kept weights; cold to warm cannot bring them back. + if (weights_kept_ && !warm) { + free_draft_weights(dw_); + free_target_weights(w_); + target_parked_ = draft_parked_ = true; + weights_kept_ = false; + } + result.status = EvictStatus::ok; + return result; + } + if (!supports_eviction(result.error)) return result; + // A deferred snapshot lives only in the device cache: copy it out first. + const int deferred_slot = live_snapshot_deferred_ ? live_snapshot_slot_ : -1; + materialize_live_snapshot(); + if (deferred_slot >= 0 && !prefix_snapshots_[deferred_slot].ctx) { + result.snapshots_lost = 1; + std::fprintf(stderr, "[evict] deferred snapshot slot %d could not be " + "copied to host; it will not restore\n", deferred_slot); + } + live_snapshot_slot_ = -1; + release_device_state(warm); + evicted_ = true; + weights_kept_ = warm; + result.status = EvictStatus::ok; + std::printf("[evict] device state released%s\n", warm ? " (weights kept)" : ""); + std::fflush(stdout); + return result; +} + +bool Qwen35Backend::reinstate(std::string & error) { + if (!evicted_) return true; + if (!supports_eviction(error)) return false; + const auto fail = [&](const std::string & what) { + error = what; + release_device_state(); // kept weights too: the model ends cold + weights_kept_ = false; + std::fprintf(stderr, "[reinstate] %s\n", what.c_str()); + return false; + }; + using clock = std::chrono::steady_clock; + auto t = clock::now(); + double stage_ms[5] = {}; // backend, target, draft, cache, mirror + const auto lap = [&](int stage) { + const auto now = clock::now(); + stage_ms[stage] = std::chrono::duration(now - t).count(); + t = now; + }; + target_backend_ = ggml_backend_cuda_init(cfg_.device.gpu); + if (!target_backend_) return fail("target backend init failed"); + draft_backend_ = cfg_.draft_path ? target_backend_ : nullptr; + lap(0); + // Fresh structs: the frees leave hyperparameters and bookkeeping behind, + // which stay readable while evicted but must not leak into the reload. + cache_ = TargetCache{}; + if (weights_kept_) { + // Warm: the weight buffers are still on the device; point them at + // the new context. + w_.backend = target_backend_; + if (cfg_.draft_path) dw_.backend = draft_backend_; + } else { + w_ = TargetWeights{}; + dw_ = DraftWeights{}; + if (!load_target_model(target_backend_, w_)) { + return fail(std::string("target load: ") + luce_last_error()); + } + target_parked_ = false; + lap(1); + if (cfg_.draft_path) { + if (!load_local_draft()) return fail(std::string("draft load: ") + luce_last_error()); + draft_parked_ = false; + } + lap(2); + } + if (!create_target_cache(w_, cfg_.device.max_ctx, layout_max_verify_tokens_, + target_backend_, cache_, /*prefill_only=*/true, + layout_ctx_alloc_, /*paged=*/false, /*n_slots=*/1, + /*fixed_chain=*/false, cfg_.cache_type_k, cfg_.cache_type_v)) { + return fail(std::string("cache: ") + luce_last_error()); + } + lap(3); + init_feature_mirror(/*fixed_chain_enabled=*/false); + lap(4); + kvflash_drafter_failed_ = false; + evicted_ = false; + weights_kept_ = false; + std::printf("[reinstate] device state restored: backend %.0f ms, target %.0f ms, " + "draft %.0f ms, cache %.0f ms, mirror %.0f ms\n", + stage_ms[0], stage_ms[1], stage_ms[2], stage_ms[3], stage_ms[4]); + std::fflush(stdout); + return true; +} + size_t Qwen35Backend::snapshot_bytes_estimate(int tokens) const { + if (evicted_) return 0; // cache tensors are freed // Paged serving keeps no single-sequence snapshots. if (cfg_.paged_attention || !snap_backend_ || tokens <= 0) return 0; return estimate_target_cache_snapshot_bytes( @@ -1116,7 +1338,22 @@ ModelBackend::MemoryReport Qwen35Backend::memory_report_for( const TargetCache & cache, const PrefixSnapshot * snapshots, int n_snapshots) { MemoryReport report; - if (!cache.base_buf) return report; + const auto add_snapshots = [&] { + for (int slot = 0; slot < n_snapshots; ++slot) { + const auto & snap = snapshots[slot]; + if (!snap.ctx || !snap.buf) continue; + report.snapshots.push_back({slot, snap.cur_pos, + ggml_backend_buffer_get_size(snap.buf), + ggml_backend_buffer_is_host(snap.buf)}); + } + }; + if (!cache.base_buf) { + // Evicted (or not yet initialized): no device cache, but host + // snapshots may still be retained. + add_snapshots(); + report.available = !report.snapshots.empty(); + return report; + } const auto bytes = [](const ggml_tensor * t) { return t ? (uint64_t) ggml_nbytes(t) : (uint64_t) 0; }; @@ -1139,13 +1376,7 @@ ModelBackend::MemoryReport Qwen35Backend::memory_report_for( const uint64_t named = report.cache.kv_bytes + report.cache.recurrent_bytes + report.cache.draft_feature_bytes; report.cache.other_bytes = allocated > named ? allocated - named : 0; - for (int slot = 0; slot < n_snapshots; ++slot) { - const auto & snap = snapshots[slot]; - if (!snap.ctx || !snap.buf) continue; - report.snapshots.push_back({slot, snap.cur_pos, - ggml_backend_buffer_get_size(snap.buf), - ggml_backend_buffer_is_host(snap.buf)}); - } + add_snapshots(); return report; } @@ -1164,7 +1395,7 @@ ModelBackend::SnapshotRef Qwen35Backend::snapshot_ref(int slot) const { bool Qwen35Backend::snapshot_adopt(int slot, ggml_context * ctx, ggml_backend_buffer_t buf, int cur_pos, int32_t last_tok) { - if (cfg_.paged_attention) return false; + if (cfg_.paged_attention || evicted_) return false; if (slot < 0 || slot >= PREFIX_SLOTS) return false; snapshot_free(slot); @@ -1482,6 +1713,7 @@ void Qwen35Backend::shutdown() { // ── Release scratch buffers between requests ──────────────────────────── void Qwen35Backend::release_scratch() { + if (evicted_) return; // Target graph allocator: grows during large prefill batches, not needed // between requests. Will be lazily recreated on next build_target_step(). if (sg_.alloc) { @@ -1519,6 +1751,11 @@ void Qwen35Backend::release_scratch() { GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, const DaemonIO & io) { + if (evicted_) { + GenerateResult evicted; + evicted.fail(GenerateErrorCode::ModelParked, "model is evicted"); + return evicted; + } materialize_live_snapshot(); live_snapshot_slot_ = -1; generating_ = true; @@ -1540,6 +1777,7 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, if (req.do_sample && sampler_.seed != 0) { sampler_rng_.seed(sampler_.seed); } + ar_logprobs_.clear(); // Zero delta-net recurrent state (SSM + conv) so a fresh prompt doesn't // inherit stale hidden state from the previous request. KV cache is @@ -1630,7 +1868,9 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, ThinkingBudget budget(req.budget_hook, ar_n_gen); // Image requests speculate too: the verify target shifts its rotary // positions by rope_delta_, and the drafter only proposes tokens. - if (cfg_.paged_attention || req.force_ar_decode) { + // Logprobs are recorded by AR decode only. + if (cfg_.paged_attention || req.force_ar_decode || + req.sampler.wants_logprobs()) { decode_ok = do_ar_decode(committed, budget, result.tokens, out_io, &result.budget_forced_close, &result.degenerate_decode_close); @@ -1663,15 +1903,93 @@ GenerateResult Qwen35Backend::generate_impl(const GenerateRequest & req, std::chrono::steady_clock::now() - t_decode_start).count(); } + result.logprobs = std::move(ar_logprobs_); result.succeed(); return result; } // ── Restore + generate ────────────────────────────────────────────────── +bool Qwen35Backend::compute_hidden_states(const std::vector & prompt, + const HiddenStatesSpec & spec, + int snap_slot, + const DaemonIO & io, + HiddenStates & out, + std::string & error) { + if (evicted_) { + error = "model is evicted"; + return false; + } + if (concurrent_slots() > 1) { + error = "hidden states run through the concurrent slot API"; + return false; + } + if (!supports_hidden_states()) { + error = "hidden states are unavailable with tensor parallelism or KVFlash"; + return false; + } + if (prompt.empty()) { + error = "hidden states need a non-empty prompt"; + return false; + } + for (int layer : spec.layers) { + if (layer < 0 || layer >= w_.n_layer) { + error = "hidden-state layer out of range"; + return false; + } + } + // The prefill below overwrites the live cache: a deferred snapshot is + // copied out first, and the live cache stops standing for any slot. + materialize_live_snapshot(); + live_snapshot_slot_ = -1; + // Same fresh-sequence setup as generate_impl, minus images and decode. + reset_recurrent_state(cache_); + if (cfg_.paged_attention) { + if (prompt.size() > (uint64_t)cfg_.device.max_ctx) { + error = "prompt exceeds max_ctx for paged attention"; + return false; + } + if (!begin_paged_sequence((uint32_t)prompt.size())) { + error = "paged KV sequence initialization failed"; + return false; + } + } + struct Guard { + Qwen35Backend * backend; + ~Guard() { + backend->hidden_accum_ = nullptr; + backend->end_paged_sequence(); + } + } guard{this}; + + HiddenStatesAccumulator accum(spec, w_.n_embd); + hidden_accum_ = &accum; + // Never restored from a snapshot. A snap_slot only receives the state at + // the prompt end: snap_pos == prompt size never falls inside a chunk, so + // do_prefill saves after its last chunk without reshaping any of them. + const int committed = do_prefill( + prompt, io, snap_slot >= 0 ? (int)prompt.size() : -1, snap_slot, + /*kv_offset=*/0, /*images=*/nullptr); + if (committed < 0) { + error = "prefill failed"; + return false; + } + if (io.is_cancelled() || accum.rows() != (int)prompt.size()) { + error = "prefill cancelled"; + return false; + } + out = accum.finish(); + return true; +} + GenerateResult Qwen35Backend::restore_and_generate_impl(int slot, const GenerateRequest & req, const DaemonIO & io) { + if (evicted_) { + GenerateResult evicted; + evicted.fail(GenerateErrorCode::ModelParked, "model is evicted"); + return evicted; + } // The live cache already is this snapshot when it was saved (or kept, // deferred) after the previous generation and nothing ran since. const bool live_restore = slot >= 0 && slot == live_snapshot_slot_ && @@ -1737,6 +2055,7 @@ GenerateResult Qwen35Backend::restore_and_generate_impl(int slot, if (req.do_sample && sampler_.seed != 0) { sampler_rng_.seed(sampler_.seed); } + ar_logprobs_.clear(); cache_.cur_pos = snap_pos; result.restored_prefix_tokens = snap_pos; @@ -1814,6 +2133,12 @@ GenerateResult Qwen35Backend::restore_and_generate_impl(int slot, // BEFORE its own build_target_step, so a null/freed graph tensor aborts // in ggml_backend_tensor_set. Build a single-token decode step graph at // the restored position now, mirroring do_ar_decode's per-step build. + if (req.sampler.wants_logprobs()) { + // No prefill ran, so there is no logits row for the first token. + result.fail(GenerateErrorCode::BackendSpecific, + "logprobs are unavailable on an exact snapshot hit"); + return result; + } const bool pool = kvflash_active(); if (!build_target_step(sg_, w_, cache_, target_backend_, /*kv_start=*/cache_.cur_pos, /*n_tokens=*/1, @@ -1841,7 +2166,7 @@ GenerateResult Qwen35Backend::restore_and_generate_impl(int slot, auto t_decode_start = std::chrono::steady_clock::now(); ThinkingBudget budget(req.budget_hook, req.n_gen); bool decode_ok = false; - if (req.force_ar_decode) { + if (req.force_ar_decode || req.sampler.wants_logprobs()) { decode_ok = do_ar_decode(committed, budget, result.tokens, out_io, &result.budget_forced_close, &result.degenerate_decode_close); @@ -1874,6 +2199,7 @@ GenerateResult Qwen35Backend::restore_and_generate_impl(int slot, std::chrono::steady_clock::now() - t_decode_start).count(); } + result.logprobs = std::move(ar_logprobs_); result.succeed(); return result; } @@ -2070,6 +2396,13 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, // decode-time windowed attention will later read. static const bool prefill_timing = std::getenv("LUCE_PREFILL_TIMING") != nullptr; const auto t_build0 = std::chrono::steady_clock::now(); + std::vector hidden_captures; + if (hidden_accum_) { + for (int layer : hidden_accum_->spec().layers) { + hidden_captures.push_back( + {layer, 0, n_tokens, hidden_accum_->spec().mean}); + } + } if (!build_target_step(sg_, w_, cache_, target_backend_, /*kv_start=*/kv_pos, /*n_tokens=*/n_tokens, with_mask, /*capture=*/true, @@ -2079,7 +2412,18 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, (start + n_tokens < prompt_len ? 1 : 0), cfg_.kq_stride_pad, should_capture_moe_router(), - /*kvflash_mask=*/kvf_paged)) { + /*kvflash_mask=*/kvf_paged, + /*capture_qk=*/false, + /*paged_attention=*/false, + /*n_seqs=*/1, /*seq_slot=*/0, + /*paged_max_kv_len=*/0, + /*n_prefill_tokens=*/0, + /*prefill_segments=*/nullptr, + /*n_prefill_segments=*/0, + /*n_logits_rows=*/0, + /*compact_slots=*/false, + hidden_captures.data(), + (int)hidden_captures.size())) { std::fprintf(stderr, "prefill build @%d\n", kv_pos); return -1; } @@ -2168,6 +2512,21 @@ int Qwen35Backend::do_prefill(const std::vector & tokens, } after_target_compute(sg_, kv_pos, n_tokens); + if (hidden_accum_) { + std::vector last((size_t)hidden), sum((size_t)hidden); + for (size_t k = 0; k < hidden_captures.size(); ++k) { + ggml_backend_tensor_get(sg_.hidden_last[k], last.data(), 0, + sizeof(float) * last.size()); + const bool has_sum = sg_.hidden_sum[k] != nullptr; + if (has_sum) { + ggml_backend_tensor_get(sg_.hidden_sum[k], sum.data(), 0, + sizeof(float) * sum.size()); + } + hidden_accum_->record(k, last.data(), has_sum ? sum.data() : nullptr); + } + hidden_accum_->add_rows(n_tokens); + } + int32_t last_tok = -1; const bool is_final_chunk = (start + n_tokens >= prompt_len); const size_t argmax_off = @@ -2421,6 +2780,14 @@ bool Qwen35Backend::do_ar_decode(int committed, ThinkingBudget & budget, first_tok = cache_.last_tok; } apply_budget(first_tok); + if (sampler_.wants_logprobs()) { + // An exact snapshot hit decodes without prefill logits. + if (!prefill_last_logits_valid_) { + set_last_error("logprobs need the prefill logits row"); + return false; + } + record_ar_logprobs(prefill_last_logits_offset_, first_tok); + } out_tokens.push_back(first_tok); io.emit(first_tok); if (kvflash_active()) kvflash_history_.push_back(first_tok); @@ -2560,6 +2927,7 @@ bool Qwen35Backend::do_ar_decode(int committed, ThinkingBudget & budget, next_tok, (int)out_tokens.size(), /*logits_row_offset=*/0); apply_budget(next_tok); + if (sampler_.wants_logprobs()) record_ar_logprobs(0, next_tok); out_tokens.push_back(next_tok); io.emit(next_tok); diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index 8562524790..345819a360 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -128,8 +128,33 @@ class Qwen35Backend : public ModelBackend { bool unpark(ParkTarget target) override; bool is_target_parked() const override { return target_parked_; } + bool supports_eviction(std::string & reason) const override; + EvictResult evict(EvictLevel level = EvictLevel::cold) override; + bool weights_resident() const override { return !evicted_ || weights_kept_; } + size_t weight_device_bytes() const override; + bool reinstate(std::string & error) override; + bool is_evicted() const override { return evicted_; } + GenerateResult generate_impl(const GenerateRequest & req, const DaemonIO & io) override; + // AR decode and the concurrent engine report logprobs; a logprobs request + // never speculates (see generate_impl). + bool supports_logprobs() const override { return true; } + + // Hidden states come from the dense target graph: single-device only, + // and not through KVFlash's pooled prefill, which evicts context. + bool supports_hidden_states() const override { + return !cfg_.device.is_tensor_parallel() && !kvflash_active(); + } + HiddenStateShape hidden_state_shape() const override { + return {w_.n_layer, w_.n_embd}; + } + bool compute_hidden_states(const std::vector & prompt, + const HiddenStatesSpec & spec, + int snap_slot, + const DaemonIO & io, + HiddenStates & out, + std::string & error) override; bool snapshot_save(int slot) override; bool snapshot_save_deferred(int slot) override; @@ -330,6 +355,23 @@ class Qwen35Backend : public ModelBackend { bool target_parked_ = false; bool draft_parked_ = false; + // ── Eviction state ─────────────────────────────────────────────── + // KV-cache arguments resolved at init. reinstate() replays them instead + // of re-deriving from free device memory, so snapshots taken before an + // eviction keep matching the rebuilt cache. + int layout_max_verify_tokens_ = 0; + int layout_ctx_alloc_ = 0; + bool evicted_ = false; + bool weights_kept_ = false; // evicted warm + // Local draft load (+ capture layers, SWA/YaRN/block-size overrides), + // shared by init() and reinstate(). + bool load_local_draft(); + // Draft feature mirror, shared by init() and reinstate(). + void init_feature_mirror(bool fixed_chain_enabled); + // Frees every device allocation; the shared teardown of evict() and of a + // failed reinstate(). Host snapshots are untouched. + void release_device_state(bool keep_weights = false); + // Vision projector, loaded next to the target weights when --mmproj is set. bool load_vision(); bool encode_images(const Qwen35ImagePrompt & prompt, Qwen35ImageRows & rows, @@ -356,6 +398,16 @@ class Qwen35Backend : public ModelBackend { std::size_t prefill_last_logits_offset_ = 0; bool prefill_last_logits_valid_ = false; + // do_ar_decode's per-token log-probabilities when sampler_ asks for + // them; moved into GenerateResult::logprobs by the generate paths. + std::vector ar_logprobs_; + // Set by compute_hidden_states for the duration of its do_prefill: each + // prefill chunk reads the spec's layers into it. + HiddenStatesAccumulator * hidden_accum_ = nullptr; + // Appends the log-probabilities of `token` from the sg_.logits row at + // byte offset `logits_offset`. + void record_ar_logprobs(std::size_t logits_offset, int32_t token); + // The single-request path owns one live sequence. The allocator and // attention op are sequence-aware; concurrent serving uses the engine // below and adds the Qwen-specific sequence-indexed DeltaNet state. diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 713799205c..5a818d528a 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -396,29 +396,11 @@ bool create_target_cache_partial(const TargetWeights & w, } // ── Zero-initialize all state tensors ───────────────────────────── - const bool meta_backend = ggml_backend_buft_is_meta( - ggml_backend_get_default_buffer_type(backend)); - if (meta_backend) { - ggml_backend_buffer_clear(out.base_buf, 0); - if (out.rollback_buf) ggml_backend_buffer_clear(out.rollback_buf, 0); - } else { - std::vector zeros(1 * 1024 * 1024, 0); - ggml_context * ctx_list[] = { out.base_ctx, out.rollback_ctx }; - for (int ci = 0; ci < 2; ci++) { - ggml_context * c = ctx_list[ci]; - if (!c) continue; - for (ggml_tensor * t = ggml_get_first_tensor(c); t != nullptr; - t = ggml_get_next_tensor(c, t)) { - size_t nb = ggml_nbytes(t); - size_t off = 0; - while (off < nb) { - size_t chunk = std::min(nb - off, zeros.size()); - ggml_backend_tensor_set(t, zeros.data(), off, chunk); - off += chunk; - } - } - } - } + // On the device (a memset per buffer): uploading host zeros moved the + // whole cache over PCIe, ~0.9 s at 64k context on a Gen3 x4 link, which + // dominated reinstating an evicted model. + ggml_backend_buffer_clear(out.base_buf, 0); + if (out.rollback_buf) ggml_backend_buffer_clear(out.rollback_buf, 0); return true; } @@ -794,6 +776,19 @@ bool migrate_prefill_cache(const TargetWeights & w, return true; } +namespace { +// SpecLA commit-graph allocator, reused across steps on the worker thread like +// the verify step graph. File scope so model eviction can release it. +thread_local ggml_gallocr_t t_specla_galloc = nullptr; +} // namespace + +void specla_release_thread_scratch() { + if (t_specla_galloc) { + ggml_gallocr_free(t_specla_galloc); + t_specla_galloc = nullptr; + } +} + // SpecLA DeltaConstruct commit (docs/SPECLA.md): one small graph advancing all // delta-net layers' durable SSM states along the accepted path, // S_A = exp(g⁺_A) S0 + Σ_{t∈path} exp(g⁺_A − g⁺_t) k_t ⊗ ṽ_t, @@ -893,7 +888,6 @@ bool specla_commit_accepted(TargetCache & cache, // Persistent metadata arena + allocator, reused across steps like the // verify step graph — avoids per-commit gallocr churn. static thread_local std::vector s_arena; - static thread_local ggml_gallocr_t s_galloc = nullptr; const size_t graph_nodes = (size_t)n_delta * 8 + 64; ggml_init_params ip{}; ip.mem_size = graph_nodes * 4 * ggml_tensor_overhead() + @@ -944,10 +938,10 @@ bool specla_commit_accepted(TargetCache & cache, ggml_build_forward_expand(gf, ggml_cpy(ctx, s_new, S)); } - if (!s_galloc) { - s_galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!t_specla_galloc) { + t_specla_galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); } - ok = s_galloc != nullptr && ggml_gallocr_alloc_graph(s_galloc, gf); + ok = t_specla_galloc != nullptr && ggml_gallocr_alloc_graph(t_specla_galloc, gf); if (ok) { ggml_backend_tensor_set(idx, accepted_idx, 0, sizeof(int32_t) * A); ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; @@ -2403,6 +2397,8 @@ QwenGraphOutputs build_qwen35_graph( if (in.capture_moe_router && w.is_moe) { og_early.moe_selected.assign((size_t)w.n_layer, nullptr); } + og_early.hidden_last.assign((size_t)in.n_hidden_captures, nullptr); + og_early.hidden_sum.assign((size_t)in.n_hidden_captures, nullptr); // DFlash target layer IDs for feature capture (from TargetWeights config). const int * CAPTURE_LAYERS = w.capture_layer_ids; @@ -2585,6 +2581,30 @@ QwenGraphOutputs build_qwen35_graph( } cur = ggml_add(ctx, ffn, ffn_residual); + // ── Residual-stream readout (/v1/hidden_states) ── + // Before the DFlash capture below, which may `continue`. + for (int k = 0; k < in.n_hidden_captures; ++k) { + const QwenHiddenCapture & hc = in.hidden_captures[k]; + if (hc.layer != il) continue; + ggml_tensor * rows = ggml_reshape_2d(ctx, cur, hidden, n_tokens); + if (rows->type != GGML_TYPE_F32) rows = ggml_cast(ctx, rows, GGML_TYPE_F32); + ggml_tensor * last = ggml_cont(ctx, ggml_view_1d(ctx, rows, hidden, + (size_t)(hc.row_begin + hc.n_rows - 1) * rows->nb[1])); + ggml_set_output(last); + ggml_build_forward_expand(gf, last); + og_early.hidden_last[(size_t)k] = last; + if (hc.sum) { + ggml_tensor * range = ggml_view_2d(ctx, rows, hidden, hc.n_rows, + rows->nb[1], (size_t)hc.row_begin * rows->nb[1]); + // [n_rows, hidden] -> [1, hidden]: the sum over the range. + ggml_tensor * sum = ggml_sum_rows( + ctx, ggml_cont(ctx, ggml_transpose(ctx, range))); + ggml_set_output(sum); + ggml_build_forward_expand(gf, sum); + og_early.hidden_sum[(size_t)k] = sum; + } + } + // ── DFlash layer feature capture ── // Write `cur` into the rolling target_feat buffer. The buffer is a // ring of `target_feat_cap` slots; position P maps to slot P%cap. diff --git a/server/src/qwen35moe/qwen35moe_backend.h b/server/src/qwen35moe/qwen35moe_backend.h index 0e27557bfa..cf56b0a8b3 100644 --- a/server/src/qwen35moe/qwen35moe_backend.h +++ b/server/src/qwen35moe/qwen35moe_backend.h @@ -32,6 +32,10 @@ class Qwen35MoeBackend : public Qwen35Backend { // These generation paths neither continue nor copy out a live state. bool snapshot_save_deferred(int slot) override { return snapshot_save(slot); } bool supports_dflash_spec_decode() const override { return true; } + // Its own generate paths do not fill GenerateResult::logprobs. + bool supports_logprobs() const override { return false; } + // Its hybrid prefill does not run the readouts of the dense graph. + bool supports_hidden_states() const override { return false; } bool set_routing_collector(MoeRoutingCollector * c) override { routing_collector_ = c; return true; } const MoeHybridRoutingStats * get_routing_stats() const override { return routing_stats_.get(); } diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index e71d34f108..3af4ec1005 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -30,6 +30,7 @@ #include "common/kv_rotation.h" #include "common/sha1.h" #include "freeze_history.h" +#include "utf8_utils.h" #ifdef LUCE_HAS_CURL #include @@ -675,6 +676,56 @@ void apply_no_thinking_sampler_defaults(const json & body, } } +int parse_request_logprobs(const json & body, ApiFormat format) { + // Only the chat-completions boolean form turns logprobs on. Every other + // shape (the legacy completions integer, a non-boolean value, a + // top_logprobs without logprobs: true, other endpoints) was ignored + // before logprobs existed and still is. + if (format != ApiFormat::OPENAI_CHAT) return -1; + const auto flag = body.find("logprobs"); + if (flag == body.end() || !flag->is_boolean() || !flag->get()) { + return -1; + } + const auto top_it = body.find("top_logprobs"); + if (top_it == body.end() || top_it->is_null()) return 0; + if (!top_it->is_number_integer()) { + throw std::invalid_argument("top_logprobs must be an integer"); + } + const int64_t top = top_it->get(); + if (top < 0 || top > kMaxTopLogprobs) { + throw std::invalid_argument( + "top_logprobs must be between 0 and " + + std::to_string(kMaxTopLogprobs)); + } + return (int) top; +} + +json build_openai_logprobs( + const std::vector & logprobs, + const std::function & token_bytes) { + // OpenAI reports -9999.0 for tokens it considers impossible; JSON has no + // -infinity (nlohmann would write null). + auto entry = [&](const TokenLogprob & lp) { + const std::string raw = token_bytes(lp.token); + json bytes = json::array(); + for (unsigned char c : raw) bytes.push_back((int) c); + return json{ + {"token", utf8_sanitize(raw)}, + {"logprob", std::isfinite(lp.logprob) ? lp.logprob : -9999.0f}, + {"bytes", std::move(bytes)}, + }; + }; + json content = json::array(); + for (const TokenLogprobs & position : logprobs) { + json item = entry(position.chosen); + json top = json::array(); + for (const TokenLogprob & alt : position.top) top.push_back(entry(alt)); + item["top_logprobs"] = std::move(top); + content.push_back(std::move(item)); + } + return {{"content", std::move(content)}}; +} + json require_messages_array(const json & body) { if (!body.contains("messages") || !body["messages"].is_array() || body["messages"].empty()) { @@ -683,6 +734,136 @@ json require_messages_array(const json & body) { return body["messages"]; } +namespace { +constexpr size_t kMaxHiddenStateLayers = 256; +constexpr const char * kChatmlAssistantHeader = "<|im_start|>assistant\n"; + +const char * pooling_name(HiddenStatesRequest::Pooling pooling) { + switch (pooling) { + case HiddenStatesRequest::Pooling::mean: return "mean"; + case HiddenStatesRequest::Pooling::both: return "both"; + case HiddenStatesRequest::Pooling::last: break; + } + return "last"; +} +} // namespace + +HiddenStatesRequest parse_hidden_states_request(const json & body) { + HiddenStatesRequest request; + if (!body.contains("layers") || !body["layers"].is_array() || + body["layers"].empty()) { + throw std::invalid_argument("layers must be a non-empty array of integers"); + } + if (body["layers"].size() > kMaxHiddenStateLayers) { + throw std::invalid_argument("layers may name at most " + + std::to_string(kMaxHiddenStateLayers) + " blocks"); + } + for (const json & layer : body["layers"]) { + if (!layer.is_number_integer()) { + throw std::invalid_argument("layers must be a non-empty array of integers"); + } + // JSON also stores uint64 values: check before converting so a large + // positive index cannot wrap into a valid negative layer index. + if (layer.is_number_unsigned() && layer.get() > INT32_MAX) { + throw std::invalid_argument("layer " + layer.dump() + " is out of range"); + } + const int64_t value = layer.get(); + if (value < INT32_MIN || value > INT32_MAX) { + throw std::invalid_argument("layer " + std::to_string(value) + " is out of range"); + } + if (std::find(request.layers.begin(), request.layers.end(), (int)value) != + request.layers.end()) { + throw std::invalid_argument("layer " + std::to_string(value) + " is repeated"); + } + request.layers.push_back((int)value); + } + if (body.contains("pooling") && !body["pooling"].is_null()) { + const json & pooling = body["pooling"]; + const std::string value = pooling.is_string() ? pooling.get() : ""; + if (value == "last") request.pooling = HiddenStatesRequest::Pooling::last; + else if (value == "mean") request.pooling = HiddenStatesRequest::Pooling::mean; + else if (value == "both") request.pooling = HiddenStatesRequest::Pooling::both; + else throw std::invalid_argument("pooling must be \"last\", \"mean\" or \"both\""); + } + if (body.contains("generation_prompt") && !body["generation_prompt"].is_null()) { + const json & prompt = body["generation_prompt"]; + const std::string value = prompt.is_string() ? prompt.get() : ""; + if (value == "bare") request.bare_generation_prompt = true; + else if (value != "template") { + throw std::invalid_argument( + "generation_prompt must be \"template\" or \"bare\""); + } + } + if (body.contains("cache_prefix") && !body["cache_prefix"].is_null()) { + if (!body["cache_prefix"].is_boolean()) { + throw std::invalid_argument("cache_prefix must be a boolean"); + } + request.cache_prefix = body["cache_prefix"].get(); + } + return request; +} + +HiddenStatesSpec resolve_hidden_states_spec(const HiddenStatesRequest & request, + int n_layers) { + HiddenStatesSpec spec; + spec.mean = request.wants_mean(); + for (int layer : request.layers) { + const int64_t block = layer < 0 ? (int64_t)n_layers + layer : layer; + if (block < 0 || block >= n_layers) { + throw std::invalid_argument( + "layer " + std::to_string(layer) + " is out of range for a " + + std::to_string(n_layers) + "-layer model"); + } + spec.layers.push_back((int)block); + } + return spec; +} + +bool trim_to_assistant_header(std::string & rendered) { + const std::string header = kChatmlAssistantHeader; + const size_t at = rendered.rfind(header); + if (at == std::string::npos) return false; + const size_t end = at + header.size(); + // A closed assistant turn is history, not the generation prompt. + if (rendered.find("<|im_end|>", end) != std::string::npos) return false; + rendered.resize(end); + return true; +} + +json build_hidden_states_response(const std::string & model, + const HiddenStatesRequest & request, + int n_layers, int n_embd, int prompt_tokens, + const HiddenStates & states, + const std::string & last_token, + double prefill_ms, + int cached_prefix_tokens) { + json layers = json::object(); + for (size_t k = 0; k < request.layers.size(); ++k) { + json entry = json::object(); + entry["last"] = k < states.last.size() ? states.last[k] : std::vector{}; + if (request.wants_mean()) { + entry["mean"] = k < states.mean.size() ? states.mean[k] : std::vector{}; + } + layers[std::to_string(request.layers[k])] = std::move(entry); + } + json body = { + {"object", "hidden_states"}, + {"model", model}, + {"n_layers", n_layers}, + {"n_embd", n_embd}, + {"prompt_tokens", prompt_tokens}, + {"pooling", pooling_name(request.pooling)}, + {"generation_prompt", request.bare_generation_prompt ? "bare" : "template"}, + {"last_token", last_token}, + {"layers", std::move(layers)}, + {"timings", {{"prefill_ms", prefill_ms}}}, + }; + if (request.cache_prefix) { + body["cache_prefix"] = {{"cached_tokens", cached_prefix_tokens}}; + } + return body; +} + static bool env_flag_enabled(const char * name) { const char * raw = std::getenv(name); if (!raw || !*raw) return false; @@ -1757,6 +1938,51 @@ int HttpServer::run(const std::vector & models) { // CORS belongs to the one listener even when a peer formats output. model->config_.enable_cors = config_.enable_cors; } + if (config_.swap_residency) { + for (auto & model : models_) { + HttpServer & server = *model.server; + std::string reason; + if (!server.backend_.supports_eviction(reason)) { + std::fprintf(stderr, "[server] --swap-residency: model %s: %s\n", + server.config_.model_name.c_str(), reason.c_str()); + return 2; + } + const std::string & target = server.config_.target_device; + const std::string & draft = server.config_.draft_device; + if (target.empty() || target == "auto") { + std::fprintf(stderr, "[server] --swap-residency: model %s needs an explicit --target-device\n", + server.config_.model_name.c_str()); + return 2; + } + model.devices = server.config_.residency_devices; + if (model.devices.empty()) { + model.devices = {target}; + if (!draft.empty() && draft != "auto" && draft != target) model.devices.push_back(draft); + } + // No worker runs yet, so the backend may be queried here. + model.residency = server.backend_.is_evicted() ? Residency::evicted + : Residency::resident; + model.weights_resident = server.backend_.weights_resident(); + model.weight_device_bytes = server.backend_.weight_device_bytes(); + // Startup evicts cold, or partially when the model has + // secondary devices; neither keeps warm weights. + if (model.residency == Residency::evicted && model.weights_resident) { + model.retained_devices = server.config_.residency_secondary_devices; + } + } + for (size_t i = 0; i < models_.size(); ++i) { + for (size_t j = i + 1; j < models_.size(); ++j) { + if (models_[i].residency == Residency::resident && + models_[j].residency == Residency::resident && + devices_overlap(models_[i], models_[j])) { + std::fprintf(stderr, "[server] --swap-residency: %s and %s are both resident on a shared device\n", + models_[i].server->config_.model_name.c_str(), + models_[j].server->config_.model_name.c_str()); + return 2; + } + } + } + } } // Every model is loaded and no worker has started: resolve one shared @@ -2074,6 +2300,9 @@ void HttpServer::handle_client(SocketHandle fd) { // Endpoint structure is parsed once by the listener. Model defaults, templates // and token IDs are resolved only by the selected model's handler. +// A failed swap-in answers 503 for this long before it may be retried. +static constexpr auto kSwapRetryDelay = std::chrono::seconds(30); + bool HttpServer::route_model_request(SocketHandle fd, ParsedRequest & req, bool count_only) { json & body = req.raw_body; @@ -2083,9 +2312,35 @@ bool HttpServer::route_model_request(SocketHandle fd, ParsedRequest & req, send_error(fd, 400, "token counting requires an explicit model name"); return true; } - // Generation always follows operator priority. Only token counting needs - // an explicit tokenizer, since no generating model has been selected yet. - const bool automatic = !count_only; + // Activations are only meaningful for the model that produced them, so + // a hidden-states request always names its model and never falls back. + const bool hidden = req.hidden_states.has_value(); + if (hidden && unnamed) { + send_error(fd, 400, "hidden states require an explicit model name"); + return true; + } + // Token counting needs the named model's tokenizer. Generation follows + // operator priority unless name routing pins a request to the model it + // names; a pinned request waits for that model and never spills over. + const bool by_name = config_.model_routing == ServerConfig::ModelRouting::name; + RoutedModel * pinned = nullptr; + if (count_only || hidden || (by_name && !unnamed)) { + for (auto & model : models_) { + if (model.server->config_.model_name == requested) pinned = &model; + } + if (!pinned && by_name && !hidden && config_.unknown_model_to_primary) { + pinned = &models_.front(); + } + if (!pinned) { + std::string names; + for (const auto & model : models_) { + names += (names.empty() ? "" : ", ") + model.server->config_.model_name; + } + send_error(fd, 404, "unknown model '" + requested + "'; available models: " + names); + return true; + } + } + const bool automatic = pinned == nullptr; // A routing reservation owns a backend slot through retirement and output // draining. Engine admission, on its worker, remains the KV authority. @@ -2111,21 +2366,43 @@ bool HttpServer::route_model_request(SocketHandle fd, ParsedRequest & req, --owner.routing_waiters_; } } waiting{*this}; + // Swap residency: a pinned request waiting for its model to come in. + // Token counting needs only the tokenizer, never residency. + const bool swapping = config_.swap_residency && !count_only; + struct ModelWaiting { + HttpServer & owner; + RoutedModel * model = nullptr; + ~ModelWaiting() { + if (!model) return; + std::lock_guard lock(owner.routing_mu_); + --model->waiters; + } + } model_waiting{*this}; std::unordered_set unfit; for (;;) { if (http_detail::inspect_peer_socket(fd) == http_detail::PeerSocketState::Disconnected) return true; - bool known = automatic; + bool needs_residency = false; for (auto & model : models_) { - if (!automatic && requested != model.server->config_.model_name) continue; - known = true; + if (pinned && &model != pinned) continue; if (unfit.count(model.server)) continue; { std::lock_guard lock(routing_mu_); if (stopping_.load()) break; + // Only a resident model that is not draining admits. Auto + // requests use residents and never trigger a swap. + if (swapping && (model.residency != Residency::resident || model.closing)) { + if (&model == pinned) needs_residency = true; + continue; + } if (!count_only && model.in_flight >= model.capacity) continue; if (!count_only) ++model.in_flight; + if (model_waiting.model == &model) { + --model.waiters; + if (model.cohort_remaining > 0) --model.cohort_remaining; + model_waiting.model = nullptr; + } if (waiting.active) { --routing_waiters_; waiting.active = false; @@ -2147,15 +2424,59 @@ bool HttpServer::route_model_request(SocketHandle fd, ParsedRequest & req, unfit.insert(model.server); } } + if (needs_residency) { + std::string error; + { + std::lock_guard lock(routing_mu_); + if (pinned->residency == Residency::failed && + std::chrono::steady_clock::now() - pinned->failed_at < kSwapRetryDelay) { + error = pinned->last_error.empty() ? "swap failed" : pinned->last_error; + } + if (!model_waiting.model) { + model_waiting.model = pinned; + ++pinned->waiters; + } + } + if (!error.empty() || !drive_swap(*pinned, error)) { + send_error(fd, 503, "model " + pinned->server->config_.model_name + + " is unavailable: " + error); + return true; + } + // Swapped in (by this request or a concurrent one): admit now + // instead of sleeping out the monitor interval. + { + std::lock_guard lock(routing_mu_); + if (pinned->residency == Residency::resident && !pinned->closing) continue; + } + } else if (swapping && automatic) { + // Unnamed requests never swap a model in, so they would wait + // forever once no resident model can take them (none resident + // after a failed swap, or the request exceeds every resident's + // context); a swap in progress may bring another one in. + bool any_resident = false; + { + std::lock_guard lock(routing_mu_); + for (const auto & model : models_) { + any_resident = any_resident || (model.residency == Residency::resident && + !unfit.count(model.server)); + } + any_resident = any_resident || swap_active_; + } + if (!any_resident && !unfit.empty()) { + send_error(fd, 400, "request exceeds the context or KV pool capacity of every resident model; " + "name a model to swap it in"); + return true; + } + if (!any_resident) { + send_error(fd, 503, "no model is resident; name a model to swap it in"); + return true; + } + } if (unfit.size() == models_.size()) { send_error(fd, 400, "request exceeds the context or KV pool capacity of every model"); return true; } - if (!known) { - send_error(fd, 404, "unknown model '" + requested + "'"); - return true; - } - if (!automatic || stopping_.load()) { + if (count_only || stopping_.load()) { send_error(fd, 503, "no available model capacity or server stopping"); return true; } @@ -2179,6 +2500,226 @@ bool HttpServer::route_model_request(SocketHandle fd, ParsedRequest & req, } } +static bool any_shared(const std::vector & a, const std::vector & b) { + for (const auto & device : a) { + if (std::find(b.begin(), b.end(), device) != b.end()) return true; + } + return false; +} + +bool HttpServer::devices_overlap(const RoutedModel & a, const RoutedModel & b) { + return any_shared(a.devices, b.devices); +} + +bool HttpServer::retained_conflict(const RoutedModel & model, const RoutedModel & target) { + if (&model == &target || model.residency == Residency::resident) return false; + if (!any_shared(model.retained_devices, target.devices)) return false; + return !model.retained_warm || target.devices.size() > 1; +} + +ModelBackend::EvictLevel HttpServer::eviction_level(const RoutedModel & model, + const RoutedModel & target) { + const auto & secondary = model.server->config_.residency_secondary_devices; + if (!secondary.empty() && !any_shared(secondary, target.devices)) { + return ModelBackend::EvictLevel::partial; + } + if (model.server->config_.swap_keep_weights && + model.devices.size() == 1 && target.devices.size() == 1) { + return ModelBackend::EvictLevel::warm; + } + return ModelBackend::EvictLevel::cold; +} + +void HttpServer::record_control(RoutedModel & model, const ControlOp & op) { + if (!op.ran) return; + (op.evict ? model.last_evict_ms : model.last_reinstate_ms) = op.ms; + model.weights_resident = op.weights_resident; + model.weight_device_bytes = op.weight_device_bytes; + if (!op.weights_resident || (!op.evict && op.ok)) { + model.retained_devices.clear(); + model.retained_warm = false; + } else if (op.evict && op.ok) { + const bool partial = op.level == ModelBackend::EvictLevel::partial; + model.retained_devices = partial + ? model.server->config_.residency_secondary_devices : model.devices; + model.retained_warm = !partial; + } +} + +bool HttpServer::drive_swap(RoutedModel & target, std::string & error) { + std::vector conflicts; + // Already evicted, but still holding memory the target needs. + std::vector releasing; + { + std::lock_guard lock(routing_mu_); + if (swap_active_ || stopping_.load() || target.residency == Residency::resident) { + return true; + } + // Check under the same lock that claims the swap: another request + // may have failed this target since the caller last checked it. + if (target.residency == Residency::failed && + std::chrono::steady_clock::now() - target.failed_at < kSwapRetryDelay) { + error = target.last_error; + return false; + } + for (auto & model : models_) { + if (retained_conflict(model, target)) releasing.push_back(&model); + if (&model == &target || model.residency != Residency::resident) continue; + if (!devices_overlap(model, target)) continue; + // Requests that were waiting when it came in are served before + // it is swapped away; departed waiters no longer count. + if (std::min(model.cohort_remaining, model.waiters) > 0) return true; + conflicts.push_back(&model); + } + swap_active_ = true; + for (auto * model : conflicts) model->closing = true; + } + const std::string & name = target.server->config_.model_name; + std::fprintf(stderr, "[residency] swapping in %s (evicting %zu model%s)\n", + name.c_str(), conflicts.size(), conflicts.size() == 1 ? "" : "s"); + + const auto finish = [&] { + std::lock_guard lock(routing_mu_); + for (auto * model : conflicts) model->closing = false; + swap_active_ = false; + routing_cv_.notify_all(); + }; + // Reinstate models already evicted by this swap after a failure. + const auto restore = [&](const std::vector & evicted) { + for (auto * model : evicted) { + auto op = model->server->post_control(/*evict=*/false); + op->wait(); + std::lock_guard lock(routing_mu_); + record_control(*model, *op); + if (op->ok) { + model->residency = Residency::resident; + } else { + model->residency = Residency::failed; + model->failed_at = std::chrono::steady_clock::now(); + model->last_error = op->error; + } + } + }; + + // Draining models admit nothing new; wait out their reservations. The + // last release may come from a probe that does not notify, so poll too. + { + std::unique_lock lock(routing_mu_); + for (;;) { + bool drained = true; + for (auto * model : conflicts) drained = drained && model->in_flight == 0; + if (drained) break; + routing_cv_.wait_for(lock, std::chrono::milliseconds(50)); + } + } + + std::vector evicted; + for (auto * model : conflicts) { + auto op = model->server->post_control(/*evict=*/true, eviction_level(*model, target)); + op->wait(); + bool ok = false; + { + std::lock_guard lock(routing_mu_); + record_control(*model, *op); + ok = op->ok; + if (ok) { + model->residency = Residency::evicted; + } else { + // Rejected before any teardown: it stays resident. + model->last_error = op->error; + } + } + if (!ok) { + error = "evicting " + model->server->config_.model_name + ": " + op->error; + { + std::lock_guard lock(routing_mu_); + target.residency = Residency::failed; + target.failed_at = std::chrono::steady_clock::now(); + target.last_error = error; + } + restore(evicted); + finish(); + return false; + } + evicted.push_back(model); + } + // A target spanning several devices loads with nothing else's weights on + // them, and a retained tier yields to a model that uses its device. + for (auto * model : releasing) { + auto cold = model->server->post_control(/*evict=*/true, ModelBackend::EvictLevel::cold); + cold->wait(); + { + std::lock_guard lock(routing_mu_); + record_control(*model, *cold); + } + if (!cold->ok) { + error = "releasing " + model->server->config_.model_name + ": " + cold->error; + { + std::lock_guard lock(routing_mu_); + target.residency = Residency::failed; + target.failed_at = std::chrono::steady_clock::now(); + target.last_error = error; + } + restore(evicted); + finish(); + return false; + } + } + + auto op = target.server->post_control(/*evict=*/false); + op->wait(); + // Evicted models that still keep weights on the target's devices. + std::vector kept; + { + std::lock_guard lock(routing_mu_); + record_control(target, *op); + for (auto & model : models_) { + if (&model != &target && model.residency != Residency::resident && + any_shared(model.retained_devices, target.devices)) { + kept.push_back(&model); + } + } + } + if (!op->ok && !kept.empty()) { + // Not enough device memory beside the kept weights: drop every other + // model's weights sharing the target's devices, then retry. + std::fprintf(stderr, "[residency] %s reinstate failed with weights kept (%s); " + "evicting others cold and retrying\n", name.c_str(), op->error.c_str()); + for (auto * model : kept) { + auto cold = model->server->post_control(/*evict=*/true, ModelBackend::EvictLevel::cold); + cold->wait(); + std::lock_guard lock(routing_mu_); + record_control(*model, *cold); + } + op = target.server->post_control(/*evict=*/false); + op->wait(); + std::lock_guard lock(routing_mu_); + record_control(target, *op); + } + if (!op->ok) { + { + std::lock_guard lock(routing_mu_); + target.residency = Residency::failed; + target.failed_at = std::chrono::steady_clock::now(); + target.last_error = op->error; + } + error = op->error; + restore(evicted); + finish(); + return false; + } + { + std::lock_guard lock(routing_mu_); + target.residency = Residency::resident; + target.last_error.clear(); + ++target.swaps_in; + target.cohort_remaining = target.waiters; + } + std::fprintf(stderr, "[residency] %s resident (reinstate %.0f ms)\n", name.c_str(), op->ms); + finish(); + return true; +} + json HttpServer::model_routing_status() { json models = json::array(); std::lock_guard lock(routing_mu_); @@ -2193,10 +2734,24 @@ json HttpServer::model_routing_status() { {"target_device", server.config_.target_device}, {"draft_device", server.config_.draft_device}, {"max_context", server.config_.max_ctx}, + {"residency", model.residency == Residency::resident ? "resident" + : model.residency == Residency::evicted ? "evicted" : "failed"}, + {"devices", model.devices}, + {"weights_resident", model.weights_resident}, + {"weight_device_bytes", model.weight_device_bytes}, + {"retained_devices", model.retained_devices}, + {"draining", model.closing}, + {"waiting_for_residency", model.waiters}, + {"swaps_in", model.swaps_in}, + {"last_evict_ms", model.last_evict_ms}, + {"last_reinstate_ms", model.last_reinstate_ms}, + {"last_residency_error", model.last_error}, {"props", std::move(props)}, {"status", server.status_.to_json()}}); } - return {{"routing", "primary-first"}, {"waiting", routing_waiters_}, + const bool by_name = config_.model_routing == ServerConfig::ModelRouting::name; + return {{"routing", by_name ? "by-name" : "primary-first"}, {"waiting", routing_waiters_}, + {"swap_residency", config_.swap_residency}, {"swap_active", swap_active_}, {"queue_limit", config_.routing_queue_limit}, {"models", models}}; } @@ -2225,6 +2780,14 @@ bool HttpServer::parse_common_request_fields( } req.sampler = parse_request_sampler(body, config_.sampler_defaults); + req.sampler.logprobs_top_n = parse_request_logprobs(body, req.format); + // Logprobs are read per token on the plain decode path; the SSE emitter + // has no per-token logprobs channel yet. Where they cannot be produced + // the request is served without them, as before logprobs existed. + if (req.sampler.wants_logprobs() && + (req.stream || !backend_.supports_logprobs())) { + req.sampler.logprobs_top_n = -1; + } if (body.contains("tools")) req.tools = body["tools"]; // Tool choice constraint for hint generation. if (body.contains("tool_choice")) req.tool_choice = body["tool_choice"]; @@ -2297,6 +2860,15 @@ bool HttpServer::parse_endpoint_request( } return true; } + if (path == "/v1/hidden_states") { + // Chat-style messages rendered through the model's template, but the + // response is the prefill's residual stream, not a completion. + req.format = ApiFormat::OPENAI_CHAT; + req.response_id = generate_id("hs"); + req.messages = require_messages_array(body); + req.hidden_states = parse_hidden_states_request(body); + return true; + } if (path == "/v1/messages/count_tokens") { // Shares Anthropic's message parsing; the flag makes route_request // short-circuit before enqueueing a generation job. @@ -2567,8 +3139,19 @@ bool HttpServer::render_and_tokenize_request( send_error(fd, 400, error); return false; } + if (req.hidden_states && req.hidden_states->bare_generation_prompt && + !trim_to_assistant_header(req.rendered_prompt)) { + send_error(fd, 400, + "generation_prompt=\"bare\" needs a chat template that ends in an " + "open <|im_start|>assistant header"); + return false; + } req.started_in_thinking = prompt_ends_in_open_think(req.rendered_prompt); req.prompt_tokens = tokenizer_.encode(req.rendered_prompt); + if (req.hidden_states && req.prompt_tokens.empty()) { + send_error(fd, 400, "hidden states need a non-empty prompt"); + return false; + } return true; } @@ -2578,7 +3161,8 @@ bool HttpServer::validate_request_context( SocketHandle fd, const ParsedRequest & req, bool send_failure) { const int prompt_tokens = (int) req.prompt_tokens.size(); const bool pflash_will_run = - !req.images && config_.pflash_mode != ServerConfig::PflashMode::OFF && + !req.images && !req.hidden_states && + config_.pflash_mode != ServerConfig::PflashMode::OFF && drafter_tokenizer_ != nullptr && (config_.pflash_mode == ServerConfig::PflashMode::ALWAYS || prompt_tokens >= config_.pflash_threshold); @@ -2659,6 +3243,9 @@ bool HttpServer::route_request(SocketHandle fd, const HttpRequest & hr) { return models_.empty() ? handle_model_request(fd, req, count_tokens_only) : route_model_request(fd, req, count_tokens_only); + } catch (const std::invalid_argument & e) { + send_error(fd, 400, std::string("Invalid request: ") + e.what()); + return true; } catch (const std::exception & e) { send_error(fd, 400, std::string("JSON parse error: ") + e.what()); return true; @@ -2671,6 +3258,18 @@ bool HttpServer::handle_model_request(SocketHandle fd, ParsedRequest & req, try { const json & body = req.raw_body; if (!parse_common_request_fields(fd, body, req)) return true; + if (req.hidden_states) { + if (!backend_.supports_hidden_states()) { + send_error(fd, 400, "hidden states are not supported by this model backend"); + return true; + } + const auto shape = backend_.hidden_state_shape(); + req.hidden_spec = resolve_hidden_states_spec(*req.hidden_states, + shape.n_layers); + // Prefill only: nothing is generated or streamed. + req.stream = false; + req.max_output = 0; + } // Image extraction and redaction apply only to an image-capable // backend; every other backend sees the request exactly as before. std::vector encoded_images; @@ -2686,6 +3285,10 @@ bool HttpServer::handle_model_request(SocketHandle fd, ParsedRequest & req, send_error(fd, 400, extraction_error); return true; } + if (req.hidden_states && !encoded_images.empty()) { + send_error(fd, 400, "hidden states do not accept image input"); + return true; + } req.messages = std::move(normalized); redact_image_urls(req.raw_body); redact_image_urls(req.messages); @@ -2714,7 +3317,8 @@ bool HttpServer::handle_model_request(SocketHandle fd, ParsedRequest & req, // PPP rearrange (optional): peel ephemeral system banners into a // following system message so the first chat boundary is stable. std::vector render_messages = chat_messages; - if (encoded_images.empty() && config_.ppp_enabled && config_.ppp_rearrange && !req.tools.empty()) { + if (encoded_images.empty() && !req.hidden_states && config_.ppp_enabled && + config_.ppp_rearrange && !req.tools.empty()) { auto layout = PinFriendlyPrompt::rearrange(chat_messages, true); if (layout.rearranged) { render_messages = std::move(layout.messages); @@ -2945,6 +3549,11 @@ json build_openai_completion_response( {"message", message}, {"finish_reason", finish_reason}, }; + if (req.sampler.wants_logprobs() && tokenizer) { + choice["logprobs"] = build_openai_logprobs( + result.logprobs, + [tokenizer](int32_t id) { return tokenizer->token_text(id); }); + } if (req.thinking_opt_in) { // finish_details mirrors ds4_eval.c's eval_think_close_info. // close_kind is "natural" when the model closed its own thinking @@ -3688,7 +4297,10 @@ std::vector HttpServer::request_restore_points(const std::vector & HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( const ParsedRequest & req, PreparedPrompt & prepared, GenerateRequest & generate_request) { - if (req.images) return {}; + // Logprobs need the prefill logits row, which an exact snapshot hit + // skips; logprobs requests are short classification prompts, so they + // bypass the prefix caches entirely. + if (req.images || req.sampler.wants_logprobs()) return {}; auto & effective_prompt = prepared.tokens; // Tool-heavy requests prefer the reusable system/tool boundary under eviction. const bool prefer_inline_snap = !req.tools.empty(); @@ -4139,7 +4751,7 @@ void HttpServer::finalize_generation_cache( GenerationCacheState & cache, const GenerateResult & result, int completion_tokens, bool visible_output_seen, bool client_disconnected) { - if (req.images) return; + if (req.images || req.sampler.wants_logprobs()) return; const auto & effective_prompt = prepared.tokens; const bool generation_produced_output = result.ok() && completion_tokens > 0 && visible_output_seen && !client_disconnected; @@ -4739,7 +5351,8 @@ void HttpServer::send_nonstream_response( const std::vector & gen_tokens, int n_gen_cap, bool budget_forced_close, bool degenerate_decode_close, const GenTimings & gen_timings, - ClientSendBuffer * send_buffer) { + ClientSendBuffer * send_buffer, + std::vector logprobs) { CompletionTokenCounts counts; counts.total = (int) gen_tokens.size(); const bool is_eos = !gen_tokens.empty() && @@ -4755,6 +5368,7 @@ void HttpServer::send_nonstream_response( result.tokens = gen_tokens; result.budget_forced_close = budget_forced_close; result.degenerate_decode_close = degenerate_decode_close; + result.logprobs = std::move(logprobs); const json response = build_non_streaming_response( req, result, n_gen_cap, gen_timings, counts, emitter, &tokenizer_); @@ -4771,15 +5385,197 @@ void HttpServer::send_nonstream_response( void HttpServer::worker_loop() { publish_memory_report(); while (true) { - ServerJob * job = dequeue(); - if (!job) break; // stopping + run_pending_controls(); + ServerJob * job = dequeue(/*wake_for_control=*/true); + if (!job) { + if (stopping_.load()) break; + continue; // a control transition is pending + } process_job(job); publish_memory_report(); } + fail_pending_controls("server stopping"); +} + +std::shared_ptr HttpServer::post_control(bool evict, + ModelBackend::EvictLevel level) { + auto op = std::make_shared(); + op->evict = evict; + op->level = level; + if (backend_.seq_engine()) { + op->finish(false, "residency control requires the serial worker"); + return op; + } + { + std::lock_guard lock(queue_mu_); + if (!stopping_.load()) { + control_queue_.push_back(op); + queue_cv_.notify_all(); + return op; + } + } + op->finish(false, "server stopping"); + return op; +} + +void HttpServer::run_pending_controls() { + for (;;) { + std::shared_ptr op; + { + std::lock_guard lock(queue_mu_); + if (control_queue_.empty()) return; + op = std::move(control_queue_.front()); + control_queue_.pop_front(); + } + const auto start = std::chrono::steady_clock::now(); + bool ok = false; + std::string error; + if (op->evict) { + const auto result = backend_.evict(op->level); + ok = result.status == ModelBackend::EvictStatus::ok; + error = result.error; + if (result.snapshots_lost > 0) { + std::fprintf(stderr, "[residency] %s: %d deferred prefix snapshot%s lost on eviction\n", + config_.model_name.c_str(), result.snapshots_lost, + result.snapshots_lost == 1 ? "" : "s"); + } + } else { + ok = backend_.reinstate(error); + } + op->ms = std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); + op->ran = true; + op->weights_resident = backend_.weights_resident(); + op->weight_device_bytes = backend_.weight_device_bytes(); + const char * what = !op->evict ? "reinstate" + : op->level == ModelBackend::EvictLevel::warm ? "evict (warm)" + : op->level == ModelBackend::EvictLevel::partial ? "evict (partial)" : "evict"; + std::fprintf(stderr, "[residency] %s %s %s in %.0f ms%s%s\n", + config_.model_name.c_str(), what, + ok ? "done" : "FAILED", op->ms, error.empty() ? "" : ": ", + error.c_str()); + publish_memory_report(); + op->finish(ok, std::move(error)); + } +} + +void HttpServer::fail_pending_controls(const std::string & reason) { + std::deque> pending; + { + std::lock_guard lock(queue_mu_); + pending.swap(control_queue_); + } + for (auto & op : pending) op->finish(false, reason); +} + +json HttpServer::hidden_states_body(const ParsedRequest & req, + const HiddenStates & states, + double prefill_s) { + const auto shape = backend_.hidden_state_shape(); + const std::string last_token = req.prompt_tokens.empty() + ? std::string() + : utf8_sanitize(tokenizer_.token_text(req.prompt_tokens.back())); + return build_hidden_states_response( + config_.model_name, *req.hidden_states, shape.n_layers, shape.n_embd, + (int)req.prompt_tokens.size(), states, last_token, prefill_s * 1000.0, + cached_hidden_prefix_tokens(req)); +} + +int HttpServer::cached_hidden_prefix_tokens(const ParsedRequest & req) { + if (!req.hidden_states || !req.hidden_states->cache_prefix || + prefix_cache_.disabled()) return 0; + return prefix_cache_.lookup_candidate( + req.prompt_tokens, (int)req.prompt_tokens.size()).second; +} + +void HttpServer::process_hidden_states_job(ServerJob * job) { + const ParsedRequest & req = job->req; + { + ServerStatus::RequestInfo info; + info.model = req.model; + info.format = "hidden_states"; + status_.set_running("", (int)req.prompt_tokens.size(), false, info); + } + broadcast_status(); + StatusGuard status_guard{status_}; + + DaemonIO io; + io.should_cancel = [job]() { + return job->client_disconnected.load(std::memory_order_acquire); + }; + // cache_prefix: save the state after the whole prompt as an inline + // prefix snapshot, keyed by the prompt, so a chat request extending it + // restores instead of prefilling. Nothing is restored here. + const int prompt_len = (int)req.prompt_tokens.size(); + PrefixCache::InlineReservation reservation; + if (req.hidden_states->cache_prefix && !prefix_cache_.disabled()) { + reservation = prefix_cache_.reserve_inline_snap( + req.prompt_tokens, /*restored_prefix_len=*/0, + /*prefer_tools_boundary=*/false, /*forced_cut=*/prompt_len, + /*restore_source_slot=*/-1, + [this](int target_cut) { + return backend_.snapshot_bytes_estimate(target_cut); + }); + } + const int snap_slot = reservation.active() ? reservation.slot() : -1; + if (snap_slot >= 0) { + forget_inline_slot_metadata(snap_slot); + backend_.snapshot_free(snap_slot); + } + + HiddenStates states; + std::string error; + const auto started_at = std::chrono::steady_clock::now(); + const bool ok = backend_.compute_hidden_states( + req.prompt_tokens, req.hidden_spec, snap_slot, io, states, error); + const double prefill_s = std::chrono::duration( + std::chrono::steady_clock::now() - started_at).count(); + if (snap_slot >= 0) { + if (ok && backend_.snapshot_used(snap_slot) && + backend_.snapshot_cur_pos(snap_slot) == prompt_len) { + reservation.commit( + req.prompt_tokens, + backend_.snapshot_bytes_estimate(prompt_len)); + trim_snapshots_after_commit(snap_slot); + slot_tokens_[snap_slot] = req.prompt_tokens; + } else { + backend_.snapshot_free(snap_slot); + reservation.abort(); + } + } + backend_.release_scratch(); + std::fprintf(stderr, + "[server] hidden_states DONE %s ok=%s in=%zu layers=%zu prefill=%.1fms%s%s\n", + req.response_id.c_str(), ok ? "true" : "false", req.prompt_tokens.size(), + req.hidden_spec.layers.size(), prefill_s * 1000.0, + ok ? "" : " error=", ok ? "" : error.c_str()); + + if (!job->client_disconnected.load(std::memory_order_acquire)) { + sock_set_block(job->fd); + if (ok) { + send_response(job->fd, 200, "application/json", + hidden_states_body(req, states, prefill_s).dump() + "\n"); + } else { + const ResponseError response_error = + ResponseError::internal("hidden_states_failed", error); + send_response(job->fd, response_error_http_status(response_error), + "application/json", + build_error_response(req.format, response_error, + req.response_id).dump() + "\n"); + } + } + stop_job_stream(job); + std::lock_guard lk(job->mu); + job->done = true; + job->cv.notify_one(); } void HttpServer::process_job(ServerJob * job) { + if (job->req.hidden_states) { + process_hidden_states_job(job); + return; + } SocketHandle fd = job->fd; const auto & req = job->req; auto started_at = std::chrono::steady_clock::now(); @@ -5167,10 +5963,13 @@ bool HttpServer::has_pending_jobs() { return queue_head_ != nullptr; } -ServerJob * HttpServer::dequeue() { +ServerJob * HttpServer::dequeue(bool wake_for_control) { std::unique_lock lk(queue_mu_); + const auto control_pending = [&] { + return wake_for_control && !control_queue_.empty(); + }; // Use timed wait so the worker periodically wakes to send SSE heartbeats. - while (!queue_head_ && !stopping_.load()) { + while (!queue_head_ && !control_pending() && !stopping_.load()) { if (queue_cv_.wait_for(lk, std::chrono::seconds(30)) == std::cv_status::timeout) { // Send SSE heartbeat (comment line) to detect disconnected clients. lk.unlock(); @@ -5178,7 +5977,9 @@ ServerJob * HttpServer::dequeue() { lk.lock(); } } - if (!queue_head_) return nullptr; + // Transitions first: the swap manager posts one only once the model is + // drained. + if (control_pending() || !queue_head_) return nullptr; ServerJob * j = queue_head_; queue_head_ = j->next; if (!queue_head_) queue_tail_ = nullptr; diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 0e9d9a029b..5a4111ad69 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -38,10 +38,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include #if !defined(_WIN32) @@ -92,12 +94,60 @@ enum class HeartbeatSendResult { HeartbeatSendResult try_send_sse_heartbeat(SocketHandle fd, size_t & offset); } +// ─── Residency control ────────────────────────────────────────────────── +// An evict/reinstate transition that runs on a model's worker thread between +// requests. Owned by shared_ptr so it outlives whoever posted it. +struct ControlOp { + bool evict = true; // false: reinstate + ModelBackend::EvictLevel level = ModelBackend::EvictLevel::cold; // evictions only + std::mutex mu; + std::condition_variable cv; + bool done = false; + bool ok = false; + std::string error; + double ms = 0.0; + // The backend's weight residency once the transition ran, read on the + // worker so other threads never query a backend that may be mid-swap. + bool ran = false; // the worker executed it (the fields below are set) + bool weights_resident = false; + size_t weight_device_bytes = 0; + + void finish(bool success, std::string err) { + std::lock_guard lock(mu); + ok = success; + error = std::move(err); + done = true; + cv.notify_all(); + } + void wait() { + std::unique_lock lock(mu); + cv.wait(lock, [&] { return done; }); + } +}; + // ─── Server configuration ─────────────────────────────────────────────── struct ServerConfig { std::string host = "0.0.0.0"; int port = 8080; int max_tokens = 4096; // default max output tokens (legacy alias for default_max_tokens) int routing_queue_limit = 32; // waiting auto requests across the listener + // How a multi-model listener picks a model for a generation request. + // balance: operator priority with capacity fallback; the request's + // `model` is ignored. name: a request naming a loaded model runs only on + // it (waiting for its capacity); omitted/"auto" still balances. + enum class ModelRouting { balance, name }; + ModelRouting model_routing = ModelRouting::balance; + // Name routing: serve an unrecognized model name on the primary instead + // of answering 404. + bool unknown_model_to_primary = false; + // Swap residency (name routing): models whose devices overlap are not + // resident together; a request naming an evicted model swaps it in. + bool swap_residency = false; + // Swap residency: this model is evicted warm (its weights stay in device + // memory) so switching back skips the weight upload; cold when memory runs + // short. Meant for the main model: keeping an occasional model warm + // shrinks what the active model can use. + bool swap_keep_weights = false; int max_ctx = 0; // 0 = use backend's DevicePlacement default (8192) bool enable_cors = true; std::string model_name = "luce"; @@ -226,6 +276,12 @@ struct ServerConfig { // server_main after CLI parse. std::string target_device; std::string draft_device; + // Every device the model occupies (each layer-split shard, draft, expert + // device), for swap residency's overlap and warm-eviction decisions. + std::vector residency_devices; + // The devices a partial eviction keeps (DeepSeek V4's expert device): + // a model that does not use them can swap in without releasing them. + std::vector residency_secondary_devices; // Idle-to-busy batching window. It is ignored by single-slot engines and // never delays an already decoding request. int admission_coalesce_ms = 20; @@ -327,6 +383,24 @@ PflashQueryWindow find_pflash_query_window( // ─── Parsed request ───────────────────────────────────────────────────── +// /v1/hidden_states fields beyond `model`, `messages` and +// `chat_template_kwargs` (see docs/API.md). +struct HiddenStatesRequest { + enum class Pooling { last, mean, both }; + // As requested: 0-based block index, negative counts from the end. + std::vector layers; + Pooling pooling = Pooling::last; + // "bare": the prompt ends at the assistant header; whatever the template + // appends after it (an empty or open think block) is dropped. + bool bare_generation_prompt = false; + // Keep this prefill in the prefix cache so a chat request whose prompt + // extends this one restores it instead of prefilling it again. The + // readout itself still comes from a full prefill: this request may store + // a checkpoint, never restore one. + bool cache_prefix = false; + bool wants_mean() const { return pooling != Pooling::last; } +}; + struct ParsedRequest { ApiFormat format; std::vector prompt_tokens; // tokenized prompt @@ -381,6 +455,11 @@ struct ParsedRequest { // The prompt ends with tool results: an agent's step, which clients only // append to. bool ends_with_tool_result = false; + // Set for /v1/hidden_states: a prefill-only request whose response is + // the residual stream, not a completion. hidden_spec holds the layers + // resolved against the serving model. + std::optional hidden_states; + HiddenStatesSpec hidden_spec; }; // Resident budget the PrefixCache enforces: the concurrent limit in paged @@ -420,6 +499,53 @@ void apply_no_thinking_sampler_defaults(const json & body, const SamplingDefaults & no_thinking, SamplerCfg & sampler); +// Resolve OpenAI chat `logprobs: true` / `top_logprobs` into +// SamplerCfg::logprobs_top_n (-1 when not requested). Other shapes and +// endpoints are ignored (-1). With logprobs: true, a top_logprobs that is not +// an integer in [0, 20] throws std::invalid_argument; route_request's catch +// turns that into a 400. +int parse_request_logprobs(const json & body, ApiFormat format); + +// OpenAI chat `choices[].logprobs` object for per-token log-probabilities: +// {"content": [{token, logprob, bytes, top_logprobs: [{token, logprob, +// bytes}]}]}. `token_bytes` maps an id to its raw bytes, which may be a +// partial UTF-8 sequence; `bytes` keeps them exact and `token` carries a +// sanitized copy so the response stays valid JSON. +json build_openai_logprobs( + const std::vector & logprobs, + const std::function & token_bytes); + +// Parse the /v1/hidden_states fields `layers`, `pooling`, +// `generation_prompt` and `cache_prefix`. Throws std::invalid_argument for a +// missing, empty or non-integer `layers`, a repeated layer, an unknown +// pooling / generation prompt, or a non-boolean `cache_prefix`; the caller +// answers 400. +HiddenStatesRequest parse_hidden_states_request(const json & body); + +// Resolve requested layers against a model with `n_layers` blocks (negative +// indices count from the end). Throws std::invalid_argument when one falls +// outside the model. +HiddenStatesSpec resolve_hidden_states_spec(const HiddenStatesRequest & request, + int n_layers); + +// generation_prompt="bare": cut a rendered ChatML prompt right after its final, +// still-open `<|im_start|>assistant\n` header, dropping whatever the template +// appended (e.g. `\n\n\n\n`). False when the prompt does not +// end in an open assistant turn. +bool trim_to_assistant_header(std::string & rendered); + +// The /v1/hidden_states response body. `last_token` is the decoded text of +// the final prompt token. `cached_prefix_tokens` is reported only when the +// request asked for cache_prefix: how many leading prompt tokens the prefix +// cache holds once the request is done (0 when nothing could be stored). +json build_hidden_states_response(const std::string & model, + const HiddenStatesRequest & request, + int n_layers, int n_embd, int prompt_tokens, + const HiddenStates & states, + const std::string & last_token, + double prefill_ms, + int cached_prefix_tokens = 0); + // Read the required `messages` field. Throws std::invalid_argument when // it is missing or not a non-empty array; route_request's catch turns // that into a 400. @@ -494,11 +620,21 @@ class HttpServer { bool handle_model_request(SocketHandle fd, ParsedRequest & req, bool count_only, RoutingAdmission * admission = nullptr); json model_routing_status(); + // Swap residency: queue a transition on this model's worker. Completes + // with an error at once when the worker is stopping or not serial. + std::shared_ptr post_control( + bool evict, ModelBackend::EvictLevel level = ModelBackend::EvictLevel::cold); // Worker thread: process jobs sequentially. process_job owns the // lifecycle of one dequeued request, including signaling completion. void worker_loop(); void process_job(ServerJob * job); + // process_job for /v1/hidden_states on the single-request worker. + void process_hidden_states_job(ServerJob * job); + json hidden_states_body(const ParsedRequest & req, + const HiddenStates & states, double prefill_s); + // cache_prefix: the leading prompt tokens the prefix cache now holds. + int cached_hidden_prefix_tokens(const ParsedRequest & req); struct PreparedPrompt { std::vector tokens; @@ -616,7 +752,8 @@ class HttpServer { const std::vector & gen_tokens, int n_gen_cap, bool budget_forced_close, bool degenerate_decode_close, const GenTimings & gen_timings, - ClientSendBuffer * send_buffer = nullptr); + ClientSendBuffer * send_buffer = nullptr, + std::vector logprobs = {}); std::string format_http_response( int status, const std::string & content_type, const std::string & body); @@ -668,7 +805,12 @@ class HttpServer { // Job queue. void enqueue(ServerJob * job); - ServerJob * dequeue(); + // Blocks for the next job. With `wake_for_control` (the serial worker) it + // also returns nullptr when a control transition is pending, which + // run_pending_controls() then executes. + ServerJob * dequeue(bool wake_for_control = false); + void run_pending_controls(); + void fail_pending_controls(const std::string & reason); bool has_pending_jobs(); // Members. @@ -744,15 +886,56 @@ class HttpServer { // Immutable model table after run() starts; only reservations mutate under // routing_mu_. A reservation spans parsing through job retirement/output // draining, so disconnects never make still-running engine work invisible. + enum class Residency { resident, evicted, failed }; struct RoutedModel { HttpServer * server; int capacity; int in_flight = 0; + // Swap residency (guarded by routing_mu_). + std::vector devices; + Residency residency = Residency::resident; + bool closing = false; // draining for an eviction: admits nothing + int waiters = 0; // pinned requests waiting for residency + int cohort_remaining = 0; // waiters to serve before swapping away + int swaps_in = 0; + double last_evict_ms = 0.0; + double last_reinstate_ms = 0.0; + std::string last_error; + std::chrono::steady_clock::time_point failed_at{}; + // Weight residency as of the last transition (or startup). + bool weights_resident = true; + size_t weight_device_bytes = 0; + // While evicted, the devices still holding its memory: its own device + // for weights kept warm, the secondary devices for a partial eviction. + std::vector retained_devices; + bool retained_warm = false; // retained_devices hold warm weights }; std::vector models_; std::mutex routing_mu_; std::condition_variable routing_cv_; int routing_waiters_ = 0; + bool swap_active_ = false; // one transition at a time (routing_mu_) + // Brings `target` in, evicting every resident model whose devices + // overlap. Called without routing_mu_; returns false with `error` when + // the swap failed. Returns true without work when another swap runs or + // the policy defers it (the caller keeps waiting). + bool drive_swap(RoutedModel & target, std::string & error); + static bool devices_overlap(const RoutedModel & a, const RoutedModel & b); + // Whether evicted `model` must release what it retains before `target` + // runs: a retained secondary tier on the target's devices, or warm + // weights there when the target spans several devices. + static bool retained_conflict(const RoutedModel & model, const RoutedModel & target); + // How resident `model` gives way to `target`: partial keeps a secondary + // device the target does not use; warm keeps weights between two + // single-device models; cold otherwise. + static ModelBackend::EvictLevel eviction_level(const RoutedModel & model, + const RoutedModel & target); + // Records a finished transition's timing and weight residency on `model` + // (call under routing_mu_). + static void record_control(RoutedModel & model, const ControlOp & op); + + // Control transitions for this model's serial worker (queue_mu_). + std::deque> control_queue_; // Request queue consumed by the serving loop owned by LuceEngine. std::mutex queue_mu_; diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 0d228e306a..99094cac62 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -55,6 +55,10 @@ struct SchedSlot { std::optional error; bool finished = false; std::vector gen_tokens; // committed + pending, in order + // Parallel to gen_tokens when the request asked for logprobs. + std::vector logprobs; + // /v1/hidden_states: the prefill's readout. Such a slot never decodes. + std::optional hidden_states; int32_t pending_tok = -1; // sampled, fed back next step // Buffered client output (see client_send_buffer.h): chunks append here and // a non-blocking flush runs every scheduler iteration, so one slow @@ -213,9 +217,22 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // delta into send_buffer, and parks the token in pending_tok as the next // step's input for this slot. Sets s.finished — but never retires the // slot — on EOS, gen cap, stop-sequence hit, or degenerate repetition. - auto advance_slot = [&](SchedSlot & s, int32_t tok) { + // `logprobs` is the engine's report for the sampled token, if any. + auto advance_slot = [&](int slot, SchedSlot & s, int32_t tok, + const std::optional & logprobs) { tok = s.budget.apply(tok).token; s.gen_tokens.push_back(tok); + if (logprobs) { + const auto actual = logprobs->chosen.token == tok + ? logprobs : engine.token_logprobs(slot, tok); + if (!actual) { + s.error = to_response_error({GenerateErrorCode::DecodeFailed, + "logprobs unavailable for the substituted token"}); + s.finished = true; + return; + } + s.logprobs.push_back(*actual); + } const bool cont = deliver_generation_token( s.job, s.job->req, *s.emitter, tok, s.completion_tokens, s.send_buffer); @@ -307,6 +324,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { "application/json", body.dump() + "\n")); } } + } else if (req.hidden_states && !s.client_disconnected) { + s.send_buffer.append(format_http_response( + 200, "application/json", + hidden_states_body(req, s.hidden_states ? *s.hidden_states + : HiddenStates{}, + s.prefill_s).dump() + "\n")); } else if (req.stream && !s.client_disconnected) { const bool is_eos = !s.gen_tokens.empty() && engine.token_is_eos(s.gen_tokens.back()); @@ -321,7 +344,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { send_nonstream_response(req, s.fd, *s.emitter, s.gen_tokens, s.n_gen_cap, s.budget.forced_close(), s.degenerate_close, gen_timings, - &s.send_buffer); + &s.send_buffer, std::move(s.logprobs)); } const double elapsed_s = std::chrono::duration( @@ -399,7 +422,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { req.max_output) : req.max_output; - if (n_gen_cap < 1) { + if (n_gen_cap < 1 && !req.hidden_states) { // 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, @@ -452,8 +475,17 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { int restore_policy_slot = -1; // Tokens alone do not identify an image: image requests never touch // the prefix cache. - const bool prefix_supported = !req.images && + // Hidden-states requests must prefill every token themselves: they + // never restore. With cache_prefix they may still leave a checkpoint + // of the whole prompt behind for a chat request that extends it. + const bool prefix_supported = !req.images && !req.hidden_states && engine.supports_prefix_store() && !prefix_cache_.disabled(); + const bool hidden_capture = req.hidden_states && + req.hidden_states->cache_prefix && + engine.supports_prefix_store() && !prefix_cache_.disabled(); + const auto estimate_bytes = [&engine](int target_cut) { + return engine.estimate_prefix_store_bytes(target_cut); + }; if (prefix_supported) { const auto hit = prefix_cache_.lookup_candidate( req.prompt_tokens, @@ -476,23 +508,30 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { /*prefer_tools_boundary=*/!req.tools.empty(), req.pin_end_token, restore_policy_slot, - [&engine](int target_cut) { - return engine.estimate_prefix_store_bytes(target_cut); - }, - /*reachable_from=*/0, backend_.restore_point_spacing()); - if (capture_reservation.active()) { - const uint64_t capture_id = next_prefix_capture_id++; - if (next_prefix_capture_id == 0) - next_prefix_capture_id = 1; - prefix_plan.capture.id = capture_id; - prefix_plan.capture.checkpoint = { - (uint64_t)capture_reservation.slot() + 1, - capture_reservation.target_cut()}; - if (prefix_plan.restore.valid() && - prefix_plan.capture.checkpoint == prefix_plan.restore) { - capture_reservation.cancel(); - prefix_plan.capture = {}; - } + estimate_bytes, + /*reachable_from=*/0, + backend_.restore_point_spacing()); + } else if (hidden_capture) { + // The cut is the prompt end, where the readout's last chunk + // already stops, so capturing reshapes no prefill chunk. + capture_reservation = prefix_cache_.reserve_inline_snap( + req.prompt_tokens, /*restored_prefix_len=*/0, + /*prefer_tools_boundary=*/false, + /*forced_cut=*/(int)req.prompt_tokens.size(), + /*restore_source_slot=*/-1, estimate_bytes); + } + if (capture_reservation.active()) { + const uint64_t capture_id = next_prefix_capture_id++; + if (next_prefix_capture_id == 0) + next_prefix_capture_id = 1; + prefix_plan.capture.id = capture_id; + prefix_plan.capture.checkpoint = { + (uint64_t)capture_reservation.slot() + 1, + capture_reservation.target_cut()}; + if (prefix_plan.restore.valid() && + prefix_plan.capture.checkpoint == prefix_plan.restore) { + capture_reservation.cancel(); + prefix_plan.capture = {}; } } if (prefix_supported) { @@ -511,7 +550,11 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // Admission only claims the slot and queues the prompt. Prefill // advances one chunk per engine step alongside live decode. const PrefixStorePlan requested_plan = prefix_plan; - auto ar = req.images + auto ar = req.hidden_states + ? engine.admit_hidden_states( + next_request_id, req.prompt_tokens, req.hidden_spec, + requested_plan.capture) + : req.images ? engine.admit_images( next_request_id, req.prompt_tokens, req.sampler, req.images) : prefix_supported @@ -1031,9 +1074,14 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.finished = true; continue; } + // Only the final token carries logprobs; slots that asked for + // them never speculate, so it is the only token. + size_t consumed = 0; consume_decode_output_tokens(out, [&](int32_t token) { if (s.finished) return false; - advance_slot(s, token); + const bool last = consumed++ == out.committed_tokens.size(); + advance_slot(out.slot, s, token, + last ? out.logprobs : std::optional{}); return !s.finished; }); } @@ -1084,6 +1132,23 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.finished = true; continue; } + if (out.status == PrefillStatus::completed && + s.job->req.hidden_states) { + s.prefilling = false; + s.prefill_s = std::chrono::duration( + std::chrono::steady_clock::now() - s.started_at).count(); + s.decode_started_at = std::chrono::steady_clock::now(); + s.hidden_states = out.hidden_states; + if (!s.hidden_states) { + s.error = to_response_error({GenerateErrorCode::PrefillFailed, + "engine completed a hidden-states prefill without a readout"}); + } + // Retired in this iteration's reap, before any decode plan. + s.finished = true; + publish_live_count(); + deferred_retry_at = {}; + continue; + } if (out.status == PrefillStatus::completed) { s.prefilling = false; publish_live_count(); @@ -1098,7 +1163,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.prefill_s = std::chrono::duration( s.decode_started_at - s.started_at).count(); } - advance_slot(s, out.token); + advance_slot(out.slot, s, out.token, out.logprobs); continue; } } diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 905c31edbb..14d8aab0d6 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -118,6 +118,21 @@ static void print_usage(const char * prog) { " --load-balancing-primary-gpu Select the primary model by its target device.\n" " Defaults to the first block; request model names\n" " do not change generation routing.\n" + " --model-routing How several model blocks share requests\n" + " (default: balance). name loads every block and runs a\n" + " request naming a block's --model-name only on that\n" + " block; omitted or auto requests still balance.\n" + " --unknown-model Name routing: answer an unrecognized\n" + " model name with 404 (default) or serve it on the primary.\n" + " --swap-residency With --model-routing name: models sharing a device\n" + " are not resident together; a request naming an\n" + " evicted model swaps it in (prefix snapshots kept).\n" + " The first model block starts resident.\n" + " --swap-keep-weights With --swap-residency: these models keep\n" + " their weights in device memory while evicted, so a\n" + " swap back rebuilds only the KV cache (cold when memory\n" + " runs short). List the main model, not occasional ones:\n" + " their weights would shrink the active model's room.\n" " --draft Draft model for speculative decode (DFlash for Qwen,\n" " Gemma and Laguna; DSpark for DeepSeek4; MTP for Qwen4Exp).\n" " Qwen4Exp auto-discovers /MTP/mtp-*.gguf.\n" @@ -430,8 +445,7 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, std::fprintf(stderr, "[server] %s belongs in the first model block: there is one listener\n", argv[i]); return 2; } - if (load_balancing && (option == "--peer-access" || option == "--expert-device" || - option == "--no-fast-rollback" || option == "--target-split-fast-rollback" || + if (load_balancing && (option == "--no-fast-rollback" || option == "--target-split-fast-rollback" || option == "--adaptive-experts" || option == "--specla" || option == "--specla-top-k" || option.rfind("--kvflash", 0) == 0 || option.rfind("--spark", 0) == 0)) { @@ -1194,6 +1208,38 @@ static void apply_launch_profile_env(const luce::server::LaunchProfile & profile kept_env.c_str()); } +// Residency compares device names, so auto:N names the compiled backend's +// device N, the one it runs on. +static std::string residency_device_name(DevicePlacement device) { + if (device.backend == PlacementBackend::Auto) device.backend = compiled_placement_backend(); + return placement_device_name(device); +} + +// Every device a model block occupies, for swap residency: each layer-split +// shard (or the single target), the draft device and the expert device. +static std::vector residency_devices(const DevicePlacement & target, + const DevicePlacement * draft, + const DevicePlacement * expert) { + std::vector list; + const auto add = [&](const DevicePlacement & placement) { + const std::string device = residency_device_name(placement); + if (std::find(list.begin(), list.end(), device) == list.end()) list.push_back(device); + }; + if (target.is_multi_device()) { + for (size_t i = 0; i < target.layer_split_gpus.size(); ++i) { + DevicePlacement shard; + shard.backend = target.layer_split_backend(i); + shard.gpu = target.layer_split_gpus[i]; + add(shard); + } + } else { + add(target); + } + if (draft) add(*draft); + if (expert) add(*expert); + return list; +} + static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_model) { if (model.profile) apply_launch_profile_env(*model.profile); if (!model.adaptive_experts_tau.empty()) @@ -1487,8 +1533,9 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod } } - // Create backend. - g_peer_access_opt_in = backend_placement.target.peer_access; + // Create backend. Peer access is process-wide: one block opts in for the + // process, and blocks loaded later must not turn it off. + if (backend_placement.target.peer_access) g_peer_access_opt_in = true; std::fprintf(stderr, "[server] creating backend...\n"); if (spark_autotune) { // Self-tuning hot/cold MoE residency: enable the bounded expert cache @@ -1543,6 +1590,16 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod if (model.expert_device && !apply_expert_device(*model.expert_device, backend_plan)) { return 2; } + // DeepSeek4 reads its expert device from the process environment, which + // another block's --expert-device sets; residency would not list it. + if (multi_model && !model.expert_device && + luce::common::arch_is_deepseek4_family(backend_plan.arch()) && + env_flag_enabled("LUCE_DS4_MOE_TP")) { + std::fprintf(stderr, "[server] model %s: DeepSeek4 without --expert-device would inherit the " + "process-wide LUCE_DS4_MOE_TP expert device; give it --expert-device or " + "unset LUCE_DS4_MOE_TP\n", model.sconfig.model_name.c_str()); + return 2; + } auto backend_owner = create_backend(backend_plan); if (!backend_owner) { std::fprintf(stderr, "[server] backend creation failed\n"); @@ -1933,6 +1990,15 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod sconfig.draft_device = backend_speculation.draft_path ? placement_device_name(backend_placement.draft) : std::string(); + sconfig.residency_devices = residency_devices( + backend_placement.target, + backend_speculation.draft_path ? &backend_placement.draft : nullptr, + model.expert_device ? &*model.expert_device : nullptr); + if (model.expert_device) { + const std::string expert = residency_device_name(*model.expert_device); + if (expert != residency_device_name(backend_placement.target)) + sconfig.residency_secondary_devices = {expert}; + } // Tokenizer ID: best-effort. The Tokenizer class doesn't currently // expose the GGUF metadata key it was loaded from, so leave empty // and let /props report null. (Add a getter on Tokenizer later.) @@ -2162,10 +2228,28 @@ int main(int argc, char ** argv) { // main's argv and outlive every backend, including factories borrowing paths. std::vector> model_args(1, {argv[0]}); bool load_balancing = false; + auto model_routing = ServerConfig::ModelRouting::balance; + bool unknown_model_to_primary = false; + bool unknown_model_set = false; + bool swap_residency = false; + std::set swap_keep_weights; // model names evicted warm std::string primary_gpu; for (int i = 1; i < argc; ++i) { if (std::strcmp(argv[i], "--load-balancing") == 0) { load_balancing = true; + } else if (std::strcmp(argv[i], "--swap-residency") == 0) { + swap_residency = true; + } else if (std::strcmp(argv[i], "--swap-keep-weights") == 0) { + if (i + 1 >= argc || argv[i + 1][0] == '-') { + std::fprintf(stderr, "[server] --swap-keep-weights requires model names (comma-separated)\n"); + return 2; + } + std::string names = argv[++i]; + for (size_t start = 0; start <= names.size();) { + const size_t end = std::min(names.find(',', start), names.size()); + if (end > start) swap_keep_weights.insert(names.substr(start, end - start)); + start = end + 1; + } } else if (std::strcmp(argv[i], "--load-balancing-primary-gpu") == 0) { DevicePlacement device; if (!primary_gpu.empty() || i + 1 >= argc || @@ -2176,6 +2260,25 @@ int main(int argc, char ** argv) { } primary_gpu = placement_device_name(device); ++i; + } else if (std::strcmp(argv[i], "--model-routing") == 0) { + const char * value = i + 1 < argc ? argv[++i] : ""; + if (std::strcmp(value, "balance") == 0) { + model_routing = ServerConfig::ModelRouting::balance; + } else if (std::strcmp(value, "name") == 0) { + model_routing = ServerConfig::ModelRouting::name; + } else { + std::fprintf(stderr, "[server] --model-routing expects balance or name\n"); + return 2; + } + } else if (std::strcmp(argv[i], "--unknown-model") == 0) { + const char * value = i + 1 < argc ? argv[++i] : ""; + if (std::strcmp(value, "reject") == 0 || std::strcmp(value, "primary") == 0) { + unknown_model_to_primary = std::strcmp(value, "primary") == 0; + unknown_model_set = true; + } else { + std::fprintf(stderr, "[server] --unknown-model expects reject or primary\n"); + return 2; + } } else if (std::strcmp(argv[i], "--model") == 0) { if (i + 1 >= argc || argv[i + 1][0] == '-') { std::fprintf(stderr, "[server] --model requires a model path\n"); @@ -2188,10 +2291,27 @@ int main(int argc, char ** argv) { } } const bool multi_model = model_args.size() > 1; - if (load_balancing && !multi_model) { - std::fprintf(stderr, "[server] --load-balancing requires at least two model blocks\n"); + const bool name_routing = model_routing == ServerConfig::ModelRouting::name; + if (unknown_model_set && !name_routing) { + std::fprintf(stderr, "[server] --unknown-model requires --model-routing name\n"); + return 2; + } + if (!swap_keep_weights.empty() && !swap_residency) { + std::fprintf(stderr, "[server] --swap-keep-weights requires --swap-residency\n"); + return 2; + } + if (swap_residency && !name_routing) { + std::fprintf(stderr, "[server] --swap-residency requires --model-routing name\n"); + return 2; + } + if ((load_balancing || name_routing) && !multi_model) { + std::fprintf(stderr, "[server] %s requires at least two model blocks\n", + name_routing ? "--model-routing name" : "--load-balancing"); return 2; } + // Name routing serves every block, so it loads them all under the same + // per-model scoping rules as load balancing. + if (name_routing) load_balancing = true; // Profile tokens are referenced by model_args for the process lifetime. static std::vector> profile_storage; std::vector profiles(model_args.size()); @@ -2220,6 +2340,18 @@ int main(int argc, char ** argv) { return 2; } } + // --expert-device and --peer-access set process-wide state (the DeepSeek4 + // expert environment, the peer-access opt-in), so one block may own them. + if (load_balancing) { + const auto owners = std::count_if(options.begin(), options.end(), [](const ModelOptions & o) { + return o.expert_device.has_value() || o.bargs.device.peer_access; + }); + if (owners > 1) { + std::fprintf(stderr, "[server] --expert-device and --peer-access change process-wide policy; " + "only one model block may set them\n"); + return 2; + } + } // Listener policy belongs to the first CLI block, independently of priority. const ServerConfig listener_config = options.front().sconfig; @@ -2245,8 +2377,11 @@ int main(int argc, char ** argv) { // Several GPU models in one process capture HIP/CUDA graphs from different // worker threads. Under relaxed capture, a blocking call from one worker // (prefix-cache checkpoint copies, host reads) invalidates the capture in - // flight on the other, so replicas fail intermittently. Eager launches cost - // ~4% on one R9700 and nothing measurable on the balanced aggregate. + // flight on the other, so replicas fail intermittently. This was observed + // with the models on different GPUs (R9700 + Strix Halo): separate devices + // or graph caches do not make capture safe, only keeping other workers' + // blocking calls out of an open capture would. Eager launches cost ~4% on + // one R9700 and nothing measurable on the balanced aggregate. // LUCE_MULTI_MODEL_GRAPHS=1 keeps graphs on (ggml reads any value of // GGML_CUDA_DISABLE_GRAPHS as disabled, so it cannot be the opt-out). if (options.size() > 1) { @@ -2267,9 +2402,24 @@ int main(int argc, char ** argv) { option.sconfig.port = listener_config.port; option.sconfig.enable_cors = listener_config.enable_cors; option.sconfig.routing_queue_limit = listener_config.routing_queue_limit; + option.sconfig.model_routing = model_routing; + option.sconfig.swap_residency = swap_residency; + option.sconfig.swap_keep_weights = swap_keep_weights.count(option.sconfig.model_name) > 0; + option.sconfig.unknown_model_to_primary = unknown_model_to_primary; + } + for (const auto & name : swap_keep_weights) { + const bool known = std::any_of(options.begin(), options.end(), [&](const ModelOptions & o) { + return o.sconfig.model_name == name; + }); + if (!known) { + std::fprintf(stderr, "[server] --swap-keep-weights: no model block named '%s'\n", name.c_str()); + return 2; + } } std::fprintf(stderr, "[server] load balancing %s; primary=%s target=%s\n", - load_balancing ? "enabled" : "disabled", options.front().sconfig.model_name.c_str(), + name_routing ? "by request model name (auto balances)" + : load_balancing ? "enabled" : "disabled", + options.front().sconfig.model_name.c_str(), options.front().target_device_auto ? "auto" : placement_device_name(options.front().bargs.device).c_str()); @@ -2286,17 +2436,110 @@ int main(int argc, char ** argv) { } } - std::vector> loaded; + std::vector> loaded(options.size()); std::vector servers; // All initialization (including backend environment defaults) finishes // before starting any scheduler. No worker observes model-loading mutations. - for (auto & option : options) { + // Swap residency never needs the models to fit together: each later block + // whose devices overlap a model that stays resident loads and is evicted + // on its own, before loading the blocks that start resident. + // The device sets here come from the command line; the listener checks + // the resolved ones again once every block is loaded. + std::vector evict_at_startup(options.size(), false); + std::vector evict_partially(options.size(), false); + if (swap_residency) { + std::vector> planned(options.size()); + for (size_t m = 0; m < options.size(); ++m) { + const ModelOptions & option = options[m]; + const std::string & name = option.sconfig.model_name; + if (option.target_device_auto) { + std::fprintf(stderr, "[server] --swap-residency: model %s needs an explicit --target-device\n", + name.c_str()); + return 2; + } + const bool has_draft = option.bargs.draft_path.has_value(); + // An unplaced draft goes where its backend decides (Qwen: GPU 0, + // DeepSeek4: the target), which is unknown before the model loads. + if (has_draft && !option.draft_device_set) { + std::fprintf(stderr, "[server] --swap-residency: model %s needs an explicit --draft-device\n", + name.c_str()); + return 2; + } + const DevicePlacement draft = resolve_draft_placement( + option.bargs.draft_device, option.draft_device_set, option.bargs.device, + /*target_auto=*/false); + planned[m] = residency_devices(option.bargs.device, has_draft ? &draft : nullptr, + option.expert_device ? &*option.expert_device : nullptr); + } + const auto overlap = [&](size_t a, size_t b) { + for (const auto & device : planned[a]) { + if (std::find(planned[b].begin(), planned[b].end(), device) != planned[b].end()) return true; + } + return false; + }; + std::vector resident = {0}; + for (size_t m = 1; m < options.size(); ++m) { + evict_at_startup[m] = std::any_of(resident.begin(), resident.end(), + [&](size_t r) { return overlap(m, r); }); + if (!evict_at_startup[m]) resident.push_back(m); + } + // An evicted block keeps its expert device (a partial eviction) when + // no resident block and no other kept tier uses it. + std::vector kept; + for (size_t m = 1; m < options.size(); ++m) { + const auto & expert = options[m].expert_device; + if (!evict_at_startup[m] || !expert) continue; + const std::string device = residency_device_name(*expert); + if (device == residency_device_name(options[m].bargs.device)) continue; + const bool used = std::find(kept.begin(), kept.end(), device) != kept.end() || + std::any_of(resident.begin(), resident.end(), [&](size_t r) { + return std::find(planned[r].begin(), planned[r].end(), device) != planned[r].end(); + }); + if (used) continue; + evict_partially[m] = true; + kept.push_back(device); + } + } + std::vector load_order; + // Load every initially evicted block before any resident block: later + // blocks can overlap a non-primary resident as well as the primary. + for (size_t m = 0; m < options.size(); ++m) { + if (evict_at_startup[m]) load_order.push_back(m); + } + for (size_t m = 0; m < options.size(); ++m) { + if (!evict_at_startup[m]) load_order.push_back(m); + } + for (size_t m : load_order) { auto model = std::make_unique(); - const int ret = load_model(option, *model, load_balancing); + const int ret = load_model(options[m], *model, load_balancing); if (ret != 0) return ret; - servers.push_back(model->server.get()); - loaded.push_back(std::move(model)); + if (swap_residency) { + ModelBackend & backend = model->engine->backend(); + std::string reason; + if (!backend.supports_eviction(reason)) { + std::fprintf(stderr, "[server] --swap-residency: model %s: %s\n", + options[m].sconfig.model_name.c_str(), reason.c_str()); + return 2; + } + } + if (evict_at_startup[m]) { + ModelBackend & backend = model->engine->backend(); + // Startup must fit one model at a time. Keep-weights applies to + // runtime swaps, where a failed warm swap can retry cold. + const auto evicted = backend.evict(evict_partially[m] + ? ModelBackend::EvictLevel::partial : ModelBackend::EvictLevel::cold); + if (evicted.status != ModelBackend::EvictStatus::ok) { + std::fprintf(stderr, "[server] --swap-residency: evicting %s at startup: %s\n", + options[m].sconfig.model_name.c_str(), evicted.error.c_str()); + return 2; + } + std::fprintf(stderr, "[server] %s loaded and evicted%s (swap residency)\n", + options[m].sconfig.model_name.c_str(), + evict_partially[m] ? " partially" : ""); + } + loaded[m] = std::move(model); } + for (auto & model : loaded) servers.push_back(model->server.get()); g_server = servers.front(); std::signal(SIGTERM, signal_handler); std::signal(SIGINT, signal_handler); diff --git a/server/test/test_model_names.cmake b/server/test/test_model_names.cmake index 94c188e613..c43ea27527 100644 --- a/server/test/test_model_names.cmake +++ b/server/test/test_model_names.cmake @@ -44,3 +44,26 @@ check_launch(disabled_second_policy 1 "failed to detect architecture from ${seco check_launch(balanced_policy_rejected 2 "changes process-wide policy" "${primary}" --model-name qwen --no-fast-rollback --model "${secondary}" --model-name ds4 --load-balancing) +# One block may own the expert device and the peer-access opt-in. +check_launch(balanced_expert_owner 1 "failed to detect architecture from ${primary}" + "${primary}" --model-name qwen + --model "${secondary}" --model-name ds4 --expert-device ${BACKEND}:1 --peer-access --load-balancing) +check_launch(balanced_expert_owners_rejected 2 "only one model block may set them" + "${primary}" --model-name qwen --peer-access + --model "${secondary}" --model-name ds4 --expert-device ${BACKEND}:1 --load-balancing) + +# Swap residency plans device sets before loading, so a draft needs a device. +check_launch(swap_draft_needs_device 2 "needs an explicit --draft-device" + "${primary}" --model-name qwen --target-device ${BACKEND}:0 + --model "${secondary}" --model-name ds4 --target-device ${BACKEND}:0 --draft "${secondary}" + --model-routing name --swap-residency) + +# Name routing loads every block, so it applies the balanced-mode name rules. +check_launch(name_routing_duplicate 2 "--model-name must be unique" + "${primary}" --model "${secondary}" --model-routing name) +check_launch(name_routing_unique 1 "failed to detect architecture from ${primary}" + "${primary}" --model-name qwen --model "${secondary}" --model-name ds4 + --model-routing name --unknown-model primary) +check_launch(unknown_model_needs_name_routing 2 "--unknown-model requires --model-routing name" + "${primary}" --model-name qwen --model "${secondary}" --model-name ds4 + --load-balancing --unknown-model primary) diff --git a/server/test/unit/test_dflash_capture_layers.cpp b/server/test/unit/test_dflash_capture_layers.cpp new file mode 100644 index 0000000000..85aa49b28a --- /dev/null +++ b/server/test/unit/test_dflash_capture_layers.cpp @@ -0,0 +1,42 @@ +// Unit tests for luce::common::adopt_drafter_capture_layers — no GPU, no +// model files. + +#include "CppUnitTestFramework.hpp" +#include "common/dflash_capture.h" +#include "luce.h" + +#include + +using namespace luce::common; + +namespace { +struct DFlashCaptureLayersFixture {}; +} + +// z-lab Qwen3.5-9B DFlash captures 8 target layers: more than the 5-slot +// default, so the table must take the drafter's count, not keep 5. +TEST_CASE(DFlashCaptureLayersFixture, adopts_drafter_count_beyond_default) { + int ids[LUCE_DRAFT_MAX_TARGET_LAYERS] = {1, 8, 15, 22, 29}; + int n = LUCE_DRAFT_N_TARGET_LAYERS; + const std::vector drafter = {1, 5, 9, 13, 17, 21, 25, 29}; + CHECK(adopt_drafter_capture_layers(drafter, /*target_n_layer=*/32, + LUCE_DRAFT_MAX_TARGET_LAYERS, ids, n)); + CHECK(n == 8); + CHECK(std::vector(ids, ids + n) == drafter); + CHECK(target_capture_index(ids, n, 29) == 7); +} + +TEST_CASE(DFlashCaptureLayersFixture, rejects_invalid_lists_untouched) { + int ids[LUCE_DRAFT_MAX_TARGET_LAYERS] = {1, 8, 15, 22, 29}; + int n = LUCE_DRAFT_N_TARGET_LAYERS; + // Layer 32 does not exist in a 32-layer target. + CHECK(!adopt_drafter_capture_layers({1, 5, 32}, 32, + LUCE_DRAFT_MAX_TARGET_LAYERS, ids, n)); + // More layers than the table holds. + CHECK(!adopt_drafter_capture_layers(std::vector(3, 1), 32, + /*max_slots=*/2, ids, n)); + CHECK(!adopt_drafter_capture_layers({}, 32, + LUCE_DRAFT_MAX_TARGET_LAYERS, ids, n)); + CHECK(n == LUCE_DRAFT_N_TARGET_LAYERS); + CHECK(ids[4] == 29); +} diff --git a/server/test/unit/test_hidden_states.cpp b/server/test/unit/test_hidden_states.cpp new file mode 100644 index 0000000000..4680319377 --- /dev/null +++ b/server/test/unit/test_hidden_states.cpp @@ -0,0 +1,216 @@ +// /v1/hidden_states: request validation, layer resolution, the "bare" +// generation prompt, chunked pooling, and the response shape. + +#include "CppUnitTestFramework.hpp" +#include "common/hidden_states.h" +#include "server/chat_template.h" +#include "server/http_server.h" + +#include +#include +#include +#include + +using luce::common::ChatFormat; +using luce::common::ChatMessage; +using luce::common::HiddenStates; +using luce::common::HiddenStatesAccumulator; +using luce::common::HiddenStatesRequest; +using luce::common::HiddenStatesSpec; +using luce::common::build_hidden_states_response; +using luce::common::json; +using luce::common::parse_hidden_states_request; +using luce::common::render_chat_template; +using luce::common::render_chat_template_jinja; +using luce::common::resolve_hidden_states_spec; +using luce::common::trim_to_assistant_header; + +namespace { +struct HiddenStatesFixture {}; + +const std::vector kConversation = { + {"system", "Route the request.", ""}, + {"user", "What is 2+2?", ""}, + {"assistant", "4", ""}, + {"user", "And 3+3?", ""}, +}; + +// A plain ChatML template whose generation prompt adds no think block. +const char * kPlainChatml = + "{% for message in messages %}" + "{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>\\n' }}" + "{% endfor %}" + "{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}"; + +bool ends_with(const std::string & s, const std::string & suffix) { + return s.size() >= suffix.size() && + s.compare(s.size() - suffix.size(), suffix.size(), suffix) == 0; +} +} // namespace + +TEST_CASE(HiddenStatesFixture, request_defaults_and_options) { + const HiddenStatesRequest plain = + parse_hidden_states_request({{"layers", {3, -1}}}); + CHECK(plain.layers == std::vector({3, -1})); + CHECK(plain.pooling == HiddenStatesRequest::Pooling::last); + CHECK(!plain.bare_generation_prompt); + CHECK(!plain.cache_prefix); + + const HiddenStatesRequest both = parse_hidden_states_request( + {{"layers", {0}}, {"pooling", "both"}, {"generation_prompt", "bare"}}); + CHECK(both.pooling == HiddenStatesRequest::Pooling::both); + CHECK(both.wants_mean()); + CHECK(both.bare_generation_prompt); + CHECK(parse_hidden_states_request( + {{"layers", {0}}, {"generation_prompt", "template"}}) + .bare_generation_prompt == false); + CHECK(parse_hidden_states_request( + {{"layers", {0}}, {"cache_prefix", true}}).cache_prefix); + CHECK(!parse_hidden_states_request( + {{"layers", {0}}, {"cache_prefix", nullptr}}).cache_prefix); +} + +TEST_CASE(HiddenStatesFixture, request_rejects_invalid_fields) { + const std::vector invalid = { + json::object(), + {{"layers", json::array()}}, + {{"layers", 3}}, + {{"layers", {1.5}}}, + {{"layers", {"3"}}}, + {{"layers", {true}}}, + {{"layers", {2, 2}}}, + json::parse(R"({"layers":[2147483648]})"), + json::parse(R"({"layers":[-2147483649]})"), + json::parse(R"({"layers":[18446744073709551615]})"), + {{"layers", {0}}, {"pooling", "max"}}, + {{"layers", {0}}, {"pooling", 1}}, + {{"layers", {0}}, {"generation_prompt", "none"}}, + {{"layers", {0}}, {"cache_prefix", "yes"}}, + {{"layers", {0}}, {"cache_prefix", 1}}, + }; + for (const json & body : invalid) { + CHECK_THROW(std::invalid_argument, + UNUSED_RETURN(parse_hidden_states_request(body))); + } +} + +TEST_CASE(HiddenStatesFixture, layers_resolve_from_the_end_and_stay_in_range) { + HiddenStatesRequest request; + request.layers = {0, 23, -1, -24}; + request.pooling = HiddenStatesRequest::Pooling::mean; + const HiddenStatesSpec spec = resolve_hidden_states_spec(request, 24); + CHECK(spec.layers == std::vector({0, 23, 23, 0})); + CHECK(spec.mean); + + for (int layer : {24, -25}) { + request.layers = {layer}; + CHECK_THROW(std::invalid_argument, + UNUSED_RETURN(resolve_hidden_states_spec(request, 24))); + } +} + +TEST_CASE(HiddenStatesFixture, bare_prompt_drops_the_template_think_block) { + for (bool thinking : {false, true}) { + std::string rendered = render_chat_template( + kConversation, ChatFormat::QWEN3, /*add_generation_prompt=*/true, + thinking); + // The Qwen template appends a think block either way. + CHECK(ends_with(rendered, thinking ? "\n" : "\n\n")); + CHECK(trim_to_assistant_header(rendered)); + CHECK(ends_with(rendered, "<|im_start|>user\nAnd 3+3?<|im_end|>\n" + "<|im_start|>assistant\n")); + // The earlier assistant turn is history and stays intact. + CHECK(rendered.find("<|im_start|>assistant\n4<|im_end|>") != + std::string::npos); + } +} + +TEST_CASE(HiddenStatesFixture, bare_prompt_keeps_a_template_without_think_block) { + const std::string rendered = render_chat_template_jinja( + kPlainChatml, kConversation, "", "", /*add_generation_prompt=*/true); + CHECK(ends_with(rendered, "And 3+3?<|im_end|>\n<|im_start|>assistant\n")); + std::string trimmed = rendered; + CHECK(trim_to_assistant_header(trimmed)); + CHECK(trimmed == rendered); +} + +TEST_CASE(HiddenStatesFixture, bare_prompt_needs_an_open_assistant_turn) { + // Without a generation prompt the last assistant header belongs to a + // closed turn; trimming there would cut real content. + std::string closed = render_chat_template_jinja( + kPlainChatml, kConversation, "", "", /*add_generation_prompt=*/false); + const std::string before = closed; + CHECK(!trim_to_assistant_header(closed)); + CHECK(closed == before); + std::string other = "<|User|>hi<|Assistant|>"; + CHECK(!trim_to_assistant_header(other)); +} + +TEST_CASE(HiddenStatesFixture, accumulator_pools_over_every_chunk) { + HiddenStatesSpec spec; + spec.layers = {5, 9}; + spec.mean = true; + HiddenStatesAccumulator accum(spec, 2); + // Chunk 1: three rows; layer entry 0 sums to {3, 6}, entry 1 to {30, 60}. + const float last_a0[] = {1, 2}, sum_a0[] = {3, 6}; + const float last_a1[] = {10, 20}, sum_a1[] = {30, 60}; + accum.record(0, last_a0, sum_a0); + accum.record(1, last_a1, sum_a1); + accum.add_rows(3); + // Chunk 2: one row, which is the prompt's last token. + const float last_b0[] = {7, 8}, last_b1[] = {70, 80}; + accum.record(0, last_b0, last_b0); + accum.record(1, last_b1, last_b1); + accum.add_rows(1); + + const HiddenStates out = accum.finish(); + CHECK(out.n_tokens == 4); + CHECK(out.last[0] == std::vector({7, 8})); + CHECK(out.last[1] == std::vector({70, 80})); + CHECK(out.mean[0] == std::vector({2.5f, 3.5f})); + CHECK(out.mean[1] == std::vector({25.0f, 35.0f})); + + spec.mean = false; + HiddenStatesAccumulator last_only(spec, 2); + last_only.record(0, last_a0, sum_a0); + last_only.add_rows(1); + CHECK(last_only.finish().mean.empty()); +} + +TEST_CASE(HiddenStatesFixture, response_keys_layers_as_requested) { + HiddenStatesRequest request; + request.layers = {-1, 3}; + HiddenStates states; + states.last = {{1.0f, 2.0f}, {3.0f, 4.0f}}; + states.mean = {{0.5f, 0.5f}, {0.25f, 0.25f}}; + states.n_tokens = 7; + + const json last = build_hidden_states_response( + "qwen35-2b", request, 24, 2, 7, states, "\n", 12.5); + CHECK(last.at("model") == "qwen35-2b"); + CHECK(last.at("n_layers") == 24); + CHECK(last.at("n_embd") == 2); + CHECK(last.at("prompt_tokens") == 7); + CHECK(last.at("generation_prompt") == "template"); + CHECK(last.at("last_token") == "\n"); + CHECK(last.at("timings").at("prefill_ms").get() == 12.5); + CHECK(last.at("layers").size() == 2u); + CHECK(last.at("layers").at("-1").at("last") == json::array({1.0f, 2.0f})); + CHECK(last.at("layers").at("3").at("last") == json::array({3.0f, 4.0f})); + CHECK(!last.at("layers").at("-1").contains("mean")); + CHECK(!last.contains("cache_prefix")); + + request.pooling = HiddenStatesRequest::Pooling::mean; + request.bare_generation_prompt = true; + const json mean = build_hidden_states_response( + "qwen35-2b", request, 24, 2, 7, states, "\n", 12.5); + CHECK(mean.at("generation_prompt") == "bare"); + CHECK(mean.at("layers").at("3").at("mean") == json::array({0.25f, 0.25f})); + CHECK(mean.at("layers").at("3").contains("last")); + + request.cache_prefix = true; + const json cached = build_hidden_states_response( + "qwen35-2b", request, 24, 2, 7, states, "\n", 12.5, + /*cached_prefix_tokens=*/7); + CHECK(cached.at("cache_prefix").at("cached_tokens") == 7); +} diff --git a/server/test/unit/test_model_routing.cpp b/server/test/unit/test_model_routing.cpp index e440ba6b77..c6ae041325 100644 --- a/server/test/unit/test_model_routing.cpp +++ b/server/test/unit/test_model_routing.cpp @@ -97,6 +97,7 @@ class HeldEngine final : public SeqEngine { for (int i = 0; i < 2; ++i) { if (active[i]) continue; active[i] = true; + logprobs_requested[i] = sampler.wants_logprobs(); prompts.push_back(prompt); temperatures.push_back(sampler.temp); ++admissions; @@ -135,12 +136,25 @@ class HeldEngine final : public SeqEngine { ++steps; result.decode.push_back({input.slot, token, false, {}}); } + for (auto & out : result.prefills) { + if (out.token >= 0) out.logprobs = token_logprobs(out.slot, out.token); + } + for (auto & out : result.decode) { + out.logprobs = token_logprobs(out.slot, out.token); + } // An unfinished fake prefill yields control to the real scheduler, // allowing admission and cancellation without timing-based completion. lock.unlock(); std::this_thread::yield(); return result; } + std::optional token_logprobs(int slot, int32_t token) override { + if (!logprobs_requested[slot]) return std::nullopt; + TokenLogprobs result; + result.chosen = {token, -float(token + 1)}; + return result; + } + std::array logprobs_requested{}; void retire(int slot) override { std::lock_guard lock(mu); active.at((size_t)slot) = false; @@ -280,7 +294,9 @@ class HeldEngine final : public SeqEngine { struct RoutedBackend : ModelBackend { HeldEngine engine; + bool logprobs = false; SeqEngine * seq_engine() override { return &engine; } + bool supports_logprobs() const override { return logprobs; } void print_ready_banner() const override {} bool park(ParkTarget) override { return true; } bool unpark(ParkTarget) override { return true; } @@ -300,6 +316,38 @@ struct RoutedBackend : ModelBackend { // The gate models work boundaries where real backends poll DaemonIO cancellation. struct HeldSingleBackend final : RoutedBackend { SeqEngine * seq_engine() override { return nullptr; } + bool supports_eviction(std::string & reason) const override { reason.clear(); return true; } + bool is_evicted() const override { return evicted; } + bool weights_resident() const override { return !evicted || kept_level; } + EvictResult evict(EvictLevel level) override { + ++evict_calls; + EvictResult result; + if (reject_eviction) { + result.error = "injected eviction rejection"; + } else { + evicted = true; + kept_level = level != EvictLevel::cold; + std::lock_guard lock(mu); + levels.push_back(level); + result.status = EvictStatus::ok; + } + return result; + } + bool reinstate(std::string & error) override { + ++reinstate_calls; + if (fail_reinstate) { + error = "injected reinstatement failure"; + return false; + } + evicted = false; + kept_level = false; + return true; + } + bool evicted = false; + bool kept_level = false; // evicted warm or partially + std::vector levels; // guarded by mu + std::atomic reject_eviction{false}, fail_reinstate{false}; + std::atomic evict_calls{0}, reinstate_calls{0}; GenerateResult generate_impl(const GenerateRequest & req, const DaemonIO & io) override { std::unique_lock lock(mu); ++calls; @@ -388,6 +436,9 @@ class RunningModels { explicit RunningModels(bool single_peer = false, bool single_listener = false, bool load_balancing = true, int queue_limit = 0, bool reverse_priority = false, size_t offload_bytes = 0, + ServerConfig::ModelRouting routing = ServerConfig::ModelRouting::balance, + bool unknown_to_primary = false, bool force_close = false, + bool swap_residency = false, bool two_device_peer = false, const char * first_output = "q") { load_tokenizer(first_tok, false, first_output); load_tokenizer(second_tok, true, "s"); @@ -406,8 +457,23 @@ class RunningModels { config.model_name = "qwen"; config.max_ctx = 64; config.routing_queue_limit = queue_limit; + config.model_routing = routing; + config.swap_residency = swap_residency; + if (swap_residency) { + config.target_device = "hip:0"; + single.evicted = true; + } + // The listener keeps weights warm; the peer spans hip:0 and an + // expert device hip:1 that a partial eviction keeps. + config.swap_keep_weights = two_device_peer; + config.unknown_model_to_primary = unknown_to_primary; config.decode_kv_offload_bytes = offload_bytes; config.default_max_tokens = 4; + if (force_close) { + config.think_close_token_ids = {3}; + config.think_max_tokens = 1; + config.hard_limit_reply_budget = 1; + } config.prefix_cache_cap = 0; config.ppp_enabled = false; config.admission_coalesce_ms = 0; @@ -422,6 +488,11 @@ class RunningModels { config.model_name = "ds4"; config.chat_template_src = "y{{ messages[0]['content'] }}"; config.sampler_defaults.temperature = 0.7f; + config.swap_keep_weights = false; + if (two_device_peer) { + config.residency_devices = {"hip:0", "hip:1"}; + config.residency_secondary_devices = {"hip:1"}; + } peer_engine = std::make_unique( single_peer ? std::move(single_owned) : std::move(second_owned)); peer = std::make_unique(*peer_engine, second_tok, config); @@ -611,7 +682,8 @@ TEST_CASE(ModelRoutingFixture, test_responses_nonstream_preserves_prose_and_mult "Listing.\n\n\n\n.\n\n\n\n" "\n\n\nREADME.md\n\n\n"; for (bool single : {false, true}) { - RunningModels models(false, single, false, 0, false, 0, generated); + RunningModels models(false, single, false, 0, false, 0, + ServerConfig::ModelRouting::balance, false, false, false, false, generated); models.first.engine.finish(); models.first_single.finish(); const json request = json::parse(R"({ @@ -987,6 +1059,158 @@ TEST_CASE(ModelRoutingFixture, test_reversed_priority_ignores_generation_model_n ROUTING_CHECK(response_body(models.post(chat("ds4")).read())["model"] == "qwen"); } +// Name routing pins a named request to its model: it waits for that model's +// capacity instead of spilling onto an idle one, while auto still balances. +TEST_CASE(ModelRoutingFixture, test_name_routing_pins_named_requests_without_spilling) { + using Routing = ServerConfig::ModelRouting; + RunningModels models(false, false, true, 1, false, 0, Routing::name); + ROUTING_CHECK(models.get("/v1/models")["data"].size() == 2); + ROUTING_CHECK(models.get("/status/json")["routing"] == "by-name"); + auto a = models.post(chat("ds4")); + auto b = models.post(chat("ds4")); + models.second.engine.wait_admissions(2); + auto waiting = models.post(chat("ds4")); + const auto deadline = Clock::now() + 5s; + while (models.get("/status/json")["waiting"] != 1) { + ROUTING_CHECK(Clock::now() < deadline); + std::this_thread::yield(); + } + ROUTING_CHECK(models.first.engine.admissions == 0); + const json unknown = response_body(models.post(chat("gpt-4o")).read(), 404); + ROUTING_CHECK(unknown["error"]["message"].get().find( + "available models: qwen, ds4") != std::string::npos); + response_body(models.post(chat("gpt-4o"), "/v1/messages/count_tokens").read(), 404); + models.first.engine.finish(); + ROUTING_CHECK(response_body(models.post(chat("auto")).read())["model"] == "qwen"); + ROUTING_CHECK(response_body(models.post(chat("qwen")).read())["model"] == "qwen"); + models.second.engine.finish(); + for (Socket * client : {&a, &b, &waiting}) { + ROUTING_CHECK(response_body(client->read())["model"] == "ds4"); + } + ROUTING_CHECK(models.second.engine.admissions == 3); +} + +TEST_CASE(ModelRoutingFixture, test_name_routing_can_serve_unknown_names_on_primary) { + RunningModels models(false, false, true, 0, false, 0, + ServerConfig::ModelRouting::name, true); + models.first.engine.finish(); + models.second.engine.finish(); + ROUTING_CHECK(response_body(models.post(chat("gpt-4o")).read())["model"] == "qwen"); + ROUTING_CHECK(response_body(models.post(chat("gpt-4o"), "/v1/messages/count_tokens").read())["input_tokens"] == 1); + ROUTING_CHECK(models.second.engine.admissions == 0); +} + +// Logprobs support belongs to the model that serves the request, not to the +// listener's primary. +TEST_CASE(ModelRoutingFixture, test_swap_failure_restores_resident_and_cools_down) { + for (bool reject_eviction : {false, true}) { + RunningModels models(true, true, true, 4, false, 0, + ServerConfig::ModelRouting::name, false, false, true); + models.first_single.finish(); + models.single.finish(); + models.first_single.reject_eviction = reject_eviction; + models.single.fail_reinstate = !reject_eviction; + response_body(models.post(chat("ds4")).read(), 503); + // Every waiter sees the same failed target, without repeating the + // drain/eviction or reinstatement during its cooldown. + response_body(models.post(chat("ds4")).read(), 503); + ROUTING_CHECK(models.first_single.evict_calls == 1); + ROUTING_CHECK(models.single.reinstate_calls == (reject_eviction ? 0 : 1)); + ROUTING_CHECK(models.first_single.reinstate_calls == (reject_eviction ? 0 : 1)); + const auto status = models.get("/status/json"); + ROUTING_CHECK(status["models"][0]["residency"] == "resident"); + ROUTING_CHECK(status["models"][1]["residency"] == "failed"); + ROUTING_CHECK(status["swap_active"] == false); + ROUTING_CHECK(response_body(models.post(chat("auto")).read())["model"] == "qwen"); + } +} + +TEST_CASE(ModelRoutingFixture, test_swap_cycles_admit_only_the_resident_model) { + RunningModels models(true, true, true, 4, false, 0, + ServerConfig::ModelRouting::name, false, false, true); + models.first_single.finish(); + models.single.finish(); + for (const char * name : {"ds4", "qwen", "ds4"}) { + ROUTING_CHECK(response_body(models.post(chat(name)).read())["model"] == name); + ROUTING_CHECK(response_body(models.post(chat("auto")).read())["model"] == name); + } + ROUTING_CHECK(models.first_single.evict_calls == 2); + ROUTING_CHECK(models.single.evict_calls == 1); + ROUTING_CHECK(models.single.reinstate_calls == 2); +} + +// An automatic request that exceeds the resident model's context never swaps, +// so it is rejected instead of waiting for a model that will not come in. +TEST_CASE(ModelRoutingFixture, test_swap_oversized_automatic_request_is_rejected) { + RunningModels models(true, true, true, 4, false, 0, + ServerConfig::ModelRouting::name, false, false, true); + models.first_single.finish(); + models.single.finish(); + auto oversized = chat("auto"); + oversized["messages"][0]["content"] = std::string(100, 'x'); + response_body(models.post(oversized).read(), 400); + const auto status = models.get("/status/json"); + ROUTING_CHECK(status["waiting"] == 0); + ROUTING_CHECK(status["models"][1]["residency"] == "evicted"); +} + +// A model with an expert device gives way to a hip:0-only model partially and +// keeps hip:1; the model coming back over both devices never leaves the other +// model's weights warm on hip:0, even when that model keeps weights. +TEST_CASE(ModelRoutingFixture, test_swap_keeps_secondary_tier_and_needs_primary_whole) { + using Level = ModelBackend::EvictLevel; + RunningModels models(true, true, true, 4, false, 0, + ServerConfig::ModelRouting::name, false, false, true, true); + models.first_single.finish(); + models.single.finish(); + for (const char * name : {"ds4", "qwen", "ds4", "qwen"}) { + ROUTING_CHECK(response_body(models.post(chat(name)).read())["model"] == name); + } + ROUTING_CHECK((models.first_single.levels == std::vector{Level::cold, Level::cold})); + ROUTING_CHECK((models.single.levels == std::vector{Level::partial, Level::partial})); + const auto status = models.get("/status/json"); + ROUTING_CHECK(status["models"][1]["residency"] == "evicted"); + ROUTING_CHECK(status["models"][1]["weights_resident"] == true); + ROUTING_CHECK((status["models"][1]["retained_devices"] == json::array({"hip:1"}))); + ROUTING_CHECK(status["models"][0]["retained_devices"].empty()); +} + +TEST_CASE(ModelRoutingFixture, test_logprobs_follow_budget_substituted_tokens) { + RunningModels models(false, false, true, 0, false, 0, + ServerConfig::ModelRouting::name, false, true); + models.first.logprobs = true; + models.first.engine.finish(); + auto request = chat("qwen"); + request["logprobs"] = true; + request["max_tokens"] = 2; + request["thinking"] = {{"type", "enabled"}, {"budget_tokens", 1}}; + const json body = response_body(models.post(request).read()); + const auto & content = body["choices"][0]["logprobs"]["content"]; + ROUTING_CHECK(content.size() == 2); + ROUTING_CHECK(content[1]["token"] == "y"); + ROUTING_CHECK(content[1]["logprob"] == -4.0f); +} + +TEST_CASE(ModelRoutingFixture, test_logprobs_support_is_decided_by_the_routed_model) { + RunningModels models(false, false, true, 0, false, 0, ServerConfig::ModelRouting::name); + models.second.logprobs = true; + models.first.engine.finish(); + models.second.engine.finish(); + auto request = chat("ds4"); + request["logprobs"] = true; + request["top_logprobs"] = 2; + const json served = response_body(models.post(request).read()); + ROUTING_CHECK(served["model"] == "ds4"); + ROUTING_CHECK(served["choices"][0].contains("logprobs")); + // A backend without logprobs serves the request without them, as + // before logprobs existed. + request["model"] = "qwen"; + const json plain = response_body(models.post(request).read()); + ROUTING_CHECK(plain["model"] == "qwen"); + ROUTING_CHECK(!plain["choices"][0].contains("logprobs")); + ROUTING_CHECK(models.first.engine.admissions == 1); +} + // These tests exercise live HTTP ownership during decode growth. Earlier // routing tests stop at admission; they cannot expose lost suspended streams. TEST_CASE(ModelRoutingFixture, test_decode_pressure_preserves_stream_and_routes_new_work) { diff --git a/server/test/unit/test_model_swap_gpu.cpp b/server/test/unit/test_model_swap_gpu.cpp new file mode 100644 index 0000000000..f71de9fc56 --- /dev/null +++ b/server/test/unit/test_model_swap_gpu.cpp @@ -0,0 +1,463 @@ +// Model-level proof for runtime model swapping: evict() returns the model's +// device memory, reinstate() rebuilds it, and a host prefix snapshot taken +// before the eviction restores afterwards with an identical greedy +// continuation. The two-model case swaps a second model through the freed +// device in between, which is the serving scenario. +// +// Needs LUCE_TEST_MODEL_QWEN35 + LUCE_TEST_MODEL_DRAFT (model A) and, for the +// swap case, LUCE_TEST_MODEL_QWEN35_B + LUCE_TEST_MODEL_DRAFT_B (model B). +// LUCE_TEST_SWAP_CYCLES (default 3) sets the evict/reinstate repetitions. +// +// DeepSeek V4.1 partial eviction runs the ds41-lucebox profile (target hip:0, +// experts hip:1, locked host memory for the secondary tier) and needs +// LUCE_TEST_MODEL_DEEPSEEK4 (the V4.1 GGUF) + LUCE_TEST_MODEL_DEEPSEEK4_DRAFT +// (its DSpark drafter); LUCE_TEST_DS4_CYCLES (default 10) sets its cycles. +#include "CppUnitTestFramework.hpp" +#include "model_test_paths.h" +#include "common/backend_factory.h" +#include "common/model_backend.h" +#include "common/peer_access.h" +#include "deepseek4/deepseek4_backend.h" +#include "server/launch_profiles.h" +#include "server/tokenizer.h" +#include "ggml-cuda.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +using namespace luce::common; +struct ModelSwapGpuFixture {}; +#define SWAP_CHECK(x) do { if (!(x)) throw std::runtime_error( \ + std::string("model swap line ") + std::to_string(__LINE__) + ": " #x); } while (false) + +constexpr std::string_view kQwen35ModelBEnv = "LUCE_TEST_MODEL_QWEN35_B"; +constexpr std::string_view kDraftModelBEnv = "LUCE_TEST_MODEL_DRAFT_B"; +constexpr double kMiB = 1024.0 * 1024.0; + +// Free device memory once released buffers are back: the amdgpu driver frees +// a buffer whose last fence is still pending in the background, so right +// after an eviction free memory can read up to ~10 GiB low for a few ms. +size_t device_free(int device = 0) { + size_t free = 0, total = 0; + ggml_backend_cuda_get_device_memory(device, &free, &total); + const auto start = std::chrono::steady_clock::now(); + for (;;) { + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + size_t now = 0; + ggml_backend_cuda_get_device_memory(device, &now, &total); + if (now <= free || std::chrono::steady_clock::now() - start > std::chrono::seconds(2)) { + return std::max(free, now); + } + free = now; + } +} + +double ms_since(std::chrono::steady_clock::time_point start) { + return std::chrono::duration( + std::chrono::steady_clock::now() - start).count(); +} + +std::unique_ptr make_backend(const std::string & model, + const std::string & draft) { + BackendArgs args; + args.model_path = model; + args.draft_path = draft; + args.device.backend = compiled_placement_backend(); + args.device.max_ctx = 4096; + args.draft_device = args.device; + args.max_concurrency = 1; + args.cache_type_k = args.cache_type_v = GGML_TYPE_Q8_0; + auto prepared = prepare_backend(args); + if (const auto * failure = std::get_if(&prepared)) { + throw std::runtime_error(failure->message); + } + auto backend = create_backend(std::get(prepared)); + SWAP_CHECK(backend); + std::string reason; + SWAP_CHECK(backend->supports_eviction(reason)); + return backend; +} + +GenerateRequest greedy(std::vector prompt, int n_gen) { + GenerateRequest req; + req.prompt = std::move(prompt); + req.n_gen = n_gen; + req.sampler.temp = 0.0f; + return req; +} + +// A session on one model: a long prompt, a snapshot after the first answer, +// and the follow-up request that restores from it. +struct Session { + std::vector prompt; + std::vector first_answer; + std::vector follow_up; + std::vector control; // follow-up continuation without any eviction + static constexpr int slot = 0; +}; + +Session start_session(ModelBackend & backend, const std::string & model_path) { + Tokenizer tokenizer; + SWAP_CHECK(tokenizer.load_from_gguf(model_path.c_str())); + std::string text; + for (int i = 0; text.size() < 6000; ++i) { + text += "Record " + std::to_string(i) + ": the courier left depot " + + std::to_string(i * 7 % 13) + " at " + std::to_string(6 + i % 12) + + " o'clock carrying " + std::to_string(i * 3 + 1) + " parcels. "; + } + Session session; + session.prompt = tokenizer.encode(text + "\nWhich depot did record 17 leave from?"); + DaemonIO io; + const auto first = backend.generate(greedy(session.prompt, 32), io); + SWAP_CHECK(first.ok()); + session.first_answer = first.tokens; + SWAP_CHECK(backend.snapshot_save(Session::slot)); + const int pos = backend.snapshot_cur_pos(Session::slot); + std::vector history = session.prompt; + history.insert(history.end(), first.tokens.begin(), first.tokens.end()); + SWAP_CHECK(pos > 0 && pos <= (int)history.size()); + history.resize((size_t)pos); + const auto question = tokenizer.encode("\nAnd how many parcels did record 23 carry?"); + history.insert(history.end(), question.begin(), question.end()); + session.follow_up = history; + const auto control = backend.restore_and_generate(Session::slot, greedy(session.follow_up, 32), io); + SWAP_CHECK(control.ok()); + SWAP_CHECK(control.restored_prefix_tokens > 0); + session.control = control.tokens; + std::printf("session: prompt %zu tokens, snapshot at %d, follow-up %zu tokens\n", + session.prompt.size(), pos, session.follow_up.size()); + return session; +} + +// After a reinstatement the snapshot must restore and reproduce the control. +void check_session(ModelBackend & backend, const Session & session) { + DaemonIO io; + const auto resumed = backend.restore_and_generate(Session::slot, greedy(session.follow_up, 32), io); + SWAP_CHECK(resumed.ok()); + SWAP_CHECK(resumed.restored_prefix_tokens > 0); + SWAP_CHECK(resumed.tokens == session.control); + const auto fresh = backend.generate(greedy(session.prompt, 32), io); + SWAP_CHECK(fresh.ok()); + SWAP_CHECK(fresh.tokens == session.first_answer); +} + +void evict(ModelBackend & backend, const char * name, + ModelBackend::EvictLevel level = ModelBackend::EvictLevel::cold) { + const auto start = std::chrono::steady_clock::now(); + const auto result = backend.evict(level); + SWAP_CHECK(result.status == ModelBackend::EvictStatus::ok); + SWAP_CHECK(result.snapshots_lost == 0); + SWAP_CHECK(backend.is_evicted()); + std::printf("%s: evict %.0f ms, device free %.0f MiB\n", name, ms_since(start), + device_free() / kMiB); +} + +void reinstate(ModelBackend & backend, const char * name) { + const auto start = std::chrono::steady_clock::now(); + std::string error; + SWAP_CHECK(backend.reinstate(error)); + SWAP_CHECK(!backend.is_evicted()); + std::printf("%s: reinstate %.0f ms, device free %.0f MiB\n", name, ms_since(start), + device_free() / kMiB); +} + +int swap_cycles() { + const char * value = std::getenv("LUCE_TEST_SWAP_CYCLES"); + return value && std::atoi(value) > 0 ? std::atoi(value) : 3; +} + +void require_gpu() { + if (ggml_backend_cuda_get_device_count() == 0) { + throw CppUnitTestFramework::TestSkippedException("GPU is unavailable"); + } +} +} // namespace + +TEST_CASE(ModelSwapGpuFixture, QwenEvictReinstateKeepsSnapshots) { + const std::string model = luce_test::require_model(luce_test::kQwen35ModelEnv).string(); + const std::string draft = luce_test::require_model(luce_test::kDraftModelEnv).string(); + require_gpu(); + const size_t baseline = device_free(); + auto backend = make_backend(model, draft); + const Session session = start_session(*backend, model); + + evict(*backend, "A"); + const size_t first_evicted = device_free(); + // Everything the model allocated is returned; the slack covers the HIP + // runtime's own per-process allocations made during the first launches. + std::printf("baseline free %.0f MiB, after evict %.0f MiB\n", + baseline / kMiB, first_evicted / kMiB); + SWAP_CHECK(first_evicted + (size_t)(512 * kMiB) >= baseline); + + // A request on an evicted model fails cleanly instead of touching freed memory. + DaemonIO io; + const auto refused = backend->generate(greedy(session.prompt, 4), io); + SWAP_CHECK(!refused.ok()); + SWAP_CHECK(backend->evict().status == ModelBackend::EvictStatus::ok); // idempotent + + reinstate(*backend, "A"); + check_session(*backend, session); + + size_t last_evicted = first_evicted; + for (int cycle = 1; cycle < swap_cycles(); ++cycle) { + evict(*backend, "A"); + last_evicted = device_free(); + reinstate(*backend, "A"); + } + check_session(*backend, session); + // No device memory creeps across cycles. + std::printf("free after evict: first %.0f MiB, last %.0f MiB\n", + first_evicted / kMiB, last_evicted / kMiB); + SWAP_CHECK(last_evicted + (size_t)(64 * kMiB) >= first_evicted); + + // Destroying an evicted backend is the shutdown path of a parked model. + evict(*backend, "A"); + backend.reset(); +} + +TEST_CASE(ModelSwapGpuFixture, QwenSwapThroughSecondModelKeepsSnapshots) { + const std::string model_a = luce_test::require_model(luce_test::kQwen35ModelEnv).string(); + const std::string draft_a = luce_test::require_model(luce_test::kDraftModelEnv).string(); + const std::string model_b = luce_test::require_model(kQwen35ModelBEnv).string(); + const std::string draft_b = luce_test::require_model(kDraftModelBEnv).string(); + require_gpu(); + + auto a = make_backend(model_a, draft_a); + const Session session_a = start_session(*a, model_a); + evict(*a, "A"); + + // B is created and used while A is evicted, then evicted in turn. + auto b = make_backend(model_b, draft_b); + const Session session_b = start_session(*b, model_b); + evict(*b, "B"); + + for (int cycle = 0; cycle < swap_cycles(); ++cycle) { + reinstate(*a, "A"); + check_session(*a, session_a); + evict(*a, "A"); + reinstate(*b, "B"); + check_session(*b, session_b); + evict(*b, "B"); + } +} + +// Warm eviction keeps both models' weights on the device and swaps only the +// per-model execution state (cache, graphs, backend context), so switching +// models skips the weight upload. +TEST_CASE(ModelSwapGpuFixture, QwenWarmSwapKeepsWeightsAndSnapshots) { + const std::string model_a = luce_test::require_model(luce_test::kQwen35ModelEnv).string(); + const std::string draft_a = luce_test::require_model(luce_test::kDraftModelEnv).string(); + const std::string model_b = luce_test::require_model(kQwen35ModelBEnv).string(); + const std::string draft_b = luce_test::require_model(kDraftModelBEnv).string(); + require_gpu(); + constexpr auto warm = ModelBackend::EvictLevel::warm; + + auto a = make_backend(model_a, draft_a); + const Session session_a = start_session(*a, model_a); + evict(*a, "A", warm); + SWAP_CHECK(a->weights_resident() && a->weight_device_bytes() > 0); + auto b = make_backend(model_b, draft_b); + const Session session_b = start_session(*b, model_b); + evict(*b, "B", warm); + const size_t both_warm = device_free(); + + for (int cycle = 0; cycle < swap_cycles(); ++cycle) { + reinstate(*a, "A (warm)"); + check_session(*a, session_a); + evict(*a, "A", warm); + reinstate(*b, "B (warm)"); + check_session(*b, session_b); + evict(*b, "B", warm); + } + // Nothing but the two weight sets stays on the device between turns. + SWAP_CHECK(device_free() + (size_t)(64 * kMiB) >= both_warm); + std::printf("weights kept: A %.0f MiB, B %.0f MiB\n", + a->weight_device_bytes() / kMiB, b->weight_device_bytes() / kMiB); + + // Warm to cold returns the weights; the model still reinstates from disk. + evict(*a, "A"); + SWAP_CHECK(!a->weights_resident()); + reinstate(*a, "A (cold)"); + check_session(*a, session_a); +} + +namespace { + +constexpr std::string_view kDeepSeek4DraftEnv = "LUCE_TEST_MODEL_DEEPSEEK4_DRAFT"; + +int ds4_cycles() { + const char * value = std::getenv("LUCE_TEST_DS4_CYCLES"); + return value && std::atoi(value) > 0 ? std::atoi(value) : 10; +} + +// The ds41-lucebox profile as the server applies it: its environment as +// defaults, --expert-device hip:1 (server_main apply_expert_device), +// --peer-access, and its flags. +std::unique_ptr make_ds41_backend(const std::string & model, + const std::string & draft) { + const auto * profile = luce::server::find_launch_profile("ds41-lucebox"); + SWAP_CHECK(profile); + for (const auto & env : profile->env) setenv(env.name, env.value, /*overwrite=*/0); + setenv("LUCE_DS4_MOE_TP", "1", 1); + setenv("LUCE_DS4_MOE_TP_INPROC", "1", 1); + setenv("LUCE_DS4_MOE_TP_GPU", "1", 1); + setenv("LUCE_DS4_MOE_TP_BACKEND", "hip", 1); + g_peer_access_opt_in = true; + const std::string share = LUCE_TEST_SHARE_DIR "/deepseek41/"; + BackendArgs args; + args.model_path = model; + args.draft_path = draft; + args.device.backend = compiled_placement_backend(); + args.device.gpu = 0; + args.device.max_ctx = 16384; + args.chunk = 4096; + args.ds4_prefill_mode = PrefillAttentionMode::Dense; + args.ds4_prefill_mode_set = true; + args.expert_placement = share + "placement_lucebox.json"; + args.ds4_router_bias = share + "router_bias_lucebox_40x384_f32.bin"; + args.ds4_protected_experts = share + "massive_experts.json"; + auto prepared = prepare_backend(args); + if (const auto * failure = std::get_if(&prepared)) { + throw std::runtime_error(failure->message); + } + auto backend = create_backend(std::get(prepared)); + SWAP_CHECK(backend); + std::string reason; + if (!backend->supports_eviction(reason)) throw std::runtime_error("DS4 eviction: " + reason); + return backend; +} + +// DS4 snapshots at a prefill boundary (after a decode, DSpark refuses one): +// the snapshot is taken at the end of the prompt, and the follow-up restores +// it and prefills its own suffix. +Session start_ds4_session(ModelBackend & backend, const std::string & model_path) { + Tokenizer tokenizer; + SWAP_CHECK(tokenizer.load_from_gguf(model_path.c_str())); + std::string text; + for (int i = 0; text.size() < 6000; ++i) { + text += "Record " + std::to_string(i) + ": the courier left depot " + + std::to_string(i * 7 % 13) + " at " + std::to_string(6 + i % 12) + + " o'clock carrying " + std::to_string(i * 3 + 1) + " parcels. "; + } + Session session; + session.prompt = tokenizer.encode(text + "\nWhich depot did record 17 leave from?"); + GenerateRequest first = greedy(session.prompt, 32); + first.snap_slot = Session::slot; + first.snap_pos = (int) session.prompt.size(); + first.restore_points = {first.snap_pos}; + DaemonIO io; + const auto answer = backend.generate(first, io); + SWAP_CHECK(answer.ok()); + session.first_answer = answer.tokens; + SWAP_CHECK(backend.snapshot_used(Session::slot)); + SWAP_CHECK(backend.snapshot_cur_pos(Session::slot) == first.snap_pos); + session.follow_up = session.prompt; + const auto question = tokenizer.encode("\nAnd how many parcels did record 23 carry?"); + session.follow_up.insert(session.follow_up.end(), question.begin(), question.end()); + const auto control = backend.restore_and_generate(Session::slot, greedy(session.follow_up, 32), io); + SWAP_CHECK(control.ok()); + SWAP_CHECK(control.restored_prefix_tokens == first.snap_pos); + session.control = control.tokens; + std::printf("ds4 session: prompt %zu tokens, snapshot at %d, follow-up %zu tokens\n", + session.prompt.size(), first.snap_pos, session.follow_up.size()); + return session; +} + +void check_ds4_session(ModelBackend & backend, const Session & session) { + DaemonIO io; + const auto resumed = backend.restore_and_generate(Session::slot, greedy(session.follow_up, 32), io); + SWAP_CHECK(resumed.ok()); + SWAP_CHECK(resumed.restored_prefix_tokens > 0); + SWAP_CHECK(resumed.tokens == session.control); +} + +} // namespace + +// Partial eviction returns the primary device while the secondary expert +// stack, the streamed expert cache and the host snapshots stay; reinstate() +// rebuilds the primary around them without reloading a secondary expert. +TEST_CASE(ModelSwapGpuFixture, DeepSeek4PartialEvictionKeepsSecondaryTierAndSnapshots) { + const std::string model = luce_test::require_model(luce_test::kDeepSeek4ModelEnv).string(); + const std::string draft = luce_test::require_model(kDeepSeek4DraftEnv).string(); + require_gpu(); + if (ggml_backend_cuda_get_device_count() < 2) { + throw CppUnitTestFramework::TestSkippedException("needs two GPUs (hip:0 target, hip:1 experts)"); + } + constexpr auto partial = ModelBackend::EvictLevel::partial; + const size_t baseline = device_free(0); + auto backend = make_ds41_backend(model, draft); + auto * ds4 = dynamic_cast(backend.get()); + SWAP_CHECK(ds4); + const Session session = start_ds4_session(*backend, model); + const auto loaded = ds4->expert_tier_report(); + SWAP_CHECK(loaded.secondary_upload_bytes > 0 && loaded.secondary_base); + + // Warm is not a DS4 level, and a refusal changes nothing. + SWAP_CHECK(backend->evict(ModelBackend::EvictLevel::warm).status == ModelBackend::EvictStatus::rejected); + SWAP_CHECK(!backend->is_evicted()); + + evict(*backend, "DS4", partial); + const size_t first_evicted = device_free(0); + const size_t secondary_evicted = device_free(1); + std::printf("hip:0 baseline free %.0f MiB, after partial evict %.0f MiB\n", + baseline / kMiB, first_evicted / kMiB); + // The primary device comes back to what it was before the model loaded, + // within the HIP runtime's own per-process allocations. + SWAP_CHECK(first_evicted + (size_t) (512 * kMiB) >= baseline); + SWAP_CHECK(backend->weights_resident()); + const auto evicted = ds4->expert_tier_report(); + SWAP_CHECK(evicted.partially_evicted && evicted.stream_cache_suspended); + SWAP_CHECK(evicted.storage == loaded.storage && evicted.secondary_base == loaded.secondary_base); + // Snapshots stay accounted for while no live cache exists. + const auto report = backend->memory_report(); + SWAP_CHECK(report.available && !report.snapshots.empty()); + DaemonIO io; + SWAP_CHECK(!backend->generate(greedy(session.prompt, 4), io).ok()); + SWAP_CHECK(backend->evict(partial).status == ModelBackend::EvictStatus::ok); // idempotent + + // The settings were resolved at load: a changed environment does not + // move the experts (reinstate verifies the placement and biases). + setenv("LUCE_DS4_MOE_TP_GPU", "0", 1); + setenv("LUCE_EXPERT_BUDGET_MB", "1", 1); + reinstate(*backend, "DS4 (partial)"); + const auto back = ds4->expert_tier_report(); + SWAP_CHECK(back.secondary_upload_bytes == loaded.secondary_upload_bytes); // no secondary re-read + SWAP_CHECK(back.secondary_base == loaded.secondary_base); + SWAP_CHECK(back.primary_upload_bytes > loaded.primary_upload_bytes); + SWAP_CHECK(!back.stream_cache_suspended); + check_ds4_session(*backend, session); + + size_t last_evicted = first_evicted; + for (int cycle = 1; cycle < ds4_cycles(); ++cycle) { + evict(*backend, "DS4", partial); + last_evicted = device_free(0); + reinstate(*backend, "DS4 (partial)"); + } + check_ds4_session(*backend, session); + SWAP_CHECK(ds4->expert_tier_report().secondary_upload_bytes == loaded.secondary_upload_bytes); + std::printf("hip:0 free after partial evict: first %.0f MiB, last %.0f MiB; hip:1 %.0f MiB\n", + first_evicted / kMiB, last_evicted / kMiB, secondary_evicted / kMiB); + // Nothing creeps on either device across cycles. + SWAP_CHECK(last_evicted + (size_t) (64 * kMiB) >= first_evicted); + evict(*backend, "DS4", partial); + SWAP_CHECK(device_free(1) + (size_t) (64 * kMiB) >= secondary_evicted); + + // Partial to cold releases the retained tier; a cold reinstate reloads it + // with the same placement and the snapshot still restores. + evict(*backend, "DS4"); + SWAP_CHECK(!backend->weights_resident()); + reinstate(*backend, "DS4 (cold)"); + check_ds4_session(*backend, session); + + // Destroying an evicted backend is the shutdown path. + evict(*backend, "DS4", partial); + backend.reset(); +} diff --git a/server/test/unit/test_moe_hybrid_storage.cpp b/server/test/unit/test_moe_hybrid_storage.cpp index 4d8245d2a9..3662468827 100644 --- a/server/test/unit/test_moe_hybrid_storage.cpp +++ b/server/test/unit/test_moe_hybrid_storage.cpp @@ -6,10 +6,17 @@ #include #include +#include #include #include #include +#if !defined(_WIN32) +#include +#include +#include +#endif + using namespace luce::common; namespace { @@ -240,3 +247,203 @@ TEST_CASE(MoeHybridStorageFixture, fractional_route_quota_rounds_over_the_batch) ggml_free(ctx); } + +#if !defined(_WIN32) +namespace { + +// A two-layer expert file, mapped read-only from offset zero as the loaders +// map a GGUF: per layer the gate, up and down stacks of six F32 experts. +// Every float encodes (layer, tensor, expert, element), so a stack's bytes +// say which experts it holds and in which order. +struct ExpertFile { + static constexpr int n_layer = 2, n_expert = 6, n_embd = 4, n_ff = 8; + static constexpr size_t expert_floats = (size_t) n_embd * n_ff; + static constexpr size_t expert_bytes = expert_floats * sizeof(float); + static constexpr size_t tensor_bytes = expert_bytes * n_expert; + + bool ok = false; + std::string path; + ggml_context * meta = nullptr; + std::vector descs; + + static float value(int layer, int tensor, int expert, size_t i) { + return (float) (layer * 1000 + tensor * 100 + expert) + (float) i / 64.0f; + } + static size_t offset(int layer, int tensor) { return ((size_t) layer * 3 + (size_t) tensor) * tensor_bytes; } + + ExpertFile() { + char name[] = "/tmp/moe-hybrid-primary-XXXXXX"; + const int fd = mkstemp(name); + if (fd < 0) return; + path = name; + std::vector data((size_t) n_layer * 3 * n_expert * expert_floats); + for (int il = 0; il < n_layer; ++il) + for (int t = 0; t < 3; ++t) + for (int e = 0; e < n_expert; ++e) + for (size_t i = 0; i < expert_floats; ++i) + data[(offset(il, t) + (size_t) e * expert_bytes) / sizeof(float) + i] = value(il, t, e, i); + const bool written = + ::write(fd, data.data(), data.size() * sizeof(float)) == (ssize_t) (data.size() * sizeof(float)); + ::close(fd); + meta = ggml_init({64 * ggml_tensor_overhead(), nullptr, true}); + if (!written || !meta) return; + for (int il = 0; il < n_layer; ++il) { + MoeLayerDesc d; + d.ffn_gate_exps = ggml_new_tensor_3d(meta, GGML_TYPE_F32, n_embd, n_ff, n_expert); + d.ffn_up_exps = ggml_new_tensor_3d(meta, GGML_TYPE_F32, n_embd, n_ff, n_expert); + d.ffn_down_exps = ggml_new_tensor_3d(meta, GGML_TYPE_F32, n_ff, n_embd, n_expert); + descs.push_back(d); + } + ok = true; + } + ~ExpertFile() { + if (meta) ggml_free(meta); + if (!path.empty()) std::remove(path.c_str()); + } + + // Builds storage over a fresh mapping (the storage unmaps it). + bool build(const std::vector> & hot, + const std::vector> & secondary, + ggml_backend_t primary, ggml_backend_t peer, int cache_slots, + MoeHybridStorage & out, std::string * err) const { + const int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) return false; + const size_t size = offset(n_layer, 0); + void * map = ::mmap(nullptr, size, PROT_READ, MAP_SHARED, fd, 0); + ::close(fd); + if (map == MAP_FAILED) return false; + MoeHybridConfig cfg; + cfg.n_layer = n_layer; + cfg.n_expert = n_expert; + cfg.n_expert_used = 2; + cfg.n_embd = n_embd; + cfg.n_ff_exp = n_ff; + cfg.cold_expert_backend = MoeHybridColdBackend::Gpu; + cfg.cold_expert_ids = secondary; + MoeHybridPlacement placement; + placement.n_layer = n_layer; + placement.n_expert = n_expert; + placement.n_expert_used = 2; + placement.hot_expert_ids = hot; + for (const auto & ids : hot) { + placement.hot_counts.push_back((int) ids.size()); + placement.total_hot += (int) ids.size(); + } + std::vector files((size_t) n_layer); + const auto * base = static_cast(map); + for (int il = 0; il < n_layer; ++il) { + files[(size_t) il].gate_exps = {base + offset(il, 0), tensor_bytes}; + files[(size_t) il].up_exps = {base + offset(il, 1), tensor_bytes}; + files[(size_t) il].down_exps = {base + offset(il, 2), tensor_bytes}; + } + const bool ok = build_moe_hybrid_storage_from_file_with_mmap( + cfg, primary, placement, descs, files, map, size, out, err, cache_slots, peer); + if (!ok) ::munmap(map, size); + return ok; + } + + // What a stack holding `experts` of (layer, tensor), in that order, must contain. + static std::vector expected(int layer, int tensor, const std::vector & experts) { + std::vector out; + for (int32_t e : experts) + for (size_t i = 0; i < expert_floats; ++i) out.push_back(value(layer, tensor, e, i)); + return out; + } +}; + +std::vector tensor_floats(const ggml_tensor * t) { + std::vector out(ggml_nbytes(t) / sizeof(float)); + ggml_backend_tensor_get(t, out.data(), 0, ggml_nbytes(t)); + return out; +} + +} // namespace + +// A partial eviction frees the primary expert stacks and later rebuilds them +// around the retained secondary stack: same rows in the same order, every +// owner map unchanged, the secondary stack neither reallocated nor reloaded. +TEST_CASE(MoeHybridStorageFixture, primary_experts_rebuild_in_place_around_the_secondary_stack) { + const ExpertFile file; + REQUIRE(file.ok); + auto primary = std::unique_ptr( + ggml_backend_cpu_init(), ggml_backend_free); + auto peer = std::unique_ptr( + ggml_backend_cpu_init(), ggml_backend_free); + REQUIRE(primary && peer); + // Unsorted hot ids: the stack row order is the list order. + const std::vector> hot = {{4, 1}, {0, 5, 2}}; + const std::vector> secondary = {{0, 2}, {3}}; + MoeHybridStorage storage; + std::string err; + REQUIRE(file.build(hot, secondary, primary.get(), peer.get(), 0, storage, &err)); + const MoeHybridOwnershipRecord placed = MoeHybridOwnershipRecord::capture(storage); + const uint64_t primary_bytes = storage.primary_upload_bytes; + const uint64_t secondary_bytes = storage.secondary_upload_bytes; + REQUIRE(primary_bytes == 5 * 3 * ExpertFile::expert_bytes); + REQUIRE(secondary_bytes == 3 * 3 * ExpertFile::expert_bytes); + std::vector cold_bufs; + for (const auto & layer : storage.layers) cold_bufs.push_back(layer.cold_buf); + + REQUIRE(storage.release_primary_experts(&err)); + REQUIRE(storage.release_primary_experts(&err)); // idempotent + for (size_t il = 0; il < storage.layers.size(); ++il) { + const auto & layer = storage.layers[il]; + CHECK(!layer.hot_buf && !layer.gate_hot && !layer.up_hot && !layer.down_hot); + CHECK(!layer.is_expert_hot(hot[il][0])); // not resident while released + CHECK(layer.cold_buf == cold_bufs[il]); + } + REQUIRE(placed.matches(storage, &err)); + + REQUIRE(storage.rebuild_primary_experts(primary.get(), file.path, &err)); + for (int il = 0; il < ExpertFile::n_layer; ++il) { + const auto & layer = storage.layers[(size_t) il]; + CHECK(tensor_floats(layer.gate_hot) == ExpertFile::expected(il, 0, hot[(size_t) il])); + CHECK(tensor_floats(layer.up_hot) == ExpertFile::expected(il, 1, hot[(size_t) il])); + CHECK(tensor_floats(layer.down_hot) == ExpertFile::expected(il, 2, hot[(size_t) il])); + CHECK(tensor_floats(layer.gate_cold) == ExpertFile::expected(il, 0, secondary[(size_t) il])); + CHECK(layer.cold_buf == cold_bufs[(size_t) il]); + for (int e = 0; e < ExpertFile::n_expert; ++e) { + const bool is_hot = std::count(hot[(size_t) il].begin(), hot[(size_t) il].end(), e) > 0; + CHECK(layer.is_expert_hot(e) == is_hot); + } + } + CHECK(placed.matches(storage, &err)); + CHECK(storage.primary_upload_bytes == 2 * primary_bytes); + CHECK(storage.secondary_upload_bytes == secondary_bytes); + + // A map that no longer matches its ordered ids (a recomputed placement) + // is refused; the storage stays released rather than half rebuilt. + REQUIRE(storage.release_primary_experts(&err)); + std::swap(storage.layers[0].hot_local_by_global[4], storage.layers[0].hot_local_by_global[1]); + err.clear(); + CHECK(!storage.rebuild_primary_experts(primary.get(), file.path, &err)); + CHECK(err.find("layer 0") != std::string::npos); + CHECK(storage.primary_experts_released()); + CHECK(!storage.layers[0].hot_buf && !storage.layers[1].hot_buf); + CHECK(!placed.matches(storage)); +} + +// A storage whose primary population is mutable (spare cache slots that +// swap experts in) cannot be rebuilt as it was, so it is not released, and a +// different placement never passes for the recorded one. +TEST_CASE(MoeHybridStorageFixture, primary_release_refuses_spare_cache_slots) { + const ExpertFile file; + REQUIRE(file.ok); + auto primary = std::unique_ptr( + ggml_backend_cpu_init(), ggml_backend_free); + auto peer = std::unique_ptr( + ggml_backend_cpu_init(), ggml_backend_free); + REQUIRE(primary && peer); + MoeHybridStorage fixed, cached; + std::string err; + REQUIRE(file.build({{4, 1}, {0, 5, 2}}, {}, primary.get(), peer.get(), 0, fixed, &err)); + REQUIRE(file.build({{1, 4}, {0, 5, 2}}, {}, primary.get(), peer.get(), 1, cached, &err)); + CHECK(!MoeHybridOwnershipRecord::capture(fixed).matches(cached)); + + err.clear(); + CHECK(!cached.release_primary_experts(&err)); + CHECK(err.find("spare") != std::string::npos); + CHECK(!cached.primary_experts_released()); + CHECK(cached.layers[0].hot_buf != nullptr); +} +#endif diff --git a/server/test/unit/test_seq_batch_plan.cpp b/server/test/unit/test_seq_batch_plan.cpp index 27a8549c93..bc00da34eb 100644 --- a/server/test/unit/test_seq_batch_plan.cpp +++ b/server/test/unit/test_seq_batch_plan.cpp @@ -146,8 +146,10 @@ int main() { // A selected prefill may terminate with a per-request error; the // scheduler retires only that slot. SeqEngine::StepResult prefill_failure = good; - prefill_failure.prefills[0] = { - 1, SeqEngine::PrefillOutput::Status::failed, -1, "prefill failed"}; + prefill_failure.prefills[0] = {}; + prefill_failure.prefills[0].slot = 1; + prefill_failure.prefills[0].status = SeqEngine::PrefillOutput::Status::failed; + prefill_failure.prefills[0].error = "prefill failed"; CHECK(validate_step_result(work, prefill_failure, 2).empty()); SeqEngine::StepResult bad_row_failure = prefill_failure; diff --git a/server/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index e1aa19ee81..013aef2e17 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -8132,6 +8132,7 @@ class SchedulerPrefixEngine final : public SeqEngine { result.slot = chosen; if (plan.restore.valid()) { + restores.push_back(plan.restore); result.prefix_store.restore_attempted = true; result.prefix_store.restore_elapsed_us = 2500; if (unrequested_restore.load(std::memory_order_relaxed)) { @@ -8176,6 +8177,7 @@ class SchedulerPrefixEngine final : public SeqEngine { output.slot = slice.slot; output.status = PrefillOutput::Status::completed; output.token = 2; + if (slot.hidden) output.hidden_states = HiddenStates{}; if (slot.capture.valid()) { output.prefix_store.status = PrefixStoreEvent::Status::saved; output.prefix_store.ticket = slot.capture; @@ -8188,6 +8190,23 @@ class SchedulerPrefixEngine final : public SeqEngine { return result; } + bool supports_hidden_states() const override { return true; } + AdmitResult admit_hidden_states( + uint64_t request_id, + const std::vector & prompt, + const HiddenStatesSpec &, + PrefixCaptureTicket capture) override { + hidden_captures.push_back(capture); + PrefixStorePlan plan; + plan.capture = capture; + AdmitResult result = + admit_with_prefix(request_id, prompt, SamplerCfg{}, plan); + if (result.status == AdmitResult::Status::admitted) { + slots_[(size_t)result.slot].hidden = true; + } + return result; + } + void retire(int slot) override { if (slot >= 0 && slot < (int)slots_.size()) { slots_[(size_t)slot] = Slot{}; @@ -8198,6 +8217,9 @@ class SchedulerPrefixEngine final : public SeqEngine { discarded.push_back(checkpoint); } + // Admission inputs as the scheduler planned them. + std::vector hidden_captures; + std::vector restores; bool saw_capture = false; bool saw_stale_restore = false; std::atomic defer_restore{false}; @@ -8210,6 +8232,7 @@ class SchedulerPrefixEngine final : public SeqEngine { struct Slot { bool active = false; bool prefilling = false; + bool hidden = false; PrefixCaptureTicket capture; }; std::vector slots_; @@ -8394,6 +8417,122 @@ TEST_CASE(ServerUnitFixture, TEST_ASSERT(stats.restore_stall_us_total == 2500); } +TEST_CASE(ServerUnitFixture, + test_scheduler_hidden_states_cache_prefix_feeds_next_chat) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + + auto backend_owner = std::make_unique(); + SchedulerPrefixBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config; + config.arch = "qwen35"; + config.max_ctx = 64; + config.prefix_cache_cap = 2; + config.concurrent_prefix_cache_max_bytes = 1024; + config.concurrent_prefix_cache = true; + config.admission_coalesce_ms = 0; + HttpServer server(engine, tokenizer, config); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + + // A "bare" readout ends at the assistant header; the chat request + // that follows appends the template's think block. + const std::vector bare = {1, 100, 3, 101, 4}; + const std::vector chat = {1, 100, 3, 101, 4, 0, 0}; + + int plain_sockets[2] = {-1, -1}; + int hidden_sockets[2] = {-1, -1}; + int chat_sockets[2] = {-1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, plain_sockets) == 0); + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, hidden_sockets) == 0); + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, chat_sockets) == 0); + + const auto hidden_job = [&](ServerJob & job, int fd, bool cache_prefix, + const char * id) { + job.fd = fd; + job.req.format = ApiFormat::OPENAI_CHAT; + job.req.prompt_tokens = bare; + job.req.stream = false; + job.req.model = "scheduler-test"; + job.req.response_id = id; + HiddenStatesRequest hidden; + hidden.layers = {0}; + hidden.cache_prefix = cache_prefix; + job.req.hidden_states = hidden; + job.req.hidden_spec.layers = {0}; + }; + ServerJob plain_job; + hidden_job(plain_job, plain_sockets[0], false, "plain-readout"); + ServerJob cached_job; + hidden_job(cached_job, hidden_sockets[0], true, "cached-readout"); + ServerJob chat_job; + chat_job.fd = chat_sockets[0]; + chat_job.req.format = ApiFormat::OPENAI_CHAT; + chat_job.req.prompt_tokens = chat; + chat_job.req.max_output = 1; + chat_job.req.stream = false; + chat_job.req.model = "scheduler-test"; + chat_job.req.response_id = "chat"; + + std::thread scheduler([&] { + SchedulerTestHarness::run(server, backend.engine); + }); + const auto run_job = [&](ServerJob & job) { + SchedulerTestHarness::enqueue(server, &job); + std::unique_lock lock(job.mu); + return job.cv.wait_for(lock, std::chrono::seconds(5), + [&] { return job.done; }); + }; + const bool plain_done = run_job(plain_job); + const auto after_plain = cache.lookup_candidate(chat, (int)chat.size() - 1); + const bool cached_done = run_job(cached_job); + const auto after_cached = cache.lookup_candidate(chat, (int)chat.size() - 1); + const bool chat_done = run_job(chat_job); + SchedulerTestHarness::stop(server); + scheduler.join(); + + const auto read_body = [](int fd) { + std::string body; + char buf[4096]; + ssize_t n; + while ((n = recv(fd, buf, sizeof(buf), MSG_DONTWAIT)) > 0) { + body.append(buf, (size_t)n); + } + const size_t at = body.find("\r\n\r\n"); + return at == std::string::npos ? json() + : json::parse(body.substr(at + 4), nullptr, false); + }; + const json plain_body = read_body(plain_sockets[1]); + const json cached_body = read_body(hidden_sockets[1]); + for (int fd : {plain_sockets[0], plain_sockets[1], hidden_sockets[0], + hidden_sockets[1], chat_sockets[0], chat_sockets[1]}) { + close(fd); + } + unlink(path.c_str()); + + TEST_ASSERT(plain_done && cached_done && chat_done); + // Default readouts leave nothing behind and say nothing about caching. + TEST_ASSERT(backend.engine.hidden_captures.size() == 2); + TEST_ASSERT(!backend.engine.hidden_captures[0].valid()); + TEST_ASSERT(after_plain.first == -1); + TEST_ASSERT(plain_body.is_object() && !plain_body.contains("cache_prefix")); + // cache_prefix captures the whole readout prompt... + TEST_ASSERT(backend.engine.hidden_captures[1].valid()); + TEST_ASSERT(backend.engine.hidden_captures[1].checkpoint.tokens == + (int)bare.size()); + TEST_ASSERT(after_cached.first >= 0); + TEST_ASSERT(after_cached.second == (int)bare.size()); + TEST_ASSERT(cached_body.is_object()); + TEST_ASSERT(cached_body["cache_prefix"]["cached_tokens"] == + (int)bare.size()); + // ...and the chat request extending it restores that checkpoint. The + // readouts themselves never restored anything. + TEST_ASSERT(backend.engine.restores == + std::vector({{ + (uint64_t)after_cached.first + 1, (int)bare.size()}})); +} + TEST_CASE(ServerUnitFixture, test_scheduler_rejects_malformed_restore_and_drops_stale_entry) { const std::string path = write_deepseek_marker_tokenizer_fixture(); diff --git a/server/test/unit/test_token_logprobs.cpp b/server/test/unit/test_token_logprobs.cpp new file mode 100644 index 0000000000..eb94f0ba15 --- /dev/null +++ b/server/test/unit/test_token_logprobs.cpp @@ -0,0 +1,129 @@ +// Logprobs for /v1/chat/completions: the log-softmax/top-k the backends run +// on a logits row, request validation, and the OpenAI response shape. + +#include "CppUnitTestFramework.hpp" +#include "common/token_logprobs.h" +#include "server/http_server.h" + +#include +#include +#include +#include +#include + +using luce::common::ApiFormat; +using luce::common::TokenLogprob; +using luce::common::TokenLogprobs; +using luce::common::build_openai_logprobs; +using luce::common::compute_token_logprobs; +using luce::common::json; +using luce::common::kMaxTopLogprobs; +using luce::common::parse_request_logprobs; + +namespace { +struct TokenLogprobsFixture {}; + +bool near(float a, double b) { return std::fabs((double)a - b) < 1e-5; } +} // namespace + +TEST_CASE(TokenLogprobsFixture, logprobs_normalize_and_rank_ties_by_id) { + // Probabilities 1/8, 2/8, 3/8, 2/8. + const std::vector logits = { + 0.0f, std::log(2.0f), std::log(3.0f), std::log(2.0f)}; + const TokenLogprobs lp = compute_token_logprobs(logits.data(), 4, 0, 3); + CHECK(lp.chosen.token == 0); + CHECK(near(lp.chosen.logprob, std::log(1.0 / 8.0))); + CHECK(lp.top.size() == 3u); + CHECK(lp.top[0].token == 2); + CHECK(near(lp.top[0].logprob, std::log(3.0 / 8.0))); + // Tokens 1 and 3 tie: the lower id ranks first. + CHECK(lp.top[1].token == 1); + CHECK(lp.top[2].token == 3); + CHECK(near(lp.top[2].logprob, std::log(2.0 / 8.0))); +} + +TEST_CASE(TokenLogprobsFixture, logprobs_stay_finite_for_large_logits) { + const std::vector logits = {1000.0f, 1000.0f, -1000.0f}; + const TokenLogprobs lp = compute_token_logprobs(logits.data(), 3, 1, 1); + CHECK(near(lp.chosen.logprob, -std::log(2.0))); + CHECK(lp.top.size() == 1u); + CHECK(lp.top[0].token == 0); +} + +TEST_CASE(TokenLogprobsFixture, logprobs_top_n_is_clamped) { + const std::vector logits(64, 0.0f); + CHECK(compute_token_logprobs(logits.data(), 64, 0, 0).top.empty()); + CHECK(compute_token_logprobs(logits.data(), 64, 0, -1).top.empty()); + CHECK(compute_token_logprobs(logits.data(), 64, 0, 100).top.size() == + (size_t)kMaxTopLogprobs); + CHECK(compute_token_logprobs(logits.data(), 3, 0, 10).top.size() == 3u); +} + +TEST_CASE(TokenLogprobsFixture, request_logprobs_resolve_to_top_n) { + const ApiFormat chat = ApiFormat::OPENAI_CHAT; + CHECK(parse_request_logprobs(json::object(), chat) == -1); + CHECK(parse_request_logprobs({{"logprobs", false}}, chat) == -1); + CHECK(parse_request_logprobs({{"logprobs", true}}, chat) == 0); + CHECK(parse_request_logprobs( + {{"logprobs", true}, {"top_logprobs", 20}}, chat) == 20); + CHECK(parse_request_logprobs( + {{"logprobs", true}, {"top_logprobs", nullptr}}, chat) == 0); +} + +TEST_CASE(TokenLogprobsFixture, request_logprobs_ignore_shapes_that_predate_them) { + const ApiFormat chat = ApiFormat::OPENAI_CHAT; + const std::vector ignored = { + {{"logprobs", 1}}, // legacy completions integer form + {{"logprobs", nullptr}}, + {{"top_logprobs", 5}}, + {{"logprobs", false}, {"top_logprobs", 5}}, + {{"logprobs", "true"}, {"top_logprobs", 99}}, + }; + for (const json & body : ignored) { + CHECK(parse_request_logprobs(body, chat) == -1); + } + for (const ApiFormat format : + {ApiFormat::ANTHROPIC, ApiFormat::RESPONSES, ApiFormat::COMPLETIONS}) { + CHECK(parse_request_logprobs( + {{"logprobs", true}, {"top_logprobs", 99}}, format) == -1); + } +} + +TEST_CASE(TokenLogprobsFixture, request_logprobs_rejects_invalid_top_logprobs) { + const ApiFormat chat = ApiFormat::OPENAI_CHAT; + const std::vector invalid = { + {{"logprobs", true}, {"top_logprobs", 21}}, + {{"logprobs", true}, {"top_logprobs", -1}}, + {{"logprobs", true}, {"top_logprobs", "5"}}, + }; + for (const json & body : invalid) { + CHECK_THROW(std::invalid_argument, + UNUSED_RETURN(parse_request_logprobs(body, chat))); + } +} + +TEST_CASE(TokenLogprobsFixture, openai_logprobs_keep_exact_bytes_and_valid_json) { + TokenLogprobs position; + position.chosen = {7, -0.25f}; + position.top = { + {7, -0.25f}, + {9, -std::numeric_limits::infinity()}, + }; + // Token 9 is the first two bytes of a three-byte UTF-8 sequence. + const json out = build_openai_logprobs( + {position}, [](int32_t id) { + return id == 7 ? std::string("easy") : std::string("\xE2\x82"); + }); + + const json & content = out.at("content"); + CHECK(content.size() == 1u); + const json & first = content[0]; + CHECK(first.at("token") == "easy"); + CHECK(first.at("logprob").get() == -0.25f); + CHECK(first.at("bytes") == json::array({101, 97, 115, 121})); + const json & partial = first.at("top_logprobs")[1]; + CHECK(partial.at("bytes") == json::array({0xE2, 0x82})); + CHECK(partial.at("token") == "\xEF\xBF\xBD\xEF\xBF\xBD"); + CHECK(partial.at("logprob").get() == -9999.0f); + CHECK_NO_THROW(UNUSED_RETURN(out.dump())); +} From fd4bb6a0bbe7098c91841dbf7de6fe9f580d3b62 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 12 Oct 2026 07:20:14 +0000 Subject: [PATCH 2/2] feat(server): supervise a local routing gateway --- server/CMakeLists.txt | 1 + server/README.md | 58 +++++++++++ server/src/deepseek4/deepseek4_backend.cpp | 9 +- server/src/deepseek4/deepseek4_backend.h | 1 + server/src/server/routing_gateway.h | 107 ++++++++++++++++++++ server/src/server/server_main.cpp | 110 ++++++++++++++++++++- server/test/test_model_names.cmake | 37 +++++++ server/test/unit/test_model_swap_gpu.cpp | 4 + server/test/unit/test_routing_gateway.cpp | 73 ++++++++++++++ 9 files changed, 396 insertions(+), 4 deletions(-) create mode 100644 server/src/server/routing_gateway.h create mode 100644 server/test/unit/test_routing_gateway.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 28402c146f..4db1d70335 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -2326,6 +2326,7 @@ if(LUCE_TESTS) test/unit/test_unit_main.cpp test/unit/test_server_unit.cpp test/unit/test_model_routing.cpp + test/unit/test_routing_gateway.cpp test/unit/test_luce_engine.cpp test/unit/test_paged_kv_offload.cpp test/unit/test_anchor_params.cpp diff --git a/server/README.md b/server/README.md index f264586abb..107e7f0b15 100644 --- a/server/README.md +++ b/server/README.md @@ -179,6 +179,64 @@ current monolithic Qwen3.5/Qwen3.6 path uses DDTree with tested defaults. Hardware- or checkpoint-specific overrides remain available through `--ddtree-budget`, `--ddtree-tau`, `--specla-top-k`, and `--draft-swa`. +### Managed local routing gateway + +`--luce-mix` enables an installed pair/topology launcher before model loading: + +```bash +LUCE_MIX_LAUNCHER=/path/to/luce-mix ./build/luce_server --luce-mix --help +``` + +The launcher defaults to `luce-mix` on `PATH`; `LUCE_MIX_LAUNCHER` selects an +explicit executable. It receives `--engine ` and the remaining +arguments unchanged, resolves its model-pair and topology options, then execs +the engine with normal model blocks and the managed gateway below. There is no +extra resident launcher process. Exactly which pair presets are available is +owned by the installed launcher. `--print-profile ` exports the engine's +resolved profile flags and environment defaults as JSON so launchers can reuse +qualified tuning without copying it. + +On POSIX hosts, `--routing-gateway /absolute/path/to/gateway` starts an installed +routing gateway alongside the engine. Combine it with `--model-routing name` +and explicit `--model-name` values. For example: + +```bash +./build/luce_server models/default.gguf --model-name default --target-device hip:0 \ + --host 127.0.0.1 --port 8080 \ + --model models/pro.gguf --model-name pro --target-device hip:0 \ + --model-routing name --swap-residency \ + --routing-gateway /opt/local-routing/start-gateway +``` + +Clients use port 8080, served by the gateway. The engine binds only to +`127.0.0.1:8491`; `--routing-engine-port` changes this internal port, which must +differ from the public port. The gateway owns model selection and request +policy. The engine owns scheduling, residency and model inference. + +The executable receives these arguments (without a shell): + +```text +--engine-url http://127.0.0.1:8491 --host 127.0.0.1 --port 8080 +``` + +Install the gateway and its dependencies separately. Its wrapper should `exec` +the gateway, honor these listener arguments, and use the engine URL for model +backends. Startup happens after model loading. The engine terminates the gateway +on shutdown, waits up to five seconds before killing it, and exits with an error +if the gateway exits unexpectedly. A missing executable fails before loading +models. The public API surface in this mode is whichever endpoints the installed +gateway supports; engine-only endpoints remain on the internal listener. + +Requests from the gateway must name the configured engine model. A named request +for an evicted model waits for conflicting requests to drain and triggers its +reinstatement; `auto` or a missing model only selects among resident models. +`/props` and `/status/json` expose model IDs, residency, retained devices, capacity, +queue counts and swap counters for a gateway that needs runtime state. Thinking +arms can set `thinking.budget_tokens` and `thinking.reply_budget` within the +request's `max_tokens` limit. Non-streaming Chat Completions expose +`choices[].finish_details.close_kind` for forced thinking closure; streamed text +alone does not provide an exact thinking-budget boundary. + ### Compression proxy mode `luce_server` can run as a **PFlash compression proxy** in front of any diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 2f63126f26..bb1d763e3e 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -693,6 +693,7 @@ static DeepSeek4ResolvedSettings resolve_deepseek4_settings(const DeepSeek4Backe out.force_full_load = env_flag_enabled("LUCE_DS4_FORCE_FULL_LOAD"); out.dense_tp_mask = env_flag_enabled("LUCE_DS4_DENSE_TP_MASK"); out.fused_decode_env = env_flag_enabled("LUCE_DS4_FUSED_DECODE"); + out.hybrid_f16_verify_kv = env_flag_enabled("LUCE_DS4_HYBRID_F16_VERIFY_KV"); out.decode_all_cold = env_flag_enabled("LUCE_DS4_DECODE_ALL_COLD"); out.spec_requested = !cfg.draft_path.empty() || env_flag_enabled("LUCE_DS4_SPEC"); if (const char * path = std::getenv("LUCE_DS4_DRAFT")) out.spec_env_draft_path = path; @@ -2230,7 +2231,7 @@ bool DeepSeek4Backend::load_model() { "backend default"); w_.fused_decode = cfg_.fused_decode && !moe_hybrid_; w_.fused_verify_f16_kv = (cfg_.fused_verify_f16_kv && !moe_hybrid_) || - (moe_hybrid_ && env_flag_enabled("LUCE_DS4_HYBRID_F16_VERIFY_KV")); + (moe_hybrid_ && settings_.hybrid_f16_verify_kv); if (moe_hybrid_ && w_.fused_verify_f16_kv) { std::fprintf(stderr, "[deepseek4] hybrid tier: F16 verify K/V on (LUCE_DS4_HYBRID_F16_VERIFY_KV)\n"); } @@ -4106,6 +4107,10 @@ bool DeepSeek4Backend::verify_frozen_state(std::string & error) const { error = "expert placement differs from the one the model was loaded with: " + err; return false; } + if (w_.fused_verify_f16_kv != settings_.hybrid_f16_verify_kv) { + error = "hybrid F16 verification mode differs from the one the model was loaded with"; + return false; + } if (w_.selection_bias_host != frozen_.selection_bias_host) { error = "routing biases differ from the ones the model was loaded with"; return false; @@ -4148,7 +4153,7 @@ bool DeepSeek4Backend::reinstate_primary(std::string & error) { w_.mixed_mmq_policy = moe_hybrid_->mixed_mmq_policy; w_.routed_expert_top_k = cfg_.expert_top_k; w_.fused_decode = false; - w_.fused_verify_f16_kv = false; + w_.fused_verify_f16_kv = settings_.hybrid_f16_verify_kv; lap(1); if (!moe_hybrid_->rebuild_primary_experts(backend_, cfg_.model_path, &err)) { return fail("primary experts: " + err); diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 6f2742c9b5..f269a672e4 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -96,6 +96,7 @@ struct DeepSeek4ResolvedSettings { bool force_full_load = false; // LUCE_DS4_FORCE_FULL_LOAD bool dense_tp_mask = false; // LUCE_DS4_DENSE_TP_MASK bool fused_decode_env = false; // LUCE_DS4_FUSED_DECODE + bool hybrid_f16_verify_kv = false; // LUCE_DS4_HYBRID_F16_VERIFY_KV bool decode_all_cold = false; // LUCE_DS4_DECODE_ALL_COLD // DSpark: --draft, or LUCE_DS4_SPEC with LUCE_DS4_DRAFT. bool spec_requested = false; diff --git a/server/src/server/routing_gateway.h b/server/src/server/routing_gateway.h new file mode 100644 index 0000000000..da7da9dc27 --- /dev/null +++ b/server/src/server/routing_gateway.h @@ -0,0 +1,107 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef _WIN32 +#include +#include +#include +#include +#include +extern char ** environ; +#endif + +namespace luce::server { + +// One optional local gateway. Only the monitor thread reaps or signals the child. +// Keep policy/configuration in the executable; no shell interpretation here. +class RoutingGateway { +public: + ~RoutingGateway() { stop(); } + RoutingGateway() = default; + RoutingGateway(const RoutingGateway &) = delete; + RoutingGateway & operator=(const RoutingGateway &) = delete; + + bool start(const std::string & executable, const std::string & host, int port, + int engine_port, std::function on_exit) { +#ifdef _WIN32 + (void)executable; (void)host; (void)port; (void)engine_port; (void)on_exit; + std::fprintf(stderr, "[server] --routing-gateway requires a POSIX host\n"); + return false; +#else + std::vector args = {executable, "--engine-url", + "http://127.0.0.1:" + std::to_string(engine_port), + "--host", host, "--port", std::to_string(port)}; + std::vector argv; + for (auto & arg : args) argv.push_back(arg.data()); + argv.push_back(nullptr); + posix_spawnattr_t attr; + posix_spawnattr_init(&attr); + posix_spawnattr_setflags(&attr, POSIX_SPAWN_SETPGROUP); + posix_spawnattr_setpgroup(&attr, 0); + posix_spawn_file_actions_t files; + posix_spawn_file_actions_init(&files); +#if defined(__GLIBC__) +#if __GLIBC_PREREQ(2, 34) + // Do not inherit model files, device handles, or server sockets. + posix_spawn_file_actions_addclosefrom_np(&files, 3); +#endif +#endif + pid_t pid; + const int error = posix_spawn(&pid, executable.c_str(), &files, &attr, argv.data(), environ); + posix_spawn_file_actions_destroy(&files); + posix_spawnattr_destroy(&attr); + if (error) { + std::fprintf(stderr, "[server] cannot start routing gateway %s: %s\n", + executable.c_str(), std::strerror(error)); + return false; + } + monitor_ = std::thread([this, pid, on_exit = std::move(on_exit)] { + int status = 0; + while (!stopping_.load()) { + const pid_t result = waitpid(pid, &status, WNOHANG); + if (result == pid || (result < 0 && errno != EINTR)) { + // A wrapper must exec its gateway, but also stop descendants + // if a broken wrapper leaves them behind when it exits. + kill(-pid, SIGTERM); + failed_.store(true); + std::fprintf(stderr, "[server] routing gateway exited unexpectedly (wait status %d)\n", status); + on_exit(); + return; + } + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + } + kill(-pid, SIGTERM); + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (std::chrono::steady_clock::now() < deadline) { + const pid_t result = waitpid(pid, &status, WNOHANG); + if (result == pid || (result < 0 && errno != EINTR)) return; + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + } + kill(-pid, SIGKILL); + while (waitpid(pid, &status, 0) < 0 && errno == EINTR) {} + }); + return true; +#endif + } + + void stop() { + stopping_.store(true); + if (monitor_.joinable()) monitor_.join(); + } + bool failed() const { return failed_.load(); } + +private: + std::atomic stopping_{false}; + std::atomic failed_{false}; + std::thread monitor_; +}; + +} // namespace luce::server diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 14d8aab0d6..ac65e1a9a9 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -14,6 +14,7 @@ #include "http_server.h" #include "chat_template.h" #include "launch_profiles.h" +#include "routing_gateway.h" #include "model_card.h" #include "common/backend_factory.h" #include "common/chain_rollback_policy.h" @@ -118,6 +119,12 @@ static void print_usage(const char * prog) { " --load-balancing-primary-gpu Select the primary model by its target device.\n" " Defaults to the first block; request model names\n" " do not change generation routing.\n" + " --luce-mix Resolve a model pair/topology with the installed luce-mix launcher\n" + " (override its path with LUCE_MIX_LAUNCHER; --luce-mix --help).\n" + " --print-profile Print resolved profile flags and environment as JSON\n" + " --routing-gateway Start a local gateway on --host/--port (POSIX).\n" + " Requires name routing; backend moves to loopback.\n" + " --routing-engine-port Internal backend port (default: 8491).\n" " --model-routing How several model blocks share requests\n" " (default: balance). name loads every block and runs a\n" " request naming a block's --model-name only on that\n" @@ -2216,6 +2223,53 @@ static bool expand_launch_profile(std::vector & block, } int main(int argc, char ** argv) { + // The private launcher resolves pair/topology before any device or model + // initialization, then execs this binary with ordinary model blocks. + if (std::any_of(argv + 1, argv + argc, [](const char * arg) { + return std::strcmp(arg, "--luce-mix") == 0; + })) { +#ifdef _WIN32 + std::fprintf(stderr, "[server] --luce-mix requires a POSIX host\n"); + return 2; +#else + const char * launcher = std::getenv("LUCE_MIX_LAUNCHER"); + if (!launcher || !*launcher) launcher = "luce-mix"; + std::error_code error; + const auto self = std::filesystem::read_symlink("/proc/self/exe", error); + std::string executable = error ? argv[0] : self.string(); + std::vector args = {const_cast(launcher), + const_cast("--engine"), executable.data()}; + for (int i = 1; i < argc; ++i) + if (std::strcmp(argv[i], "--luce-mix") != 0) args.push_back(argv[i]); + args.push_back(nullptr); + execvp(launcher, args.data()); + std::fprintf(stderr, "[server] cannot start Luce Mix launcher '%s': %s; " + "install it on PATH or set LUCE_MIX_LAUNCHER\n", + launcher, std::strerror(errno)); + return 2; +#endif + } + if (argc >= 2 && std::strcmp(argv[1], "--print-profile") == 0) { + const auto * profile = argc == 3 ? luce::server::find_launch_profile(argv[2]) : nullptr; + if (!profile) { + std::fprintf(stderr, "[server] --print-profile requires a profile name (%s)\n", + luce::server::launch_profile_names().c_str()); + return 2; + } + nlohmann::json flags = nlohmann::json::array(), env = nlohmann::json::object(); + for (const auto & flag : profile->flags) { + flags.push_back(flag.flag); + if (flag.value) { + const std::string value = flag.value; + flags.push_back(value.rfind("share/", 0) == 0 + ? resolve_profile_data_path(value, argv[0]) : value); + } + } + for (const auto & entry : profile->env) env[entry.name] = entry.value; + const nlohmann::json result = {{"name", profile->name}, {"flags", flags}, {"env", env}}; + std::puts(result.dump().c_str()); + return 0; + } if (argc >= 2 && std::strcmp(argv[1], "--list-devices") == 0) { if (argc > 3) { std::fprintf(stderr, "Usage: %s --list-devices [model.gguf]\n", argv[0]); @@ -2234,9 +2288,28 @@ int main(int argc, char ** argv) { bool swap_residency = false; std::set swap_keep_weights; // model names evicted warm std::string primary_gpu; + std::string routing_gateway; + int routing_engine_port = 8491; + bool routing_engine_port_set = false; for (int i = 1; i < argc; ++i) { if (std::strcmp(argv[i], "--load-balancing") == 0) { load_balancing = true; + } else if (std::strcmp(argv[i], "--routing-gateway") == 0) { + if (i + 1 >= argc || argv[i + 1][0] == '-') { + std::fprintf(stderr, "[server] --routing-gateway requires an executable path\n"); + return 2; + } + routing_gateway = argv[++i]; + } else if (std::strcmp(argv[i], "--routing-engine-port") == 0) { + const char * value = i + 1 < argc ? argv[++i] : ""; + const char * end = value + std::strlen(value); + const auto parsed = std::from_chars(value, end, routing_engine_port); + if (parsed.ec != std::errc{} || parsed.ptr != end || + routing_engine_port < 1 || routing_engine_port > 65535) { + std::fprintf(stderr, "[server] --routing-engine-port requires a port from 1 to 65535\n"); + return 2; + } + routing_engine_port_set = true; } else if (std::strcmp(argv[i], "--swap-residency") == 0) { swap_residency = true; } else if (std::strcmp(argv[i], "--swap-keep-weights") == 0) { @@ -2292,6 +2365,21 @@ int main(int argc, char ** argv) { } const bool multi_model = model_args.size() > 1; const bool name_routing = model_routing == ServerConfig::ModelRouting::name; + if (routing_engine_port_set && routing_gateway.empty()) { + std::fprintf(stderr, "[server] --routing-engine-port requires --routing-gateway\n"); + return 2; + } + if (!routing_gateway.empty()) { +#ifdef _WIN32 + std::fprintf(stderr, "[server] --routing-gateway requires a POSIX host\n"); + return 2; +#else + if (!name_routing || access(routing_gateway.c_str(), X_OK) != 0) { + std::fprintf(stderr, "[server] --routing-gateway requires name routing and an executable path\n"); + return 2; + } +#endif + } if (unknown_model_set && !name_routing) { std::fprintf(stderr, "[server] --unknown-model requires --model-routing name\n"); return 2; @@ -2354,7 +2442,17 @@ int main(int argc, char ** argv) { } // Listener policy belongs to the first CLI block, independently of priority. - const ServerConfig listener_config = options.front().sconfig; + ServerConfig listener_config = options.front().sconfig; + const std::string gateway_host = listener_config.host; + const int gateway_port = listener_config.port; + if (!routing_gateway.empty()) { + if (gateway_port == routing_engine_port) { + std::fprintf(stderr, "[server] gateway and engine ports must differ\n"); + return 2; + } + listener_config.host = "127.0.0.1"; + listener_config.port = routing_engine_port; + } if (!primary_gpu.empty()) { size_t selected = options.size(); for (size_t m = 0; m < options.size(); ++m) { @@ -2543,7 +2641,15 @@ int main(int argc, char ** argv) { g_server = servers.front(); std::signal(SIGTERM, signal_handler); std::signal(SIGINT, signal_handler); + luce::server::RoutingGateway gateway; + if (!routing_gateway.empty() && + !gateway.start(routing_gateway, gateway_host, gateway_port, routing_engine_port, + [server = g_server] { server->request_stop(); })) { + g_server = nullptr; + return 1; + } const int ret = load_balancing ? g_server->run(servers) : g_server->run(); + gateway.stop(); g_server = nullptr; - return ret; + return gateway.failed() ? 1 : ret; } diff --git a/server/test/test_model_names.cmake b/server/test/test_model_names.cmake index c43ea27527..308a9c4165 100644 --- a/server/test/test_model_names.cmake +++ b/server/test/test_model_names.cmake @@ -67,3 +67,40 @@ check_launch(name_routing_unique 1 "failed to detect architecture from ${primary check_launch(unknown_model_needs_name_routing 2 "--unknown-model requires --model-routing name" "${primary}" --model-name qwen --model "${secondary}" --model-name ds4 --load-balancing --unknown-model primary) + +check_launch(gateway_missing_path 2 "--routing-gateway requires an executable path" + "${primary}" --routing-gateway) +check_launch(gateway_bad_port 2 "--routing-engine-port requires a port from 1 to 65535" + "${primary}" --routing-engine-port 70000) +check_launch(gateway_port_without_gateway 2 "--routing-engine-port requires --routing-gateway" + "${primary}" --routing-engine-port 8492) +if(UNIX) + check_launch(gateway_missing_executable 2 "requires name routing and an executable path" + "${primary}" --model-routing name --routing-gateway "${primary}") + check_launch(gateway_port_collision 2 "gateway and engine ports must differ" + "${primary}" --model-name default --port 8491 + --model "${secondary}" --model-name pro --model-routing name + --routing-gateway /bin/true) +endif() + +check_launch(profile_unknown 2 "--print-profile requires a profile name" + --print-profile missing-profile) +execute_process(COMMAND "${SERVER}" --print-profile ds41-lucebox + RESULT_VARIABLE profile_status OUTPUT_VARIABLE profile_json ERROR_VARIABLE profile_error) +string(JSON profile_name GET "${profile_json}" name) +string(JSON profile_budget GET "${profile_json}" env LUCE_EXPERT_BUDGET_MB) +if(NOT profile_status EQUAL 0 OR NOT profile_name STREQUAL "ds41-lucebox" OR NOT profile_budget STREQUAL "10500") + message(FATAL_ERROR "Profile export failed: ${profile_error} ${profile_json}") +endif() +if(UNIX) + set(mix_launcher "${CMAKE_CURRENT_BINARY_DIR}/test-luce-mix-launcher") + file(WRITE "${mix_launcher}" "#!/bin/sh\nprintf '%s\\n' \"$@\"\n") + file(CHMOD "${mix_launcher}" PERMISSIONS OWNER_READ OWNER_WRITE OWNER_EXECUTE) + execute_process(COMMAND "${CMAKE_COMMAND}" -E env "LUCE_MIX_LAUNCHER=${mix_launcher}" + "${SERVER}" --luce-mix --mix-pair "small+big" --mix-big-topology strix-only --mix-big-model "path with spaces.gguf" + RESULT_VARIABLE mix_status OUTPUT_VARIABLE mix_args ERROR_VARIABLE mix_error TIMEOUT 10) + if(NOT mix_status EQUAL 0 OR NOT mix_args MATCHES "--engine\n[^\n]+\n--mix-pair\nsmall[+]big\n--mix-big-topology\nstrix-only\n--mix-big-model\npath with spaces.gguf\n") + message(FATAL_ERROR "Luce Mix handoff failed (${mix_status}): ${mix_error} ${mix_args}") + endif() + file(REMOVE "${mix_launcher}") +endif() diff --git a/server/test/unit/test_model_swap_gpu.cpp b/server/test/unit/test_model_swap_gpu.cpp index f71de9fc56..efb079b4a3 100644 --- a/server/test/unit/test_model_swap_gpu.cpp +++ b/server/test/unit/test_model_swap_gpu.cpp @@ -393,6 +393,7 @@ TEST_CASE(ModelSwapGpuFixture, DeepSeek4PartialEvictionKeepsSecondaryTierAndSnap } constexpr auto partial = ModelBackend::EvictLevel::partial; const size_t baseline = device_free(0); + setenv("LUCE_DS4_HYBRID_F16_VERIFY_KV", "1", 1); auto backend = make_ds41_backend(model, draft); auto * ds4 = dynamic_cast(backend.get()); SWAP_CHECK(ds4); @@ -427,6 +428,9 @@ TEST_CASE(ModelSwapGpuFixture, DeepSeek4PartialEvictionKeepsSecondaryTierAndSnap // move the experts (reinstate verifies the placement and biases). setenv("LUCE_DS4_MOE_TP_GPU", "0", 1); setenv("LUCE_EXPERT_BUDGET_MB", "1", 1); + // Reinstatement must retain the captured verification mode too, for + // both partial and cold restores below (checked by verify_frozen_state). + setenv("LUCE_DS4_HYBRID_F16_VERIFY_KV", "0", 1); reinstate(*backend, "DS4 (partial)"); const auto back = ds4->expert_tier_report(); SWAP_CHECK(back.secondary_upload_bytes == loaded.secondary_upload_bytes); // no secondary re-read diff --git a/server/test/unit/test_routing_gateway.cpp b/server/test/unit/test_routing_gateway.cpp new file mode 100644 index 0000000000..444df4a32d --- /dev/null +++ b/server/test/unit/test_routing_gateway.cpp @@ -0,0 +1,73 @@ +#include "CppUnitTestFramework.hpp" +#include "server/routing_gateway.h" + +#ifndef _WIN32 +#include +#include +#include +#include + +namespace { +void require_gateway(bool condition) { + if (!condition) throw std::runtime_error("routing gateway lifecycle assertion failed"); +} +struct RoutingGatewayFixture { + std::string dir; + RoutingGatewayFixture() { + char pattern[] = "/tmp/luce-gateway-XXXXXX"; + const char * path = mkdtemp(pattern); + require_gateway(path != nullptr); + dir = path; + } + ~RoutingGatewayFixture() { std::filesystem::remove_all(dir); } + std::string script(const std::string & body) { + const std::string path = dir + "/gateway"; + std::ofstream(path) << "#!/bin/sh\n" << body; + require_gateway(chmod(path.c_str(), 0700) == 0); + return path; + } +}; +template void wait_for(Predicate ready) { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3); + while (!ready() && std::chrono::steady_clock::now() < deadline) + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + require_gateway(ready()); +} +} + +TEST_CASE(RoutingGatewayFixture, missing_executable_fails_without_child) { + luce::server::RoutingGateway gateway; + require_gateway(!gateway.start(dir + "/missing", "127.0.0.1", 8080, 8491, [] {})); +} + +TEST_CASE(RoutingGatewayFixture, unexpected_exit_stops_owner) { + std::atomic stopped{false}; + luce::server::RoutingGateway gateway; + require_gateway(gateway.start(script("exit 17\n"), "127.0.0.1", 8080, 8491, + [&] { stopped.store(true); })); + wait_for([&] { return stopped.load(); }); + gateway.stop(); + require_gateway(gateway.failed()); +} + +TEST_CASE(RoutingGatewayFixture, forwards_listener_and_reaps_on_shutdown) { + std::atomic stopped{false}; + const std::string path = script( + "trap 'exit 0' TERM\n" + "printf '%s\\n' \"$@\" > \"$0.args\"\n" + "echo $$ > \"$0.pid\"\n" + "while :; do sleep 1; done\n"); + luce::server::RoutingGateway gateway; + require_gateway(gateway.start(path, "0.0.0.0", 8087, 8497, [&] { stopped.store(true); })); + int pid = 0; + wait_for([&] { std::ifstream(path + ".pid") >> pid; return pid > 0; }); + std::ifstream args(path + ".args"); + const std::string text((std::istreambuf_iterator(args)), {}); + require_gateway(text == "--engine-url\nhttp://127.0.0.1:8497\n--host\n0.0.0.0\n--port\n8087\n"); + gateway.stop(); + require_gateway(!gateway.failed() && !stopped.load()); + require_gateway(kill(pid, 0) == -1 && errno == ESRCH); + int status; + require_gateway(waitpid(pid, &status, WNOHANG) == -1 && errno == ECHILD); +} +#endif