From 905f443327e47ee407de113d8dc2c34136de1e9b Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 11 Aug 2026 13:11:00 +0000 Subject: [PATCH 1/2] feat(ds4): add monolithic Strix Halo concurrent serving --- .github/workflows/ci.yml | 8 +- README.md | 5 +- server/CMakeLists.txt | 55 ++ server/docs/DS4.md | 37 + server/share/status.html | 4 +- server/src/common/backend_args.h | 11 + server/src/common/backend_factory.cpp | 6 + server/src/common/concurrency/seq_engine.h | 333 ++++++++ .../common/concurrency/seq_slot_manager.cpp | 271 +++++++ .../src/common/concurrency/seq_slot_manager.h | 131 ++++ server/src/common/feature_gate.cpp | 55 +- server/src/common/kvflash_pager.h | 9 + server/src/common/model_backend.h | 14 + server/src/common/model_capabilities.h | 2 +- server/src/common/paged_attention_config.h | 40 + server/src/common/paged_kv_pool.cpp | 68 +- server/src/common/paged_kv_pool.h | 41 +- server/src/deepseek4/deepseek4_backend.cpp | 169 +++- server/src/deepseek4/deepseek4_backend.h | 5 + .../src/deepseek4/deepseek4_fused_verify.inc | 145 +++- server/src/deepseek4/deepseek4_graph.cpp | 510 ++++++++++-- server/src/deepseek4/deepseek4_internal.h | 46 ++ server/src/deepseek4/deepseek4_page_layout.h | 55 ++ .../src/deepseek4/deepseek4_paged_cache.cpp | 219 ++++++ server/src/deepseek4/deepseek4_paged_cache.h | 63 ++ server/src/deepseek4/deepseek4_seq_engine.cpp | 272 +++++++ server/src/deepseek4/deepseek4_seq_engine.h | 41 + server/src/server/api_types.h | 12 + server/src/server/client_send_buffer.h | 139 ++++ server/src/server/http_server.cpp | 164 +++- server/src/server/http_server.h | 45 +- server/src/server/scheduler.cpp | 739 ++++++++++++++++++ server/src/server/server_main.cpp | 48 +- server/src/server/server_status.h | 22 +- server/src/server/tool_memory.cpp | 2 + server/src/server/tool_memory.h | 6 + server/test/host_check.h | 14 + server/test/seq_engine_contract.h | 349 +++++++++ server/test/test_client_send_buffer.cpp | 149 ++++ server/test/test_deepseek4_page_layout.cpp | 53 ++ server/test/test_deepseek4_paged_cache.cpp | 86 ++ server/test/test_feature_gate.cpp | 332 +++++--- server/test/test_kvflash_pool_sizing.cpp | 5 + server/test/test_paged_kv_pool.cpp | 163 +++- server/test/test_seq_batch_plan.cpp | 172 ++++ server/test/test_seq_engine_contract.cpp | 323 ++++++++ server/test/test_seq_slot_manager.cpp | 426 ++++++++++ server/test/test_server_unit.cpp | 36 +- server/tests/test_server_parallel.py | 589 ++++++++++++++ 49 files changed, 6246 insertions(+), 243 deletions(-) create mode 100644 server/src/common/concurrency/seq_engine.h create mode 100644 server/src/common/concurrency/seq_slot_manager.cpp create mode 100644 server/src/common/concurrency/seq_slot_manager.h create mode 100644 server/src/deepseek4/deepseek4_page_layout.h create mode 100644 server/src/deepseek4/deepseek4_paged_cache.cpp create mode 100644 server/src/deepseek4/deepseek4_paged_cache.h create mode 100644 server/src/deepseek4/deepseek4_seq_engine.cpp create mode 100644 server/src/deepseek4/deepseek4_seq_engine.h create mode 100644 server/src/server/client_send_buffer.h create mode 100644 server/src/server/scheduler.cpp create mode 100644 server/test/host_check.h create mode 100644 server/test/seq_engine_contract.h create mode 100644 server/test/test_client_send_buffer.cpp create mode 100644 server/test/test_deepseek4_page_layout.cpp create mode 100644 server/test/test_deepseek4_paged_cache.cpp create mode 100644 server/test/test_seq_batch_plan.cpp create mode 100644 server/test/test_seq_engine_contract.cpp create mode 100644 server/test/test_seq_slot_manager.cpp create mode 100644 server/tests/test_server_parallel.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 58d75ce3e..7144a409f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -79,12 +79,16 @@ jobs: -DCMAKE_BUILD_TYPE=Release cmake --build build --target \ test_dflash test_generate test_flash_attn_sparse test_server_unit \ - test_deepseek4_unit -j$(nproc) + test_deepseek4_unit test_feature_gate test_seq_slot_manager \ + test_seq_engine_contract test_seq_batch_plan test_client_send_buffer \ + test_deepseek4_page_layout test_deepseek4_paged_cache -j$(nproc) - name: Run C++ server unit tests run: | cd server/build - ctest --output-on-failure -R "server_unit|deepseek4_unit" --no-tests=error + ctest --output-on-failure \ + -R "server_unit|deepseek4_unit|feature_gate|seq_slot_manager|seq_engine_contract|seq_batch_plan|client_send_buffer|deepseek4_page_layout|deepseek4_paged_cache" \ + --no-tests=error - name: Populate venv with cu128 torch + setuptools # First pass: install the workspace's default deps. dflash declares diff --git a/README.md b/README.md index d96b49592..f48894332 100644 --- a/README.md +++ b/README.md @@ -350,7 +350,9 @@ When compression is on, the request path picks one of three modes automatically, | `DFLASH_PREFILL_CACHE_SLOTS=N` | `0` | Container-entrypoint equivalent of `--prefill-cache-slots`; the native binary itself uses the CLI flag. | | `--kv-cache-dir ` | — | Persist prefix cache to disk | | `--kv-cache-budget N` | — | On-disk cache size cap | -| `--paged-attention` | off | Exact 16-token block-table decode for single-device Qwen3.6-27B AR; see [paged attention](optimizations/paged_attention/README.md) | +| `--paged-attention` | off | Exact block-table attention for single-device Qwen3.6-27B and monolithic DeepSeek4 on Strix Halo; see [paged attention](optimizations/paged_attention/README.md) and [DeepSeek4](server/docs/DS4.md#strix-halo-concurrent-serving) | +| `--max-concurrency N` | `1` | DeepSeek4 Strix Halo sequence slots. Values 2–16 enable paged attention automatically; Qwen remains single-sequence. | +| `--kv-pool-tokens N` | `0` (auto) | Shared physical K/V capacity for concurrent paged serving. Requires `--max-concurrency` greater than 1. Zero derives capacity from available device memory; explicit values are rounded to whole 16-token blocks. | **Bounded KV residency (KVFlash)** @@ -388,6 +390,7 @@ Pages the attention KV cache through a fixed pool of GPU slots; cold 64-token ch | `--draft-ipc-bin ` | — | Out-of-process draft binary (mixed CUDA/HIP) | | `--peer-access` | off | Enable P2P between target GPUs | | `--chunk N` | backend default | Prefill ubatch size | +| `--admission-coalesce-ms N` | `20` | Idle-to-busy batching window for concurrent serving, from 0 to 1000 ms; `0` disables it. | | `--no-cors` | CORS on | Disable CORS headers | | `DFLASH_TARGET_GPU=N` | `0` | Env var equivalent of `--target-gpu` | | `DFLASH_DRAFT_GPU=N` | same as target | Env var equivalent of `--draft-gpu` | diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 57b663156..23c9be404 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -147,6 +147,7 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip") endif() set(GGML_HIP ON CACHE BOOL "" FORCE) set(GGML_HIP_RCCL OFF CACHE BOOL "" FORCE) + set(GGML_HIP_GRAPHS OFF CACHE BOOL "Enable experimental HIP graph replay") set(DFLASH27B_GGML_BACKEND_TARGET ggml-hip) set(DFLASH27B_HIP_ARCHITECTURES "" CACHE STRING "HIP GPU targets, e.g. gfx906;gfx1100") set(GGML_BACKEND_DL OFF CACHE BOOL "" FORCE) @@ -391,6 +392,8 @@ add_library(dflash_common STATIC src/deepseek4/deepseek4_loader.cpp src/deepseek4/deepseek4_graph.cpp src/deepseek4/deepseek4_roctx.cpp + src/deepseek4/deepseek4_paged_cache.cpp + src/deepseek4/deepseek4_seq_engine.cpp src/deepseek4/deepseek4_backend.cpp src/deepseek4/deepseek4_daemon.cpp src/deepseek4/deepseek4_layer_split_adapter.cpp @@ -423,6 +426,7 @@ add_library(dflash_common STATIC src/common/dflash_draft_kv.cpp src/common/dflash_spec_decode.cpp src/common/paged_kv_pool.cpp + src/common/concurrency/seq_slot_manager.cpp src/common/layer_split_backend.cpp src/common/layer_split_runtime.cpp src/qwen35/graph_builders.cpp @@ -1365,6 +1369,55 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/src) list(APPEND _raw_unit_test_targets test_paged_kv_pool) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_seq_slot_manager.cpp") + # Host-side slot bookkeeping test (concurrent serving): no GPU. + add_executable(test_seq_slot_manager + test/test_seq_slot_manager.cpp + src/common/concurrency/seq_slot_manager.cpp + src/common/paged_kv_pool.cpp) + target_include_directories(test_seq_slot_manager PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_seq_slot_manager) + endif() + + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_deepseek4_page_layout.cpp") + add_executable(test_deepseek4_page_layout + test/test_deepseek4_page_layout.cpp) + target_include_directories(test_deepseek4_page_layout PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src ${CMAKE_CURRENT_SOURCE_DIR}/include) + list(APPEND _raw_unit_test_targets test_deepseek4_page_layout) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_deepseek4_paged_cache.cpp") + add_executable(test_deepseek4_paged_cache + test/test_deepseek4_paged_cache.cpp + src/deepseek4/deepseek4_paged_cache.cpp) + target_compile_definitions(test_deepseek4_paged_cache PRIVATE DFLASH_DS4_PLAN_ONLY=1) + target_include_directories(test_deepseek4_paged_cache PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_deepseek4_paged_cache) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_seq_engine_contract.cpp") + # SeqEngine conformance checker + the fakes that prove it bites: no GPU. + add_executable(test_seq_engine_contract test/test_seq_engine_contract.cpp) + target_include_directories(test_seq_engine_contract PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_seq_engine_contract) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_seq_batch_plan.cpp") + # Pure-host tests for model-neutral token-budget/FIFO planning. + add_executable(test_seq_batch_plan test/test_seq_batch_plan.cpp) + target_include_directories(test_seq_batch_plan PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_seq_batch_plan) + endif() + if(UNIX AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_client_send_buffer.cpp") + # Buffered non-blocking client writer test (socketpair): no GPU. + add_executable(test_client_send_buffer test/test_client_send_buffer.cpp) + target_include_directories(test_client_send_buffer PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_client_send_buffer) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_dflash.cpp") add_executable(test_dflash test/test_dflash.cpp) target_include_directories(test_dflash PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) @@ -1494,6 +1547,7 @@ if(DFLASH27B_TESTS) add_executable(test_server_unit ${_server_unit_sources}) target_sources(test_server_unit PRIVATE src/server/http_server.cpp + src/server/scheduler.cpp src/server/model_card.cpp src/server/prompt_normalize.cpp src/qwen3/anchor_scan.cpp) @@ -1743,6 +1797,7 @@ if(DFLASH27B_SERVER) add_executable(dflash_server src/server/server_main.cpp src/server/http_server.cpp + src/server/scheduler.cpp src/server/model_card.cpp src/server/prompt_normalize.cpp ) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index be5b1ab28..5b6a5394f 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -207,6 +207,43 @@ performance profile held 48.1 tok/s median on the deterministic 128-token workload. The all-6-expert reference-exact mode is a correctness profile, not a throughput profile. +### Strix Halo concurrent serving + +DeepSeek4 paged concurrency is deliberately a single-device path: one local +HIP target on Strix Halo (`gfx1151`), with the complete model and every expert +resident on that device. It does not use layer splitting, CUDA/HIP expert +ownership, host-streamed experts, or DSpark. + +The backend keeps raw MLA rows, compressed rows, indexer state, sequence +lengths, and block tables in a persistent 128-token paged cache. The shared +HTTP scheduler performs admission, cancellation, slow-client isolation, and +fair continuous batching. DeepSeek4 lowers each scheduler plan into one exact +gathered graph with up to 16 independent lanes. Decode rows share the weight +pass; each selected prompt advances by one exact token because the graph must +not contain two rows from the same sequence. + +```bash +cmake -S . -B build-hip \ + -DDFLASH27B_GPU_BACKEND=hip \ + -DDFLASH27B_HIP_ARCHITECTURES=gfx1151 \ + -DDFLASH27B_SERVER=ON +cmake --build build-hip -j + +./build-hip/dflash_server /path/to/deepseek4-target.gguf \ + --target-device hip:0 \ + --paged-attention \ + --max-concurrency 16 \ + --kv-pool-tokens 8192 \ + --max-ctx 4096 \ + --ds4-prefill exact \ + --prefix-cache-slots 0 +``` + +This mode fails closed for non-gfx1151 devices, CUDA, layer or remote target +splits, `DFLASH_DS4_MOE_TP`, drafts/DSpark, DDTree, PFlash/KVFlash, fused +decode, approximate prefill, windowed attention, and prefix-cache parking. +There is no automatic fallback to a slower or asymmetric execution mode. + ### Local single-shard If the adapter decides all 43 layers fit on one CUDA GPU, it loads a single shard locally and no IPC daemon is involved. diff --git a/server/share/status.html b/server/share/status.html index dbfe6f3aa..6ba70fc3e 100644 --- a/server/share/status.html +++ b/server/share/status.html @@ -197,7 +197,9 @@

Decode Performance

function update(data) { const badge = document.getElementById('phase-badge'); - badge.textContent = data.phase; + badge.textContent = data.active_requests + ? data.phase + ' (' + data.active_requests + ' active)' + : data.phase; badge.className = 'badge badge-' + data.phase; document.getElementById('total-req').textContent = data.total_requests; diff --git a/server/src/common/backend_args.h b/server/src/common/backend_args.h index 5b765ce58..ee7b55d33 100644 --- a/server/src/common/backend_args.h +++ b/server/src/common/backend_args.h @@ -25,6 +25,11 @@ struct BackendFeatureConfig { // time rather than through BackendArgs. bool routing_stats_requested = false; // --freq / --collect-routing bool adaptive_experts_requested = false; // --adaptive-experts + + // A fixed KVFlash pool requested through DFLASH_KVFLASH. "auto" is + // resolved later by the backend because only it has the VRAM budget needed + // to know whether a pool will actually be active. + bool kvflash_enabled = false; }; // A superset of all per-architecture config fields. The factory reads only @@ -59,6 +64,12 @@ struct BackendArgs { // only the fields they support. int fa_window = 0; // 0 = full attention. qwen3.6 full-attn layers must see the whole context; a finite window drops the system prompt/tools -> breaks tool calls. bool paged_attention = false; // 16-token paged K/V blocks for AR decode + // Concurrent decode slots (--max-concurrency). > 1 requires paged_attention; + // the backend serves that many sequences through the seq_* slot API. + int max_concurrency = 1; + // Total paged K/V pool in tokens shared by all slots (--kv-pool-tokens; + // block-rounded). 0 = derive capacity from available device memory. + long long kv_pool_tokens = 0; int kq_stride_pad = 32; int draft_swa_window = 0; int draft_ctx_max = 4096; diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 1514d476f..daf7dddb2 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -104,6 +104,9 @@ DFLASH_CHECK_ARCH("deepseek4", DeepSeek4BackendConfig, DeepSeek4LayerSplitAdapte // never reads a block table — which is why its capability row is Never.) DFLASH_CHECK_ARCH_OPTION("qwen35", Qwen35Config, Qwen35LayerSplitAdapterConfig, has_paged_attention, paged_attn); +DFLASH_CHECK_ARCH_OPTION("deepseek4", DeepSeek4BackendConfig, + DeepSeek4LayerSplitAdapterConfig, + has_paged_attention, paged_attn); #undef DFLASH_CHECK_ARCH #undef DFLASH_CHECK_ARCH_OPTION @@ -423,6 +426,9 @@ std::unique_ptr create_backend( cfg.expert_top_k = args.ds4_expert_top_k; cfg.fused_decode = args.ds4_fused_decode; cfg.prefill_mode = args.ds4_prefill_mode; + cfg.paged_attention = args.paged_attention; + cfg.max_concurrency = args.max_concurrency; + cfg.kv_pool_tokens = args.kv_pool_tokens; auto backend = std::make_unique(cfg); if (!backend->init()) { diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h new file mode 100644 index 000000000..4beeaf116 --- /dev/null +++ b/server/src/common/concurrency/seq_engine.h @@ -0,0 +1,333 @@ +// SeqEngine — concurrent serving over decode slots (iteration-level +// scheduling), the boundary between the HTTP scheduler and a backend that can +// hold several live sequences at once. +// +// A backend qualifies when it can keep several independent sequences in a +// paged KV cache and execute a batched decode step. Any additional +// per-sequence model state is owned by the concrete engine, not by this +// interface. admit() claims a slot and queues its prompt without compute. +// Each step() then advances a scheduler-selected cohort of prompt slices +// alongside the complete live decode batch. Once a prefill completes, the +// scheduler advances that slot one token per step(), feeding each sampled +// token back as the next step's input — which is what lets it override a token +// (thinking-budget force-close) before it is committed to the cache. +// +// The split of duties is deliberate and is the reason this interface exists +// apart from ModelBackend: +// scheduler policy — who gets admitted, when a slot stops, what reaches +// the client, how a slow reader is handled +// engine mechanism — KV blocks, backend-owned per-sequence state, the +// batched forward, sampling +// Nothing above this interface knows about block tables, and nothing below it +// knows about sockets. +// +// Threading contract: every call comes from ONE thread — the same one that +// calls ModelBackend::generate() — so implementations need no locking. +// +// ── Adding a backend ──────────────────────────────────────────────────── +// Implement this interface next to the model; DeepSeek4's sequence engine is +// the worked example. Return it +// from ModelBackend::seq_engine(). Nothing else in the server changes: +// scheduler_loop() takes over the worker thread +// as soon as seq_engine() returns non-null. +// +// Reuse only genuinely model-neutral pieces such as PagedKvPool. Prompt and +// slot lifecycle state belongs beside the backend that interprets it; a new +// engine supplies that host record together with its device-side prefill, +// batched forward, and metadata uploads. +// +// What must never reach this interface, or anything above it: block tables, +// graph shapes, and per-sequence model state (recurrent/SSM/conv tensors, +// cache layout). A model that seems to need SeqEngine widened to serve +// concurrently is the signal that the split has slipped — the model-specific +// part belongs inside the engine, not in the contract. +// +// test/seq_engine_contract.h drives an engine through exactly the call +// sequence the scheduler makes and returns the violations it finds; run a new +// engine through it before wiring it up. + +#pragma once + +#include +#include +#include +#include +#include + +#include "common/sampler.h" + +namespace dflash::common { + +// Model-neutral prefill planning for continuous batching. The scheduler owns +// arrival order and fairness; the engine advertises the useful work envelope +// and lowers the selected slices into its model-specific graph. +struct PrefillCandidate { + int slot = -1; + uint64_t order = 0; +}; + +struct PrefillSlice { + int slot = -1; + int max_tokens = 0; +}; + +struct StepPlanLimits { + int max_prefill_sequences = 1; + int max_prefill_tokens_per_sequence = 512; + int max_prefill_tokens_total = 512; + int prefill_allocation_quantum = 512; +}; + +// Select the oldest eligible sequences, then distribute the step's token +// capacity in engine-owned quanta. Strict FIFO determines cohort membership; +// a rotating cursor prevents the oldest member from always winning a partial +// final round. +inline std::vector plan_prefill_slices( + const std::vector & candidates, + const StepPlanLimits & limits, + size_t round_robin_start = 0) { + std::vector slices; + if (candidates.empty() || limits.max_prefill_sequences <= 0 || + limits.max_prefill_tokens_per_sequence <= 0 || + limits.max_prefill_tokens_total <= 0 || + limits.prefill_allocation_quantum <= 0) { + return slices; + } + + std::vector ordered = candidates; + std::stable_sort(ordered.begin(), ordered.end(), + [](const PrefillCandidate & a, const PrefillCandidate & b) { + if (a.order != b.order) return a.order < b.order; + return a.slot < b.slot; + }); + + const int selected = std::min( + (int)ordered.size(), limits.max_prefill_sequences); + int budget = std::min( + limits.max_prefill_tokens_total, + selected * limits.max_prefill_tokens_per_sequence); + slices.reserve((size_t)selected); + for (int i = 0; i < selected; ++i) { + slices.push_back({ordered[(size_t)i].slot, 0}); + } + + while (budget > 0) { + bool granted = false; + for (size_t offset = 0; offset < slices.size(); ++offset) { + const size_t idx = + (round_robin_start + offset) % slices.size(); + PrefillSlice & slice = slices[idx]; + const int room = + limits.max_prefill_tokens_per_sequence - slice.max_tokens; + if (room <= 0) continue; + const int grant = std::min({ + limits.prefill_allocation_quantum, room, budget}); + if (grant <= 0) continue; + slice.max_tokens += grant; + budget -= grant; + granted = true; + if (budget == 0) break; + } + if (!granted) break; + } + + slices.erase( + std::remove_if(slices.begin(), slices.end(), + [](const PrefillSlice & slice) { return slice.max_tokens <= 0; }), + slices.end()); + return slices; +} + +class SeqEngine { +public: + virtual ~SeqEngine() = default; + + // Number of decode slots served concurrently. Fixed for the engine's + // lifetime: the scheduler sizes its own per-slot array from this once and + // then indexes it by the slot id admit() returns. + virtual int slot_count() const = 0; + // Per-sequence logical context bound. The scheduler owns generation + // policy and clamps its output cap against this value after admission. + virtual int max_context() const = 0; + + struct AdmitResult { + enum class Status { + admitted, + busy, + failed, + }; + + Status status = Status::failed; + int slot = -1; + std::string error; + }; + + // Admit one request into a free slot and queue its prompt for chunked + // prefill. No model compute is performed here. Implementations may reserve + // persistent capacity atomically so that every admitted prompt can finish; + // subsequent step() calls advance scheduler-selected prompt slices. + // + // `sampler` is the only source of truth for how the slot samples: + // sampler.needs_logit_processing() selects CPU sampling over GPU argmax + // AND decides whether sampler.seed is honoured. There is deliberately no + // separate do_sample flag — a second copy of that one fact is something + // an engine can disagree with, and the failure mode is silent (a seeded + // request sampling nondeterministically, with no error anywhere). + virtual AdmitResult admit(uint64_t request_id, + const std::vector & prompt, + const SamplerCfg & sampler) = 0; + + struct StepInput { + int slot = -1; + int32_t token = -1; // token to commit at this slot's next position + }; + struct DecodeOutput { + int slot = -1; + int32_t token = -1; // newly sampled token (pending until next step) + bool failed = false; + // Present when failed=true so the scheduler can report an honest + // per-request error instead of silently truncating generation. + std::string error; + }; + + struct PrefillOutput { + enum class Status { + advanced, + completed, + failed, + }; + + int slot = -1; + Status status = Status::advanced; + // 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; + // Present only for failed. + std::string error; + }; + + // One scheduler iteration owns both kinds of logical work. `decode` must + // contain every currently decoding slot exactly once; `prefills` is the + // bounded subset of pending prompt work selected by scheduler policy. + // Device graph shapes, staging indices, cache blocks, and model state stay + // behind the engine boundary. + struct StepPlan { + std::vector decode; + std::vector prefills; + }; + + struct StepResult { + std::vector decode; + std::vector prefills; + // A non-empty error is fatal for the whole live cohort and carries + // no usable row output. Validation + // failures occur before mutation; a device/build/compute failure may + // leave backend state partially advanced, so the caller must retire + // every live sequence before invoking step() again. + std::string error; + + bool ok() const { return error.empty(); } + }; + + // Useful per-step work envelope at the requested live decode width. The + // scheduler fills this capacity and distributes it fairly; an engine may + // advertise different sequence, per-sequence, and total-token limits for + // idle, mixed, or larger decode buckets. + virtual StepPlanLimits step_plan_limits(int decode_rows) const = 0; + + // A successful result returns one decode output for every decode input and + // one explicit advanced/completed/failed result for every selected + // prefill. Invalid plans return a fatal error without advancing state. + // Runtime failures are terminal for the live cohort and may follow partial + // backend mutation, but expose no consumable payload. + virtual StepResult step(const StepPlan & plan) = 0; + + // Release a slot's KV blocks and mark it free. Safe on failed slots. + virtual void retire(int slot) = 0; + + // EOS check for scheduler-side stop decisions. + virtual bool token_is_eos(int32_t token) const = 0; +}; + +// Validate the model-neutral step protocol before the scheduler consumes any +// output. Malformed row ownership is fatal because re-feeding a token after an +// omitted output would silently corrupt that sequence. +inline std::string validate_step_result( + const SeqEngine::StepPlan & plan, + const SeqEngine::StepResult & result, + int slot_count) { + if (slot_count < 1) return "slot count must be positive"; + const bool payload_empty = result.decode.empty() && result.prefills.empty(); + if (!result.error.empty()) { + return payload_empty ? std::string{} + : "fatal result exposes partial payload"; + } + + std::vector decode_planned((size_t)slot_count, 0); + std::vector prefill_planned((size_t)slot_count, 0); + for (const SeqEngine::StepInput & input : plan.decode) { + if (input.slot < 0 || input.slot >= slot_count || input.token < 0) + return "decode plan contains an invalid row"; + if (decode_planned[(size_t)input.slot]) + return "decode plan contains a duplicate slot"; + decode_planned[(size_t)input.slot] = 1; + } + for (const PrefillSlice & slice : plan.prefills) { + if (slice.slot < 0 || slice.slot >= slot_count || + slice.max_tokens <= 0) + return "prefill plan contains an invalid slice"; + if (decode_planned[(size_t)slice.slot] || + prefill_planned[(size_t)slice.slot]) + return "step plan assigns a slot more than once"; + prefill_planned[(size_t)slice.slot] = 1; + } + + std::vector decode_seen((size_t)slot_count, 0); + for (const SeqEngine::DecodeOutput & output : result.decode) { + if (output.slot < 0 || output.slot >= slot_count || + !decode_planned[(size_t)output.slot]) + return "decode output names an unplanned slot"; + if (decode_seen[(size_t)output.slot]) + return "step returned duplicate decode outputs"; + if (output.failed ? output.error.empty() : output.token < 0) + return output.failed ? "failed decode has no diagnostic" + : "successful decode has no token"; + decode_seen[(size_t)output.slot] = 1; + } + + using PrefillStatus = SeqEngine::PrefillOutput::Status; + std::vector prefill_seen((size_t)slot_count, 0); + for (const SeqEngine::PrefillOutput & output : result.prefills) { + if (output.slot < 0 || output.slot >= slot_count || + !prefill_planned[(size_t)output.slot]) + return "prefill output names an unselected slot"; + if (prefill_seen[(size_t)output.slot]) + return "step returned duplicate prefill outputs"; + if (output.status != PrefillStatus::advanced && + output.status != PrefillStatus::completed && + output.status != PrefillStatus::failed) + return "prefill output has an unknown status"; + if (output.status == PrefillStatus::advanced && + (output.token >= 0 || !output.error.empty())) + return "advanced prefill carries completion payload"; + if (output.status == PrefillStatus::completed && + (output.token < 0 || !output.error.empty())) + return "completed prefill has invalid payload"; + if (output.status == PrefillStatus::failed && + (output.token >= 0 || output.error.empty())) + return "failed prefill has invalid payload"; + prefill_seen[(size_t)output.slot] = 1; + } + + for (const SeqEngine::StepInput & input : plan.decode) + if (!decode_seen[(size_t)input.slot]) + return "step omitted an output for a decode slot"; + for (const PrefillSlice & slice : plan.prefills) + if (!prefill_seen[(size_t)slice.slot]) + return "step omitted an output for a selected prefill"; + if (payload_empty && (!plan.decode.empty() || !plan.prefills.empty())) + return "valid planned work returned no output"; + return {}; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/seq_slot_manager.cpp b/server/src/common/concurrency/seq_slot_manager.cpp new file mode 100644 index 000000000..1ff91ed8d --- /dev/null +++ b/server/src/common/concurrency/seq_slot_manager.cpp @@ -0,0 +1,271 @@ +#include "common/concurrency/seq_slot_manager.h" + +#include +#include + +namespace dflash::common { + +SeqSlotManager::SeqSlotManager(PagedKvPool & pool, int max_ctx) + : pool_(pool), max_ctx_(max_ctx) { + slots_.assign(pool.max_sequences(), SeqSlot{}); +} + +int SeqSlotManager::decoding_count() const { + int n = 0; + for (const SeqSlot & s : slots_) { + n += s.decoding() ? 1 : 0; + } + return n; +} + +uint32_t SeqSlotManager::decode_headroom_capacity(int logical_tokens) const { + const uint64_t extended = + static_cast(std::max(0, logical_tokens)) + + pool_.block_size(); + return static_cast(std::min( + static_cast(max_ctx_), extended)); +} + +bool SeqSlotManager::capacity_fits_pool(uint32_t token_capacity) const { + const uint64_t blocks = token_capacity == 0 ? 0 : + 1 + (static_cast(token_capacity) - 1) / + pool_.block_size(); + return blocks <= pool_.physical_block_count(); +} + +PagedKvStatus SeqSlotManager::protect_decode_headroom() { + struct TopUp { + PagedKvSequenceHandle handle; + uint32_t token_capacity = 0; + }; + + std::vector topups; + topups.reserve(slots_.size()); + uint64_t total_additional = 0; + const uint64_t block_size = pool_.block_size(); + for (const SeqSlot & slot : slots_) { + if (!slot.decoding()) continue; + const uint32_t capacity = decode_headroom_capacity(slot.cur_pos); + if (!capacity_fits_pool(capacity)) continue; + + uint32_t owned_blocks = 0; + const PagedKvStatus status = + pool_.owned_block_count(slot.handle, owned_blocks); + if (status != PagedKvStatus::Ok) return status; + const uint64_t target_blocks = capacity == 0 ? 0 : + 1 + (static_cast(capacity) - 1) / block_size; + if (target_blocks <= owned_blocks) continue; + const uint32_t additional = + static_cast(target_blocks - owned_blocks); + total_additional += additional; + topups.push_back({slot.handle, capacity}); + } + + // Preflight the whole cohort before moving a block, so a failed admission + // attempt cannot protect only whichever decoder happened to be visited + // first. + if (total_additional > pool_.free_block_count()) { + return PagedKvStatus::BlocksExhausted; + } + for (const TopUp & topup : topups) { + const PagedKvStatus status = + pool_.reserve_capacity(topup.handle, topup.token_capacity); + if (status != PagedKvStatus::Ok) return status; + } + return PagedKvStatus::Ok; +} + +bool SeqSlotManager::is_active(int slot) const { + return slot >= 0 && slot < (int)slots_.size() && + slots_[(size_t)slot].active(); +} + +bool SeqSlotManager::is_prefilling(int slot) const { + return is_active(slot) && slots_[(size_t)slot].prefilling(); +} + +SeqEngine::AdmitResult SeqSlotManager::admit( + uint64_t request_id, const std::vector & prompt, + const SamplerCfg & sampler) { + using AdmitStatus = SeqEngine::AdmitResult::Status; + SeqEngine::AdmitResult r; + if (prompt.empty()) { + r.error = "empty prompt"; + return r; + } + if (prompt.size() > static_cast(max_ctx_)) { + r.error = "prompt exceeds max_ctx"; + return r; + } + const int prompt_len = static_cast(prompt.size()); + + // A prompt larger than the whole pool can NEVER be admitted; waiting + // for other sequences to drain would stall the queue forever and then + // fail anyway. Hard-fail it up front instead of reporting busy. + const uint64_t pool_capacity = + (uint64_t)pool_.physical_block_count() * pool_.block_size(); + if ((uint64_t)prompt_len > pool_capacity) { + r.error = "prompt needs " + std::to_string(prompt_len) + + " KV tokens but the pool holds " + + std::to_string(pool_capacity) + + "; raise --kv-pool-tokens or shorten the prompt"; + return r; + } + + int slot = -1; + for (int i = 0; i < (int)slots_.size(); i++) { + if (!slots_[(size_t)i].active()) { slot = i; break; } + } + if (slot < 0) { + r.status = AdmitStatus::busy; + r.error = "all decode slots are busy"; + return r; + } + + // A newly freed block belongs to any older decoder missing its rolling + // next-page reserve before it can belong to this admission. + const PagedKvStatus headroom_status = protect_decode_headroom(); + if (headroom_status != PagedKvStatus::Ok) { + r.status = headroom_status == PagedKvStatus::BlocksExhausted + ? AdmitStatus::busy : AdmitStatus::failed; + r.error = r.status == AdmitStatus::busy + ? "existing decoders need the available KV headroom" + : paged_kv_status_string(headroom_status); + return r; + } + + PagedKvSequenceHandle handle; + uint32_t reservation_capacity = + decode_headroom_capacity(prompt_len); + if (!capacity_fits_pool(reservation_capacity)) { + // The prompt itself fits, but this physical pool can never hold its + // following page. Preserve useful prompt-only behavior and report + // decode exhaustion later if the sequence reaches that boundary. + reservation_capacity = static_cast(prompt_len); + } + const PagedKvStatus status = pool_.acquire_reserved( + request_id, reservation_capacity, handle); + if (status != PagedKvStatus::Ok) { + r.status = status == PagedKvStatus::SequenceSlotsExhausted || + status == PagedKvStatus::BlocksExhausted + ? AdmitStatus::busy : AdmitStatus::failed; + r.error = status == PagedKvStatus::BlocksExhausted + ? "not enough unreserved KV blocks for the prompt and decode headroom" + : paged_kv_status_string(status); + return r; + } + + SeqSlot & s = slots_[(size_t)slot]; + s.phase = SeqSlotPhase::prefill; + s.handle = handle; + s.cur_pos = 0; + s.prompt = prompt; + s.sampler = sampler; + s.sample_history = prompt; + // Same predicate the engine uses to pick CPU sampling over GPU argmax: + // a seed only means anything when the sampler actually draws. + if (sampler.needs_logit_processing() && sampler.seed != 0) { + s.rng.seed(sampler.seed); + } else { + s.rng.seed(std::random_device{}()); + } + + r.status = AdmitStatus::admitted; + r.slot = slot; + return r; +} + +SeqSlotManager::PrefillChunk SeqSlotManager::append_prefill( + int slot, int n_tokens) { + PrefillChunk out; + if (!is_prefilling(slot) || n_tokens < 1) return out; + + SeqSlot & s = slots_[(size_t)slot]; + if (s.cur_pos > (int)s.prompt.size() || + n_tokens > (int)s.prompt.size() - s.cur_pos) { + return out; + } + + PagedKvAppendResult app = pool_.append(s.handle, (uint32_t)n_tokens); + if (!app) { + // Admission reserved the whole prompt. Treat exhaustion here as a + // broken invariant, not a retryable condition: retrying a batch of + // all-prefill slots without any decoder able to retire would livelock. + if (app.status == PagedKvStatus::BlocksExhausted) { + std::fprintf(stderr, + "[parallel] reserved prefill capacity missing for slot %d\n", + slot); + } + out.busy = false; + return out; + } + + out.rows.reserve(app.write_slots.size()); + for (const PagedKvWriteSlot & write : app.write_slots) { + out.rows.push_back((int64_t)write.physical_token_index); + if (write.block_offset == 0) { + if (out.first_new_block < 0) { + out.first_new_block = + (int)(write.logical_position / pool_.block_size()); + } + out.new_blocks.push_back((int32_t)write.physical_block); + } + } + s.cur_pos += n_tokens; + out.ok = true; + return out; +} + +void SeqSlotManager::commit_prefill(int slot) { + if (!is_prefilling(slot)) return; + SeqSlot & s = slots_[(size_t)slot]; + if (s.cur_pos != (int)s.prompt.size()) return; + s.phase = SeqSlotPhase::decode; +} + +SeqSlotManager::StepAppend SeqSlotManager::append_token(int slot, + int32_t fed_token) { + StepAppend out; + if (!is_active(slot) || !slots_[(size_t)slot].decoding()) return out; + SeqSlot & s = slots_[(size_t)slot]; + if (s.cur_pos >= max_ctx_) { + // No context left; the scheduler should have stopped this slot. + return out; + } + PagedKvAppendResult app = pool_.append( + s.handle, 1, /*only_first_last_slots=*/true); + if (!app || app.token_count != 1 || + app.last.logical_position != (uint32_t)s.cur_pos) { + out.busy = app.status == PagedKvStatus::BlocksExhausted; + return out; + } + s.sample_history.push_back(fed_token); + + out.ok = true; + out.physical_row = (int64_t)app.last.physical_token_index; + out.position = s.cur_pos; + if ((uint32_t)s.cur_pos % pool_.block_size() == 0) { + out.new_block = (int32_t)app.last.physical_block; + out.new_block_index = s.cur_pos / (int)pool_.block_size(); + } + return out; +} + +void SeqSlotManager::commit_step(int slot) { + if (!is_active(slot)) return; + slots_[(size_t)slot].cur_pos += 1; +} + +void SeqSlotManager::retire(int slot) { + if (slot < 0 || slot >= (int)slots_.size()) return; + SeqSlot & s = slots_[(size_t)slot]; + if (!s.active()) return; + const PagedKvStatus status = pool_.release(s.handle); + if (status != PagedKvStatus::Ok && status != PagedKvStatus::StaleHandle) { + std::fprintf(stderr, "[parallel] slot %d release failed: %s\n", + slot, paged_kv_status_string(status)); + } + s = SeqSlot{}; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/seq_slot_manager.h b/server/src/common/concurrency/seq_slot_manager.h new file mode 100644 index 000000000..2699ebade --- /dev/null +++ b/server/src/common/concurrency/seq_slot_manager.h @@ -0,0 +1,131 @@ +// SeqSlotManager — complete host-side state for each concurrent serving slot. +// +// Companion of PagedKvPool: the pool hands out sequence handles and physical +// blocks; this class owns everything else a slot needs between admission and +// retirement — the pool-handle lifecycle (including every error path), the +// admission arithmetic (context clamp, prompt reservation, and rolling decode +// headroom), on-demand block allocation, per-slot sampler/RNG/penalty-history +// state, and the position counters. +// +// It deliberately owns NO device state. Prefill/decode allocation returns +// physical rows and block-table deltas as plain vectors. Prompt, KV ownership, +// sampler, and progress live together here; the scheduler keeps +// only its coarse request phase. +// +// Not thread-safe; the single scheduler thread is the only caller. + +#pragma once + +#include "common/paged_kv_pool.h" +#include "common/sampler.h" +#include "common/concurrency/seq_engine.h" + +#include +#include +#include +#include + +namespace dflash::common { + +enum class SeqSlotPhase { + free, + prefill, + decode, +}; + +struct SeqSlot { + SeqSlotPhase phase = SeqSlotPhase::free; + PagedKvSequenceHandle handle; + std::vector prompt; + int cur_pos = 0; + SamplerCfg sampler; + std::mt19937_64 rng{0x9E3779B97F4A7C15ull}; + // Penalty history is recorded as fed rather than sampled: the scheduler + // may override a sample before the model consumes it. + std::vector sample_history; + + bool active() const { return phase != SeqSlotPhase::free; } + bool prefilling() const { return phase == SeqSlotPhase::prefill; } + bool decoding() const { return phase == SeqSlotPhase::decode; } +}; + +class SeqSlotManager { +public: + // `max_ctx` is the per-sequence logical bound; slot count comes from the + // pool's max_sequences. The pool must outlive the manager. + SeqSlotManager(PagedKvPool & pool, int max_ctx); + + // Claim a free slot and atomically reserve all K/V blocks needed by the + // known prompt plus its next logical decode page when that page can exist + // in both max_ctx and the physical pool. Existing decoders are topped up + // first, so a younger admission cannot steal their next-page headroom. + // Prompts larger than the whole pool hard-fail; temporary capacity pressure + // reports busy. Seeds the slot RNG from sampler.seed only when the sampler + // actually draws, else nondeterministically. + SeqEngine::AdmitResult admit(uint64_t request_id, + const std::vector & prompt, + const SamplerCfg & sampler); + + struct PrefillChunk { + bool ok = false; + // The pool is temporarily out of blocks; retrying after another slot + // retires can succeed. BlocksExhausted leaves the pool unchanged. + bool busy = false; + std::vector rows; + // Delta to patch into the slot's device block-table column. + std::vector new_blocks; + int first_new_block = -1; + }; + + // Append `n_tokens` more prompt rows for a prefilling slot. Physical block + // ids come from the slot's admission reservation, so any append within the + // admitted prompt is guaranteed not to wait on another sequence. + PrefillChunk append_prefill(int slot, int n_tokens); + + // Record a finished prefill and expose the slot to decode. + void commit_prefill(int slot); + + struct StepAppend { + bool ok = false; + bool busy = false; // no physical block available right now + int64_t physical_row = -1; + int position = -1; // logical position the fed token is written at + int32_t new_block = -1; + int new_block_index = -1; + }; + + // Allocate the next decode token's cache row, report any new block-table + // entry, and log it to sample_history. cur_pos waits for commit_step(). + StepAppend append_token(int slot, int32_t fed_token); + + // The batched step's compute succeeded: cur_pos++. + void commit_step(int slot); + + // Release the slot's blocks and clear its state. Safe on inactive slots + // and after a failed admission/prefill. + void retire(int slot); + + int slot_count() const { return (int)slots_.size(); } + int max_context() const { return max_ctx_; } + int decoding_count() const; + bool is_active(int slot) const; + bool is_prefilling(int slot) const; + SeqSlot & slot(int i) { return slots_[(size_t)i]; } + const SeqSlot & slot(int i) const { return slots_[(size_t)i]; } + +private: + // Logical extent whose block count includes the sequence's current pages + // plus one future page, capped at max_ctx. + uint32_t decode_headroom_capacity(int logical_tokens) const; + bool capacity_fits_pool(uint32_t token_capacity) const; + + // Atomically preflight and top up every decoding slot as one cohort before + // a younger sequence may reserve capacity. + PagedKvStatus protect_decode_headroom(); + + PagedKvPool & pool_; + int max_ctx_ = 0; + std::vector slots_; +}; + +} // namespace dflash::common diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index e6ec4adc9..3bc305a8c 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -165,14 +165,19 @@ std::string check_feature_compatibility( // ── --paged-attention × architecture, placement, and decode features // Paged decode swaps the contiguous K/V cache for a block table owned by - // the monolithic qwen35 backend, so every rule below is about reaching - // that one code path. All are errors rather than warnings: running dense + // a monolithic model backend, so every rule below is about reaching + // that code path. All are errors rather than warnings: running dense // instead would hide the memory behavior the flag was chosen for. if (args.paged_attention) { if (!arch_supports_paged_attention(arch, /*is_layer_split=*/false)) { - return "--paged-attention requires a Qwen3.5/Qwen3.6 dense target " + return "--paged-attention requires a supported monolithic target " "(architecture '" + arch + "' has no paged decode path)"; } + if (arch == "deepseek4" && + target_backend != PlacementBackend::Hip) { + return "DeepSeek4 --paged-attention currently requires one " + "local HIP target (Strix Halo/gfx1151)"; + } // No rule for "requires a CUDA or HIP build": those are the only two // backends this binary can be configured with, and GGML_OP_PAGED_ATTN // is compiled into both. @@ -192,6 +197,17 @@ std::string check_feature_compatibility( return "--paged-attention cannot be combined with PFlash prefill " "compression"; } + if (features.kvflash_enabled) { + return "--paged-attention cannot be combined with KVFlash"; + } + if (arch == "deepseek4") { + if (args.ds4_prefill_mode != PrefillAttentionMode::Exact) { + return "DeepSeek4 --paged-attention requires --ds4-prefill exact"; + } + if (args.ds4_fused_decode) { + return "DeepSeek4 --paged-attention uses the gathered reference graph and cannot combine with --ds4-fused-decode"; + } + } // The pool rounds max_ctx up to a whole number of blocks, so the top // of the range is what can be rounded without overflowing int. if (args.device.max_ctx <= 0 || @@ -201,6 +217,39 @@ std::string check_feature_compatibility( } } + // ── --max-concurrency × DeepSeek4 paged attention + // Multiple live sequences are intentionally a DeepSeek4-only path. Qwen35 + // keeps its existing single-sequence paged implementation unchanged. + if (args.max_concurrency < 1) { + return "--max-concurrency must be at least 1"; + } + if (args.max_concurrency > 1) { + if (!args.paged_attention) { + return "--max-concurrency requires --paged-attention"; + } + if (arch != "deepseek4") { + return "--max-concurrency greater than 1 is currently supported only for DeepSeek4 on Strix Halo"; + } + if (args.max_concurrency > 16) { + return "DeepSeek4 --max-concurrency must be at most 16"; + } + // Physical capacity is memory-derived and capped independently of the + // logical slot count. The DeepSeek4 backend performs the final gfx1151 + // runtime check after HIP has selected the device. + } + if (args.kv_pool_tokens != 0) { + if (args.max_concurrency <= 1) { + return "--kv-pool-tokens requires --max-concurrency greater than 1"; + } + const int64_t max_pool_tokens = paged_kv_address_cap(); + if (args.kv_pool_tokens < PAGED_BLOCK_SIZE || + args.kv_pool_tokens > max_pool_tokens) { + return "--kv-pool-tokens must be in [" + + std::to_string(PAGED_BLOCK_SIZE) + ", " + + std::to_string(max_pool_tokens) + "]"; + } + } + // ── --ds4-prefill × architecture if (args.ds4_prefill_mode_set && arch != "deepseek4") { return "--ds4-prefill is only valid for deepseek4 models (detected '" + diff --git a/server/src/common/kvflash_pager.h b/server/src/common/kvflash_pager.h index d61707d96..c9af59a38 100644 --- a/server/src/common/kvflash_pager.h +++ b/server/src/common/kvflash_pager.h @@ -599,6 +599,15 @@ struct KvFlashAutoBudget { int speed_cap_tokens = 16384; }; +// The compatibility gate can reject a fixed KVFlash pool before model setup. +// "auto" is deliberately excluded: only the backend's VRAM-aware sizing can +// determine whether an automatic pool will actually be active. +inline bool kvflash_fixed_pool_requested(const char * value) { + return value != nullptr && + std::strcmp(value, "auto") != 0 && + std::atoi(value) > 0; +} + // Pool size from DFLASH_KVFLASH for a backend with `cfg` protections: // 0 = off; otherwise rounded to a 256 multiple, floored at // min_pool_tokens(cfg) (eviction must keep a victim) and clamped to diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 0a4d6e00d..c78180646 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -22,6 +22,7 @@ #include "ggml.h" #include "ggml-backend.h" #include "sampler.h" +#include "concurrency/seq_engine.h" #include "placement/draft_residency.h" namespace dflash::common { @@ -321,6 +322,19 @@ struct ModelBackend { virtual GenerateResult generate_impl(const GenerateRequest & req, const DaemonIO & io) = 0; + // ── 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 + // SeqEngine (common/concurrency/seq_engine.h). Any additional + // per-sequence model state is an implementation detail of that engine. + // nullptr — the + // default — means this backend serves one request at a time and the + // server drives it through generate(). + // + // The engine is owned by the backend; the returned pointer is borrowed + // and stays valid until shutdown(). + virtual SeqEngine * seq_engine() { return nullptr; } + // ── Snapshots ──────────────────────────────────────────────────── // With right-sized CPU-resident snapshots, each slot costs only // ~(cur_pos × 5 KB) of system RAM, so we can afford many slots. diff --git a/server/src/common/model_capabilities.h b/server/src/common/model_capabilities.h index cf14f414b..5c5f0410c 100644 --- a/server/src/common/model_capabilities.h +++ b/server/src/common/model_capabilities.h @@ -75,7 +75,7 @@ inline constexpr ArchCapabilities kArchCapabilities[] = { {"laguna", true, false, false, true, kMono, kMono, kMono, kNever, kNever, kNever}, {"qwen3", false, false, true, false, kNever, kNever, kNever, kNever, kNever, kNever}, {"gemma4", true, false, false, false, kMono, kNever, kNever, kBoth, kNever, kNever}, - {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever, kNever}, + {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever, kMono}, }; inline constexpr std::size_t kArchCount = diff --git a/server/src/common/paged_attention_config.h b/server/src/common/paged_attention_config.h index e78e0c346..73f5d278c 100644 --- a/server/src/common/paged_attention_config.h +++ b/server/src/common/paged_attention_config.h @@ -8,6 +8,9 @@ #pragma once +#include +#include + namespace dflash::common { constexpr int PAGED_BLOCK_SIZE = 16; @@ -20,4 +23,41 @@ constexpr int paged_token_capacity(int max_ctx) { return paged_block_count(max_ctx) * PAGED_BLOCK_SIZE; } +constexpr int64_t paged_kv_address_cap() { + return ((int64_t)INT32_MAX - PAGED_BLOCK_SIZE) / PAGED_BLOCK_SIZE * + PAGED_BLOCK_SIZE; +} + +// Inputs for sizing a concurrent paged pool from memory that remains after +// model weights have loaded. `fixed_cache_bytes` covers cache allocations that +// do not shrink with the pool (one prefill staging K/V, per-slot recurrent +// state, and metadata); `reserve_bytes` leaves room for runtime graph buffers. +struct PagedKvAutoBudget { + int64_t free_bytes = 0; + int64_t fixed_cache_bytes = 0; + int64_t reserve_bytes = 0; + int64_t bytes_per_token = 0; +}; + +// Return a whole-block physical capacity, capped by the old +// n_slots * max_ctx policy and by the pool's signed-int tensor address space. +// Zero means the supplied memory budget cannot hold even one block. +inline int64_t paged_kv_auto_pool_tokens(int max_ctx, int n_slots, + const PagedKvAutoBudget & budget) { + if (max_ctx < 1 || n_slots < 1 || budget.free_bytes <= 0 || + budget.bytes_per_token <= 0) { + return 0; + } + const int64_t usable = std::max( + 0, budget.free_bytes - budget.fixed_cache_bytes - + budget.reserve_bytes); + int64_t tokens = usable / budget.bytes_per_token; + tokens = (tokens / PAGED_BLOCK_SIZE) * PAGED_BLOCK_SIZE; + const int64_t logical_cap = + (int64_t)n_slots * (int64_t)paged_token_capacity(max_ctx); + const int64_t address_cap = paged_kv_address_cap(); + return std::max( + 0, std::min(tokens, std::min(logical_cap, address_cap))); +} + } // namespace dflash::common diff --git a/server/src/common/paged_kv_pool.cpp b/server/src/common/paged_kv_pool.cpp index 19a03f751..6bdf0f3fe 100644 --- a/server/src/common/paged_kv_pool.cpp +++ b/server/src/common/paged_kv_pool.cpp @@ -65,6 +65,12 @@ PagedKvPool::PagedKvPool(uint32_t physical_block_count, PagedKvStatus PagedKvPool::acquire(PagedKvRequestId request_id, PagedKvSequenceHandle & out_handle) { + return acquire_reserved(request_id, 0, out_handle); +} + +PagedKvStatus PagedKvPool::acquire_reserved( + PagedKvRequestId request_id, uint32_t token_capacity, + PagedKvSequenceHandle & out_handle) { if (request_to_slot_.find(request_id) != request_to_slot_.end()) { return PagedKvStatus::DuplicateRequest; } @@ -72,6 +78,11 @@ PagedKvStatus PagedKvPool::acquire(PagedKvRequestId request_id, return PagedKvStatus::SequenceSlotsExhausted; } + const uint32_t reserve_blocks = blocks_for_tokens(token_capacity); + if (reserve_blocks > free_blocks_.size()) { + return PagedKvStatus::BlocksExhausted; + } + const uint32_t slot = take_lowest(free_sequence_slots_); SequenceState & sequence = sequences_[slot]; uint64_t generation = sequence.generation + 1; @@ -82,12 +93,34 @@ PagedKvStatus PagedKvPool::acquire(PagedKvRequestId request_id, sequence.request_id = request_id; sequence.generation = generation; sequence.active = true; + take_reserved_blocks(sequence, reserve_blocks); request_to_slot_.emplace(request_id, slot); out_handle = {slot, generation}; return PagedKvStatus::Ok; } +PagedKvStatus PagedKvPool::reserve_capacity( + PagedKvSequenceHandle handle, uint32_t token_capacity) { + const PagedKvStatus status = validate(handle); + if (status != PagedKvStatus::Ok) return status; + + SequenceState & sequence = sequences_[handle.slot]; + const uint32_t target_blocks = blocks_for_tokens(token_capacity); + const uint64_t owned_blocks = + (uint64_t)sequence.block_table.size() + + sequence.reserved_blocks.size(); + if (target_blocks <= owned_blocks) return PagedKvStatus::Ok; + + const uint32_t additional_blocks = + target_blocks - static_cast(owned_blocks); + if (additional_blocks > free_blocks_.size()) { + return PagedKvStatus::BlocksExhausted; + } + take_reserved_blocks(sequence, additional_blocks); + return PagedKvStatus::Ok; +} + PagedKvAppendResult PagedKvPool::append(PagedKvSequenceHandle handle, uint32_t token_count, bool only_first_last_slots) { @@ -145,10 +178,14 @@ PagedKvStatus PagedKvPool::release(PagedKvSequenceHandle handle) { for (uint32_t block : sequence.block_table) { give_back(free_blocks_, block); } + for (uint32_t block : sequence.reserved_blocks) { + give_back(free_blocks_, block); + } sequence.request_id = 0; sequence.kv_seq_len = 0; sequence.active = false; sequence.block_table.clear(); + sequence.reserved_blocks.clear(); give_back(free_sequence_slots_, handle.slot); return PagedKvStatus::Ok; } @@ -161,6 +198,7 @@ void PagedKvPool::reset() { sequence.kv_seq_len = 0; sequence.active = false; sequence.block_table.clear(); + sequence.reserved_blocks.clear(); } refill(free_sequence_slots_, static_cast(sequences_.size())); refill(free_blocks_, physical_block_count_); @@ -176,10 +214,23 @@ PagedKvStatus PagedKvPool::sequence( PagedKvSequenceSnapshot snapshot; snapshot.kv_seq_len = sequence.kv_seq_len; snapshot.block_table = sequence.block_table; + snapshot.reserved_block_count = + static_cast(sequence.reserved_blocks.size()); out_sequence = std::move(snapshot); return PagedKvStatus::Ok; } +PagedKvStatus PagedKvPool::owned_block_count( + PagedKvSequenceHandle handle, uint32_t & out_count) const { + const PagedKvStatus status = validate(handle); + if (status != PagedKvStatus::Ok) return status; + + const SequenceState & sequence = sequences_[handle.slot]; + out_count = static_cast( + sequence.block_table.size() + sequence.reserved_blocks.size()); + return PagedKvStatus::Ok; +} + uint32_t PagedKvPool::blocks_for_tokens(uint32_t token_count) const { if (token_count == 0) return 0; return 1 + (token_count - 1) / block_size_; @@ -203,15 +254,28 @@ PagedKvStatus PagedKvPool::extend_block_table(SequenceState & sequence, if (required_blocks <= current_blocks) return PagedKvStatus::Ok; const uint32_t additional_blocks = required_blocks - current_blocks; - if (additional_blocks > free_blocks_.size()) { + const uint64_t available_blocks = + (uint64_t)sequence.reserved_blocks.size() + free_blocks_.size(); + if (additional_blocks > available_blocks) { return PagedKvStatus::BlocksExhausted; } sequence.block_table.reserve(required_blocks); for (uint32_t i = 0; i < additional_blocks; ++i) { - sequence.block_table.push_back(take_lowest(free_blocks_)); + std::vector & source = sequence.reserved_blocks.empty() + ? free_blocks_ : sequence.reserved_blocks; + sequence.block_table.push_back(take_lowest(source)); } return PagedKvStatus::Ok; } +void PagedKvPool::take_reserved_blocks(SequenceState & sequence, + uint32_t additional_blocks) { + sequence.reserved_blocks.reserve( + sequence.reserved_blocks.size() + additional_blocks); + for (uint32_t i = 0; i < additional_blocks; ++i) { + give_back(sequence.reserved_blocks, take_lowest(free_blocks_)); + } +} + } // namespace dflash::common diff --git a/server/src/common/paged_kv_pool.h b/server/src/common/paged_kv_pool.h index b8cca2499..9d745739f 100644 --- a/server/src/common/paged_kv_pool.h +++ b/server/src/common/paged_kv_pool.h @@ -7,8 +7,7 @@ // block/slot indices into offsets within their own pooled K/V tensors. // // The single-request backend consumes sequence() directly; the concurrent -// scheduler drives the multi-sequence slot, capacity, and handle-generation -// machinery through SeqSlotManager. +// engines compose this allocator with their own per-slot lifecycle records. // // Not thread-safe; callers must serialize access. Pool state is unspecified // if a std::bad_alloc escapes any call. @@ -77,6 +76,9 @@ struct PagedKvAppendResult { struct PagedKvSequenceSnapshot { uint32_t kv_seq_len = 0; std::vector block_table; + // Physical blocks held for future append() calls by this sequence. They + // are not visible in block_table until append consumes them. + uint32_t reserved_block_count = 0; }; // Allocator front-end: hands out sequence slots and physical block indices. @@ -91,13 +93,14 @@ class PagedKvPool { uint32_t block_size() const { return block_size_; } uint32_t physical_block_count() const { return physical_block_count_; } - // Admission capacity: SeqSlotManager sizes its slot table from this. + // Admission capacity: a concurrent engine sizes its slot table from this. uint32_t max_sequences() const { return static_cast(sequences_.size()); } uint32_t active_sequence_count() const { return static_cast(request_to_slot_.size()); } + // Blocks that are neither appended nor reserved by an active sequence. uint32_t free_block_count() const { return static_cast(free_blocks_.size()); } @@ -107,6 +110,25 @@ class PagedKvPool { PagedKvStatus acquire(PagedKvRequestId request_id, PagedKvSequenceHandle & out_handle); + // Claim a sequence slot and atomically reserve enough physical blocks for + // `token_capacity` future tokens. The logical sequence still starts empty; + // append() moves reserved blocks into its visible block table as needed. + // On any status failure, no slot or block is consumed and `out_handle` is + // unchanged. This is the admission primitive for chunked prompt prefill: + // once it succeeds, another sequence cannot strand this prompt halfway + // through by consuming the remainder of its capacity. + PagedKvStatus acquire_reserved(PagedKvRequestId request_id, + uint32_t token_capacity, + PagedKvSequenceHandle & out_handle); + + // Atomically ensure an active sequence owns enough appended plus reserved + // blocks to cover `token_capacity` logical tokens. This does not advance + // kv_seq_len or expose new block-table entries; append() consumes the + // private reservation later. Existing excess capacity is retained. On a + // status failure no block moves and the sequence is unchanged. + PagedKvStatus reserve_capacity(PagedKvSequenceHandle handle, + uint32_t token_capacity); + // Advance kv_seq_len by `token_count`, allocating new blocks as needed. // By default return every appended token's cache destination; when // `only_first_last_slots` is true, return only the range endpoints to @@ -129,6 +151,10 @@ class PagedKvPool { PagedKvStatus sequence(PagedKvSequenceHandle handle, PagedKvSequenceSnapshot & out_sequence) const; + // Return appended plus reserved blocks without copying the block table. + PagedKvStatus owned_block_count(PagedKvSequenceHandle handle, + uint32_t & out_count) const; + private: // Bookkeeping for one sequence slot. `generation` survives release so // the next acquire on this slot invalidates old handles. @@ -138,6 +164,10 @@ class PagedKvPool { uint32_t kv_seq_len = 0; bool active = false; std::vector block_table; + // Min-heap of blocks promised to this sequence but not yet made + // visible by append(). Reserved blocks are excluded from the global + // free list and returned by release()/reset(). + std::vector reserved_blocks; }; // Blocks needed to hold `token_count` tokens (ceiling division). @@ -152,6 +182,11 @@ class PagedKvPool { PagedKvStatus extend_block_table(SequenceState & sequence, uint32_t required_blocks); + // Move exactly `additional_blocks` globally free blocks into a sequence's + // private reservation. Caller must preflight availability. + void take_reserved_blocks(SequenceState & sequence, + uint32_t additional_blocks); + // Take the lowest free index off `free_list`, which must be non-empty. static uint32_t take_lowest(std::vector & free_list); // Return one index to `free_list`. diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index d4489a4f4..2f58d4909 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -4,6 +4,7 @@ #include "deepseek4_backend.h" #include "deepseek4_budget_hook.h" #include "deepseek4_internal.h" +#include "deepseek4_page_layout.h" #include "common/dynamic_backend.h" #include "common/peer_access.h" #include "common/sampler.h" @@ -44,6 +45,17 @@ static bool env_flag_enabled(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } +static bool is_gfx1151_device(int gpu) { +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + cudaDeviceProp prop{}; + return cudaGetDeviceProperties(&prop, gpu) == cudaSuccess && + std::strncmp(prop.gcnArchName, "gfx1151", 7) == 0; +#else + (void) gpu; + return false; +#endif +} + static void configure_gfx1151_dspark_mmvq_default(int gpu) { #if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) if (!env_flag_enabled("DFLASH_DS4_SPEC") || @@ -335,13 +347,14 @@ static uint64_t estimate_ds4_cache_bytes(const DeepSeek4Weights & w, int max_ctx const size_t comp_cap = (size_t) (max_ctx / (int) ratio) + 16; total_bytes += comp_cap * head_dim * sizeof(uint16_t); - const size_t window = (ratio == 4) ? 8 : ratio; - total_bytes += window * head_dim * sizeof(float) * 2; + const size_t state_rows = (ratio == 4) ? 8 : ratio; + const size_t comp_width = head_dim * (ratio == 4 ? 2 : 1); + total_bytes += state_rows * comp_width * sizeof(float) * 2; if (ratio == 4) { - const size_t index_comp_width = (size_t) w.n_indexer_head * (size_t) w.n_indexer_head_dim; - total_bytes += comp_cap * index_comp_width * sizeof(uint16_t); - total_bytes += window * index_comp_width * sizeof(float) * 2; + const size_t index_dim = (size_t) w.n_indexer_head_dim; + total_bytes += comp_cap * index_dim * sizeof(uint16_t); + total_bytes += 8 * (2 * index_dim) * sizeof(float) * 2; } } @@ -479,6 +492,7 @@ static bool fill_profiled_hot_placement(const DeepSeek4Weights & w, static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, ggml_backend_t backend, int max_ctx, + const DeepSeek4BackendConfig * cfg, Ds4HybridBudgetInfo & out, std::string * err) { out = {}; @@ -500,6 +514,41 @@ static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, out.core_bytes = moe_hybrid_core_bytes_from_memory( "deepseek4", out.gpu_free, out.gpu_total); out.kv_bytes = estimate_ds4_cache_bytes(w, max_ctx); + if (cfg && cfg->paged_attention) { + const uint64_t requested = cfg->kv_pool_tokens > 0 + ? (uint64_t) cfg->kv_pool_tokens + : (uint64_t) max_ctx * (uint64_t) cfg->max_concurrency; + const uint64_t physical_blocks = + (std::max(requested, (uint64_t) max_ctx) + + DS4_PAGE_TOKENS - 1) / DS4_PAGE_TOKENS; + DeepSeek4PagedCachePlan paged; + if (physical_blocks > UINT32_MAX || + !plan_deepseek4_paged_cache( + (uint32_t) w.head_dim, (uint32_t) w.n_indexer_head_dim, + (uint32_t) cfg->max_concurrency, (uint32_t) max_ctx, + (uint32_t) physical_blocks, w.compress_ratios, paged)) { + if (err) *err = "could not plan paged cache for expert budget"; + return false; + } + if (out.kv_bytes > UINT64_MAX - paged.total_persistent_bytes) { + if (err) *err = "paged cache budget overflow"; + return false; + } + out.kv_bytes += paged.total_persistent_bytes; + // The gathered reference graph materializes an independent MLA lane + // per slot and retains a two-shape mixed-backend cache. Reserve the + // measured 3090 worst-case scratch envelope before expert placement; + // otherwise an expert-heavy automatic placement can load correctly + // and then OOM when the first wide graph is built. + constexpr uint64_t graph_bytes_per_lane = 320ULL * 1024 * 1024; + const uint64_t graph_bytes = + (uint64_t) cfg->max_concurrency * graph_bytes_per_lane; + if (out.kv_bytes > UINT64_MAX - graph_bytes) { + if (err) *err = "paged graph budget overflow"; + return false; + } + out.kv_bytes += graph_bytes; + } if (out.gpu_total > out.core_bytes + out.kv_bytes + out.warm_bytes + out.safety_bytes) { out.expert_budget = out.gpu_total - out.core_bytes - out.kv_bytes - out.warm_bytes - out.safety_bytes; @@ -569,7 +618,7 @@ DeepSeek4Backend::~DeepSeek4Backend() { } bool DeepSeek4Backend::requires_monolithic_model() const { - return cfg_.fused_decode || + return cfg_.paged_attention || cfg_.fused_decode || prefill_attention_mode_is_approximate(cfg_.prefill_mode); } @@ -602,8 +651,9 @@ bool DeepSeek4Backend::load_model() { ? compiled_placement_backend() : cfg_.device.backend; - // Fused decode and layer-major prefill normally require monolithic expert - // residency. Heterogeneous TP is the exception: its fused graph owns the + // Paged concurrency, fused decode, and layer-major prefill require + // monolithic expert residency. Heterogeneous TP is the exception for + // non-paged modes: its fused graph owns the // routed experts across two local GPU backends, so forcing a full load would // disable the requested split before the TP runtime can initialize. const bool force_full = env_flag_enabled("DFLASH_DS4_FORCE_FULL_LOAD"); @@ -614,8 +664,9 @@ bool DeepSeek4Backend::load_model() { (force_full || need_monolithic)) { std::fprintf(stderr, "[deepseek4] monolithic execution requested " - "(forced=%s, fused_decode=%s, prefill=%s)\n", + "(forced=%s, paged=%s, fused_decode=%s, prefill=%s)\n", force_full ? "yes" : "no", + cfg_.paged_attention ? "on" : "off", cfg_.fused_decode ? "on" : "off", prefill_attention_mode_name(cfg_.prefill_mode)); if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { @@ -826,12 +877,43 @@ int DeepSeek4Backend::capture_safe_prefill_tokens( } bool DeepSeek4Backend::init() { + if (cfg_.paged_attention) { + const PlacementBackend target_backend = + cfg_.device.backend == PlacementBackend::Auto + ? compiled_placement_backend() : cfg_.device.backend; + if (target_backend != PlacementBackend::Hip || + !is_gfx1151_device(cfg_.device.gpu)) { + std::fprintf(stderr, + "[deepseek4] paged concurrency currently requires one " + "local Strix Halo (gfx1151) HIP target\n"); + return false; + } + if (env_flag_enabled("DFLASH_DS4_MOE_TP")) { + std::fprintf(stderr, + "[deepseek4] paged concurrency keeps all experts resident " + "on Strix Halo and cannot use DFLASH_DS4_MOE_TP\n"); + 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. configure_gfx1151_dspark_mmvq_default(cfg_.device.gpu); configure_gfx1201_hybrid_sub_batch_default(cfg_.device.gpu); + if (cfg_.paged_attention && + (cfg_.max_concurrency < 1 || cfg_.max_concurrency > 16 || + cfg_.device.is_layer_split() || + cfg_.prefill_mode != PrefillAttentionMode::Exact || + cfg_.fused_decode || env_flag_enabled("DFLASH_DS4_FUSED_DECODE") || + env_flag_enabled("DFLASH_DS4_SPEC"))) { + std::fprintf(stderr, + "[deepseek4] paged serving requires 1..16 local slots, exact " + "prefill, and autoregressive non-fused decode\n"); + return false; + } + backend_ = ggml_backend_cuda_init(cfg_.device.gpu); if (!backend_) { std::fprintf(stderr, "[deepseek4] failed to create CUDA backend (gpu=%d)\n", @@ -855,15 +937,48 @@ bool DeepSeek4Backend::init() { } const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; - if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { - std::fprintf(stderr, "[deepseek4] failed to allocate KV cache (ctx=%d)\n", max_ctx); - return false; + if (cfg_.paged_attention) { + uint64_t requested = cfg_.kv_pool_tokens > 0 + ? (uint64_t)cfg_.kv_pool_tokens + : (uint64_t)max_ctx * (uint64_t)cfg_.max_concurrency; + requested = std::max(requested, (uint64_t)max_ctx); + const uint64_t blocks64 = + (requested + DS4_PAGE_TOKENS - 1) / DS4_PAGE_TOKENS; + if (blocks64 == 0 || blocks64 > UINT32_MAX || + !create_deepseek4_paged_cache( + backend_, w_, (uint32_t)cfg_.max_concurrency, + (uint32_t)max_ctx, (uint32_t)blocks64, paged_cache_)) { + std::fprintf(stderr, + "[deepseek4] paged cache allocation failed (ctx=%d slots=%d " + "requested_pool_tokens=%llu); reduce --max-ctx/--max-concurrency " + "or set --kv-pool-tokens\n", max_ctx, cfg_.max_concurrency, + (unsigned long long)requested); + return false; + } + } else { + if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { + std::fprintf(stderr, "[deepseek4] failed to allocate KV cache (ctx=%d)\n", max_ctx); + return false; + } + cache_.prefill_mode = cfg_.prefill_mode; } - cache_.prefill_mode = cfg_.prefill_mode; if (env_flag_enabled("DFLASH_DS4_MOE_TP") && !init_moe_tensor_parallel()) { return false; } + if (cfg_.paged_attention && expert_runtime_.compute) { + std::fprintf(stderr, + "[deepseek4] paged serving cannot use the out-of-process expert " + "compute callback; select in-process DFLASH_DS4_MOE_TP or disable paged attention\n"); + return false; + } + if (cfg_.paged_attention && moe_hybrid_ && + !moe_hybrid_->materialized_cold_experts) { + std::fprintf(stderr, + "[deepseek4] paged serving requires statically materialized " + "expert ownership; enable in-process DFLASH_DS4_MOE_TP\n"); + return false; + } if (const char * stats_path = std::getenv("DFLASH_DS4_ROUTING_STATS_OUT")) { if (*stats_path) { @@ -888,6 +1003,16 @@ bool DeepSeek4Backend::init() { std::fprintf(stderr, "[deepseek4-moe-tp] in-memory routing stats enabled\n"); } + if (cfg_.paged_attention) { + seq_engine_ = std::make_unique( + *this, *paged_cache_.pool, max_ctx, + paged_cache_.plan.max_blocks_per_sequence); + std::fprintf(stderr, + "[deepseek4-parallel] enabled %d slots, %u x %d-token physical " + "blocks; prefill is exact reference mode at one prompt token per slot per scheduler iteration\n", + cfg_.max_concurrency, paged_cache_.plan.physical_blocks, + DS4_PAGE_TOKENS); + } const int active_experts = w_.routed_expert_top_k > 0 ? w_.routed_expert_top_k : w_.n_expert_used; std::fprintf(stderr, @@ -898,7 +1023,7 @@ bool DeepSeek4Backend::init() { prefill_attention_mode_name(cfg_.prefill_mode), moe_hybrid_ ? " [hybrid]" : ""); - if (env_flag_enabled("DFLASH_DS4_SPEC")) { + if (!cfg_.paged_attention && env_flag_enabled("DFLASH_DS4_SPEC")) { const char * dp = std::getenv("DFLASH_DS4_DRAFT"); if (dp && *dp) { spec_draft_path_ = dp; @@ -981,7 +1106,8 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & MoeHybridPlacement & out, std::string * err) const { Ds4HybridBudgetInfo budget; - if (!compute_ds4_hybrid_budget_info(w, backend_, max_ctx, budget, err)) { + if (!compute_ds4_hybrid_budget_info( + w, backend_, max_ctx, &cfg_, budget, err)) { return false; } @@ -1236,7 +1362,10 @@ bool DeepSeek4Backend::init_hybrid_model() { void DeepSeek4Backend::print_ready_banner() const { std::printf("[deepseek4-daemon] ready layers=%d ctx=%d experts=%d/%d\n", - w_.n_layer, cache_.max_ctx, w_.n_expert_used, w_.n_expert); + w_.n_layer, + cfg_.paged_attention ? (int)paged_cache_.plan.max_ctx + : cache_.max_ctx, + w_.n_expert_used, w_.n_expert); std::fflush(stdout); } @@ -1250,6 +1379,12 @@ bool DeepSeek4Backend::park(ParkTarget target) { std::fflush(stdout); } if (!want_target_model || parked_) return true; + if (cfg_.paged_attention) { + std::fprintf(stderr, + "[deepseek4] target park is unavailable while paged serving owns " + "live graph and slot state\n"); + return false; + } maybe_save_routing_stats(); for (int i = 0; i < PREFIX_SLOTS; ++i) { @@ -1992,6 +2127,8 @@ void DeepSeek4Backend::shutdown() { for (int i = 0; i < PREFIX_SLOTS; i++) { snapshot_free(i); } + seq_engine_.reset(); + free_deepseek4_paged_cache(paged_cache_); free_deepseek4_cache(cache_); expert_runtime_.reset(); stream_engine_.destroy(); diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 32f7230ce..da272624d 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -15,6 +15,7 @@ #include "../common/moe_hybrid_stream.h" #include "deepseek4_internal.h" #include "deepseek4_dspark.h" +#include "deepseek4_seq_engine.h" #include "ggml.h" #include "ggml-backend.h" @@ -60,6 +61,7 @@ class DeepSeek4Backend : public ModelBackend { void free_drafter() override; void shutdown() override; + SeqEngine * seq_engine() override { return seq_engine_.get(); } const MoeHybridRoutingStats * get_routing_stats() const override { return routing_stats_.get(); @@ -72,6 +74,8 @@ class DeepSeek4Backend : public ModelBackend { ggml_backend_t expert_backend_ = nullptr; DeepSeek4Weights w_; DeepSeek4Cache cache_; + DeepSeek4PagedCache paged_cache_; + std::unique_ptr seq_engine_; bool parked_ = false; // Sampler @@ -152,6 +156,7 @@ class DeepSeek4Backend : public ModelBackend { MoeExpertComputeRuntime expert_runtime_; std::shared_ptr routing_stats_; std::string routing_stats_out_path_; + friend class DeepSeek4SeqEngine; }; } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index ac583ad65..637bc6668 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -220,7 +220,8 @@ static bool ds4_fused_read_route_matrix( static void ds4_fused_consume_route_diagnostics( DeepSeek4FusedDecodeGraph & fg, const MoeHybridStorage * hybrid, - MoeHybridRoutingStats * routing_stats) { + MoeHybridRoutingStats * routing_stats, + const int32_t * active_slots = nullptr) { const bool cache_audit = ds4_env_flag("DFLASH_DS4_TP_CACHE_AUDIT"); if (!routing_stats && !cache_audit) return; @@ -242,6 +243,7 @@ static void ds4_fused_consume_route_diagnostics( continue; } for (int token = 0; token < route.n_tokens; ++token) { + if (active_slots && active_slots[route.lane_start + token] < 0) continue; const int32_t * token_ids = ids.data() + (size_t) token * route.width; const float * token_weights = weights.data() + @@ -344,7 +346,10 @@ static bool ds4_build_fused_verify_graph( bool have_token_ids, const std::vector & capture_ids, MoeHybridStorage * hybrid, - std::vector && shape_key) { + std::vector && shape_key, + DeepSeek4PagedCache * paged_cache = nullptr, + const std::vector> * paged_rows = nullptr) { + const bool paged_mode = paged_cache && paged_rows; if (fg.sched) { ggml_backend_sched_free(fg.sched); fg.sched = nullptr; @@ -372,7 +377,11 @@ static bool ds4_build_fused_verify_graph( fg.sg.ctx = ggml_init(params); if (!fg.sg.ctx) return false; ggml_context * ctx = fg.sg.ctx; - constexpr size_t graph_capacity = 65536u; + // Gathered paged attention builds one q=1 MLA lane per concurrent + // sequence. Above eight lanes the resulting whole-model graph exceeds + // the verifier-era 64K scheduler hash table even though the metadata + // arena still has ample room. + const size_t graph_capacity = q > 8 ? 131072u : 65536u; fg.sg.gf = ggml_new_graph_custom(ctx, graph_capacity, false); ggml_cgraph * gf = fg.sg.gf; @@ -412,7 +421,8 @@ static bool ds4_build_fused_verify_graph( const int preserved_rows = q > 1 ? q : 0; mask_total += (int64_t) (w.n_swa + padded + preserved_rows) * q; } - fg.mask_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mask_total); + fg.mask_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, + std::max(mask_total, 1)); ggml_set_input(fg.mask_bundle); int64_t mask_off = 0; @@ -454,7 +464,115 @@ static bool ds4_build_fused_verify_graph( attn_in = attn_in ? ggml_concat(ctx, attn_in, w2, 1) : w2; } - // ── Batched attention ── + // ── Attention ── + ggml_tensor * attn_out = nullptr; + if (paged_mode) { + // q denotes independent decode lanes here, not a causal sequence. + // Gather each lane's immutable chronological history and run the + // established MLA lane core at q=1; all surrounding HC/MoE/output + // machinery remains q-wide and unchanged. + if (q < 1 || q > 16 || paged_rows->size() != (size_t) w.n_layer || + (*paged_rows)[(size_t) il].size() != (size_t) q) return false; + DeepSeek4PagedLayerCache & plc = paged_cache->layers[(size_t) il]; + ggml_tensor * raw_flat = ggml_reshape_2d( + ctx, plc.raw_kv, w.head_dim, + (int64_t) DS4_PAGE_TOKENS * paged_cache->plan.slots); + for (int t = 0; t < q; ++t) { + const auto & rows = (*paged_rows)[(size_t) il][(size_t) t]; + auto & px = ex.paged.emplace_back(); + px.pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(px.pos); + px.neg_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(px.neg_pos); + px.raw_gather = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, + std::max(rows.raw_history.size(), 1)); ggml_set_input(px.raw_gather); + ggml_tensor * raw_history = rows.raw_history.empty() ? nullptr + : ggml_get_rows(ctx, raw_flat, px.raw_gather); + ggml_tensor * comp_history = nullptr; + ggml_tensor * index_history = nullptr; + if (ratio > 0) { + px.comp_gather = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, + std::max(rows.compressed_history.size(), 1)); + ggml_set_input(px.comp_gather); + if (!rows.compressed_history.empty()) + comp_history = ggml_get_rows(ctx, plc.comp_kv, px.comp_gather); + if (ratio == 4) { + px.index_gather = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, + std::max(rows.compressed_history.size(), 1)); + ggml_set_input(px.index_gather); + if (!rows.compressed_history.empty()) + index_history = ggml_get_rows(ctx, plc.index_comp_kv, px.index_gather); + } + } + px.raw_write = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); ggml_set_input(px.raw_write); + if (ratio > 0) { + px.comp_write = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); ggml_set_input(px.comp_write); + px.comp_read = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(px.comp_read); + px.ape = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(px.ape); + px.state_row = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); ggml_set_input(px.state_row); + px.comp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(px.comp_pos); + } + DeepSeek4CompressorState attn_state{}, index_state{}; + if (ratio > 0 && rows.slot >= 0) { + const size_t off = (size_t) rows.slot * plc.attn_compressor.state_kv->nb[2]; + attn_state.state_kv = ggml_view_2d(ctx, plc.attn_compressor.state_kv, + plc.attn_compressor.state_kv->ne[0], plc.attn_compressor.state_kv->ne[1], + plc.attn_compressor.state_kv->nb[1], off); + attn_state.state_score = ggml_view_2d(ctx, plc.attn_compressor.state_score, + plc.attn_compressor.state_score->ne[0], plc.attn_compressor.state_score->ne[1], + plc.attn_compressor.state_score->nb[1], off); + if (ratio == 4) { + const size_t io = (size_t) rows.slot * plc.indexer_compressor.state_kv->nb[2]; + index_state.state_kv = ggml_view_2d(ctx, plc.indexer_compressor.state_kv, + plc.indexer_compressor.state_kv->ne[0], plc.indexer_compressor.state_kv->ne[1], + plc.indexer_compressor.state_kv->nb[1], io); + index_state.state_score = ggml_view_2d(ctx, plc.indexer_compressor.state_score, + plc.indexer_compressor.state_score->ne[0], plc.indexer_compressor.state_score->ne[1], + plc.indexer_compressor.state_score->nb[1], io); + } + } + DeepSeek4MlaLaneBindings lane{}; + lane.history_mode = DeepSeek4MlaLaneBindings::HistoryMode::ChronologicalGathered; + lane.raw_history = raw_history; lane.n_raw_history = (int) rows.raw_history.size(); + lane.comp_history = comp_history; lane.n_comp_history = (int) rows.compressed_history.size(); + lane.index_comp_history = index_history; + lane.n_index_comp_history = (int) rows.compressed_history.size(); + // Writes use the same flattened physical-row geometry as the + // gather indices (not the slot-local 128-row second axis). + lane.raw_kv = raw_flat; lane.comp_kv = plc.comp_kv; + lane.index_comp_kv = plc.index_comp_kv; + lane.raw_write_rows = px.raw_write; lane.comp_write_rows = px.comp_write; + lane.index_comp_write_rows = px.comp_write; + lane.comp_read_rows = px.comp_read; + lane.index_comp_read_rows = px.comp_read; + lane.write_enabled = rows.slot >= 0; + lane.attn_compressor = &attn_state; lane.indexer_compressor = &index_state; + DeepSeek4AttentionGraphInputs ci{}; + ci.rope_pos = px.pos; ci.neg_pos = px.neg_pos; + ci.attn_ape_row = px.ape; ci.attn_state_rows = px.state_row; + ci.attn_comp_pos = px.comp_pos; ci.index_ape_row = px.ape; + ci.index_state_rows = px.state_row; ci.index_comp_pos = px.comp_pos; + std::vector ib; + std::vector iab; + std::vector lab; + ggml_tensor * col = ggml_view_2d(ctx, attn_in, n_embd, 1, + attn_in->nb[1], (size_t) t * attn_in->nb[1]); + ggml_tensor * one = build_mla_attention_lane_core( + ctx, gf, build_rms_norm(ctx, col, L.attn_norm, w.rms_eps), + w, L, lane, il, (int) rows.position, 1, &ci, ib, iab, lab, + nullptr, DeepSeek4AttentionImpl::Explicit); + if (!one || !ib.empty() || !iab.empty() || !lab.empty()) { + std::fprintf(stderr, + "[deepseek4-paged] layer %d lane %d attention build " + "failed (graph=%d i32=%zu arrays=%zu i64=%zu, " + "first_array=%d)\n", + il, t, one != nullptr, ib.size(), iab.size(), lab.size(), + iab.empty() || iab[0].values.empty() + ? INT32_MIN : iab[0].values[0]); + return false; + } + attn_out = attn_out ? ggml_concat(ctx, attn_out, one, 1) : one; + } + } else { + // ── Batched speculative attention ── DeepSeek4AttentionGraphInputs ain{}; ain.rope_pos = ex.pos_q; ain.neg_pos = ex.neg_q; @@ -492,7 +610,7 @@ static bool ds4_build_fused_verify_graph( std::vector i32ab; std::vector i64ab; ggml_tensor * normed = build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); - ggml_tensor * attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, + attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, lane_kv_start, lane_q, &ain, i32b, i32ab, i64ab); if (!attn_out) return false; @@ -500,6 +618,7 @@ static bool ds4_build_fused_verify_graph( std::fprintf(stderr, "[ds4-fused-verify] layer %d dynamic bindings; cannot fuse\n", il); return false; } + } // ── HC post (attention) + HC pre (FFN), per token ── ggml_tensor * ffn_in = nullptr; @@ -713,7 +832,11 @@ static bool ds4_build_fused_verify_graph( } } } - if (!ffn_out) return false; + if (!ffn_out) { + std::fprintf(stderr, + "[deepseek4-paged] layer %d FFN graph build failed\n", il); + return false; + } // ── HC post (FFN), per token; capture at drafter layers ── for (int t = 0; t < lane_q; ++t) { @@ -777,7 +900,7 @@ static bool ds4_build_fused_verify_graph( fg.logits = ggml_mul_mat(ctx, w.output, out_normed); // [n_vocab, q] ggml_set_output(fg.logits); ggml_build_forward_expand(gf, fg.logits); - if (ds4_env_flag("DFLASH_DS4_GPU_ARGMAX_VERIFY")) { + if (paged_mode || ds4_env_flag("DFLASH_DS4_GPU_ARGMAX_VERIFY")) { ex.argmax = ggml_argmax(ctx, fg.logits); ggml_set_output(ex.argmax); ggml_build_forward_expand(gf, ex.argmax); @@ -901,6 +1024,12 @@ static bool ds4_build_fused_verify_graph( pin_main(fg.i32_bundle); pin_main(fg.i64_bundle); pin_main(fg.mask_bundle); + for (const auto & px : ex.paged) { + pin_main(px.pos); pin_main(px.neg_pos); pin_main(px.raw_gather); pin_main(px.comp_gather); + pin_main(px.index_gather); + pin_main(px.raw_write); pin_main(px.comp_write); pin_main(px.comp_read); pin_main(px.ape); + pin_main(px.state_row); pin_main(px.comp_pos); + } for (ggml_tensor * hids : fg.hash_ids) pin_main(hids); for (const MoeHybridGraphInputs & inputs : fg.hybrid_inputs) { if (mixed_policy.pin_route_weights) { diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 7014fafd4..f196dc7ef 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -11,6 +11,7 @@ #include "deepseek4_internal.h" #include "deepseek4_hc_cuda.h" #include "deepseek4_roctx.h" +#include "deepseek4_page_layout.h" #include "internal.h" #include "../common/step_graph.h" #include "../common/cuda_graph_overrides.h" @@ -994,7 +995,9 @@ static void build_compressor_step( ggml_tensor * cur_all = nullptr, int n_tokens_all = 1, int kv_start_all = -1, - bool indexer_qat = false) { + bool indexer_qat = false, + ggml_tensor ** current_comp_out = nullptr, + bool paged_physical_row = false) { if (!gf || !cur_last || !ape || !kv_proj || !gate_proj || !norm_weight || !state.state_kv || !state.state_score || !comp_cache || ratio <= 0) { return; @@ -1202,10 +1205,14 @@ static void build_compressor_step( if (indexer_qat) { pooled = ggml_ds4_indexer_qat(ctx, ggml_cont(ctx, pooled)); } + if (current_comp_out) { + *current_comp_out = pooled; + } ggml_tensor * pooled_f16 = ggml_cast(ctx, pooled, GGML_TYPE_F16); const int comp_row = token_pos / ratio; - if (comp_row >= (int) comp_cache->ne[1]) { + if ((!comp_rows_inp || !paged_physical_row) && + comp_row >= (int) comp_cache->ne[1]) { return; } @@ -1315,7 +1322,8 @@ static void build_indexer_compressor_step( ggml_tensor * cur_last, const DeepSeek4Weights & w, const DeepSeek4Layer & L, - DeepSeek4LayerCache & lc, + DeepSeek4CompressorState & indexer_compressor, + ggml_tensor * index_comp_kv, int token_pos, ggml_tensor * ape_row_inp, ggml_tensor * state_rows_inp, @@ -1328,14 +1336,16 @@ static void build_indexer_compressor_step( ggml_tensor * cur_all = nullptr, int n_tokens_all = 1, int kv_start_all = -1, - bool indexer_qat = false) { + bool indexer_qat = false, + ggml_tensor ** current_comp_out = nullptr, + bool paged_physical_row = false) { build_compressor_step(ctx, gf, cur_last, L.indexer_compressor_ape, L.indexer_compressor_kv, L.indexer_compressor_gate, L.indexer_compressor_norm, - lc.indexer_compressor, - lc.index_comp_kv, + indexer_compressor, + index_comp_kv, 4, w.n_indexer_head_dim, // indexer head_dim = 128 token_pos, @@ -1357,7 +1367,9 @@ static void build_indexer_compressor_step( cur_all, n_tokens_all, kv_start_all, - indexer_qat); + indexer_qat, + current_comp_out, + paged_physical_row); } static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int ratio, int token_pos) { @@ -1506,13 +1518,86 @@ static ggml_tensor * build_indexer_topk( // ─── MLA Attention Block ──────────────────────────────────────────────── -static ggml_tensor * build_mla_attention( +// All persistent and live-state bindings consumed by one MLA lane. Keeping +// this internal seam tensor-based is intentional: a paged adapter can later +// supply gathered history and slot-specific compressor state without the +// graph builder consulting DeepSeek4LayerCache or host cache counters. +struct DeepSeek4MlaLaneBindings { + enum class HistoryMode { + ContiguousRing, + ChronologicalGathered, + }; + + HistoryMode history_mode = HistoryMode::ContiguousRing; + // In gathered mode these are immutable, chronological attention inputs. + // Counts are explicit so adapters may bind capacity-padded tensors. + ggml_tensor * raw_history = nullptr; + int n_raw_history = 0; + ggml_tensor * comp_history = nullptr; + int n_comp_history = 0; + ggml_tensor * index_comp_history = nullptr; + int n_index_comp_history = 0; + + // Persistent mutation targets are deliberately independent of history. + ggml_tensor * raw_kv = nullptr; + ggml_tensor * comp_kv = nullptr; + ggml_tensor * index_comp_kv = nullptr; + ggml_tensor * raw_write_rows = nullptr; + ggml_tensor * comp_write_rows = nullptr; + ggml_tensor * index_comp_write_rows = nullptr; + ggml_tensor * comp_read_rows = nullptr; // GET_ROWS requires I32 + ggml_tensor * index_comp_read_rows = nullptr; + + // Optional passive outputs let a future adapter scatter current products. + ggml_tensor ** current_raw_out = nullptr; + ggml_tensor ** current_comp_out = nullptr; + ggml_tensor ** current_index_comp_out = nullptr; + // False is the padding/inactive-lane contract: build attention against the + // supplied padded history, but emit no persistent current-row mutations. + bool write_enabled = true; + DeepSeek4CompressorState * attn_compressor = nullptr; + DeepSeek4CompressorState * indexer_compressor = nullptr; + int n_comp_live = 0; + int n_index_comp_live = 0; + int n_comp_committed = 0; +}; + +// Projection/RoPE products handed to the history/update portion of a lane. +// This is deliberately a passive bundle: introducing graph operations in a +// separate builder would risk changing decode graph ordering. +struct DeepSeek4PreparedProjectedLane { + ggml_tensor * normalized_q_lora = nullptr; + ggml_tensor * q = nullptr; + ggml_tensor * kv = nullptr; + ggml_tensor * rope_pos = nullptr; +}; + +static DeepSeek4MlaLaneBindings deepseek4_contiguous_lane_bindings( + DeepSeek4LayerCache & lc, + int ratio, + int token_pos) { + DeepSeek4MlaLaneBindings lane; + lane.history_mode = DeepSeek4MlaLaneBindings::HistoryMode::ContiguousRing; + lane.raw_kv = lc.raw_kv; + lane.comp_kv = lc.comp_kv; + lane.index_comp_kv = lc.index_comp_kv; + lane.attn_compressor = &lc.attn_compressor; + lane.indexer_compressor = &lc.indexer_compressor; + lane.n_comp_live = ratio > 0 + ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; + lane.n_index_comp_live = ratio == 4 + ? ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos) : 0; + lane.n_comp_committed = lc.n_comp; + return lane; +} + +static ggml_tensor * build_mla_attention_lane_core( ggml_context * ctx, ggml_cgraph * gf, ggml_tensor * cur, // [n_embd, n_tokens] const DeepSeek4Weights & w, const DeepSeek4Layer & L, - DeepSeek4LayerCache & lc, + const DeepSeek4MlaLaneBindings & lane, int layer_idx, int kv_start, int n_tokens, @@ -1530,6 +1615,8 @@ static ggml_tensor * build_mla_attention( const int n_out_group = w.n_out_group; const int n_lora_o = w.n_lora_o; const int ratio = w.compress_ratios[layer_idx]; + const bool gathered_history = lane.history_mode == + DeepSeek4MlaLaneBindings::HistoryMode::ChronologicalGathered; // ── Q path: cur → q_a → norm → q_b → per-head norm ───────────── // q_a: [n_embd, n_tokens] → [n_lora_q, n_tokens] @@ -1587,6 +1674,11 @@ static ggml_tensor * build_mla_attention( rope_freq, rope_scale, rope_ext, rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + const DeepSeek4PreparedProjectedLane projected = {qr, q, kv, rope_pos}; + // Keep the established local names below to make the no-topology-change + // property obvious; the bundle is the handoff seam for a future adapter. + (void) projected; + // ── Causal batched step (exact multi-token target semantics) ─── // The target model is causal: token i must not attend to batch tokens // j > i, must see the compressed-row count as of its own position, and — @@ -1601,15 +1693,15 @@ static ggml_tensor * build_mla_attention( ggml_tensor * old_rows_scratch = nullptr; int n_old_rows = 0; ggml_tensor * prior_rows_scratch = nullptr; - int n_prior_rows = 0; + int n_prior_rows = gathered_history ? lane.n_raw_history : 0; const bool fused_causal = cached_inputs && cached_inputs->attn_row_mask && n_tokens > 1; - if (fused_causal) { + if (!gathered_history && fused_causal) { // Fused verify: ALWAYS q preserved rows so the topology is stable; // unwrapped/garbage rows are masked by the host-filled mask values. for (int ti = 0; ti < n_tokens; ti++) { ggml_tensor * slot = ggml_view_2d( - ctx, lc.raw_kv, head_dim, 1, lc.raw_kv->nb[1], - (size_t)((kv_start + ti) % w.n_swa) * lc.raw_kv->nb[1]); + ctx, lane.raw_kv, head_dim, 1, lane.raw_kv->nb[1], + (size_t)((kv_start + ti) % w.n_swa) * lane.raw_kv->nb[1]); ggml_tensor * saved = ggml_cont(ctx, slot); ggml_build_forward_expand(gf, saved); old_rows_scratch = old_rows_scratch @@ -1617,14 +1709,14 @@ static ggml_tensor * build_mla_attention( n_old_rows++; } old_rows_scratch = ds4_cast_if_needed(ctx, old_rows_scratch, GGML_TYPE_F32); - } else if (causal_batch && !layer_major_batch) { + } else if (!gathered_history && causal_batch && !layer_major_batch) { // Copy the to-be-overwritten rows FIRST; same-stream build order runs // these before the ring writes below. for (int ti = 0; ti < n_tokens; ti++) { if (kv_start + ti < w.n_swa) continue; // slot never held an older pos ggml_tensor * slot = ggml_view_2d( - ctx, lc.raw_kv, head_dim, 1, lc.raw_kv->nb[1], - (size_t)((kv_start + ti) % w.n_swa) * lc.raw_kv->nb[1]); + ctx, lane.raw_kv, head_dim, 1, lane.raw_kv->nb[1], + (size_t)((kv_start + ti) % w.n_swa) * lane.raw_kv->nb[1]); ggml_tensor * saved = ggml_cont(ctx, slot); ggml_build_forward_expand(gf, saved); old_rows_scratch = old_rows_scratch @@ -1634,7 +1726,7 @@ static ggml_tensor * build_mla_attention( if (old_rows_scratch) { old_rows_scratch = ds4_cast_if_needed(ctx, old_rows_scratch, GGML_TYPE_F32); } - } else if (layer_major_batch) { + } else if (!gathered_history && layer_major_batch) { // Snapshot the chronological pre-chunk window before any ring writes. // Attention then consumes [prior F16 rows | current F32 rows], matching // the single-token path and avoiding an F16 round-trip for this chunk. @@ -1644,8 +1736,8 @@ static ggml_tensor * build_mla_attention( const int tail = std::min(n_prior_rows, w.n_swa - first); auto snapshot_span = [&](int row, int count) { ggml_tensor * span = ggml_view_2d( - ctx, lc.raw_kv, head_dim, count, lc.raw_kv->nb[1], - (size_t) row * lc.raw_kv->nb[1]); + ctx, lane.raw_kv, head_dim, count, lane.raw_kv->nb[1], + (size_t) row * lane.raw_kv->nb[1]); return ggml_cont(ctx, span); }; prior_rows_scratch = snapshot_span(first, tail); @@ -1663,13 +1755,18 @@ static ggml_tensor * build_mla_attention( // ── Store ALL KV rows in the raw SWA ring ───────────────────── // For decode (n_tokens=1): write single row. For prefill: write all rows. - ggml_tensor * raw_kv_source = lc.raw_kv; - ggml_tensor * raw_kv_rows = cached_inputs - ? cached_inputs->raw_kv_rows - : nullptr; - if (raw_kv_rows) { + ggml_tensor * raw_kv_source = lane.raw_kv; + ggml_tensor * raw_kv_rows = lane.raw_write_rows + ? lane.raw_write_rows + : (cached_inputs ? cached_inputs->raw_kv_rows : nullptr); + if (lane.current_raw_out) { + *lane.current_raw_out = kv; + } + if (!lane.write_enabled) { + // Inactive/padding lanes intentionally have no cache mutation. + } else if (raw_kv_rows) { ggml_tensor * kv_f32 = ggml_is_contiguous(kv) ? kv : ggml_cont(ctx, kv); - raw_kv_source = ggml_set_rows(ctx, lc.raw_kv, kv_f32, raw_kv_rows); + raw_kv_source = ggml_set_rows(ctx, lane.raw_kv, kv_f32, raw_kv_rows); ggml_build_forward_expand(gf, raw_kv_source); } else { // The attention graph consumes the whole current ubatch directly. @@ -1681,8 +1778,8 @@ static ggml_tensor * build_mla_attention( ggml_tensor * kv_row = ggml_view_2d( ctx, kv, head_dim, 1, kv->nb[1], (size_t)ti * kv->nb[1]); ggml_tensor * kv_slot = ggml_view_2d( - ctx, lc.raw_kv, head_dim, 1, lc.raw_kv->nb[1], - (size_t)(pos_ti % w.n_swa) * lc.raw_kv->nb[1]); + ctx, lane.raw_kv, head_dim, 1, lane.raw_kv->nb[1], + (size_t)(pos_ti % w.n_swa) * lane.raw_kv->nb[1]); ggml_build_forward_expand(gf, ggml_cpy(ctx, ggml_cast(ctx, kv_row, GGML_TYPE_F16), kv_slot)); } } @@ -1691,15 +1788,15 @@ static ggml_tensor * build_mla_attention( // ── Learned compression update ────────────────────────────────── ggml_tensor * cur_last = ggml_view_2d( ctx, cur, n_embd, 1, cur->nb[1], (size_t)(n_tokens - 1) * cur->nb[1]); - ggml_tensor * comp_kv_source = lc.comp_kv; - if (ratio > 0 && L.attn_compressor_kv) { + ggml_tensor * comp_kv_source = lane.comp_kv; + if (lane.write_enabled && ratio > 0 && L.attn_compressor_kv) { build_compressor_step(ctx, gf, cur_last, L.attn_compressor_ape, L.attn_compressor_kv, L.attn_compressor_gate, L.attn_compressor_norm, - lc.attn_compressor, - lc.comp_kv, + *lane.attn_compressor, + lane.comp_kv, ratio, head_dim, token_pos, @@ -1712,7 +1809,8 @@ static ggml_tensor * build_mla_attention( (int)w.rope_orig_ctx, cached_inputs ? cached_inputs->attn_ape_row : nullptr, cached_inputs ? cached_inputs->attn_state_rows : nullptr, - cached_inputs ? cached_inputs->attn_comp_rows : nullptr, + lane.comp_write_rows ? lane.comp_write_rows : + (cached_inputs ? cached_inputs->attn_comp_rows : nullptr), cached_inputs ? cached_inputs->attn_comp_pos : nullptr, i64_array_inputs, i32_array_inputs, @@ -1720,15 +1818,20 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->flush_rows : nullptr, (causal_batch || fused_causal) ? cur : nullptr, n_tokens, - kv_start); + kv_start, + false, + lane.current_comp_out, + gathered_history); } - ggml_tensor * index_comp_kv_source = lc.index_comp_kv; - if (ratio == 4 && L.indexer_compressor_kv) { - build_indexer_compressor_step(ctx, gf, cur_last, w, L, lc, token_pos, + ggml_tensor * index_comp_kv_source = lane.index_comp_kv; + if (lane.write_enabled && ratio == 4 && L.indexer_compressor_kv) { + build_indexer_compressor_step(ctx, gf, cur_last, w, L, + *lane.indexer_compressor, lane.index_comp_kv, token_pos, cached_inputs ? cached_inputs->index_ape_row : nullptr, cached_inputs ? cached_inputs->index_state_rows : nullptr, - cached_inputs ? cached_inputs->index_comp_rows : nullptr, + lane.index_comp_write_rows ? lane.index_comp_write_rows : + (cached_inputs ? cached_inputs->index_comp_rows : nullptr), cached_inputs ? cached_inputs->index_comp_pos : nullptr, i64_array_inputs, i32_array_inputs, @@ -1738,7 +1841,9 @@ static ggml_tensor * build_mla_attention( n_tokens, kv_start, attention_impl == - DeepSeek4AttentionImpl::SparseFlash); + DeepSeek4AttentionImpl::SparseFlash, + lane.current_index_comp_out, + gathered_history); } // ── MLA Dot-Product Attention (SWA + compressed KV) ──────────── @@ -1747,20 +1852,45 @@ static ggml_tensor * build_mla_attention( // comp_kv: [head_dim, comp_cap] F16 compressed rows. // n_raw = min(kv_start + n_tokens, n_swa) const bool masked_kv = cached_inputs && cached_inputs->attn_row_mask; - const int n_comp_live = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; + const bool gathered_emits_comp = gathered_history && lane.write_enabled && + ratio > 0 && ((token_pos + 1) % ratio) == 0; + const int n_comp_live = gathered_history + ? lane.n_comp_history + (gathered_emits_comp ? 1 : 0) : lane.n_comp_live; + ggml_tensor * comp_history_source = gathered_history + ? lane.comp_history : comp_kv_source; + ggml_tensor * index_comp_history_source = gathered_history + ? lane.index_comp_history : index_comp_kv_source; + if (gathered_emits_comp) { + // Gather through the post-update source to make the compressor write a + // graph dependency. Reading the F16 cache row preserves ordinary q=1 + // rounding at a boundary instead of feeding the transient F32 pool. + ggml_tensor * emitted = ggml_get_rows( + ctx, comp_kv_source, lane.comp_read_rows); + comp_history_source = lane.comp_history + ? ggml_concat(ctx, lane.comp_history, emitted, 1) : emitted; + if (ratio == 4) { + ggml_tensor * index_emitted = ggml_get_rows( + ctx, index_comp_kv_source, lane.index_comp_read_rows); + index_comp_history_source = lane.index_comp_history + ? ggml_concat(ctx, lane.index_comp_history, index_emitted, 1) + : index_emitted; + } + } ggml_tensor * indexer_topk = nullptr; if (attention_impl == DeepSeek4AttentionImpl::SparseFlash && ratio == 4 && f32_array_inputs) { - const int n_index_comp = ds4_comp_rows_used( - lc.index_comp_kv, lc.n_index_comp, 4, token_pos); + const int n_index_comp = gathered_history + ? lane.n_index_comp_history + (gathered_emits_comp ? 1 : 0) + : lane.n_index_comp_live; indexer_topk = build_indexer_topk( - ctx, qr, cur, w, L, index_comp_kv_source, + ctx, qr, cur, w, L, index_comp_history_source, n_index_comp, kv_start, n_tokens, rope_pos, i32_array_inputs); } // Stable path reads the full physical ring (masking not-yet-written slots) // and a padded compressed-row span; the plain path reads only valid rows. - const int n_raw = masked_kv ? w.n_swa + const int n_raw = gathered_history ? lane.n_raw_history + n_tokens + : masked_kv ? w.n_swa : layer_major_batch ? n_prior_rows + n_tokens : std::min(kv_start + n_tokens, w.n_swa); const int n_comp_attn = masked_kv ? cached_inputs->padded_comp : n_comp_live; @@ -1772,13 +1902,24 @@ static ggml_tensor * build_mla_attention( // write and see the previous contents of the raw KV slot. auto raw_kv_view = [&](int row, int count) -> ggml_tensor * { ggml_tensor * view = ggml_view_2d( - ctx, lc.raw_kv, head_dim, count, lc.raw_kv->nb[1], - (size_t)row * lc.raw_kv->nb[1]); + ctx, lane.raw_kv, head_dim, count, lane.raw_kv->nb[1], + (size_t)row * lane.raw_kv->nb[1]); return ds4_cast_if_needed(ctx, view, GGML_TYPE_F32); }; ggml_tensor * kv_attn = nullptr; - if (masked_kv) { + if (gathered_history) { + ggml_tensor * current = ds4_cast_if_needed(ctx, kv, GGML_TYPE_F32); + if (lane.n_raw_history > 0 && lane.raw_history) { + ggml_tensor * history = ggml_view_2d( + ctx, lane.raw_history, head_dim, lane.n_raw_history, + lane.raw_history->nb[1], 0); + history = ds4_cast_if_needed(ctx, history, GGML_TYPE_F32); + kv_attn = ggml_concat(ctx, history, current, 1); + } else { + kv_attn = current; + } + } else if (masked_kv) { // Fused stable-KV path: read the full physical ring; rows not yet // written are masked to -1e30 in the score matrix (exact 0 after // softmax). Only the fused decode graph sets attn_row_mask. Read @@ -1802,7 +1943,7 @@ static ggml_tensor * build_mla_attention( // KV at its runtime row in an F32 snapshot instead. The tokenwise // prefill helper takes the same branch and row ordering. ggml_tensor * ring = ggml_view_2d( - ctx, lc.raw_kv, head_dim, w.n_swa, lc.raw_kv->nb[1], 0); + ctx, lane.raw_kv, head_dim, w.n_swa, lane.raw_kv->nb[1], 0); ring = ds4_cast_if_needed(ctx, ring, GGML_TYPE_F32); kv_attn = ggml_set_rows(ctx, ring, cur_kv, raw_kv_rows); ggml_build_forward_expand(gf, kv_attn); @@ -1822,8 +1963,8 @@ static ggml_tensor * build_mla_attention( } else { kv_attn = raw_kv_view(0, n_raw); } - if (n_comp_attn > 0 && comp_kv_source) { - ggml_tensor * comp = ggml_view_2d(ctx, comp_kv_source, head_dim, n_comp_attn, comp_kv_source->nb[1], 0); + if (n_comp_attn > 0 && comp_history_source) { + ggml_tensor * comp = ggml_view_2d(ctx, comp_history_source, head_dim, n_comp_attn, comp_history_source->nb[1], 0); comp = ds4_cast_if_needed(ctx, comp, GGML_TYPE_F32); kv_attn = ggml_concat(ctx, kv_attn, comp, 1); } @@ -1869,7 +2010,9 @@ static ggml_tensor * build_mla_attention( } } if (n_comp_attn > 0) { - const int vis = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, pos_i); + const int vis = gathered_history ? n_comp_attn + : ds4_comp_rows_used( + lane.comp_kv, lane.n_comp_committed, ratio, pos_i); for (int c = vis; c < n_comp_attn; c++) col[n_raw + c] = -1e30f; } } @@ -1893,8 +2036,9 @@ static ggml_tensor * build_mla_attention( if (pos_r > pos_i) col[r] = -1e30f; } if (n_comp_attn > 0) { - const int visible = ds4_comp_rows_used( - lc.comp_kv, lc.n_comp, ratio, pos_i); + const int visible = gathered_history ? n_comp_attn + : ds4_comp_rows_used( + lane.comp_kv, lane.n_comp_committed, ratio, pos_i); for (int c = visible; c < n_comp_attn; ++c) { col[n_raw + c] = -1e30f; } @@ -1939,7 +2083,7 @@ static ggml_tensor * build_mla_attention( const int first_count = DS4_NUMERICAL_PREFILL_BAND; const int second_count = n_tokens - first_count; const int first_comp = ratio > 0 - ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, + ? ds4_comp_rows_used(lane.comp_kv, lane.n_comp_committed, ratio, kv_start + first_count - 1) : 0; const int second_comp = n_comp_live; @@ -1983,7 +2127,7 @@ static ggml_tensor * build_mla_attention( } if (comp_count > 0) { const int visible = ds4_comp_rows_used( - lc.comp_kv, lc.n_comp, ratio, pos_i); + lane.comp_kv, lane.n_comp_committed, ratio, pos_i); for (int c = visible; c < comp_count; ++c) { col[raw_count + c] = -1e30f; } @@ -2200,6 +2344,34 @@ static ggml_tensor * build_mla_attention( return out; } +// Legacy contiguous-cache adapter. Both decode and the consecutive q>1 +// verifier/prefill path enter through here, so their graph construction order +// remains exactly the order in build_mla_attention_lane_core. +static ggml_tensor * build_mla_attention( + ggml_context * ctx, + ggml_cgraph * gf, + ggml_tensor * cur, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + DeepSeek4LayerCache & lc, + int layer_idx, + int kv_start, + int n_tokens, + const DeepSeek4AttentionGraphInputs * cached_inputs, + std::vector & i32_inputs, + std::vector & i32_array_inputs, + std::vector & i64_array_inputs, + std::vector * f32_array_inputs = nullptr, + DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit) { + const int ratio = w.compress_ratios[layer_idx]; + DeepSeek4MlaLaneBindings lane = deepseek4_contiguous_lane_bindings( + lc, ratio, kv_start + n_tokens - 1); + return build_mla_attention_lane_core( + ctx, gf, cur, w, L, lane, layer_idx, kv_start, n_tokens, + cached_inputs, i32_inputs, i32_array_inputs, i64_array_inputs, + f32_array_inputs, attention_impl); +} + struct DeepSeek4CachedDecodeHcPreGraph { const ggml_context * owner_ctx = nullptr; ggml_backend_t backend = nullptr; @@ -4600,6 +4772,19 @@ struct Ds4FusedVerifyCache { std::array slots; struct Extra { + struct PagedLane { + ggml_tensor * pos = nullptr; + ggml_tensor * neg_pos = nullptr; + ggml_tensor * raw_gather = nullptr; + ggml_tensor * comp_gather = nullptr; + ggml_tensor * index_gather = nullptr; + ggml_tensor * raw_write = nullptr; + ggml_tensor * comp_write = nullptr; + ggml_tensor * comp_read = nullptr; + ggml_tensor * ape = nullptr; + ggml_tensor * state_row = nullptr; + ggml_tensor * comp_pos = nullptr; + }; ggml_tensor * pos_q = nullptr; // i32 [q] ggml_tensor * neg_q = nullptr; // i32 [q] ggml_tensor * rawrows = nullptr; // i64 [1,q] @@ -4612,6 +4797,7 @@ struct Ds4FusedVerifyCache { // Reused host staging for the context-sized additive attention mask. // Keeping it per slot removes one allocation from every verify step. std::vector mask_values; + std::vector paged; // [layer*q], paged mode only int q = 0; void reset() { *this = Extra{}; } @@ -6586,6 +6772,222 @@ static bool initialize_layer_range_cache( runtime.owns_output = owns_output; return true; } + +struct Ds4PagedGatheredRuntime { + DeepSeek4LayerRangeCache model; + const MoeHybridStorage * hybrid_identity = nullptr; + ggml_backend_t hybrid_cpu_backend = nullptr; + ggml_backend_t hybrid_cold_backend = nullptr; +}; + +void deepseek4_release_paged_gathered_runtime(DeepSeek4PagedCache & cache) { + delete static_cast(cache.gathered_runtime); + cache.gathered_runtime = nullptr; +} + +bool deepseek4_paged_gathered_step( + ggml_backend_t backend, int device, const DeepSeek4Weights & w, + DeepSeek4PagedCache & cache, const float * embeddings, + const int32_t * token_ids, const int64_t * positions, + const int32_t * slots, uint32_t lanes, const int32_t * block_tables, + uint32_t block_table_stride, std::vector & out_logits, + std::vector & out_argmax, MoeHybridStorage * hybrid, + MoeHybridRoutingStats * routing_stats) { + if (!backend || !embeddings || !positions || !slots || !block_tables || + lanes < 1 || lanes > 16 || cache.layers.size() != (size_t) w.n_layer || + block_table_stride < cache.plan.max_blocks_per_sequence) return false; + for (uint32_t lane = 0; lane < lanes; ++lane) { + if (slots[lane] < 0) continue; + if ((uint32_t) slots[lane] >= cache.plan.slots || positions[lane] < 0 || + (uint64_t) positions[lane] >= cache.plan.max_ctx || + positions[lane] > INT32_MAX) return false; + for (uint32_t prior = 0; prior < lane; ++prior) + if (slots[prior] == slots[lane]) return false; + } + // Active logical pages must have valid, exclusive physical ownership. + // Aliasing would make one lane's compressor write mutate another lane's + // chronological history and is therefore malformed addressing. + std::vector physical_owner(cache.plan.physical_blocks, -1); + for (uint32_t lane = 0; lane < lanes; ++lane) { + if (slots[lane] < 0) continue; + const uint64_t last_block = (uint64_t) positions[lane] / DS4_PAGE_TOKENS; + if (last_block >= block_table_stride) return false; + for (uint64_t logical = 0; logical <= last_block; ++logical) { + const int32_t physical = block_tables[(size_t) lane * block_table_stride + logical]; + if (physical < 0 || (uint32_t) physical >= cache.plan.physical_blocks || + physical_owner[(size_t) physical] >= 0) return false; + physical_owner[(size_t) physical] = (int32_t) lane; + } + } + if (hybrid) { + for (size_t il = 0; il < hybrid->layers.size(); ++il) { + if (hybrid->layers[il].cache_slots > 0) { + std::fprintf(stderr, + "[deepseek4-paged] layer %zu uses mutable expert-cache " + "placement, which gathered serving cannot capture\n", il); + return false; + } + } + } + auto * rt = static_cast(cache.gathered_runtime); + if (!rt) { + rt = new (std::nothrow) Ds4PagedGatheredRuntime; + if (!rt) return false; + cache.gathered_runtime = rt; + } + if (rt->hybrid_identity != hybrid || + rt->hybrid_cpu_backend != (hybrid ? hybrid->cpu_backend : nullptr) || + rt->hybrid_cold_backend != (hybrid ? hybrid->cold_backend : nullptr)) { + rt->model.fused_verify_graph_cache.destroy(); + rt->hybrid_identity = hybrid; + rt->hybrid_cpu_backend = hybrid ? hybrid->cpu_backend : nullptr; + rt->hybrid_cold_backend = hybrid ? hybrid->cold_backend : nullptr; + } + if (!rt->model.matches(w, backend, device, 0, w.n_layer, true) && + !initialize_layer_range_cache(rt->model, backend, device, w, + 0, w.n_layer, true)) { + std::fprintf(stderr, + "[deepseek4-paged] failed to initialize whole-model graph cache\n"); + return false; + } + + std::vector> prepared((size_t) w.n_layer); + std::vector key = {0x5041474544LL, (int64_t) lanes, + token_ids ? 1 : 0, hybrid ? 1 : 0}; + for (uint32_t lane = 0; lane < lanes; ++lane) key.push_back(slots[lane]); + for (int il = 0; il < w.n_layer; ++il) { + const uint32_t ratio = cache.layers[(size_t) il].ratio; + if (!prepare_deepseek4_gathered_lane_rows( + slots, positions, lanes, block_tables, block_table_stride, + cache.plan.physical_blocks, ratio, prepared[(size_t) il])) return false; + for (const auto & row : prepared[(size_t) il]) { + key.push_back((int64_t) row.raw_history.size()); + key.push_back((int64_t) row.compressed_history.size()); + key.push_back(row.slot < 0 ? -1 : + (ratio ? row.position % ratio : row.position % DS4_PAGE_TOKENS)); + } + } + + auto & vc = rt->model.fused_verify_graph_cache; + auto & mc = rt->model.fused_decode_graph_cache; + if (vc.owner_ctx != w.ctx || vc.backend != backend || + vc.peer_backend != (hybrid ? hybrid->cold_backend : nullptr)) { + vc.destroy(); vc.owner_ctx = w.ctx; vc.backend = backend; + vc.peer_backend = hybrid ? hybrid->cold_backend : nullptr; + } + if (mc.owner_ctx != w.ctx || mc.backend != backend) { + mc.destroy(); mc.owner_ctx = w.ctx; mc.backend = backend; + } + if (!ds4_fused_ensure_fn_mirrors(mc, backend, w, + rt->model.hc_layer_weights, rt->model.hc_output_weights)) return false; + vc.counter++; + DeepSeek4FusedDecodeGraph * fg = nullptr; + Ds4FusedVerifyCache::Extra * ex = nullptr; + const size_t slot_limit = hybrid ? ds4_fused_verify_hybrid_slot_limit() + : vc.slots.size(); + for (size_t i = 0; i < slot_limit; ++i) { + if (vc.slots[i].built() && vc.slots[i].shape_key == key) { + fg = &vc.slots[i]; ex = &vc.extra[i]; break; + } + } + if (!fg) { + size_t pick = 0; + for (size_t i = 0; i < slot_limit; ++i) { + if (!vc.slots[i].built()) { pick = i; break; } + if (vc.slots[i].last_use < vc.slots[pick].last_use) pick = i; + } + fg = &vc.slots[pick]; ex = &vc.extra[pick]; + fg->destroy(vc.backend, vc.peer_backend); ex->reset(); + if (!ds4_build_fused_verify_graph( + mc, *fg, *ex, backend, w, cache.prefill_staging, + rt->model.hc_layer_weights, rt->model.hc_output_weights, + rt->model.hash_routing_tables, 0, (int) lanes, + token_ids != nullptr, {}, hybrid, std::move(key), + &cache, &prepared)) { + std::fprintf(stderr, + "[deepseek4-paged] failed to build gathered graph " + "(lanes=%u)\n", lanes); + return false; + } + } + fg->last_use = vc.counter; + ds4_fv_set(fg->inp_embed, embeddings, + sizeof(float) * (size_t) w.n_embd * lanes); + size_t pi = 0; + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) cache.layers[(size_t) il].ratio; + for (uint32_t lane = 0; lane < lanes; ++lane, ++pi) { + const auto & row = prepared[(size_t) il][lane]; + const auto & px = ex->paged[pi]; + const int32_t pos = (int32_t) row.position; + const int32_t neg_pos = -pos; + ds4_fv_set(px.pos, &pos, sizeof(pos)); + ds4_fv_set(px.neg_pos, &neg_pos, sizeof(neg_pos)); + std::vector idx(std::max(row.raw_history.size(), 1), 0); + for (size_t i = 0; i < row.raw_history.size(); ++i) idx[i] = (int32_t) row.raw_history[i]; + ds4_fv_set(px.raw_gather, idx.data(), idx.size() * sizeof(int32_t)); + const int64_t raw_write = std::max(row.raw_scatter, 0); + ds4_fv_set(px.raw_write, &raw_write, sizeof(raw_write)); + if (ratio > 0) { + idx.assign(std::max(row.compressed_history.size(), 1), 0); + for (size_t i = 0; i < row.compressed_history.size(); ++i) + idx[i] = (int32_t) row.compressed_history[i]; + ds4_fv_set(px.comp_gather, idx.data(), idx.size() * sizeof(int32_t)); + if (px.index_gather) + ds4_fv_set(px.index_gather, idx.data(), idx.size() * sizeof(int32_t)); + const int64_t cw = std::max(row.compressed_scatter, 0); + const int32_t cr = (int32_t) cw; + const int32_t ape = pos % ratio; + const int64_t state = ratio == 4 ? 4 + ape : ape; + const int32_t comp_pos = pos + 1 - ratio; + ds4_fv_set(px.comp_write, &cw, sizeof(cw)); + ds4_fv_set(px.comp_read, &cr, sizeof(cr)); + ds4_fv_set(px.ape, &ape, sizeof(ape)); + ds4_fv_set(px.state_row, &state, sizeof(state)); + ds4_fv_set(px.comp_pos, &comp_pos, sizeof(comp_pos)); + } + } + } + if (token_ids) { + for (int il = 0; il < w.n_layer; ++il) { + ggml_tensor * ids = fg->hash_ids[(size_t) il]; if (!ids) continue; + std::vector values((size_t) ids->ne[0] * lanes); + for (uint32_t lane = 0; lane < lanes; ++lane) { + const int32_t * src = hash_routing_row(rt->model.hash_routing_tables[(size_t) il], + slots[lane] < 0 ? 0 : token_ids[lane], + w.n_expert_used); + if (!src) return false; + std::memcpy(values.data() + lane * ids->ne[0], src, + (size_t) ids->ne[0] * sizeof(int32_t)); + } + ds4_fv_set(ids, values.data(), values.size() * sizeof(int32_t)); + } + } + const enum ggml_status status = fg->sched + ? ggml_backend_sched_graph_compute(fg->sched, fg->sg.gf) + : ggml_backend_graph_compute(backend, fg->sg.gf); + if (status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[deepseek4-paged] gathered graph compute failed: status=%d\n", + (int) status); + return false; + } + ds4_fused_consume_route_diagnostics(*fg, hybrid, routing_stats, slots); + out_logits.resize((size_t) w.n_vocab * lanes); + out_argmax.resize(lanes); + ggml_backend_tensor_get(fg->logits, out_logits.data(), 0, + out_logits.size() * sizeof(float)); + ggml_backend_tensor_get(ex->argmax, out_argmax.data(), 0, + out_argmax.size() * sizeof(int32_t)); + for (uint32_t lane = 0; lane < lanes; ++lane) { + if (slots[lane] >= 0) continue; + std::fill_n(out_logits.data() + (size_t) lane * w.n_vocab, + w.n_vocab, 0.0f); + out_argmax[lane] = -1; + } + return true; +} + bool deepseek4_step_layer_range( ggml_backend_t backend, int device, diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index b0e80ec00..b4a3d8f03 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -15,6 +15,7 @@ #pragma once #include +#include #include #include @@ -25,6 +26,8 @@ #include "internal.h" #include "common/layer_split_utils.h" #include "common/prefill_attention_mode.h" +#include "common/paged_kv_pool.h" +#include "deepseek4_paged_cache.h" namespace dflash::common { @@ -291,6 +294,26 @@ struct DeepSeek4Cache { ggml_backend_buffer_t buf = nullptr; }; +struct DeepSeek4PagedLayerCache : DeepSeek4LayerCache { + uint32_t ratio = 0; + uint64_t physical_rows = 0; +}; + +struct DeepSeek4PagedCache { + std::unique_ptr pool; + DeepSeek4PagedCachePlan plan; + ggml_tensor * block_table = nullptr; + ggml_tensor * sequence_lengths = nullptr; + ggml_tensor * active_slot_ids = nullptr; + std::vector layers; + DeepSeek4Cache prefill_staging; + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + // Dedicated bounded gathered-reference graph cache (opaque here because + // its implementation shares the fused verifier's private machinery). + void * gathered_runtime = nullptr; +}; + struct DeepSeek4Snapshot; struct DeepSeek4RawRingSpan { @@ -309,6 +332,9 @@ struct DeepSeek4BackendConfig { int max_ctx = 0; // 0 = auto from SWA + compression capacity int expert_top_k = 0; // 0 = use all model-routed experts bool fused_decode = false; // single-graph GPU decode + bool paged_attention = false; + int max_concurrency = 1; + long long kv_pool_tokens = 0; }; // ─── Function declarations ────────────────────────────────────────────── @@ -334,6 +360,26 @@ bool create_deepseek4_cache(ggml_backend_t backend, DeepSeek4Cache & out); void free_deepseek4_cache(DeepSeek4Cache & c); +bool create_deepseek4_paged_cache(ggml_backend_t backend, + const DeepSeek4Weights & w, + uint32_t slots, uint32_t max_ctx, + uint32_t physical_blocks, + DeepSeek4PagedCache & out); +void reset_deepseek4_paged_slot(DeepSeek4PagedCache & c, uint32_t slot); +void free_deepseek4_paged_cache(DeepSeek4PagedCache & c); +// Exact gathered-reference decode for 1..16 independent lanes. Inputs are +// lane-major; negative slots are inactive padding lanes. `out_logits` is +// [n_vocab, lanes] and `out_argmax` is [lanes]. +bool deepseek4_paged_gathered_step( + ggml_backend_t backend, int device, const DeepSeek4Weights & w, + DeepSeek4PagedCache & cache, const float * embeddings, + const int32_t * token_ids, const int64_t * positions, + const int32_t * slots, uint32_t lanes, const int32_t * block_tables, + uint32_t block_table_stride, std::vector & out_logits, + std::vector & out_argmax, + MoeHybridStorage * moe_hybrid = nullptr, + MoeHybridRoutingStats * routing_stats = nullptr); +void deepseek4_release_paged_gathered_runtime(DeepSeek4PagedCache & cache); void reset_deepseek4_cache(DeepSeek4Cache & c); int deepseek4_previous_raw_ring_spans( int kv_start, diff --git a/server/src/deepseek4/deepseek4_page_layout.h b/server/src/deepseek4/deepseek4_page_layout.h new file mode 100644 index 000000000..b2ee71631 --- /dev/null +++ b/server/src/deepseek4/deepseek4_page_layout.h @@ -0,0 +1,55 @@ +// Host-side address geometry for DeepSeek V4's paged raw and compressed KV. +#pragma once + +#include +#include + +namespace dflash::common { + +inline constexpr uint32_t DS4_PAGE_TOKENS = 128; + +// Raw KV remains slot-indexed: every sequence reuses this 128-row ring. +inline constexpr uint32_t ds4_raw_ring_row(uint64_t logical_token) { + return static_cast(logical_token % DS4_PAGE_TOKENS); +} + +// Number of physically paged compressed rows. Rejects unsupported ratios and +// arithmetic that cannot be represented by the row-index type. +inline bool ds4_compressed_page_capacity(uint64_t physical_blocks, + uint32_t ratio, + uint64_t & rows) { + if (ratio != 4 && ratio != 128) return false; + const uint64_t rows_per_block = DS4_PAGE_TOKENS / ratio; + if (physical_blocks > + std::numeric_limits::max() / rows_per_block) { + return false; + } + rows = physical_blocks * rows_per_block; + return true; +} + +// Computes the destination for a completed compression group. `emitted` is +// false between group boundaries and `row` is left unchanged in that case. +// Physical block IDs need not be contiguous. +inline bool ds4_compressed_page_row(uint64_t logical_token, + uint64_t physical_block, + uint32_t ratio, + uint64_t & row, + bool & emitted) { + if (ratio != 4 && ratio != 128) return false; + emitted = logical_token % ratio == ratio - 1; + if (!emitted) return true; + + const uint64_t rows_per_block = DS4_PAGE_TOKENS / ratio; + if (physical_block > + std::numeric_limits::max() / rows_per_block) { + return false; + } + const uint64_t base = physical_block * rows_per_block; + const uint64_t offset = (logical_token % DS4_PAGE_TOKENS) / ratio; + if (base > std::numeric_limits::max() - offset) return false; + row = base + offset; + return true; +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_paged_cache.cpp b/server/src/deepseek4/deepseek4_paged_cache.cpp new file mode 100644 index 000000000..cdb395dbb --- /dev/null +++ b/server/src/deepseek4/deepseek4_paged_cache.cpp @@ -0,0 +1,219 @@ +#include "deepseek4_paged_cache.h" + +#include "deepseek4_page_layout.h" + +#ifndef DFLASH_DS4_PLAN_ONLY +#include "deepseek4_internal.h" +#endif + +#include +#include +#include + +namespace dflash::common { +namespace { +bool add_mul(uint64_t & dst, uint64_t a, uint64_t b) { + if (a && b > std::numeric_limits::max() / a) return false; + const uint64_t v = a * b; + if (dst > std::numeric_limits::max() - v) return false; + dst += v; + return true; +} +} + +bool prepare_deepseek4_gathered_lane_rows( + const int32_t * slots, const int64_t * positions, uint32_t lanes, + const int32_t * block_tables, uint32_t block_table_stride, + uint32_t physical_blocks, uint32_t ratio, + std::vector & out) { + if (!slots || !positions || !block_tables || !block_table_stride || + !physical_blocks || (ratio != 0 && ratio != 4 && ratio != 128)) { + return false; + } + std::vector prepared(lanes); + for (uint32_t lane = 0; lane < lanes; ++lane) { + auto & rows = prepared[lane]; + rows.slot = slots[lane]; + rows.position = positions[lane]; + if (rows.slot < 0) continue; // Padding must remain entirely passive. + if (rows.position < 0) return false; + const uint64_t pos = static_cast(rows.position); + // The current row is appended in-graph, so retain at most the 127 + // preceding rows that can coexist with it in the 128-row SWA window. + const uint64_t first_raw = pos >= DS4_PAGE_TOKENS + ? pos - DS4_PAGE_TOKENS + 1 : 0; + rows.raw_history.reserve(static_cast(pos - first_raw)); + for (uint64_t p = first_raw; p < pos; ++p) { + rows.raw_history.push_back( + int64_t(rows.slot) * DS4_PAGE_TOKENS + ds4_raw_ring_row(p)); + } + rows.raw_scatter = int64_t(rows.slot) * DS4_PAGE_TOKENS + + ds4_raw_ring_row(pos); + if (!ratio) continue; + + // Every completed group before the current token contributes one + // chronological row. Looking up each logical page (rather than + // assuming contiguous physical pages) is the reference behaviour. + const uint64_t completed = pos / ratio; + rows.compressed_history.reserve(static_cast(completed)); + for (uint64_t group = 0; group < completed; ++group) { + const uint64_t end_token = group * ratio + ratio - 1; + const uint64_t logical_block = end_token / DS4_PAGE_TOKENS; + if (logical_block >= block_table_stride) return false; + const int32_t physical = + block_tables[size_t(lane) * block_table_stride + logical_block]; + if (physical < 0 || uint32_t(physical) >= physical_blocks) return false; + uint64_t row = 0; bool emitted = false; + if (!ds4_compressed_page_row(end_token, uint32_t(physical), ratio, + row, emitted) || !emitted || + row > uint64_t(INT64_MAX)) return false; + rows.compressed_history.push_back(static_cast(row)); + } + const uint64_t logical_block = pos / DS4_PAGE_TOKENS; + if (logical_block >= block_table_stride) return false; + const int32_t physical = + block_tables[size_t(lane) * block_table_stride + logical_block]; + if (physical < 0 || uint32_t(physical) >= physical_blocks) return false; + uint64_t scatter = 0; + if (!ds4_compressed_page_row(pos, uint32_t(physical), ratio, scatter, + rows.compressed_emitted) || + scatter > uint64_t(INT64_MAX)) return false; + if (rows.compressed_emitted) rows.compressed_scatter = int64_t(scatter); + } + out = std::move(prepared); + return true; +} + +bool plan_deepseek4_paged_cache(uint32_t head_dim, uint32_t indexer_head_dim, + uint32_t slots, uint32_t max_ctx, + uint32_t physical_blocks, + const std::vector & ratios, + DeepSeek4PagedCachePlan & out) { + DeepSeek4PagedCachePlan p; + if (!head_dim || !indexer_head_dim || !slots || !max_ctx || + !physical_blocks || ratios.empty() || + physical_blocks > UINT32_MAX / DS4_PAGE_TOKENS) return false; + p.slots = slots; p.max_ctx = max_ctx; p.physical_blocks = physical_blocks; + p.max_blocks_per_sequence = 1 + (max_ctx - 1) / DS4_PAGE_TOKENS; + p.ratios = ratios; + p.physical_rows.resize(ratios.size()); + // block table, lengths, and active IDs, all I32. + if (!add_mul(p.metadata_bytes, p.max_blocks_per_sequence, uint64_t(slots) * 4) || + !add_mul(p.metadata_bytes, slots, 8)) return false; + for (size_t i = 0; i < ratios.size(); ++i) { + const uint32_t r = ratios[i]; + if (r != 0 && r != 4 && r != 128) return false; + if (!add_mul(p.raw_bytes, uint64_t(head_dim) * DS4_PAGE_TOKENS * 2, slots)) return false; + if (!r) continue; + uint64_t rows = 0; + if (!ds4_compressed_page_capacity(physical_blocks, r, rows)) return false; + p.physical_rows[i] = rows; + if (!add_mul(p.compressed_bytes, uint64_t(head_dim) * 2, rows)) return false; + const uint64_t width = uint64_t(head_dim) * (r == 4 ? 2 : 1); + const uint64_t state_rows = r == 4 ? 8 : 128; + if (!add_mul(p.state_bytes, width * state_rows * 8, slots)) return false; // KV + score F32 + if (r == 4) { + if (!add_mul(p.compressed_bytes, uint64_t(indexer_head_dim) * 2, rows)) return false; + if (!add_mul(p.state_bytes, uint64_t(indexer_head_dim) * 2 * 8 * 8, slots)) return false; + } + } + p.total_persistent_bytes = p.metadata_bytes; + if (p.total_persistent_bytes > UINT64_MAX - p.raw_bytes) return false; + p.total_persistent_bytes += p.raw_bytes; + if (p.total_persistent_bytes > UINT64_MAX - p.compressed_bytes) return false; + p.total_persistent_bytes += p.compressed_bytes; + if (p.total_persistent_bytes > UINT64_MAX - p.state_bytes) return false; + p.total_persistent_bytes += p.state_bytes; + out = std::move(p); + return true; +} + +#ifndef DFLASH_DS4_PLAN_ONLY +bool create_deepseek4_paged_cache(ggml_backend_t backend, + const DeepSeek4Weights & w, uint32_t slots, + uint32_t max_ctx, uint32_t physical_blocks, + DeepSeek4PagedCache & out) { + free_deepseek4_paged_cache(out); + DeepSeek4PagedCachePlan plan; + if (!backend || w.n_layer <= 0 || w.compress_ratios.size() != size_t(w.n_layer) || + !plan_deepseek4_paged_cache(w.head_dim, w.n_indexer_head_dim, slots, + max_ctx, physical_blocks, w.compress_ratios, plan)) return false; + try { out.pool = std::make_unique(physical_blocks, slots, DS4_PAGE_TOKENS); } + catch (...) { free_deepseek4_paged_cache(out); return false; } + out.plan = plan; + out.layers.resize(w.n_layer); + ggml_init_params ip{ggml_tensor_overhead() * size_t(w.n_layer * 9 + 8) + 4096, nullptr, true}; + out.ctx = ggml_init(ip); + if (!out.ctx) { free_deepseek4_paged_cache(out); return false; } + out.block_table = ggml_new_tensor_2d(out.ctx, GGML_TYPE_I32, plan.max_blocks_per_sequence, slots); + out.sequence_lengths = ggml_new_tensor_1d(out.ctx, GGML_TYPE_I32, slots); + out.active_slot_ids = ggml_new_tensor_1d(out.ctx, GGML_TYPE_I32, slots); + for (int il = 0; il < w.n_layer; ++il) { + auto & l = out.layers[il]; const uint32_t r = plan.ratios[il]; + l.ratio = r; l.physical_rows = plan.physical_rows[il]; + l.raw_kv = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F16, w.head_dim, DS4_PAGE_TOKENS, slots); + if (!r) continue; + l.comp_kv = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F16, w.head_dim, l.physical_rows); + const int64_t width = int64_t(w.head_dim) * (r == 4 ? 2 : 1), sr = r == 4 ? 8 : 128; + l.attn_compressor.state_kv = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F32, width, sr, slots); + l.attn_compressor.state_score = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F32, width, sr, slots); + if (r == 4) { + l.index_comp_kv = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F16, w.n_indexer_head_dim, l.physical_rows); + const int64_t iw = int64_t(w.n_indexer_head_dim) * 2; + l.indexer_compressor.state_kv = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F32, iw, 8, slots); + l.indexer_compressor.state_score = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F32, iw, 8, slots); + } + } + out.buf = ggml_backend_alloc_ctx_tensors(out.ctx, backend); + if (!out.buf) { free_deepseek4_paged_cache(out); return false; } + ggml_backend_buffer_clear(out.buf, 0); + // One shared, contiguous legacy cache is intentionally retained for prefill. + if (!create_deepseek4_cache(backend, w, int(max_ctx), out.prefill_staging)) { + free_deepseek4_paged_cache(out); return false; + } + return true; +} + +void reset_deepseek4_paged_slot(DeepSeek4PagedCache & c, uint32_t slot) { + if (!c.buf || slot >= c.plan.slots) return; + auto clear_slot = [slot](ggml_tensor * tensor) { + if (!tensor || tensor->ne[2] <= (int64_t) slot) return; + const size_t bytes = tensor->nb[2]; + std::vector zeros(bytes, 0); + ggml_backend_tensor_set(tensor, zeros.data(), (size_t) slot * bytes, + bytes); + }; + for (DeepSeek4PagedLayerCache & layer : c.layers) { + clear_slot(layer.attn_compressor.state_kv); + clear_slot(layer.attn_compressor.state_score); + clear_slot(layer.indexer_compressor.state_kv); + clear_slot(layer.indexer_compressor.state_score); + } + if (c.block_table) { + std::vector empty(c.plan.max_blocks_per_sequence, -1); + ggml_backend_tensor_set(c.block_table, empty.data(), + (size_t) slot * c.block_table->nb[1], + empty.size() * sizeof(int32_t)); + } + const int32_t zero = 0; + const int32_t inactive = -1; + if (c.sequence_lengths) { + ggml_backend_tensor_set(c.sequence_lengths, &zero, + (size_t) slot * sizeof(int32_t), sizeof(zero)); + } + if (c.active_slot_ids) { + ggml_backend_tensor_set(c.active_slot_ids, &inactive, + (size_t) slot * sizeof(int32_t), sizeof(inactive)); + } +} + +void free_deepseek4_paged_cache(DeepSeek4PagedCache & c) { + deepseek4_release_paged_gathered_runtime(c); + free_deepseek4_cache(c.prefill_staging); + if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } + if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } + c.pool.reset(); c.layers.clear(); c.block_table = nullptr; + c.sequence_lengths = nullptr; c.active_slot_ids = nullptr; c.plan = {}; +} +#endif +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_paged_cache.h b/server/src/deepseek4/deepseek4_paged_cache.h new file mode 100644 index 000000000..d8bb3cc83 --- /dev/null +++ b/server/src/deepseek4/deepseek4_paged_cache.h @@ -0,0 +1,63 @@ +#pragma once + +#include +#include +#include + +namespace dflash::common { + +// Pure host-side allocation plan. Byte counts describe tensor payloads (ggml +// alignment/padding is deliberately excluded). +struct DeepSeek4PagedCachePlan { + uint32_t slots = 0; + uint32_t max_ctx = 0; + uint32_t physical_blocks = 0; + uint32_t max_blocks_per_sequence = 0; + uint64_t metadata_bytes = 0; + uint64_t raw_bytes = 0; + uint64_t compressed_bytes = 0; + uint64_t state_bytes = 0; + uint64_t total_persistent_bytes = 0; + std::vector ratios; + std::vector physical_rows; +}; + +// Host metadata for the gathered-reference decode graph. Rows are expressed +// in the flattened persistent tensors: raw rows are [slot, ring-row], while +// compressed rows use the physical page geometry from deepseek4_page_layout.h. +// A negative slot denotes a padding lane and consequently has no scatter rows. +struct DeepSeek4GatheredLaneRows { + int32_t slot = -1; + int64_t position = 0; + std::vector raw_history; + std::vector compressed_history; + int64_t raw_scatter = -1; + int64_t compressed_scatter = -1; + bool compressed_emitted = false; +}; + +// block_tables is lane-major with block_table_stride entries per lane. +// Physical block IDs may be fragmented and are validated against +// physical_blocks. History excludes the current token; compressed history is +// in chronological group order. Returns false for malformed active lanes. +bool prepare_deepseek4_gathered_lane_rows( + const int32_t * slots, + const int64_t * positions, + uint32_t lanes, + const int32_t * block_tables, + uint32_t block_table_stride, + uint32_t physical_blocks, + uint32_t ratio, + std::vector & out); + +// Ratios must contain only 0, 4, or 128. A ratio-zero layer has a raw ring +// but no compressed storage or compressor state. +bool plan_deepseek4_paged_cache(uint32_t head_dim, + uint32_t indexer_head_dim, + uint32_t slots, + uint32_t max_ctx, + uint32_t physical_blocks, + const std::vector & ratios, + DeepSeek4PagedCachePlan & out); + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_seq_engine.cpp b/server/src/deepseek4/deepseek4_seq_engine.cpp new file mode 100644 index 000000000..745046783 --- /dev/null +++ b/server/src/deepseek4/deepseek4_seq_engine.cpp @@ -0,0 +1,272 @@ +#include "deepseek4_seq_engine.h" + +#include "deepseek4_backend.h" +#include "common/sampler.h" + +#include +#include + +namespace dflash::common { + +DeepSeek4SeqEngine::DeepSeek4SeqEngine( + DeepSeek4Backend & backend, PagedKvPool & pool, int max_ctx, + uint32_t table_stride) + : b_(backend), slots_(pool, max_ctx), stride_(table_stride), + host_tables_((size_t)pool.max_sequences() * table_stride, -1) {} + +bool DeepSeek4SeqEngine::token_is_eos(int32_t token) const { + return deepseek4_is_eos_tok(token, b_.w_); +} + +StepPlanLimits DeepSeek4SeqEngine::step_plan_limits( + int decode_rows) const { + // The gathered graph accepts at most sixteen independent lanes and does + // not permit two rows from the same sequence. A prompt therefore advances + // by one token while every live decoder still shares the same weight pass. + const int available = std::max(0, 16 - decode_rows); + return {available, 1, available, 1}; +} + +SeqEngine::AdmitResult DeepSeek4SeqEngine::admit( + uint64_t request_id, const std::vector & prompt, + const SamplerCfg & sampler) { + using AdmitStatus = AdmitResult::Status; + AdmitResult result = slots_.admit( + request_id, prompt, sampler); + if (result.status != AdmitStatus::admitted) return result; + if (result.slot < 0 || result.slot >= slots_.slot_count()) { + result.status = AdmitStatus::failed; + result.error = "invalid DeepSeek4 serving slot"; + return result; + } + std::fill_n(host_tables_.data() + (size_t)result.slot * stride_, + stride_, -1); + reset_deepseek4_paged_slot(b_.paged_cache_, (uint32_t)result.slot); + return result; +} + +bool DeepSeek4SeqEngine::set_block(int slot, int logical, int32_t physical) { + if (slot < 0 || slot >= slots_.slot_count() || logical < 0 || + (uint32_t)logical >= stride_ || !b_.paged_cache_.block_table) { + return false; + } + host_tables_[(size_t)slot * stride_ + (size_t)logical] = physical; + ggml_tensor * table = b_.paged_cache_.block_table; + ggml_backend_tensor_set( + table, &physical, + (size_t)slot * table->nb[1] + + (size_t)logical * sizeof(int32_t), + sizeof(physical)); + return true; +} + +void DeepSeek4SeqEngine::fail_prefill( + int slot, std::vector & outputs, + const std::string & error) { + std::fprintf(stderr, "[deepseek4-parallel] prefill slot %d: %s\n", + slot, error.c_str()); + PrefillOutput out; + out.slot = slot; + out.status = PrefillOutput::Status::failed; + out.error = error; + outputs.push_back(std::move(out)); +} + +SeqEngine::StepResult DeepSeek4SeqEngine::step(const StepPlan & plan) { + StepResult result; + const std::vector & inputs = plan.decode; + const int n_slots = slots_.slot_count(); + + auto fail_step = [&](const std::string & error) { + result.decode.clear(); + result.prefills.clear(); + result.error = error; + return std::move(result); + }; + + if ((int)inputs.size() != slots_.decoding_count()) { + return fail_step("decode plan does not cover every live DeepSeek4 slot"); + } + std::vector decode_seen((size_t)n_slots, 0); + for (const StepInput & input : inputs) { + if (input.slot < 0 || input.slot >= n_slots || input.token < 0 || + decode_seen[(size_t)input.slot] || + !slots_.is_active(input.slot) || + slots_.is_prefilling(input.slot)) { + return fail_step("invalid or duplicate DeepSeek4 decode row"); + } + decode_seen[(size_t)input.slot] = 1; + } + + const StepPlanLimits limits = step_plan_limits((int)inputs.size()); + if ((int)plan.prefills.size() > limits.max_prefill_sequences) { + return fail_step("DeepSeek4 step exceeds the sixteen-lane graph"); + } + std::vector prefill_seen((size_t)n_slots, 0); + for (const PrefillSlice & slice : plan.prefills) { + if (slice.slot < 0 || slice.slot >= n_slots || + slice.max_tokens != 1 || prefill_seen[(size_t)slice.slot] || + decode_seen[(size_t)slice.slot] || + !slots_.is_prefilling(slice.slot)) { + return fail_step("invalid or duplicate DeepSeek4 prefill row"); + } + prefill_seen[(size_t)slice.slot] = 1; + } + if (inputs.empty() && plan.prefills.empty()) return result; + + std::vector lane_tokens; + std::vector lane_positions; + std::vector lane_slots; + std::vector decode_lanes; + lane_tokens.reserve(inputs.size() + plan.prefills.size()); + lane_positions.reserve(inputs.size() + plan.prefills.size()); + lane_slots.reserve(inputs.size() + plan.prefills.size()); + decode_lanes.reserve(inputs.size()); + result.decode.reserve(inputs.size()); + result.prefills.reserve(plan.prefills.size()); + + for (const StepInput & input : inputs) { + DecodeOutput out; + out.slot = input.slot; + const SeqSlotManager::StepAppend append = + slots_.append_token(input.slot, input.token); + if (!append.ok) { + out.failed = true; + out.error = append.busy + ? "paged KV pool exhausted during DeepSeek4 decode; raise " + "--kv-pool-tokens or lower --max-ctx/--max-concurrency" + : "DeepSeek4 decode K/V append failed"; + decode_lanes.push_back(-1); + result.decode.push_back(std::move(out)); + continue; + } + if (append.new_block >= 0 && + !set_block(input.slot, append.new_block_index, + append.new_block)) { + out.failed = true; + out.error = "DeepSeek4 decode block-table update failed"; + decode_lanes.push_back(-1); + result.decode.push_back(std::move(out)); + continue; + } + decode_lanes.push_back((int)lane_tokens.size()); + lane_tokens.push_back(input.token); + lane_positions.push_back(append.position); + lane_slots.push_back(input.slot); + result.decode.push_back(std::move(out)); + } + + struct PrefillLane { + int slot = -1; + int lane = -1; + bool commit = false; + }; + std::vector prefill_lanes; + prefill_lanes.reserve(plan.prefills.size()); + for (const PrefillSlice & slice : plan.prefills) { + SeqSlotManager::PrefillChunk chunk = + slots_.append_prefill(slice.slot, 1); + if (!chunk.ok || chunk.rows.size() != 1) { + fail_prefill(slice.slot, result.prefills, + "DeepSeek4 prefill K/V append failed"); + continue; + } + if (!chunk.new_blocks.empty() && + !set_block(slice.slot, chunk.first_new_block, + chunk.new_blocks.front())) { + fail_prefill(slice.slot, result.prefills, + "DeepSeek4 prefill block-table update failed"); + continue; + } + const SeqSlot & slot = slots_.slot(slice.slot); + const bool commit = slot.cur_pos == (int)slot.prompt.size(); + prefill_lanes.push_back( + {slice.slot, (int)lane_tokens.size(), commit}); + lane_tokens.push_back(slot.prompt[(size_t)slot.cur_pos - 1]); + lane_positions.push_back(slot.cur_pos - 1); + lane_slots.push_back(slice.slot); + } + + if (lane_tokens.empty()) return result; + if (lane_tokens.size() > 16) { + return fail_step("DeepSeek4 gathered step exceeds sixteen lanes"); + } + + std::vector embeddings( + (size_t)b_.w_.n_embd * lane_tokens.size()); + if (!b_.w_.embedder.embed(lane_tokens.data(), (int)lane_tokens.size(), + embeddings.data())) { + return fail_step("DeepSeek4 token embedding failed"); + } + + std::vector compact_tables( + lane_tokens.size() * stride_, -1); + for (size_t lane = 0; lane < lane_slots.size(); ++lane) { + std::copy_n( + host_tables_.data() + (size_t)lane_slots[lane] * stride_, + stride_, compact_tables.data() + lane * stride_); + } + + std::vector lengths((size_t)n_slots, 0); + std::vector active((size_t)n_slots, -1); + for (size_t lane = 0; lane < lane_slots.size(); ++lane) { + lengths[(size_t)lane_slots[lane]] = + (int32_t)lane_positions[lane] + 1; + active[lane] = lane_slots[lane]; + } + ggml_backend_tensor_set( + b_.paged_cache_.sequence_lengths, lengths.data(), 0, + lengths.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + b_.paged_cache_.active_slot_ids, active.data(), 0, + active.size() * sizeof(int32_t)); + + std::vector logits; + std::vector argmax; + if (!deepseek4_paged_gathered_step( + b_.backend_, b_.cfg_.device.gpu, b_.w_, b_.paged_cache_, + embeddings.data(), lane_tokens.data(), lane_positions.data(), + lane_slots.data(), (uint32_t)lane_tokens.size(), + compact_tables.data(), stride_, logits, argmax, + b_.moe_hybrid_.get(), b_.routing_stats_.get())) { + return fail_step("DeepSeek4 gathered paged graph failed"); + } + + auto sample_lane = [&](int slot_id, int lane) { + SeqSlot & slot = slots_.slot(slot_id); + if (!slot.sampler.needs_logit_processing()) { + return argmax[(size_t)lane]; + } + return sample_logits( + logits.data() + (size_t)lane * b_.w_.n_vocab, + b_.w_.n_vocab, slot.sampler, slot.sample_history, slot.rng); + }; + + for (size_t i = 0; i < inputs.size(); ++i) { + const int lane = decode_lanes[i]; + if (lane < 0) continue; + DecodeOutput & out = result.decode[i]; + slots_.commit_step(out.slot); + out.token = sample_lane(out.slot, lane); + } + for (const PrefillLane & prefill : prefill_lanes) { + PrefillOutput out; + out.slot = prefill.slot; + if (prefill.commit) { + out.status = PrefillOutput::Status::completed; + out.token = sample_lane(prefill.slot, prefill.lane); + slots_.commit_prefill(prefill.slot); + } + result.prefills.push_back(std::move(out)); + } + return result; +} + +void DeepSeek4SeqEngine::retire(int slot) { + if (!slots_.is_active(slot)) return; + slots_.retire(slot); + reset_deepseek4_paged_slot(b_.paged_cache_, (uint32_t)slot); + std::fill_n(host_tables_.data() + (size_t)slot * stride_, stride_, -1); +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_seq_engine.h b/server/src/deepseek4/deepseek4_seq_engine.h new file mode 100644 index 000000000..9eeadfd81 --- /dev/null +++ b/server/src/deepseek4/deepseek4_seq_engine.h @@ -0,0 +1,41 @@ +#pragma once + +#include "common/concurrency/seq_engine.h" +#include "common/concurrency/seq_slot_manager.h" + +#include +#include + +namespace dflash::common { + +class DeepSeek4Backend; + +// Exact concurrent serving path for DeepSeek4. Model state remains in +// DeepSeek4PagedCache; this class owns only scheduler-facing slot state and +// the host mirror of the model's block table. +class DeepSeek4SeqEngine final : public SeqEngine { +public: + DeepSeek4SeqEngine(DeepSeek4Backend & backend, PagedKvPool & pool, + int max_ctx, uint32_t table_stride); + + int slot_count() const override { return slots_.slot_count(); } + int max_context() const override { return slots_.max_context(); } + AdmitResult admit(uint64_t request_id, const std::vector & prompt, + const SamplerCfg & sampler) override; + StepResult step(const StepPlan & plan) override; + StepPlanLimits step_plan_limits(int decode_rows) const override; + void retire(int slot) override; + bool token_is_eos(int32_t token) const override; + +private: + bool set_block(int slot, int logical, int32_t physical); + void fail_prefill(int slot, std::vector & outputs, + const std::string & error); + + DeepSeek4Backend & b_; + SeqSlotManager slots_; + uint32_t stride_ = 0; + std::vector host_tables_; +}; + +} // namespace dflash::common diff --git a/server/src/server/api_types.h b/server/src/server/api_types.h index 9fa187997..da970080f 100644 --- a/server/src/server/api_types.h +++ b/server/src/server/api_types.h @@ -5,4 +5,16 @@ namespace dflash::common { enum class ApiFormat { OPENAI_CHAT, ANTHROPIC, RESPONSES, COMPLETIONS }; +// Log/status name of a format — shared by the request-tracing logs of the +// classic worker loop and the concurrent scheduler. +inline const char * api_format_name(ApiFormat format) { + switch (format) { + case ApiFormat::OPENAI_CHAT: return "chat"; + case ApiFormat::ANTHROPIC: return "anthropic"; + case ApiFormat::RESPONSES: return "responses"; + case ApiFormat::COMPLETIONS: return "completions"; + default: return "unknown"; + } +} + } // namespace dflash::common diff --git a/server/src/server/client_send_buffer.h b/server/src/server/client_send_buffer.h new file mode 100644 index 000000000..b4debe9d8 --- /dev/null +++ b/server/src/server/client_send_buffer.h @@ -0,0 +1,139 @@ +// ClientSendBuffer — buffered non-blocking writer for one client socket. +// +// The concurrent scheduler must never block on a client: a stalled SSE +// reader would head-of-line-block every co-scheduled request's decode. +// Token chunks and final responses append here; the scheduler drains the +// buffer with non-blocking sends once per decode iteration and consults +// the stall policy (no drain progress for a deadline, or a byte cap) to +// decide when a dead reader should be dropped. +// +// The fd is owned by the caller and must already be non-blocking (the HTTP +// server sets client sockets non-blocking at enqueue time). + +#pragma once + +#include "socket_handle.h" + +#if defined(_WIN32) +#if !defined(NOMINMAX) +#define NOMINMAX +#endif +#if !defined(WIN32_LEAN_AND_MEAN) +#define WIN32_LEAN_AND_MEAN +#endif +#include +#else +#include +#include +#include +#endif + +#include +#include +#include +#include +#include + +namespace dflash::common { + +class ClientSendBuffer { +public: + void append(std::string_view bytes) { buf_.append(bytes); } + + bool empty() const { return off_ == buf_.size(); } + size_t pending() const { return buf_.size() - off_; } + + // Drain as much as the socket accepts right now. Never blocks. Returns + // false on a hard socket error (peer gone). Records drain progress for + // the stall policy and compacts the drained prefix. + bool flush(SocketHandle fd) { + const size_t off_before = off_; + while (off_ < buf_.size()) { + const auto n = send_some(fd, buf_.data() + off_, + buf_.size() - off_); + if (n > 0) { off_ += (size_t)n; continue; } + if (n < 0) { + const int error = send_error(); + if (send_would_block(error)) break; + if (send_was_interrupted(error)) continue; + } + return false; + } + const bool made_progress = off_ != off_before; + if (off_ == buf_.size()) { + buf_.clear(); + off_ = 0; + } else if (off_ > kCompactAt) { + buf_.erase(0, off_); + off_ = 0; + } + if (made_progress) { + last_progress_ = std::chrono::steady_clock::now(); + } + return true; + } + + // True when the reader should be dropped: bytes are pending and either + // the buffer exceeds `cap` or nothing drained since `stall` ago. The + // deadline does NOT reset while the reader makes no progress — a + // trickling reader is bounded by the cap instead. + bool should_drop(std::chrono::steady_clock::time_point now, + std::chrono::seconds stall, size_t cap) const { + if (empty()) return false; + return pending() > cap || now - last_progress_ > stall; + } + + // Start (or restart) the stall clock, e.g. at admission. + void mark_progress(std::chrono::steady_clock::time_point now) { + last_progress_ = now; + } + +private: + static constexpr size_t kCompactAt = 64u << 10; + +#if defined(_WIN32) + static int send_some(SocketHandle fd, const char * data, size_t len) { + const int chunk = static_cast((std::min)( + len, static_cast((std::numeric_limits::max)()))); + return ::send(fd, data, chunk, 0); + } + + static int send_error() { + return WSAGetLastError(); + } + + static bool send_would_block(int error) { + return error == WSAEWOULDBLOCK; + } + + static bool send_was_interrupted(int error) { + return error == WSAEINTR; + } +#else + static ssize_t send_some(SocketHandle fd, const char * data, size_t len) { +#if defined(MSG_NOSIGNAL) + return ::send(fd, data, len, MSG_NOSIGNAL); +#else + return ::send(fd, data, len, 0); +#endif + } + + static int send_error() { + return errno; + } + + static bool send_would_block(int error) { + return error == EAGAIN || error == EWOULDBLOCK; + } + + static bool send_was_interrupted(int error) { + return error == EINTR; + } +#endif + + std::string buf_; + size_t off_ = 0; + std::chrono::steady_clock::time_point last_progress_{}; +}; + +} // namespace dflash::common diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 2cad0e142..4859f5db1 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -449,18 +449,6 @@ static std::string generate_id(const char * prefix) { return buf; } -// Logging helpers shared by route_request() / worker_loop(). Kept static -// (file-scope) so they don't leak into the public ABI; the chat lifecycle -// logs that use them are part of #270's request-tracing instrumentation. -static const char * api_format_name(ApiFormat format) { - switch (format) { - case ApiFormat::OPENAI_CHAT: return "chat"; - case ApiFormat::ANTHROPIC: return "anthropic"; - case ApiFormat::RESPONSES: return "responses"; - default: return "unknown"; - } -} - static size_t json_array_size(const json & value) { return value.is_array() ? value.size() : 0; } @@ -739,6 +727,9 @@ json build_props_body(const ServerConfig & config, // (dflash/scripts/bench_http_capability.py) read // /props.runtime wholesale into result.json.server_info. {"chunk", config.chunk}, + {"continuous_batching", { + {"admission_coalesce_ms", config.admission_coalesce_ms}, + }}, // Device placement strings (e.g. "auto:0", "cuda:0"). Empty // string when no draft model is loaded. {"target_device", config.target_device}, @@ -1332,8 +1323,16 @@ int HttpServer::run() { std::fprintf(stderr, "[server] listening on http://%s:%d\n", config_.host.c_str(), config_.port); - // Start worker thread. - worker_thread_ = std::thread([this]() { worker_loop(); }); + // A backend-provided sequence engine replaces the one-request worker + // with the concurrent scheduler. Upstream forwarding stays on the + // classic path even when the local backend exposes an engine. + if (SeqEngine * engine = backend_.seq_engine(); + engine && config_.pflash_upstream_base.empty()) { + worker_thread_ = + std::thread([this, engine]() { scheduler_loop(*engine); }); + } else { + worker_thread_ = std::thread([this]() { worker_loop(); }); + } // Accept loop. while (!stopping_.load()) { @@ -1916,8 +1915,7 @@ void HttpServer::log_parsed_request(const ParsedRequest & req) const { void HttpServer::enqueue_request_and_wait(SocketHandle fd, ParsedRequest req) { // Set socket non-blocking for send() stall detection during streaming. - const int flags = sock_get_flags(fd); - if (flags >= 0) sock_set_nonblock(fd); + sock_set_nonblock(fd); ServerJob job; job.fd = fd; @@ -2331,10 +2329,8 @@ json build_responses_api_response( json build_non_streaming_response( const ParsedRequest & req, const GenerateResult & result, - int generation_cap, const GenTimings & timings, Tokenizer & tokenizer, - SseEmitter & emitter) { - const CompletionTokenCounts counts = feed_non_streaming_tokens( - result.tokens, tokenizer, emitter); + int generation_cap, const GenTimings & timings, + const CompletionTokenCounts & counts, SseEmitter & emitter) { switch (req.format) { case ApiFormat::OPENAI_CHAT: return build_openai_completion_response( @@ -2350,6 +2346,16 @@ json build_non_streaming_response( } } +json build_non_streaming_response( + const ParsedRequest & req, const GenerateResult & result, + int generation_cap, const GenTimings & timings, Tokenizer & tokenizer, + SseEmitter & emitter) { + const CompletionTokenCounts counts = feed_non_streaming_tokens( + result.tokens, tokenizer, emitter); + return build_non_streaming_response( + req, result, generation_cap, timings, counts, emitter); +} + // Prompt preparation applies exactly one compression policy: FlowKV for // continuations, a verbatim turn-one anchor, or whole-prompt PFlash. bool is_continuation_request(const json & messages) { @@ -3447,6 +3453,70 @@ void HttpServer::configure_generation_io( }; } +bool HttpServer::deliver_generation_token( + const ParsedRequest & req, SseEmitter & emitter, int32_t token, + int & completion_tokens, ClientSendBuffer & send_buffer) { + ++completion_tokens; + + std::string text; + const TokenDelivery delivery = + classify_generated_token(tokenizer_, token, text); + if (delivery == TokenDelivery::kSkip) return true; + + // Non-stream replay counts every non-skipped token, including tokens + // whose decoded text is empty. Keep concurrent usage accounting aligned; + // streaming still has no frame to send for an empty string. + if (text.empty() && req.stream) return true; + + const auto chunks = emitter.emit_token(text); + if (req.stream) { + for (const auto & chunk : chunks) { + send_buffer.append(chunk); + } + } + return delivery == TokenDelivery::kThinkTag || !emitter.stop_hit(); +} + +void HttpServer::send_nonstream_response( + const ParsedRequest & req, SocketHandle fd, SseEmitter & emitter, + const std::vector & gen_tokens, int n_gen_cap, + bool budget_forced_close, bool degenerate_decode_close, + const GenTimings & gen_timings, + ClientSendBuffer * send_buffer) { + CompletionTokenCounts counts; + counts.total = (int) gen_tokens.size(); + emitter.emit_finish(counts.total); + const int first_content = emitter.first_content_token_index(); + const int emitted = emitter.emit_token_count(); + counts.reasoning = first_content < 0 ? emitted : first_content; + counts.content = first_content < 0 ? 0 : emitted - first_content; + + GenerateResult result; + result.tokens = gen_tokens; + result.budget_forced_close = budget_forced_close; + result.degenerate_decode_close = degenerate_decode_close; + + const json response = build_non_streaming_response( + req, result, n_gen_cap, gen_timings, counts, emitter); + + const std::string body = response.dump() + "\n"; + if (send_buffer) { + send_buffer->append( + format_http_response(200, "application/json", body)); + } else { + send_response(fd, 200, "application/json", body); + } +} + +std::array HttpServer::sse_error_close_chunks( + const std::string & message) { + const json err = {{"error", { + {"message", message}, + {"type", "server_error"}, + }}}; + return {"data: " + err.dump() + "\n\n", "data: [DONE]\n\n"}; +} + void HttpServer::worker_loop() { while (true) { ServerJob * job = dequeue(); @@ -3502,11 +3572,9 @@ void HttpServer::process_job(ServerJob * job) { std::fprintf(stderr, "[server] request failed: %s\n", message.c_str()); if (req.stream) { stop_job_stream(job); - json err = {{"error", {{"message", message}, {"type", "server_error"}}}}; - const std::string chunk = "data: " + err.dump() + "\n\n"; - send_job_bytes(job, chunk.data(), chunk.size()); - const char done[] = "data: [DONE]\n\n"; - send_job_bytes(job, done, sizeof(done) - 1); + for (const std::string & chunk : sse_error_close_chunks(message)) { + send_job_bytes(job, chunk.data(), chunk.size()); + } } else { send_error(fd, status, message); } @@ -3708,8 +3776,7 @@ void HttpServer::process_job(ServerJob * job) { req, result, n_gen_cap, gen_timings, tokenizer_, emitter); // Streaming uses non-blocking sends; restore blocking mode before // writing a complete JSON response on this shared socket path. - const int flags = sock_get_flags(fd); - if (flags >= 0) sock_set_block(fd); + sock_set_block(fd); send_response(fd, 200, "application/json", response.dump() + "\n"); } @@ -3794,6 +3861,33 @@ ServerJob * HttpServer::dequeue() { return j; } +ServerJob * HttpServer::try_dequeue() { + std::lock_guard lk(queue_mu_); + if (!queue_head_) return nullptr; + ServerJob * job = queue_head_; + queue_head_ = job->next; + if (!queue_head_) queue_tail_ = nullptr; + job->next = nullptr; + return job; +} + +ServerJob * HttpServer::dequeue_for( + std::chrono::steady_clock::duration timeout) { + std::unique_lock lk(queue_mu_); + if (!queue_head_ && !stopping_.load() && + timeout > std::chrono::steady_clock::duration::zero()) { + queue_cv_.wait_for(lk, timeout, [&] { + return queue_head_ != nullptr || stopping_.load(); + }); + } + if (!queue_head_) return nullptr; + ServerJob * job = queue_head_; + queue_head_ = job->next; + if (!queue_head_) queue_tail_ = nullptr; + job->next = nullptr; + return job; +} + // ─── HTTP I/O ─────────────────────────────────────────────────────────── bool HttpServer::read_http_request(SocketHandle fd, HttpRequest & out) { @@ -3980,9 +4074,9 @@ void HttpServer::maybe_send_job_heartbeat( job->last_stream_write = now; } -bool HttpServer::send_response( - SocketHandle fd, int status, const std::string & content_type, - const std::string & body) { +std::string HttpServer::format_http_response( + int status, const std::string & content_type, + const std::string & body) { const char * reason = "OK"; switch (status) { case 200: reason = "OK"; break; @@ -4006,7 +4100,15 @@ bool HttpServer::send_response( header += "Content-Length: " + std::to_string(body.size()) + "\r\n"; header += "Connection: close\r\n\r\n"; header += body; - return send_all(fd, header.data(), header.size()); + return header; +} + +bool HttpServer::send_response( + SocketHandle fd, int status, const std::string & content_type, + const std::string & body) { + const std::string payload = + format_http_response(status, content_type, body); + return send_all(fd, payload.data(), payload.size()); } bool HttpServer::send_error( diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index caf5b6946..8859b05f7 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -14,6 +14,7 @@ #pragma once #include "socket_handle.h" +#include "client_send_buffer.h" #include "common/model_backend.h" #include "tokenizer.h" #include "chat_template.h" @@ -28,9 +29,11 @@ #include "model_card.h" #include "adaptive_keep_ratio.h" #include "server_status.h" +#include "sse_emitter.h" #include #include +#include #include #include #include @@ -51,7 +54,6 @@ using json = nlohmann::json; // ─── Forward declarations ─────────────────────────────────────────────── struct ServerJob; -class SseEmitter; namespace http_detail { // Non-consuming peer-state probe used by the client-thread job monitor. @@ -176,6 +178,9 @@ struct ServerConfig { // server_main after CLI parse. std::string target_device; std::string draft_device; + // Idle-to-busy batching window. It is ignored by single-slot engines and + // never delays an already decoding request. + int admission_coalesce_ms = 20; // PFlash (speculative prefill compression) enum class PflashMode { OFF, AUTO, ALWAYS }; @@ -395,6 +400,35 @@ class HttpServer { ServerJob * job, const ParsedRequest & req, SseEmitter & emitter, GenerationOutputState & output, DaemonIO & io); + // Worker thread, concurrent mode (the backend exposes a SeqEngine): + // iteration-level scheduler. Admission is claim-only; this baseline + // drains its pending prefill between decode iterations, then advances + // active slots together in one batched step. + void scheduler_loop(SeqEngine & engine); + + // Non-blocking dequeue used for admission polling between decode steps. + ServerJob * try_dequeue(); + // Bounded wait used only during an idle-to-busy admission window. + ServerJob * dequeue_for( + std::chrono::steady_clock::duration timeout); + + // Concurrent-scheduler token delivery and shared response construction. + // A send buffer keeps slow clients off the shared decode loop. + bool deliver_generation_token( + const ParsedRequest & req, SseEmitter & emitter, int32_t token, + int & completion_tokens, ClientSendBuffer & send_buffer); + void send_nonstream_response( + const ParsedRequest & req, SocketHandle fd, SseEmitter & emitter, + 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); + std::string format_http_response( + int status, const std::string & content_type, + const std::string & body); + static std::array sse_error_close_chunks( + const std::string & message); + // Parse HTTP request from socket. struct HttpRequest { std::string method; @@ -531,6 +565,15 @@ struct ServerJob { std::chrono::steady_clock::time_point last_stream_write{}; std::atomic client_disconnected{false}; ServerJob * next = nullptr; + + // Concurrent-scheduler state that survives a pool-full admission retry. + // The classic worker leaves these fields untouched. + bool announced = false; + bool sse_started = false; + // First concurrent-scheduler attempt; retained across busy deferrals so + // server-side prefill/elapsed telemetry does not erase queueing delay. + std::chrono::steady_clock::time_point parallel_started_at{}; + std::unique_ptr emitter; }; // ─── Parse session_id from a chat-completion JSON body ────────────────── diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp new file mode 100644 index 000000000..b812a8556 --- /dev/null +++ b/server/src/server/scheduler.cpp @@ -0,0 +1,739 @@ +// Concurrent scheduler for --max-concurrency serving: the worker thread's +// iteration-level loop over a backend's SeqEngine decode slots. +// +// Split from http_server.cpp: this TU owns non-blocking admission (one +// prefill chunk per engine step, fused with the live decode batch), FIFO +// pool-full deferrals, per-slot streaming through ClientSendBuffer, and +// retirement. SSE emission, error-close chunks, and HTTP response +// formatting are shared with the classic worker so both paths emit +// matching wire formats. + +#include "http_server.h" +#include "common/concurrency/seq_engine.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +// Per-slot request state for the iteration-level scheduler. Indexed by the +// engine slot id returned from admit(), so scheduler and engine agree on +// which engine-owned state record a request owns. This remains the one +// external phase: sockets stay here, prompt/KV/sampler/progress stay in Qwen. +struct SchedSlot { + ServerJob * job = nullptr; + SocketHandle fd = kInvalidSocket; + std::unique_ptr emitter; + bool prefilling = false; + uint64_t admission_order = 0; + std::chrono::steady_clock::time_point started_at{}; + std::chrono::steady_clock::time_point decode_started_at{}; + double prefill_s = 0.0; + int n_gen_cap = 0; + int completion_tokens = 0; + bool client_disconnected = false; + bool failed = false; + std::string error; + bool finished = false; + std::vector gen_tokens; // committed + pending, in order + 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 + // reader can never head-of-line-block the shared decode loop. + dflash::common::ClientSendBuffer send_buffer; + // Thinking-budget force-close, applied scheduler-side before the token + // is fed back (mirrors do_ar_decode's maybe_force_close). + dflash::common::BudgetHook hook; + bool hook_started = false; + int hook_pos = 0; + bool budget_forced_close = false; + bool degenerate_close = false; +}; + +// Outcome of one admission attempt. The three cases differ in who owns the +// job afterwards: Admitted hands it to a slot, Deferred hands it back to the +// caller to retry at the head of the line, Retired means the job is already +// answered and holds nothing. +enum class AdmissionDisposition { + Admitted, // job owns an engine slot and its first token is on the wire + Deferred, // engine had no room; caller keeps the job and retries it first + Retired, // job finished or failed during admission; no slot was taken +}; + +} // namespace + +void HttpServer::scheduler_loop(SeqEngine & engine) { + const int n_slots = engine.slot_count(); + std::vector slots((size_t)n_slots); + uint64_t next_request_id = 1; + uint64_t next_admission_order = 0; + // Admission-deferred job (pool blocks/slots exhausted). Kept at the head + // of the line so FIFO order survives the deferral. + ServerJob * deferred = nullptr; + std::chrono::steady_clock::time_point deferred_retry_at{}; + + // Degenerate-run guard shared with do_ar_decode: explicit env override, + // else 32 when the min-tokens floor is active, else off. + static const int repeat_guard = [] { + if (const char * s = std::getenv("DFLASH_DEGENERATE_RUN_TOKENS")) { + const int v = std::atoi(s); + if (v >= 0) return v; + } + const char * f = std::getenv("DFLASH_MIN_TOKENS"); + return (f && std::atoi(f) > 0) ? 32 : 0; + }(); + + // Cached live-slot count — incremented on admit, decremented on retire. + // Replaces the O(n_slots) scan that was called 2-3× per iteration. + int live_slots = 0; + + int published_live_count = -1; + auto publish_live_count = [&]() { + if (live_slots == published_live_count) return; + published_live_count = live_slots; + if (live_slots > 0) status_.set_concurrent_requests(live_slots); + else status_.set_idle(); + broadcast_status(); + }; + + auto finish_job = [this](ServerJob * job) { + stop_job_stream(job); + std::lock_guard lk(job->mu); + job->done = true; + job->cv.notify_one(); + }; + + // A stalled reader may buffer at most this much before being dropped. + constexpr size_t kMaxSlotSendBuffer = 1u << 20; + constexpr auto kClientStallTimeout = std::chrono::seconds(30); + + // Final payloads of retired slots still draining to their sockets. The + // job is signalled done only once its bytes are out (or the deadline / + // hard error gives up) because the parked client thread closes the fd + // the moment it wakes. + struct DrainJob { + ServerJob * job = nullptr; + SocketHandle fd = kInvalidSocket; + ClientSendBuffer send_buffer; + std::chrono::steady_clock::time_point deadline{}; + }; + std::vector drains; + + auto service_drains = [&]() { + if (drains.empty()) return; + const auto now = std::chrono::steady_clock::now(); + for (size_t i = 0; i < drains.size();) { + DrainJob & d = drains[i]; + const size_t pending_before = d.send_buffer.pending(); + const bool ok = d.send_buffer.flush(d.fd); + if (d.send_buffer.pending() < pending_before) { + d.deadline = std::chrono::steady_clock::now() + + kClientStallTimeout; + } + if (!ok || d.send_buffer.empty() || now > d.deadline) { + finish_job(d.job); + drains[i] = std::move(drains.back()); + drains.pop_back(); + continue; + } + ++i; + } + }; + + auto maybe_force_close = [](SchedSlot & s, int32_t & tok) { + if (s.hook.close_token_ids.empty()) return; + if (s.hook_started) { + if (s.hook_pos < (int)s.hook.close_token_ids.size()) { + tok = s.hook.close_token_ids[(size_t)s.hook_pos++]; + } + return; + } + const int generated = (int)s.gen_tokens.size(); + const int remaining = s.n_gen_cap - generated; + if (remaining <= s.hook.hard_limit_remaining) { + const int32_t first_close = s.hook.close_token_ids.front(); + s.hook_started = true; + s.hook_pos = 1; + if (tok != first_close) { + tok = first_close; + s.budget_forced_close = true; + } + } + }; + + // Advances one slot by a single sampled token — the post-sample path + // shared by the first (prefill-logits) token and every decode-step token. + // Note `tok` is by value but not passthrough: maybe_force_close may + // *substitute* a close token for it, and that substitute is what gets + // recorded, emitted, and fed back. Appends to gen_tokens, streams the + // 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) { + maybe_force_close(s, tok); + // Heartbeats are needed while prefill is silent. Once token frames enter + // the buffered writer, stop them so a comment cannot split a partial + // SSE frame across separate non-blocking flushes. + if (s.job->req.stream) stop_job_stream(s.job); + s.gen_tokens.push_back(tok); + const bool cont = deliver_generation_token( + s.job->req, *s.emitter, tok, s.completion_tokens, s.send_buffer); + s.pending_tok = tok; + if (!cont || engine.token_is_eos(tok) || + (int)s.gen_tokens.size() >= s.n_gen_cap) { + s.finished = true; + return; + } + // Single-token run guard (matches do_ar_decode's repeat break). + if (repeat_guard > 0 && (int)s.gen_tokens.size() >= repeat_guard) { + int run = 1; + for (int j = (int)s.gen_tokens.size() - 2; j >= 0; --j) { + if (s.gen_tokens[(size_t)j] != tok) break; + run++; + } + if (run >= repeat_guard) { + std::fprintf(stderr, + "[parallel] token %d repeated %d times — stopping slot\n", + tok, run); + s.degenerate_close = true; + s.finished = true; + return; + } + } + // Post-close repetition watchdog (periods 12..80), mirrors + // do_ar_decode's sweep once the close sequence has fully injected. + if (s.hook_started && + s.hook_pos >= (int)s.hook.close_token_ids.size()) { + const auto end = s.gen_tokens.end(); + const int avail = (int)s.gen_tokens.size(); + for (int P = 12; P <= 80; P++) { + if (avail < 2 * P) break; + if (std::equal(end - 2 * P, end - P, end - P)) { + std::fprintf(stderr, + "[parallel] post-close period=%d repeated — " + "stopping slot\n", P); + s.degenerate_close = true; + s.finished = true; + return; + } + } + } + }; + + auto retire_slot = [&](int idx, bool backend_ok) { + SchedSlot & s = slots[(size_t)idx]; + if (!s.job) return; + const ParsedRequest & req = s.job->req; + // Stop monitor-thread heartbeats before queuing terminal frames. + stop_job_stream(s.job); + const double decode_s = std::chrono::duration( + std::chrono::steady_clock::now() - s.decode_started_at).count(); + const int prompt_tokens = (int)req.prompt_tokens.size(); + GenTimings gen_timings{ + s.prefill_s, + decode_s, + /*cache_hit=*/false, + /*cached_prefix_tokens=*/0, + /*prefilled_tokens=*/prompt_tokens, + /*effective_prompt_tokens=*/prompt_tokens, + }; + + if (backend_ok && !s.failed) { + PerfRecord perf; + perf.prompt_tokens = (int)req.prompt_tokens.size(); + perf.completion_tokens = s.completion_tokens; + perf.prefill_tok_s = s.prefill_s > 0.0 + ? (double)req.prompt_tokens.size() / s.prefill_s : 0.0; + perf.decode_tok_s = decode_s > 0.0 + ? (double)s.completion_tokens / decode_s : 0.0; + status_.record_perf(perf); + } + + if (s.failed || !backend_ok) { + const std::string message = + s.error.empty() ? "generation failed" : s.error; + if (!s.client_disconnected) { + if (req.stream) { + for (const std::string & chunk : + sse_error_close_chunks(message)) { + s.send_buffer.append(chunk); + } + } else { + json err = {{"error", {{"message", message}, + {"type", "invalid_request_error"}}}}; + s.send_buffer.append(format_http_response( + 500, "application/json", err.dump() + "\n")); + } + } + } else if (req.stream && !s.client_disconnected) { + auto final_chunks = + s.emitter->emit_finish(s.completion_tokens, &gen_timings); + for (const auto & chunk : final_chunks) { + s.send_buffer.append(chunk); + } + } else if (!req.stream && !s.client_disconnected) { + 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); + } + + const double elapsed_s = std::chrono::duration( + std::chrono::steady_clock::now() - s.started_at).count(); + const int out_tokens = (int)s.gen_tokens.size(); + std::fprintf(stderr, + "[server] chat DONE %s ok=%s in=%zu out=%d %.1fs %.1f tok/s " + "finish=%s slot=%d prefill=%.1fs decode=%.1fs(%.1ftok/s) parallel\n", + req.response_id.c_str(), + (!s.failed && backend_ok) ? "true" : "false", + req.prompt_tokens.size(), out_tokens, elapsed_s, + elapsed_s > 0.0 ? out_tokens / elapsed_s : 0.0, + s.client_disconnected ? "client_disconnect" + : s.emitter->finish_reason().c_str(), + idx, s.prefill_s, decode_s, + decode_s > 0.0 ? out_tokens / decode_s : 0.0); + + engine.retire(idx); + // A retirement may have released the blocks the head job needs. + deferred_retry_at = {}; + + // Hand any undrained bytes to the drain list; the job stays parked + // until they are out (or the drain gives up). + bool drained = s.client_disconnected; + if (!drained) { + drained = s.send_buffer.flush(s.fd) ? s.send_buffer.empty() : true; + } + if (drained) { + finish_job(s.job); + } else { + DrainJob d; + d.job = s.job; + d.fd = s.fd; + d.send_buffer = std::move(s.send_buffer); + d.deadline = std::chrono::steady_clock::now() + + kClientStallTimeout; + drains.push_back(std::move(d)); + } + s = SchedSlot{}; + live_slots--; + publish_live_count(); + }; + + auto admit_job = [&](ServerJob * job) -> AdmissionDisposition { + const ParsedRequest & req = job->req; + if (job->parallel_started_at == + std::chrono::steady_clock::time_point{}) { + job->parallel_started_at = std::chrono::steady_clock::now(); + } + const auto started_at = job->parallel_started_at; + + // Same thinking-budget n_gen math as the classic worker loop. + const bool budget_active = req.thinking_opt_in; + const int effective_think_ceiling = (req.per_req_phase1_cap >= 0) + ? req.per_req_phase1_cap + : config_.think_max_tokens; + const int eff_reply_for_n_gen = (req.per_req_reply_budget >= 0) + ? req.per_req_reply_budget + : config_.hard_limit_reply_budget; + const int n_gen_cap = budget_active + ? (std::min)(effective_think_ceiling + eff_reply_for_n_gen, + req.max_output) + : req.max_output; + + if (n_gen_cap < 1) { + // Degenerate ask: reply with an empty completion, no slot needed. + SseEmitter emitter(req.format, req.response_id, req.model, + (int)req.prompt_tokens.size(), req.tools, + &tool_memory_, req.stop_sequences, + req.started_in_thinking); + GenTimings t{ + 0.0, + 0.0, + /*cache_hit=*/false, + /*cached_prefix_tokens=*/0, + /*prefilled_tokens=*/0, + /*effective_prompt_tokens=*/(int)req.prompt_tokens.size(), + }; + if (req.stream) { + if (send_sse_headers(job)) { + bool ok = true; + for (const auto & c : emitter.emit_start()) { + if (!send_job_bytes(job, c.data(), c.size())) { ok = false; break; } + } + if (ok) { + for (const auto & c : emitter.emit_finish(0, &t)) { + if (!send_job_bytes(job, c.data(), c.size())) break; + } + } + } + } else { + send_nonstream_response(req, job->fd, emitter, {}, n_gen_cap, + false, false, t); + } + finish_job(job); + return AdmissionDisposition::Retired; + } + + if (!job->announced) { + job->announced = true; + std::fprintf(stderr, + "[server] chat START %s format=%s stream=%s prompt_tokens=%zu " + "max_tokens=%d live=%d parallel\n", + req.response_id.c_str(), api_format_name(req.format), + req.stream ? "true" : "false", req.prompt_tokens.size(), + req.max_output, live_slots); + } + + // Commit the SSE preamble BEFORE the (multi-second) prefill so + // streaming clients see the 200 immediately — and so a dead client + // is detected before its prefill is paid for. The socket is fresh + // (nothing sent yet), so these few hundred bytes cannot stall. + // sse_started survives a busy deferral: retries do not resend. + if (!job->emitter) { + job->emitter = std::make_unique( + req.format, req.response_id, req.model, + (int)req.prompt_tokens.size(), req.tools, &tool_memory_, + req.stop_sequences, req.started_in_thinking); + } + if (req.stream && !job->sse_started) { + job->sse_started = true; + bool ok = send_sse_headers(job); + if (ok) { + for (const auto & c : job->emitter->emit_start()) { + if (!send_job_bytes(job, c.data(), c.size())) { + ok = false; + break; + } + } + } + if (!ok) { + finish_job(job); + return AdmissionDisposition::Retired; + } + start_job_stream(job); + } + + // Admission only claims the slot and queues the prompt. Prefill + // advances one chunk per engine step alongside live decode. + auto ar = engine.admit(next_request_id, req.prompt_tokens, + req.sampler); + if (ar.status == SeqEngine::AdmitResult::Status::busy) + return AdmissionDisposition::Deferred; + if (ar.status != SeqEngine::AdmitResult::Status::admitted) { + std::fprintf(stderr, "[server] admit failed: %s\n", + ar.error.c_str()); + if (req.stream && job->sse_started) { + stop_job_stream(job); + // Headers are already on the wire: report in-stream, like + // the classic worker's fail_request after SSE start. + for (const std::string & chunk : sse_error_close_chunks( + "admission failed: " + ar.error)) { + send_job_bytes(job, chunk.data(), chunk.size()); + } + } else { + send_error(job->fd, 500, "admission failed: " + ar.error); + } + finish_job(job); + return AdmissionDisposition::Retired; + } + next_request_id++; + + SchedSlot & s = slots[(size_t)ar.slot]; + s = SchedSlot{}; + s.job = job; + s.fd = job->fd; + s.prefilling = true; + s.admission_order = next_admission_order++; + s.started_at = started_at; + s.decode_started_at = started_at; // sane on prefill failure + s.n_gen_cap = std::min( + n_gen_cap, + engine.max_context() - (int)req.prompt_tokens.size() + 1); + s.emitter = std::move(job->emitter); + s.send_buffer.mark_progress(std::chrono::steady_clock::now()); + if (budget_active && !config_.think_close_token_ids.empty() && + config_.hard_limit_reply_budget > 0) { + s.hook.close_token_ids = config_.think_close_token_ids; + s.hook.hard_limit_remaining = eff_reply_for_n_gen; + } + live_slots++; + publish_live_count(); + return AdmissionDisposition::Admitted; + }; + + // Scheduler loop. Every iteration walks the same five phases: + // + // 1. Admit — fill every available slot (deferred first, then the + // queue). Their prefills advance inside later steps. + // 2. Idle — nothing live: service drains and loop back, where the + // admission phase parks in the blocking dequeue. + // 3. Step — advance one pending prefill chunk alongside every + // decoding slot in one engine pass. + // 4. Flush — non-blocking write of the buffered chunks; readers that + // stall or overflow their buffer are dropped. + // 5. Reap — service drains, then retire whatever finished this + // iteration so its blocks are free for the next admit. + // + // Exits on stopping_ (checked after admission), leaving the teardown + // below to answer every client still parked in a slot, drain, or queue. + // Hoisted per-iteration buffers — capacity persists across iterations. + SeqEngine::StepPlan step_plan; + step_plan.decode.reserve((size_t)n_slots); + step_plan.prefills.reserve((size_t)n_slots); + std::vector prefill_candidates; + prefill_candidates.reserve((size_t)n_slots); + size_t prefill_round_robin_start = 0; + + while (true) { + // Phase 1 — Admission: deferred job first (FIFO), then the queue. + // Blocking dequeue only when idle; between decode steps only a poll. + // Engines reject atomically when their current slot, staging, or + // reserved-capacity limit is reached, so policy stays model-neutral. + auto idle_admission_deadline = + std::chrono::steady_clock::time_point{}; + while (live_slots < n_slots && !stopping_.load()) { + // A deferred job owns the front of the line, so nothing else may + // be admitted while its retry backoff is still running. + if (deferred && + std::chrono::steady_clock::now() < deferred_retry_at) { + break; + } + // Retry the deferred job first; it was already queued ahead of + // everything still in the queue. + ServerJob * job = deferred; + deferred = nullptr; + bool woke_from_idle = false; + if (!job) { + if (live_slots == 0 && drains.empty()) { + // Fully idle: no stream is waiting on us, so give the + // scratch memory back and park in a blocking dequeue. + publish_live_count(); + backend_.release_scratch(); + job = dequeue(); + woke_from_idle = job != nullptr; + } else { + // A fresh idle transition may spend its bounded batching + // window here; all ongoing decode paths only poll. + const auto now = std::chrono::steady_clock::now(); + if (idle_admission_deadline > now) { + job = dequeue_for(idle_admission_deadline - now); + if (!job) idle_admission_deadline = {}; + } else { + job = try_dequeue(); + } + } + } + if (!job) break; // queue empty: go decode what is already live + if (job->client_disconnected.load(std::memory_order_acquire)) { + finish_job(job); + continue; + } + const AdmissionDisposition outcome = admit_job(job); + if (outcome == AdmissionDisposition::Deferred) { + deferred = job; + deferred_retry_at = std::chrono::steady_clock::now() + + std::chrono::seconds(1); + break; // wait for a retire to free blocks + } + if (outcome == AdmissionDisposition::Admitted) { + if (woke_from_idle && n_slots > 1 && + config_.admission_coalesce_ms > 0) { + idle_admission_deadline = + std::chrono::steady_clock::now() + + std::chrono::milliseconds( + config_.admission_coalesce_ms); + } + continue; // fill the remaining slots before step() + } + } + if (stopping_.load()) break; + + // Phase 2 — Idle: no slot to step, so only the drains need service. + if (live_slots == 0) { + service_drains(); + if (deferred) { + // A defensive busy response with no live sequence must not + // turn the worker into a tight retry loop. Real capacity + // releases clear deferred_retry_at in retire_slot(). + const auto now = std::chrono::steady_clock::now(); + if (deferred_retry_at > now) { + std::this_thread::sleep_until(std::min( + deferred_retry_at, now + std::chrono::milliseconds(5))); + } + continue; + } + if (!drains.empty()) { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + continue; // keep draining; don't block in dequeue + } + continue; // dequeue() blocks in admission + } + + // Retire cancellations before spending another model step. + for (int i = 0; i < n_slots; ++i) { + SchedSlot & s = slots[(size_t)i]; + if (s.job && s.job->client_disconnected.load( + std::memory_order_acquire)) { + s.client_disconnected = true; + s.finished = true; + retire_slot(i, true); + } + } + if (live_slots == 0) continue; + + // Phase 3 — Build one model-neutral batch plan: every decode row plus + // a FIFO, engine-bounded subset of pending prompt work. The engine + // lowers this plan into whatever graph/state representation it owns. + step_plan.decode.clear(); + prefill_candidates.clear(); + for (int i = 0; i < n_slots; i++) { + if (slots[(size_t)i].job && !slots[(size_t)i].prefilling) { + step_plan.decode.push_back( + {i, slots[(size_t)i].pending_tok}); + } else if (slots[(size_t)i].job) { + prefill_candidates.push_back( + {i, slots[(size_t)i].admission_order}); + } + } + const StepPlanLimits step_limits = + engine.step_plan_limits((int)step_plan.decode.size()); + step_plan.prefills = plan_prefill_slices( + prefill_candidates, step_limits, prefill_round_robin_start); + if (!prefill_candidates.empty()) { + ++prefill_round_robin_start; + } + + SeqEngine::StepResult step_result = engine.step(step_plan); + const std::string protocol_error = + validate_step_result(step_plan, step_result, n_slots); + if (!protocol_error.empty()) { + step_result.decode.clear(); + step_result.prefills.clear(); + step_result.error = + "engine step protocol violation: " + protocol_error; + } + + if (!step_result.ok()) { + const std::string & error = step_result.error; + std::fprintf(stderr, + "[parallel] engine step failed: %s — " + "failing all live requests\n", error.c_str()); + for (int i = 0; i < n_slots; i++) { + if (slots[(size_t)i].job) { + slots[(size_t)i].failed = true; + slots[(size_t)i].error = error; + retire_slot(i, false); + } + } + continue; + } + for (const auto & out : step_result.decode) { + if (out.slot < 0 || out.slot >= n_slots) continue; + SchedSlot & s = slots[(size_t)out.slot]; + if (!s.job) continue; + if (out.failed) { + s.failed = true; + s.error = out.error; + s.finished = true; + continue; + } + advance_slot(s, out.token); + } + using PrefillStatus = SeqEngine::PrefillOutput::Status; + for (const auto & out : step_result.prefills) { + if (out.slot < 0 || out.slot >= n_slots) continue; + SchedSlot & s = slots[(size_t)out.slot]; + if (!s.job) continue; + if (out.status == PrefillStatus::failed) { + s.failed = true; + s.error = out.error; + s.finished = true; + continue; + } + if (out.status == PrefillStatus::completed) { + s.prefilling = false; + // A prefill lane just became reusable, so the FIFO head may + // be admissible even though no KV blocks were retired. + deferred_retry_at = {}; + s.decode_started_at = std::chrono::steady_clock::now(); + s.prefill_s = std::chrono::duration( + s.decode_started_at - s.started_at).count(); + advance_slot(s, out.token); + continue; + } + } + // Phase 4 — Non-blocking flush of every live slot's chunks. Progress + // resets the stall clock; a reader that makes no progress for 30 s + // or lets the buffer hit the cap is dropped (its slot retires). + { + const auto now = std::chrono::steady_clock::now(); + for (int i = 0; i < n_slots; i++) { + SchedSlot & s = slots[(size_t)i]; + if (!s.job || s.client_disconnected) continue; + if (!s.send_buffer.flush(s.fd)) { + s.client_disconnected = true; + s.finished = true; + continue; + } + if (s.send_buffer.should_drop(now, kClientStallTimeout, + kMaxSlotSendBuffer)) { + std::fprintf(stderr, + "[parallel] slot %d client stalled — dropping stream\n", i); + s.client_disconnected = true; + s.finished = true; + } + } + } + // Phase 5 — Reap: finish the drains, then hand back the blocks of + // every slot that ended this iteration so the next admit can use them. + service_drains(); + for (int i = 0; i < n_slots; i++) { + if (slots[(size_t)i].job && slots[(size_t)i].finished) { + retire_slot(i, true); + } + } + } + + // Shutdown: unblock every parked client thread. + for (int i = 0; i < n_slots; i++) { + if (slots[(size_t)i].job) { + slots[(size_t)i].failed = true; + retire_slot(i, false); + } + } + service_drains(); + for (DrainJob & d : drains) finish_job(d.job); + drains.clear(); + if (deferred) { + // admit_job() sends the SSE headers and opening event before asking + // the engine for a slot, so a pool-full deferred stream is already + // live on the wire. Close that protocol cleanly on shutdown instead + // of waking the client thread and letting it truncate the response. + const ParsedRequest & req = deferred->req; + if (req.stream && deferred->sse_started) { + for (const std::string & chunk : + sse_error_close_chunks("server shutting down")) { + send_all(deferred->fd, chunk.data(), chunk.size()); + } + } else { + send_error(deferred->fd, 503, "server shutting down"); + } + finish_job(deferred); + } + // Jobs that never reached admission are still parked in their client + // threads too. Drain the raw queue before returning so run() does not hit + // its client-shutdown timeout and the destructor never has to wake threads + // after the server/backend teardown has already started. + while (ServerJob * queued = try_dequeue()) { + send_error(queued->fd, 503, "server shutting down"); + finish_job(queued); + } +} + + +} // namespace dflash::common diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index adfedc300..e6cc1c79b 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -24,8 +24,10 @@ #include "common/peer_access.h" #include "placement/pflash_placement.h" #include "placement/draft_residency.h" +#include "kvflash_pager.h" #include +#include #include #include #include @@ -104,8 +106,15 @@ static void print_usage(const char * prog) { " --fa-window Flash-attention sliding window (default: 0=full).\n" " WARNING: >0 drops system prompt / tool definitions\n" " from attention at long contexts. Use 0 for tools.\n" - " --paged-attention Use 16-token paged KV blocks for Qwen3.6-27B\n" - " autoregressive decode (experimental)\n" + " --paged-attention Use model-native paged state for Qwen3.6 or\n" + " DeepSeek V4 autoregressive decode (experimental)\n" + " --max-concurrency Maximum concurrent decode sequences\n" + " (enables paged attention; default: 1)\n" + " --admission-coalesce-ms Idle-to-busy batching window\n" + " (default: 20; 0 disables)\n" + " --kv-pool-tokens Total paged K/V pool shared by all\n" + " --max-concurrency slots, in tokens\n" + " (default: sized from available device memory)\n" " --model-name Model name for /v1/models (default: dflash)\n" " --prefix-cache-slots Prefix cache slots (default: 32, 0 disables)\n" " --prefill-cache-slots Full prompt/prefill cache slots (default: 0)\n" @@ -353,6 +362,27 @@ int main(int argc, char ** argv) { bargs.fa_window = std::atoi(argv[++i]); } else if (std::strcmp(argv[i], "--paged-attention") == 0) { bargs.paged_attention = true; + } else if (std::strcmp(argv[i], "--max-concurrency") == 0 && i + 1 < argc) { + bargs.max_concurrency = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--admission-coalesce-ms") == 0 && + i + 1 < argc) { + sconfig.admission_coalesce_ms = std::atoi(argv[++i]); + if (sconfig.admission_coalesce_ms < 0 || + sconfig.admission_coalesce_ms > 1000) { + std::fprintf(stderr, + "[server] --admission-coalesce-ms must be in [0,1000]\n"); + return 2; + } + } else if (std::strcmp(argv[i], "--kv-pool-tokens") == 0 && i + 1 < argc) { + const char * value = argv[++i]; + const char * end = value + std::strlen(value); + const auto parsed = std::from_chars( + value, end, bargs.kv_pool_tokens); + if (parsed.ec != std::errc{} || parsed.ptr != end) { + std::fprintf(stderr, + "[server] --kv-pool-tokens must be an integer\n"); + return 2; + } } else if (std::strcmp(argv[i], "--model-name") == 0 && i + 1 < argc) { sconfig.model_name = argv[++i]; } else if (std::strcmp(argv[i], "--prefix-cache-slots") == 0 && i + 1 < argc) { @@ -602,6 +632,13 @@ int main(int argc, char ** argv) { } } + // Concurrent serving is implemented by the paged engine, so one user-facing + // concurrency flag selects the complete serving mode. Keep the feature gate + // strict for non-CLI callers that construct BackendArgs directly. + if (bargs.max_concurrency > 1) { + bargs.paged_attention = true; + } + // Ask the factory to resolve model/placement facts and apply its feature // admission policy before any setup work. server_main only maps the // categorized result to the existing process exit convention. @@ -613,6 +650,11 @@ int main(int argc, char ** argv) { backend_features.routing_stats_requested = sconfig.freq_tracking || !sconfig.collect_routing_path.empty(); backend_features.adaptive_experts_requested = adaptive_experts_set; + // Fixed pools are known incompatibilities before model setup. Automatic + // sizing needs the backend's real VRAM budget; if it produces a live pool, + // the backend rejects the pairing after sizing. + backend_features.kvflash_enabled = + kvflash_fixed_pool_requested(std::getenv("DFLASH_KVFLASH")); const BackendPreparation backend_preparation = prepare_backend(bargs, backend_features); if (!backend_preparation.ok()) { @@ -1029,6 +1071,8 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "[server] │ peer_access = %s\n", bargs.device.peer_access ? "ON" : "off"); std::fprintf(stderr, "[server] │ chunk = %d\n", bargs.chunk); + std::fprintf(stderr, "[server] │ admission_wait = %d ms\n", + sconfig.admission_coalesce_ms); if (arch == "deepseek4") { std::fprintf(stderr, "[server] │ ds4_fused = %s\n", bargs.ds4_fused_decode ? "ON" : "off"); diff --git a/server/src/server/server_status.h b/server/src/server/server_status.h index 98794e919..d715d19d5 100644 --- a/server/src/server/server_status.h +++ b/server/src/server/server_status.h @@ -67,6 +67,7 @@ class ServerStatus { bool is_stream, const RequestInfo & info) { std::lock_guard lk(mu_); phase_ = InferencePhase::PREFILL; + active_requests_ = 0; prompt_excerpt_ = prompt_excerpt; prompt_tokens_ = prompt_tokens; completion_tokens_ = 0; @@ -109,6 +110,21 @@ class ServerStatus { void set_idle() { std::lock_guard lk(mu_); phase_ = InferencePhase::IDLE; + active_requests_ = 0; + prompt_excerpt_.clear(); + draft_tokens_.clear(); + } + + // Concurrent serving intentionally exposes only aggregate live state. + // A single request record/token feed cannot represent multiple slots + // without mixing unrelated requests. + void set_concurrent_requests(int n) { + std::lock_guard lk(mu_); + if (active_requests_ == 0 && n > 0) { + started_at_ = std::chrono::steady_clock::now(); + } + active_requests_ = n; + phase_ = n > 0 ? InferencePhase::DECODE : InferencePhase::IDLE; prompt_excerpt_.clear(); draft_tokens_.clear(); } @@ -136,6 +152,7 @@ class ServerStatus { RequestInfo info; bool cache_hit = false, pflash = false, spec_decode = false; std::string messages_json; + int active_requests = 0; { std::lock_guard lk(mu_); @@ -152,6 +169,7 @@ class ServerStatus { pflash = pflash_; spec_decode = spec_decode_; messages_json = messages_json_; + active_requests = active_requests_; if (phase != InferencePhase::IDLE) { elapsed_s = std::chrono::duration( std::chrono::steady_clock::now() - started_at_).count(); @@ -161,8 +179,9 @@ class ServerStatus { json j; j["phase"] = phase_name(phase); j["total_requests"] = total_requests; + j["active_requests"] = active_requests; - if (phase != InferencePhase::IDLE) { + if (phase != InferencePhase::IDLE && active_requests == 0) { j["current"] = { {"prompt_excerpt", prompt_excerpt}, {"prompt_tokens", prompt_tokens}, @@ -226,6 +245,7 @@ class ServerStatus { bool pflash_ = false; bool spec_decode_ = false; std::string messages_json_; + int active_requests_ = 0; // History. std::vector perf_history_; diff --git a/server/src/server/tool_memory.cpp b/server/src/server/tool_memory.cpp index 1d54f11f7..a09376b72 100644 --- a/server/src/server/tool_memory.cpp +++ b/server/src/server/tool_memory.cpp @@ -14,6 +14,7 @@ ToolMemory::ToolMemory(size_t max_entries, size_t max_bytes) void ToolMemory::remember(const std::vector & call_ids, const std::string & raw_text) { if (disabled() || raw_text.empty()) return; + std::lock_guard lk(mu_); // Deduplicate call_ids std::vector unique_ids; @@ -58,6 +59,7 @@ void ToolMemory::remember(const std::vector & call_ids, } std::string ToolMemory::lookup(const std::vector & call_ids) { + std::lock_guard lk(mu_); std::string result_text; bool first = true; diff --git a/server/src/server/tool_memory.h b/server/src/server/tool_memory.h index 5144e229f..b2a7abe8e 100644 --- a/server/src/server/tool_memory.h +++ b/server/src/server/tool_memory.h @@ -10,12 +10,16 @@ #include #include +#include #include #include #include namespace dflash::common { +// Thread-safe: lookup() is called from HTTP client threads during request +// parsing while the worker/scheduler thread calls remember() at stream +// finish, and lookup() mutates the LRU. One mutex covers both. class ToolMemory { public: explicit ToolMemory(size_t max_entries = 50000, @@ -40,10 +44,12 @@ class ToolMemory { size_t current_bytes; }; Stats stats() const { + std::lock_guard lk(mu_); return {max_entries_, max_bytes_, by_id_.size(), total_bytes_}; } private: + mutable std::mutex mu_; struct Block { size_t size_bytes; size_t refs = 0; diff --git a/server/test/host_check.h b/server/test/host_check.h new file mode 100644 index 000000000..873dfc5d2 --- /dev/null +++ b/server/test/host_check.h @@ -0,0 +1,14 @@ +#pragma once + +#include +#include + +#define CHECK(cond) \ + do { \ + g_checks++; \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, \ + #cond); \ + std::exit(1); \ + } \ + } while (0) diff --git a/server/test/seq_engine_contract.h b/server/test/seq_engine_contract.h new file mode 100644 index 000000000..63ca035ab --- /dev/null +++ b/server/test/seq_engine_contract.h @@ -0,0 +1,349 @@ +// Host-side executable contract for common/concurrency/seq_engine.h. +// +// The checker deliberately knows only logical slots, decode inputs, and +// scheduler-selected prefill slices. It exercises the largest useful cohort +// up to two sequences allowed by the advertised limits, mixed +// decode/prefill, validation failures, retryable blocking, and slot reuse +// without importing a model graph or cache representation. + +#pragma once + +#include "common/concurrency/seq_engine.h" + +#include +#include +#include +#include + +namespace dflash::common { + +inline std::vector check_seq_engine_contract(SeqEngine & engine) { + std::vector violations; + auto require = [&violations](bool ok, const char * message) { + if (!ok) violations.emplace_back(message); + }; + + const int n_slots = engine.slot_count(); + require(n_slots >= 2, + "slot_count() must be at least 2 for the concurrency contract"); + if (n_slots < 2) return violations; + + const StepPlanLimits idle_limits = engine.step_plan_limits(0); + const StepPlanLimits mixed_limits = engine.step_plan_limits(1); + const auto supports_one_prefill = [](const StepPlanLimits & limits) { + return limits.max_prefill_sequences >= 1 && + limits.max_prefill_tokens_per_sequence >= 1 && + limits.max_prefill_tokens_total >= 1 && + limits.prefill_allocation_quantum >= 1; + }; + require(supports_one_prefill(idle_limits), + "idle step_plan_limits() must permit one prefill token"); + require(supports_one_prefill(mixed_limits), + "mixed step_plan_limits() must permit one prefill token"); + require(engine.max_context() >= 3, + "max_context() must fit the contract-check prompts"); + if (!supports_one_prefill(idle_limits) || + !supports_one_prefill(mixed_limits) || engine.max_context() < 3) { + return violations; + } + + // Exercise at most two concurrent prefills, but never infer K=2 support + // from slot_count(): sequence count and the total-token cap are separate + // engine capabilities. + const int idle_cohort_size = std::min({ + 2, + n_slots, + idle_limits.max_prefill_sequences, + idle_limits.max_prefill_tokens_total, + }); + + const SamplerCfg greedy{}; + SamplerCfg seeded{}; + seeded.temp = 0.7f; + seeded.seed = 20260809; + + std::vector active((size_t)n_slots, false); + std::vector decoding((size_t)n_slots, false); + std::vector remaining((size_t)n_slots, 0); + std::vector next_token((size_t)n_slots, -1); + + auto retire_all = [&]() { + for (int slot = 0; slot < n_slots; ++slot) { + if (active[(size_t)slot]) engine.retire(slot); + active[(size_t)slot] = false; + decoding[(size_t)slot] = false; + remaining[(size_t)slot] = 0; + } + }; + + auto record_admit = [&](uint64_t request_id, + const std::vector & prompt, + const SamplerCfg & sampler) { + const SeqEngine::AdmitResult result = + engine.admit(request_id, prompt, sampler); + const bool admitted = + result.status == SeqEngine::AdmitResult::Status::admitted; + require(admitted, "admit() must succeed while a slot is free"); + if (!admitted) return -1; + require(result.slot >= 0 && result.slot < n_slots, + "admit() returned an unknown slot"); + if (result.slot < 0 || result.slot >= n_slots) return -1; + require(!active[(size_t)result.slot], + "admit() reused a live slot"); + if (active[(size_t)result.slot]) return -1; + active[(size_t)result.slot] = true; + decoding[(size_t)result.slot] = false; + remaining[(size_t)result.slot] = (int)prompt.size(); + return result.slot; + }; + + auto slice_limit = [&](const SeqEngine::StepPlan & plan) { + return engine.step_plan_limits((int)plan.decode.size()) + .max_prefill_tokens_per_sequence; + }; + + // Validate and apply one successful logical step to the checker mirror. + // Contract prefills below request one token at a time, so advanced means + // exactly one prompt token without reintroducing a progress counter. + auto apply_progress = [&](const SeqEngine::StepPlan & plan, + const SeqEngine::StepResult & result) { + const std::string protocol_error = + validate_step_result(plan, result, n_slots); + if (!protocol_error.empty()) { + violations.emplace_back(protocol_error); + return false; + } + require(result.ok(), "valid planned work must succeed"); + if (!result.ok()) return false; + + std::vector decode_answered((size_t)n_slots, false); + std::vector prefill_answered((size_t)n_slots, false); + + for (const SeqEngine::DecodeOutput & output : result.decode) { + require(!output.failed, + "contract-check decode work must not fail"); + if (output.slot < 0 || output.slot >= n_slots || + output.failed) { + continue; + } + decode_answered[(size_t)output.slot] = true; + next_token[(size_t)output.slot] = output.token; + } + + using PrefillStatus = SeqEngine::PrefillOutput::Status; + for (const SeqEngine::PrefillOutput & output : result.prefills) { + require(output.status != PrefillStatus::failed, + "contract-check prefill work must not fail"); + if (output.slot < 0 || output.slot >= n_slots || + output.status == PrefillStatus::failed) { + continue; + } + auto selected = std::find_if( + plan.prefills.begin(), plan.prefills.end(), + [&](const PrefillSlice & slice) { + return slice.slot == output.slot; + }); + if (selected == plan.prefills.end()) continue; + require(selected->max_tokens == 1, + "contract checker must use unit prefill slices"); + prefill_answered[(size_t)output.slot] = true; + if (output.status == PrefillStatus::advanced) { + require(remaining[(size_t)output.slot] > 1, + "prefill reported advanced for its final token"); + --remaining[(size_t)output.slot]; + } else { + require(remaining[(size_t)output.slot] == 1, + "prefill reported completion before its final token"); + remaining[(size_t)output.slot] = 0; + decoding[(size_t)output.slot] = true; + next_token[(size_t)output.slot] = output.token; + } + } + + for (const SeqEngine::StepInput & input : plan.decode) { + if (input.slot >= 0 && input.slot < n_slots) { + require(decode_answered[(size_t)input.slot], + "step() left a decoding slot without an output"); + } + } + for (const PrefillSlice & slice : plan.prefills) { + require(slice.max_tokens > 0 && + slice.max_tokens <= slice_limit(plan), + "StepPlan prefill slice exceeded engine limits"); + if (slice.slot >= 0 && slice.slot < n_slots) { + require(prefill_answered[(size_t)slice.slot], + "step() left a selected prefill without an output"); + } + } + return true; + }; + + auto execute = [&](const SeqEngine::StepPlan & plan) { + return apply_progress(plan, engine.step(plan)); + }; + + auto decode_inputs = [&]() { + std::vector inputs; + for (int slot = 0; slot < n_slots; ++slot) { + if (active[(size_t)slot] && decoding[(size_t)slot]) { + inputs.push_back({slot, next_token[(size_t)slot]}); + } + } + return inputs; + }; + + // Admit the long prompt in the same idle cohort only when the engine can + // actually execute two slices. K=1 engines admit it after the short + // prompt releases its staging resource, then exercise the same mixed path. + const int short_slot = record_admit(1, {11}, greedy); + int long_slot = -1; + if (idle_cohort_size >= 2) { + long_slot = record_admit(2, {21, 22, 23}, seeded); + } + if (short_slot < 0 || (idle_cohort_size >= 2 && long_slot < 0)) { + retire_all(); + return violations; + } + if (long_slot >= 0) { + require(short_slot != long_slot, + "two pending admissions must own distinct slots"); + } + + SeqEngine::StepPlan first_plan; + first_plan.prefills.push_back({short_slot, 1}); + if (long_slot >= 0) { + first_plan.prefills.push_back({long_slot, 1}); + } + if (!execute(first_plan)) { + retire_all(); + return violations; + } + require(decoding[(size_t)short_slot], + "the short prefill must complete in its selected slice"); + + if (long_slot < 0) { + long_slot = record_admit(2, {21, 22, 23}, seeded); + if (long_slot < 0) { + retire_all(); + return violations; + } + } + require(!decoding[(size_t)long_slot] && + remaining[(size_t)long_slot] > 0, + "long member must remain pending after the short member completes"); + + for (int iteration = 0; + remaining[(size_t)long_slot] > 0 && iteration < 8; + ++iteration) { + SeqEngine::StepPlan mixed; + mixed.decode = decode_inputs(); + const StepPlanLimits limits = + engine.step_plan_limits((int)mixed.decode.size()); + require(supports_one_prefill(limits), + "mixed step_plan_limits() must permit continued prefill"); + if (!supports_one_prefill(limits)) break; + mixed.prefills.push_back({long_slot, 1}); + if (!execute(mixed)) break; + } + require(remaining[(size_t)long_slot] == 0, + "mixed prefill did not complete within bounded progress steps"); + + // Full decode coverage, including the token handoff back into the engine. + SeqEngine::StepPlan decode_plan; + decode_plan.decode = decode_inputs(); + require(decode_plan.decode.size() == 2, + "both completed admissions must enter decode"); + if (!execute(decode_plan)) { + retire_all(); + return violations; + } + + // A full engine is retryable admission pressure, not a request error. + if (n_slots == 2) { + const SeqEngine::AdmitResult full = engine.admit(3, {31}, greedy); + require(full.status == SeqEngine::AdmitResult::Status::busy && + full.slot < 0, + "a full engine must report busy without claiming a slot"); + } + + // Omitting a decoder or attaching prefill work to a decoding slot is a + // terminal plan-validation failure and must not partially advance state. + auto require_failed = [&](const SeqEngine::StepPlan & invalid, + const char * message) { + const SeqEngine::StepResult result = engine.step(invalid); + require(!result.ok(), message); + require(!result.error.empty(), + "failed step must explain the validation error"); + require(result.decode.empty() && result.prefills.empty(), + "failed step must not report partial progress"); + }; + + SeqEngine::StepPlan omitted; + omitted.decode.push_back(decode_plan.decode.front()); + require_failed(omitted, + "step() must reject a plan that omits a decoding slot"); + + SeqEngine::StepPlan invalid_prefill = decode_plan; + invalid_prefill.prefills.push_back({short_slot, 1}); + require_failed(invalid_prefill, + "step() must reject prefill work for a decoding slot"); + + // Start fresh and complete the advertised idle cohort in one step. When + // K=2 is available, this still catches engines that accidentally retain a + // scalar completion/output path. + retire_all(); + const int simultaneous_a = record_admit(10, {41}, greedy); + int simultaneous_b = -1; + if (idle_cohort_size >= 2) { + simultaneous_b = record_admit(11, {51}, seeded); + } + if (simultaneous_a >= 0 && + (idle_cohort_size < 2 || simultaneous_b >= 0)) { + SeqEngine::StepPlan simultaneous; + simultaneous.prefills.push_back({simultaneous_a, 1}); + if (simultaneous_b >= 0) { + simultaneous.prefills.push_back({simultaneous_b, 1}); + } + execute(simultaneous); + require(decoding[(size_t)simultaneous_a], + "selected prefill must report completion"); + if (simultaneous_b >= 0) { + require(decoding[(size_t)simultaneous_b], + "one K=2 step must report simultaneous completions"); + } + } + + // Retire is idempotent, frees capacity, and cancels unfinished prefill. + const int freed = simultaneous_a; + if (freed >= 0) { + engine.retire(freed); + engine.retire(freed); + active[(size_t)freed] = false; + decoding[(size_t)freed] = false; + const int replacement = record_admit(12, {61, 62}, greedy); + require(replacement >= 0, + "admit() must reuse capacity after retire()"); + if (replacement >= 0) { + engine.retire(replacement); + engine.retire(replacement); + active[(size_t)replacement] = false; + decoding[(size_t)replacement] = false; + remaining[(size_t)replacement] = 0; + } + } + + retire_all(); + const SeqEngine::StepResult idle = engine.step({}); + require(idle.ok(), "step() with no work must succeed"); + require(idle.decode.empty() && idle.prefills.empty() && + idle.error.empty(), + "idle result must not retain work or errors"); + + const int reused = record_admit(13, {71}, greedy); + require(reused >= 0, + "an engine must admit again after every slot retires"); + retire_all(); + return violations; +} + +} // namespace dflash::common diff --git a/server/test/test_client_send_buffer.cpp b/server/test/test_client_send_buffer.cpp new file mode 100644 index 000000000..b5a1f91c9 --- /dev/null +++ b/server/test/test_client_send_buffer.cpp @@ -0,0 +1,149 @@ +// Host-side unit test for ClientSendBuffer (concurrent scheduler writer). +// +// Uses a socketpair with a shrunken send buffer to exercise partial drains, +// backpressure, the stall policy, and hard-error detection. + +#include "server/client_send_buffer.h" +#include "host_check.h" + +#include +#include +#include +#include +#include +#include + +using namespace dflash::common; +using clock_t_ = std::chrono::steady_clock; +using std::chrono::seconds; + +static int g_checks = 0; + +static void make_pair(int fds[2]) { + CHECK(socketpair(AF_UNIX, SOCK_STREAM, 0, fds) == 0); + // Non-blocking writer (matches the HTTP server's client sockets) and a + // small send buffer so a few KB reliably back-pressures. + CHECK(fcntl(fds[0], F_SETFL, O_NONBLOCK) == 0); + const int sndbuf = 4096; + CHECK(setsockopt(fds[0], SOL_SOCKET, SO_SNDBUF, + &sndbuf, sizeof(sndbuf)) == 0); +} + +static std::string drain_peer(int fd) { + std::string got; + char tmp[4096]; + for (;;) { + const ssize_t n = recv(fd, tmp, sizeof(tmp), MSG_DONTWAIT); + if (n <= 0) break; + got.append(tmp, (size_t)n); + } + return got; +} + +int main() { + // Happy path: everything fits the socket buffer in one flush. + { + int fds[2]; + make_pair(fds); + ClientSendBuffer buffer; + CHECK(buffer.empty()); + buffer.append("hello "); + buffer.append("world"); + CHECK(buffer.pending() == 11); + CHECK(buffer.flush(fds[0])); + CHECK(buffer.empty()); + CHECK(drain_peer(fds[1]) == "hello world"); + close(fds[0]); + close(fds[1]); + } + + // Backpressure: a payload larger than the socket buffer drains across + // flushes as the peer reads, and never blocks. + { + int fds[2]; + make_pair(fds); + ClientSendBuffer buffer; + std::string payload(256 * 1024, '\0'); + for (size_t i = 0; i < payload.size(); ++i) { + payload[i] = static_cast((i * 131 + 17) & 0xff); + } + buffer.append(payload); + CHECK(buffer.flush(fds[0])); + CHECK(!buffer.empty()); // socket buffer full, remainder pending + std::string got; + for (int i = 0; i < 1000 && !buffer.empty(); i++) { + got += drain_peer(fds[1]); + CHECK(buffer.flush(fds[0])); + } + got += drain_peer(fds[1]); + CHECK(buffer.empty()); + CHECK(got == payload); + close(fds[0]); + close(fds[1]); + } + + // A single flush that sends past the compaction threshold but leaves a + // tail must still count as progress. Compaction normalizes off_ back to + // zero, so progress cannot be inferred from the post-compaction offset. + { + int fds[2]; + CHECK(socketpair(AF_UNIX, SOCK_STREAM, 0, fds) == 0); + CHECK(fcntl(fds[0], F_SETFL, O_NONBLOCK) == 0); + const int sndbuf = 256 * 1024; + CHECK(setsockopt(fds[0], SOL_SOCKET, SO_SNDBUF, + &sndbuf, sizeof(sndbuf)) == 0); + + ClientSendBuffer buffer; + const std::string payload(2 * 1024 * 1024, 'p'); + buffer.append(payload); + buffer.mark_progress(clock_t_::now() - seconds(3600)); + CHECK(buffer.flush(fds[0])); + CHECK(!buffer.empty()); + CHECK(payload.size() - buffer.pending() > (64u << 10)); + CHECK(!buffer.should_drop(clock_t_::now(), seconds(30), + payload.size())); + close(fds[0]); + close(fds[1]); + } + + // Stall policy: pending bytes with no progress trip the deadline and the + // cap; an empty send buffer never drops; progress resets the clock. + { + int fds[2]; + make_pair(fds); + ClientSendBuffer buffer; + const auto t0 = clock_t_::now(); + buffer.mark_progress(t0); + CHECK(!buffer.should_drop(t0 + seconds(3600), seconds(30), 1u << 20)); + + buffer.append(std::string(64 * 1024, 'y')); + CHECK(buffer.flush(fds[0])); // fills the 4 KB socket buffer, stalls + CHECK(!buffer.empty()); + CHECK(!buffer.should_drop(clock_t_::now(), seconds(30), 1u << 20)); + CHECK(buffer.should_drop(clock_t_::now() + seconds(31), seconds(30), + 1u << 20)); + CHECK(buffer.should_drop(clock_t_::now(), seconds(30), /*cap=*/1024)); + + // Reader progress resets the deadline. + drain_peer(fds[1]); + CHECK(buffer.flush(fds[0])); + CHECK(!buffer.should_drop(clock_t_::now() + seconds(29), seconds(30), + 1u << 20)); + close(fds[0]); + close(fds[1]); + } + + // Hard error: peer closed -> flush returns false. + { + int fds[2]; + make_pair(fds); + close(fds[1]); + ClientSendBuffer buffer; + buffer.append("doomed"); + CHECK(!buffer.flush(fds[0])); + close(fds[0]); + } + + std::printf("OK test_client_send_buffer (%d checks)\n", g_checks); + return 0; +} diff --git a/server/test/test_deepseek4_page_layout.cpp b/server/test/test_deepseek4_page_layout.cpp new file mode 100644 index 000000000..5b98665cf --- /dev/null +++ b/server/test/test_deepseek4_page_layout.cpp @@ -0,0 +1,53 @@ +#include "deepseek4/deepseek4_page_layout.h" +#include "host_check.h" + +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +int main() { + CHECK(ds4_raw_ring_row(0) == 0); + CHECK(ds4_raw_ring_row(127) == 127); + CHECK(ds4_raw_ring_row(128) == 0); + CHECK(ds4_raw_ring_row(255) == 127); + + uint64_t row = 999; + bool emitted = true; + CHECK(ds4_compressed_page_row(3, 7, 4, row, emitted)); + CHECK(emitted && row == 7 * 32); + CHECK(ds4_compressed_page_row(4, 7, 4, row, emitted)); + CHECK(!emitted && row == 7 * 32); + CHECK(ds4_compressed_page_row(127, 7, 4, row, emitted)); + CHECK(emitted && row == 7 * 32 + 31); + CHECK(ds4_compressed_page_row(128, 42, 4, row, emitted)); + CHECK(!emitted); + CHECK(ds4_compressed_page_row(131, 42, 4, row, emitted)); + CHECK(emitted && row == 42 * 32); + CHECK(ds4_compressed_page_row(255, 3, 4, row, emitted)); + CHECK(emitted && row == 3 * 32 + 31); + + CHECK(ds4_compressed_page_row(127, 91, 128, row, emitted)); + CHECK(emitted && row == 91); + CHECK(ds4_compressed_page_row(128, 2, 128, row, emitted)); + CHECK(!emitted); + CHECK(ds4_compressed_page_row(255, 2, 128, row, emitted)); + CHECK(emitted && row == 2); + + uint64_t capacity = 0; + CHECK(ds4_compressed_page_capacity(5, 4, capacity) && capacity == 160); + CHECK(ds4_compressed_page_capacity(5, 128, capacity) && capacity == 5); + CHECK(!ds4_compressed_page_capacity(5, 0, capacity)); + CHECK(!ds4_compressed_page_capacity(5, 16, capacity)); + CHECK(!ds4_compressed_page_capacity( + std::numeric_limits::max(), 4, capacity)); + CHECK(!ds4_compressed_page_row(3, + std::numeric_limits::max(), 4, row, emitted)); + CHECK(!ds4_compressed_page_row(3, 0, 16, row, emitted)); + + std::printf("OK test_deepseek4_page_layout (%d checks)\n", g_checks); + return 0; +} diff --git a/server/test/test_deepseek4_paged_cache.cpp b/server/test/test_deepseek4_paged_cache.cpp new file mode 100644 index 000000000..dd9681bf0 --- /dev/null +++ b/server/test/test_deepseek4_paged_cache.cpp @@ -0,0 +1,86 @@ +#include "deepseek4/deepseek4_paged_cache.h" +#include "host_check.h" +#include +#include +using namespace dflash::common; +static int g_checks = 0; +int main() { + DeepSeek4PagedCachePlan p, twice; + CHECK(plan_deepseek4_paged_cache(512, 128, 3, 4096, 40, {0, 4, 128}, p)); + CHECK(p.max_blocks_per_sequence == 32 && p.physical_rows[0] == 0); + CHECK(p.physical_rows[1] == 1280 && p.physical_rows[2] == 40); + CHECK(p.raw_bytes == uint64_t(3) * 512 * 128 * 3 * 2); + CHECK(p.metadata_bytes == uint64_t(32 * 3 + 3 + 3) * 4); + CHECK(plan_deepseek4_paged_cache(512, 128, 6, 4096, 40, {0, 4, 128}, twice)); + // Paged rows are shared; only raw rings, metadata, and compressor state scale by slots. + CHECK(twice.compressed_bytes == p.compressed_bytes); + CHECK(twice.raw_bytes == p.raw_bytes * 2 && twice.state_bytes == p.state_bytes * 2); + DeepSeek4PagedCachePlan sixteen; + CHECK(plan_deepseek4_paged_cache(512, 128, 16, 4096, 40, + {0, 4, 128}, sixteen)); + CHECK(sixteen.slots == 16 && sixteen.max_blocks_per_sequence == 32); + CHECK(sixteen.compressed_bytes == p.compressed_bytes); + CHECK(sixteen.raw_bytes == p.raw_bytes / 3 * 16); + CHECK(!plan_deepseek4_paged_cache(512, 128, 1, 4096, 40, {4, 16}, twice)); + CHECK(!plan_deepseek4_paged_cache(512, 128, 1, 4096, + std::numeric_limits::max(), {4}, twice)); + + const int32_t slots[] = {2, 5, -1}; + const int64_t positions[] = {3, 259, 999}; + const int32_t tables[] = {4, 3, 2, 6, 1, 7, 0, 0, 0}; + std::vector rows; + CHECK(prepare_deepseek4_gathered_lane_rows( + slots, positions, 3, tables, 3, 8, 4, rows)); + CHECK(rows.size() == 3); + CHECK(rows[0].raw_history == std::vector({256, 257, 258})); + CHECK(rows[0].raw_scatter == 259); + CHECK(rows[0].compressed_emitted && rows[0].compressed_scatter == 4 * 32); + CHECK(rows[0].compressed_history.empty()); + CHECK(rows[1].raw_history.size() == 127); + CHECK(rows[1].raw_history.front() == 5 * 128 + 4); + CHECK(rows[1].raw_history.back() == 5 * 128 + 2); + CHECK(rows[1].compressed_history.size() == 64); + CHECK(rows[1].compressed_history.front() == 6 * 32); + CHECK(rows[1].compressed_history[31] == 6 * 32 + 31); + CHECK(rows[1].compressed_history[32] == 1 * 32); + CHECK(rows[1].compressed_history.back() == 1 * 32 + 31); + CHECK(rows[1].compressed_emitted && rows[1].compressed_scatter == 7 * 32); + CHECK(rows[2].raw_history.empty() && rows[2].compressed_history.empty()); + CHECK(rows[2].raw_scatter == -1 && rows[2].compressed_scatter == -1); + + std::vector sixteen_slots(16); + std::vector sixteen_positions(16, 0); + std::vector sixteen_tables(16); + for (int i = 0; i < 16; ++i) { + sixteen_slots[(size_t) i] = i; + sixteen_tables[(size_t) i] = i; + } + CHECK(prepare_deepseek4_gathered_lane_rows( + sixteen_slots.data(), sixteen_positions.data(), 16, + sixteen_tables.data(), 1, 16, 4, rows)); + CHECK(rows.size() == 16); + for (int i = 0; i < 16; ++i) { + CHECK(rows[(size_t) i].slot == i); + CHECK(rows[(size_t) i].raw_history.empty()); + CHECK(rows[(size_t) i].raw_scatter == int64_t(i * 128)); + } + + const int32_t boundary_slot[] = {1}; + const int32_t boundary_table[] = {0, 1}; + for (int64_t pos : {127LL, 128LL, 129LL}) { + CHECK(prepare_deepseek4_gathered_lane_rows( + boundary_slot, &pos, 1, boundary_table, 2, 2, 0, rows)); + CHECK(rows[0].raw_history.size() == (pos == 127 ? 127u : 127u)); + CHECK(rows[0].raw_history.front() == 128 + (pos == 127 ? 0 : pos - 127)); + CHECK(rows[0].raw_history.back() == 128 + ((pos - 1) % 128)); + } + + const int64_t ratio128_pos[] = {255}; + CHECK(prepare_deepseek4_gathered_lane_rows( + slots, ratio128_pos, 1, tables, 3, 8, 128, rows)); + CHECK(rows[0].compressed_history.size() == 1); + CHECK(rows[0].compressed_history[0] == 4); + CHECK(rows[0].compressed_emitted && rows[0].compressed_scatter == 3); + std::printf("OK test_deepseek4_paged_cache (%d checks)\n", g_checks); + return 0; +} diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index 399197359..2c1442e3d 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -53,13 +53,13 @@ static BackendArgs gate_args_hip_deepseek4() { return args; } -static std::string gate_result( +static bool gate_accepts( const BackendArgs & args, const std::string & arch, PlacementBackend backend, const BackendFeatureConfig & features = {}) { return check_feature_compatibility( - args, features, arch, backend, backend); + args, features, arch, backend, backend).empty(); } static std::string gate_result_for_binary( @@ -75,15 +75,15 @@ static std::string gate_result_for_binary( static void test_feature_gate_accepts_plain_launch() { BackendArgs args; args.model_path = "/nonexistent/model.gguf"; - TEST_ASSERT(gate_result( - args, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts( + args, "qwen35", PlacementBackend::Cuda)); } static void test_feature_gate_rejects_undetected_arch() { BackendArgs args; args.model_path = "/nonexistent/model.gguf"; - TEST_ASSERT(!gate_result( - args, "", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + args, "", PlacementBackend::Cuda)); } static void test_feature_gate_requires_compiled_target_backend() { @@ -99,14 +99,14 @@ static void test_feature_gate_ipc_options_require_ipc_binary() { BackendArgs draft; draft.model_path = "/nonexistent/model.gguf"; draft.remote_draft.work_dir = "/tmp/draft"; - TEST_ASSERT(!gate_result( - draft, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + draft, "qwen35", PlacementBackend::Cuda)); BackendArgs target; target.model_path = "/nonexistent/model.gguf"; target.remote_target_shard.work_dir = "/tmp/target"; - TEST_ASSERT(!gate_result( - target, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + target, "qwen35", PlacementBackend::Cuda)); } static void test_feature_gate_mixed_draft_placement_requires_ipc() { @@ -116,16 +116,16 @@ static void test_feature_gate_mixed_draft_placement_requires_ipc() { args.device.backend = PlacementBackend::Cuda; args.draft_device.backend = PlacementBackend::Hip; - TEST_ASSERT(!gate_result( - args, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + args, "qwen35", PlacementBackend::Cuda)); args.remote_draft.ipc_bin = "/usr/bin/draft-ipc"; - TEST_ASSERT(gate_result( - args, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts( + args, "qwen35", PlacementBackend::Cuda)); args.draft_device.backend = PlacementBackend::Cuda; - TEST_ASSERT(!gate_result( - args, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + args, "qwen35", PlacementBackend::Cuda)); } static void test_feature_gate_pflash_requires_drafter_and_supported_arch() { @@ -134,39 +134,39 @@ static void test_feature_gate_pflash_requires_drafter_and_supported_arch() { BackendFeatureConfig features; features.pflash_enabled = true; - TEST_ASSERT(!gate_result( - args, "qwen35", PlacementBackend::Cuda, features).empty()); + TEST_ASSERT(!gate_accepts( + args, "qwen35", PlacementBackend::Cuda, features)); features.pflash_drafter_configured = true; - TEST_ASSERT(gate_result( - args, "gemma4", PlacementBackend::Cuda, features).empty()); + TEST_ASSERT(gate_accepts( + args, "gemma4", PlacementBackend::Cuda, features)); args.device.backend = PlacementBackend::Cuda; args.draft_device.backend = PlacementBackend::Hip; args.remote_draft.ipc_bin = "/usr/bin/draft-ipc"; - TEST_ASSERT(!gate_result( - args, "gemma4", PlacementBackend::Cuda, features).empty()); - TEST_ASSERT(gate_result( - args, "qwen35", PlacementBackend::Cuda, features).empty()); + TEST_ASSERT(!gate_accepts( + args, "gemma4", PlacementBackend::Cuda, features)); + TEST_ASSERT(gate_accepts( + args, "qwen35", PlacementBackend::Cuda, features)); } static void test_feature_gate_validates_target_split_topology() { BackendArgs weights; weights.model_path = "/nonexistent/model.gguf"; weights.device.layer_split_weights = {1.0, 1.0}; - TEST_ASSERT(!gate_result( - weights, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + weights, "qwen35", PlacementBackend::Cuda)); BackendArgs mixed; mixed.model_path = "/nonexistent/model.gguf"; TEST_ASSERT(parse_placement_device_list( "cuda:0,hip:0", mixed.device)); - TEST_ASSERT(!gate_result( - mixed, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + mixed, "qwen35", PlacementBackend::Cuda)); mixed.remote_target_shard.ipc_bin = "/usr/bin/target-shard"; - TEST_ASSERT(gate_result( - mixed, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts( + mixed, "qwen35", PlacementBackend::Cuda)); BackendArgs two_boundaries; two_boundaries.model_path = "/nonexistent/model.gguf"; @@ -174,8 +174,8 @@ static void test_feature_gate_validates_target_split_topology() { "cuda:0,hip:0,cuda:1", two_boundaries.device)); two_boundaries.remote_target_shard.ipc_bin = "/usr/bin/target-shard"; - TEST_ASSERT(!gate_result( - two_boundaries, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + two_boundaries, "qwen35", PlacementBackend::Cuda)); } static void test_feature_gate_tensor_parallel_requirements() { @@ -184,52 +184,52 @@ static void test_feature_gate_tensor_parallel_requirements() { TEST_ASSERT(parse_placement_device_list( "cuda:0,cuda:1", valid.device)); valid.device.split_mode = TargetSplitMode::Tensor; - TEST_ASSERT(gate_result( - valid, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts( + valid, "qwen35", PlacementBackend::Cuda)); BackendArgs missing_devices; missing_devices.model_path = "/nonexistent/model.gguf"; missing_devices.device.split_mode = TargetSplitMode::Tensor; - TEST_ASSERT(!gate_result( - missing_devices, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + missing_devices, "qwen35", PlacementBackend::Cuda)); - TEST_ASSERT(!gate_result( - valid, "laguna", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + valid, "laguna", PlacementBackend::Cuda)); BackendArgs hip; hip.model_path = "/nonexistent/model.gguf"; TEST_ASSERT(parse_placement_device_list("hip:0,hip:1", hip.device)); hip.device.split_mode = TargetSplitMode::Tensor; - TEST_ASSERT(!gate_result( - hip, "qwen35", PlacementBackend::Hip).empty()); + TEST_ASSERT(!gate_accepts( + hip, "qwen35", PlacementBackend::Hip)); BackendArgs mixed = valid; TEST_ASSERT(parse_placement_device_list( "cuda:0,hip:0", mixed.device)); mixed.device.split_mode = TargetSplitMode::Tensor; - TEST_ASSERT(!gate_result( - mixed, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + mixed, "qwen35", PlacementBackend::Cuda)); BackendArgs weighted = valid; weighted.device.layer_split_weights = {1.0, 1.0}; - TEST_ASSERT(!gate_result( - weighted, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + weighted, "qwen35", PlacementBackend::Cuda)); BackendArgs remote = valid; remote.remote_target_shard.ipc_bin = "/usr/bin/target-shard"; - TEST_ASSERT(!gate_result( - remote, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + remote, "qwen35", PlacementBackend::Cuda)); BackendFeatureConfig pflash; pflash.pflash_enabled = true; pflash.pflash_drafter_configured = true; - TEST_ASSERT(!gate_result( - valid, "qwen35", PlacementBackend::Cuda, pflash).empty()); + TEST_ASSERT(!gate_accepts( + valid, "qwen35", PlacementBackend::Cuda, pflash)); BackendArgs draft = valid; draft.draft_path = "/nonexistent/draft.gguf"; - TEST_ASSERT(gate_result( - draft, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts( + draft, "qwen35", PlacementBackend::Cuda)); } static void test_feature_gate_ds4_prefill_requires_deepseek4() { @@ -237,10 +237,10 @@ static void test_feature_gate_ds4_prefill_requires_deepseek4() { args.ds4_prefill_mode_set = true; args.ds4_prefill_mode = PrefillAttentionMode::Dense; - TEST_ASSERT(!gate_result( - args, "qwen35", PlacementBackend::Hip).empty()); - TEST_ASSERT(gate_result( - args, "deepseek4", PlacementBackend::Hip).empty()); + TEST_ASSERT(!gate_accepts( + args, "qwen35", PlacementBackend::Hip)); + TEST_ASSERT(gate_accepts( + args, "deepseek4", PlacementBackend::Hip)); } static void test_feature_gate_approximate_ds4_prefill_requires_local_hip() { @@ -249,60 +249,60 @@ static void test_feature_gate_approximate_ds4_prefill_requires_local_hip() { args.ds4_prefill_mode = PrefillAttentionMode::Sparse; // CUDA has no approximate prefill path. - TEST_ASSERT(!gate_result( - args, "deepseek4", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + args, "deepseek4", PlacementBackend::Cuda)); // Neither does the layer-split adapter, even on HIP. BackendArgs split = args; TEST_ASSERT(parse_placement_device_list("hip:0,hip:1", split.device)); - TEST_ASSERT(!gate_result( - split, "deepseek4", PlacementBackend::Hip).empty()); + TEST_ASSERT(!gate_accepts( + split, "deepseek4", PlacementBackend::Hip)); // Nor a remote target shard. BackendArgs remote = args; remote.remote_target_shard.ipc_bin = "/usr/bin/shard"; - TEST_ASSERT(!gate_result( - remote, "deepseek4", PlacementBackend::Hip).empty()); + TEST_ASSERT(!gate_accepts( + remote, "deepseek4", PlacementBackend::Hip)); // Single local HIP device is the supported placement. - TEST_ASSERT(gate_result( - args, "deepseek4", PlacementBackend::Hip).empty()); + TEST_ASSERT(gate_accepts( + args, "deepseek4", PlacementBackend::Hip)); // Exact prefill is unrestricted. BackendArgs exact = gate_args_hip_deepseek4(); exact.ds4_prefill_mode_set = true; exact.ds4_prefill_mode = PrefillAttentionMode::Exact; - TEST_ASSERT(gate_result( - exact, "deepseek4", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts( + exact, "deepseek4", PlacementBackend::Cuda)); } static void test_feature_gate_ds4_decode_options_require_monolithic_hip() { BackendArgs fused = gate_args_hip_deepseek4(); fused.ds4_fused_decode = true; - TEST_ASSERT(!gate_result( - fused, "deepseek4", PlacementBackend::Cuda).empty()); - TEST_ASSERT(gate_result( - fused, "deepseek4", PlacementBackend::Hip).empty()); + TEST_ASSERT(!gate_accepts( + fused, "deepseek4", PlacementBackend::Cuda)); + TEST_ASSERT(gate_accepts( + fused, "deepseek4", PlacementBackend::Hip)); BackendArgs topk = gate_args_hip_deepseek4(); topk.ds4_expert_top_k = 4; - TEST_ASSERT(!gate_result( - topk, "qwen35", PlacementBackend::Hip).empty()); - TEST_ASSERT(gate_result( - topk, "deepseek4", PlacementBackend::Hip).empty()); + TEST_ASSERT(!gate_accepts( + topk, "qwen35", PlacementBackend::Hip)); + TEST_ASSERT(gate_accepts( + topk, "deepseek4", PlacementBackend::Hip)); // Top-k is a model policy in the monolithic backend and is independent of // the GPU vendor. Unlike fused decode, mixed CUDA-primary expert // placement can therefore use it. BackendArgs cuda_topk = topk; cuda_topk.device.backend = PlacementBackend::Cuda; - TEST_ASSERT(gate_result( - cuda_topk, "deepseek4", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts( + cuda_topk, "deepseek4", PlacementBackend::Cuda)); BackendArgs split_topk = topk; split_topk.device.layer_split_gpus = {0, 1}; - TEST_ASSERT(!gate_result( - split_topk, "deepseek4", PlacementBackend::Hip).empty()); + TEST_ASSERT(!gate_accepts( + split_topk, "deepseek4", PlacementBackend::Hip)); } static void test_feature_gate_remote_draft_requires_supported_arch() { @@ -313,16 +313,16 @@ static void test_feature_gate_remote_draft_requires_supported_arch() { args.draft_device.backend = PlacementBackend::Hip; args.remote_draft.ipc_bin = "/usr/bin/draft-ipc"; - TEST_ASSERT(!gate_result( - args, "gemma4", PlacementBackend::Cuda).empty()); - TEST_ASSERT(gate_result( - args, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + args, "gemma4", PlacementBackend::Cuda)); + TEST_ASSERT(gate_accepts( + args, "qwen35", PlacementBackend::Cuda)); // Without a draft model or PFlash, remote draft IPC is unnecessary. BackendArgs no_draft = args; no_draft.draft_path = nullptr; - TEST_ASSERT(!gate_result( - no_draft, "gemma4", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + no_draft, "gemma4", PlacementBackend::Cuda)); } static void test_feature_gate_layer_split_requires_supported_arch() { @@ -332,34 +332,37 @@ static void test_feature_gate_layer_split_requires_supported_arch() { // These four have a layer-split adapter. for (const char * arch : {"qwen35", "laguna", "gemma4", "deepseek4"}) { - TEST_ASSERT(gate_result(args, arch, PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts(args, arch, PlacementBackend::Cuda)); } // These two do not: the factory would hand the split placement to a // monolithic backend, which reads only the primary GPU. for (const char * arch : {"qwen35moe", "qwen3"}) { - TEST_ASSERT(!gate_result(args, arch, PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts(args, arch, PlacementBackend::Cuda)); } // Single-device placement is unaffected for the same architectures. BackendArgs single; single.model_path = "/nonexistent/model.gguf"; - TEST_ASSERT(gate_result(single, "qwen35moe", PlacementBackend::Cuda).empty()); - TEST_ASSERT(gate_result(single, "qwen3", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts(single, "qwen35moe", PlacementBackend::Cuda)); + TEST_ASSERT(gate_accepts(single, "qwen3", PlacementBackend::Cuda)); } -static void test_feature_gate_paged_attention_requires_qwen35_monolithic() { +static void test_feature_gate_paged_attention_requires_supported_monolithic() { BackendArgs args; args.model_path = "/nonexistent/model.gguf"; args.paged_attention = true; - TEST_ASSERT(gate_result(args, "qwen35", PlacementBackend::Cuda).empty()); - TEST_ASSERT(gate_result(args, "qwen35", PlacementBackend::Hip).empty()); - - // Only qwen35 has a paged decode path. qwen35moe shares Qwen35Config, so + TEST_ASSERT(gate_accepts(args, "qwen35", PlacementBackend::Cuda)); + TEST_ASSERT(gate_accepts(args, "qwen35", PlacementBackend::Hip)); + TEST_ASSERT(!gate_accepts(args, "deepseek4", PlacementBackend::Cuda)); + TEST_ASSERT(gate_accepts(args, "deepseek4", PlacementBackend::Hip)); + + // Qwen35 supports both GPU backends; DeepSeek4 concurrent paging is + // intentionally limited to a monolithic HIP target. qwen35moe shares + // Qwen35Config, so // its rejection is this gate's job — the factory's field-presence // cross-check cannot tell the two apart. - for (const char * arch : {"qwen35moe", "laguna", "qwen3", - "gemma4", "deepseek4"}) { - TEST_ASSERT(!gate_result(args, arch, PlacementBackend::Cuda).empty()); + for (const char * arch : {"qwen35moe", "laguna", "qwen3", "gemma4"}) { + TEST_ASSERT(!gate_accepts(args, arch, PlacementBackend::Cuda)); } // Only the monolithic qwen35 backend owns a paged K/V pool. Both @@ -367,20 +370,44 @@ static void test_feature_gate_paged_attention_requires_qwen35_monolithic() { // rejection has to come from the paged rule. BackendArgs split = args; TEST_ASSERT(parse_placement_device_list("cuda:0,cuda:1", split.device)); - TEST_ASSERT(!gate_result(split, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts(split, "qwen35", PlacementBackend::Cuda)); + TEST_ASSERT(!gate_accepts(split, "deepseek4", PlacementBackend::Cuda)); BackendArgs remote_shard = args; remote_shard.remote_target_shard.ipc_bin = "/usr/bin/target-shard"; - TEST_ASSERT(!gate_result( - remote_shard, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + remote_shard, "qwen35", PlacementBackend::Cuda)); for (BackendArgs * relaxed : {&split, &remote_shard}) { relaxed->paged_attention = false; - TEST_ASSERT(gate_result( - *relaxed, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts( + *relaxed, "qwen35", PlacementBackend::Cuda)); } } +static void test_feature_gate_deepseek4_paged_reference_constraints() { + BackendArgs args = gate_args_hip_deepseek4(); + args.paged_attention = true; + args.max_concurrency = 16; + TEST_ASSERT(gate_accepts(args, "deepseek4", PlacementBackend::Hip)); + + BackendArgs split = args; + TEST_ASSERT(parse_placement_device_list("hip:0,hip:1", split.device)); + TEST_ASSERT(!gate_accepts(split, "deepseek4", PlacementBackend::Hip)); + BackendArgs remote = args; + remote.remote_target_shard.ipc_bin = "/usr/bin/target-shard"; + TEST_ASSERT(!gate_accepts(remote, "deepseek4", PlacementBackend::Hip)); + + args.max_concurrency = 17; + TEST_ASSERT(!gate_accepts(args, "deepseek4", PlacementBackend::Hip)); + args.max_concurrency = 16; + args.ds4_fused_decode = true; + TEST_ASSERT(!gate_accepts(args, "deepseek4", PlacementBackend::Hip)); + args.ds4_fused_decode = false; + args.ds4_prefill_mode = PrefillAttentionMode::Dense; + TEST_ASSERT(!gate_accepts(args, "deepseek4", PlacementBackend::Hip)); +} + static void test_feature_gate_paged_attention_requires_plain_ar_decode() { BackendArgs base; base.model_path = "/nonexistent/model.gguf"; @@ -388,49 +415,122 @@ static void test_feature_gate_paged_attention_requires_plain_ar_decode() { BackendArgs draft = base; draft.draft_path = "/nonexistent/draft.gguf"; - TEST_ASSERT(!gate_result(draft, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts(draft, "qwen35", PlacementBackend::Cuda)); BackendArgs ddtree = base; ddtree.ddtree_mode = true; - TEST_ASSERT(!gate_result(ddtree, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts(ddtree, "qwen35", PlacementBackend::Cuda)); BackendArgs windowed = base; windowed.fa_window = 4096; - TEST_ASSERT(!gate_result( - windowed, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + windowed, "qwen35", PlacementBackend::Cuda)); BackendFeatureConfig pflash; pflash.pflash_enabled = true; pflash.pflash_drafter_configured = true; - TEST_ASSERT(!gate_result( - base, "qwen35", PlacementBackend::Cuda, pflash).empty()); + TEST_ASSERT(!gate_accepts( + base, "qwen35", PlacementBackend::Cuda, pflash)); + + BackendFeatureConfig kvflash; + kvflash.kvflash_enabled = true; + TEST_ASSERT(!gate_accepts( + base, "qwen35", PlacementBackend::Cuda, kvflash)); // The pool rounds max_ctx up to whole blocks, so both ends of the range // are rejected: nothing to allocate, and rounding that overflows int. BackendArgs empty_ctx = base; empty_ctx.device.max_ctx = 0; - TEST_ASSERT(!gate_result( - empty_ctx, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + empty_ctx, "qwen35", PlacementBackend::Cuda)); BackendArgs huge_ctx = base; huge_ctx.device.max_ctx = INT_MAX; - TEST_ASSERT(!gate_result( - huge_ctx, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(!gate_accepts( + huge_ctx, "qwen35", PlacementBackend::Cuda)); BackendArgs max_ctx = base; max_ctx.device.max_ctx = INT_MAX - PAGED_BLOCK_SIZE + 1; - TEST_ASSERT(gate_result( - max_ctx, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts( + max_ctx, "qwen35", PlacementBackend::Cuda)); // None of these are rules about paged attention itself: without the flag // every one of them is a supported qwen35 launch. for (BackendArgs * args : {&draft, &ddtree, &windowed, &empty_ctx, &huge_ctx}) { args->paged_attention = false; - TEST_ASSERT(gate_result(*args, "qwen35", PlacementBackend::Cuda).empty()); + TEST_ASSERT(gate_accepts(*args, "qwen35", PlacementBackend::Cuda)); } } +static void test_feature_gate_parallel_and_kv_pool_rules() { + // The zero/one values are validated for every architecture. + BackendArgs plain; + plain.model_path = "/nonexistent/model.gguf"; + plain.max_concurrency = 0; + TEST_ASSERT(!gate_accepts(plain, "qwen35", PlacementBackend::Cuda)); + plain.max_concurrency = 1; + TEST_ASSERT(gate_accepts(plain, "qwen35", PlacementBackend::Cuda)); + + // Qwen keeps its existing single-sequence paged path; concurrent slots are + // intentionally limited to monolithic HIP DeepSeek4. + BackendArgs qwen = plain; + qwen.paged_attention = true; + qwen.max_concurrency = 2; + TEST_ASSERT(!gate_accepts(qwen, "qwen35", PlacementBackend::Cuda)); + + BackendArgs parallel = gate_args_hip_deepseek4(); + parallel.max_concurrency = 2; + TEST_ASSERT(!gate_accepts( + parallel, "deepseek4", PlacementBackend::Hip)); + parallel.paged_attention = true; + TEST_ASSERT(gate_accepts( + parallel, "deepseek4", PlacementBackend::Hip)); + + // Lane counts need not be powers of two; the gathered graph supports 1..16. + parallel.max_concurrency = 3; + TEST_ASSERT(gate_accepts( + parallel, "deepseek4", PlacementBackend::Hip)); + parallel.max_concurrency = 16; + TEST_ASSERT(gate_accepts( + parallel, "deepseek4", PlacementBackend::Hip)); + parallel.max_concurrency = 17; + TEST_ASSERT(!gate_accepts( + parallel, "deepseek4", PlacementBackend::Hip)); + + // --kv-pool-tokens sizes the shared pool, so it needs slots to share. + BackendArgs pool = gate_args_hip_deepseek4(); + pool.paged_attention = true; + pool.kv_pool_tokens = 4096; + TEST_ASSERT(!gate_accepts( + pool, "deepseek4", PlacementBackend::Hip)); + pool.max_concurrency = 2; + TEST_ASSERT(gate_accepts( + pool, "deepseek4", PlacementBackend::Hip)); + + pool.kv_pool_tokens = PAGED_BLOCK_SIZE - 1; + TEST_ASSERT(!gate_accepts( + pool, "deepseek4", PlacementBackend::Hip)); + pool.kv_pool_tokens = PAGED_BLOCK_SIZE; + TEST_ASSERT(gate_accepts( + pool, "deepseek4", PlacementBackend::Hip)); + const long long max_pool_tokens = paged_kv_address_cap(); + pool.kv_pool_tokens = max_pool_tokens + 1; + TEST_ASSERT(!gate_accepts( + pool, "deepseek4", PlacementBackend::Hip)); + pool.kv_pool_tokens = max_pool_tokens; + TEST_ASSERT(gate_accepts( + pool, "deepseek4", PlacementBackend::Hip)); + + // Automatic capacity remains memory-derived instead of max_ctx * slots. + BackendArgs large = gate_args_hip_deepseek4(); + large.paged_attention = true; + large.max_concurrency = 2; + large.device.max_ctx = 1 << 30; + TEST_ASSERT(gate_accepts( + large, "deepseek4", PlacementBackend::Hip)); +} + // ── Inert-flag warnings ───────────────────────────────────────────────── // Warnings must never gate admission, so each case also asserts the same // configuration passes check_feature_compatibility(). @@ -558,10 +658,12 @@ static void test_model_capability_tables() { TEST_ASSERT(!arch_supports_draft_swa("qwen36", false)); TEST_ASSERT(!arch_supports_paged_attention("qwen36", false)); - // Paged decode lives in the monolithic qwen35 backend alone. + // Paged decode is monolithic for both supported architectures. TEST_ASSERT(arch_supports_paged_attention("qwen35", false)); TEST_ASSERT(!arch_supports_paged_attention("qwen35", true)); TEST_ASSERT(!arch_supports_paged_attention("qwen35moe", false)); + TEST_ASSERT(arch_supports_paged_attention("deepseek4", false)); + TEST_ASSERT(!arch_supports_paged_attention("deepseek4", true)); } int main() { @@ -579,8 +681,10 @@ int main() { RUN_TEST(test_feature_gate_ds4_decode_options_require_monolithic_hip); RUN_TEST(test_feature_gate_remote_draft_requires_supported_arch); RUN_TEST(test_feature_gate_layer_split_requires_supported_arch); - RUN_TEST(test_feature_gate_paged_attention_requires_qwen35_monolithic); + RUN_TEST(test_feature_gate_paged_attention_requires_supported_monolithic); + RUN_TEST(test_feature_gate_deepseek4_paged_reference_constraints); RUN_TEST(test_feature_gate_paged_attention_requires_plain_ar_decode); + RUN_TEST(test_feature_gate_parallel_and_kv_pool_rules); RUN_TEST(test_feature_warnings_silent_when_supported); RUN_TEST(test_feature_warnings_report_inert_draft); RUN_TEST(test_feature_warnings_report_inert_decode_tunables); diff --git a/server/test/test_kvflash_pool_sizing.cpp b/server/test/test_kvflash_pool_sizing.cpp index 37b1fbbf0..07091be97 100644 --- a/server/test/test_kvflash_pool_sizing.cpp +++ b/server/test/test_kvflash_pool_sizing.cpp @@ -20,6 +20,11 @@ struct KvflashPoolSizingFixture {}; } TEST_CASE(KvflashPoolSizingFixture, kvflash_pool_sizing_suite) { + REQUIRE(!kvflash_fixed_pool_requested(nullptr)); + REQUIRE(!kvflash_fixed_pool_requested("0")); + REQUIRE(!kvflash_fixed_pool_requested("auto")); + REQUIRE(kvflash_fixed_pool_requested("4096")); + { const luce_test::ScopedEnvVar kvflash_off("DFLASH_KVFLASH", "0"); REQUIRE(kvflash_pool_from_env(131072) == 0); diff --git a/server/test/test_paged_kv_pool.cpp b/server/test/test_paged_kv_pool.cpp index 76367185d..71fe183c4 100644 --- a/server/test/test_paged_kv_pool.cpp +++ b/server/test/test_paged_kv_pool.cpp @@ -2,7 +2,8 @@ #define GENERATE_UNIT_TEST_MAIN #include "CppUnitTestFramework.hpp" -#include "../src/common/paged_kv_pool.h" +#include "common/paged_kv_pool.h" +#include "../src/common/paged_attention_config.h" #include #include @@ -50,17 +51,23 @@ static bool state_unchanged(PagedKvPool & pool, const auto after = sequence(pool, handle); return after.kv_seq_len == before.kv_seq_len && after.block_table == before.block_table && + after.reserved_block_count == before.reserved_block_count && pool.free_block_count() == free_before; } // True when every handle-taking operation rejects `handle` as stale. static bool all_ops_stale(PagedKvPool & pool, PagedKvSequenceHandle handle) { PagedKvSequenceSnapshot snapshot; + uint32_t owned_blocks = 0; return pool.append(handle, 1).status == PagedKvStatus::StaleHandle && pool.append(handle, 1, /*only_first_last_slots=*/true).status == PagedKvStatus::StaleHandle && + pool.reserve_capacity(handle, 16) == + PagedKvStatus::StaleHandle && pool.release(handle) == PagedKvStatus::StaleHandle && - pool.sequence(handle, snapshot) == PagedKvStatus::StaleHandle; + pool.sequence(handle, snapshot) == PagedKvStatus::StaleHandle && + pool.owned_block_count(handle, owned_blocks) == + PagedKvStatus::StaleHandle; } TEST_CASE(PagedKvPoolFixture, block_boundaries) { @@ -103,6 +110,126 @@ TEST_CASE(PagedKvPoolFixture, nondefault_block_size) { CHECK(append.write_slots[14].physical_token_index == 14); } +TEST_CASE(PagedKvPoolFixture, reserved_acquire_feeds_chunked_append) { + PagedKvPool pool(/*physical_block_count=*/6, + /*max_sequences=*/3, /*block_size=*/16); + PagedKvSequenceHandle handle; + CHECK(pool.acquire_reserved(77, /*token_capacity=*/40, handle) == + PagedKvStatus::Ok); + CHECK(pool.active_sequence_count() == 1); + CHECK(pool.free_block_count() == 3); + + auto snapshot = sequence(pool, handle); + CHECK(snapshot.kv_seq_len == 0); + CHECK(snapshot.block_table.empty()); + CHECK(snapshot.reserved_block_count == 3); + uint32_t owned_blocks = 0; + CHECK(pool.owned_block_count(handle, owned_blocks) == PagedKvStatus::Ok); + CHECK(owned_blocks == 3); + + const auto first = pool.append(handle, 17); + CHECK(first.status == PagedKvStatus::Ok); + CHECK(equals(sequence(pool, handle).block_table, {0, 1})); + CHECK(sequence(pool, handle).reserved_block_count == 1); + CHECK(pool.owned_block_count(handle, owned_blocks) == PagedKvStatus::Ok); + CHECK(owned_blocks == 3); + // Consuming a reservation never changes globally available capacity. + CHECK(pool.free_block_count() == 3); + + const auto tail = pool.append(handle, 23); + CHECK(tail.status == PagedKvStatus::Ok); + snapshot = sequence(pool, handle); + CHECK(snapshot.kv_seq_len == 40); + CHECK(equals(snapshot.block_table, {0, 1, 2})); + CHECK(snapshot.reserved_block_count == 0); + CHECK(pool.free_block_count() == 3); + CHECK(pool.owned_block_count(handle, owned_blocks) == PagedKvStatus::Ok); + CHECK(owned_blocks == 3); + + CHECK(pool.release(handle) == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 6); +} + +TEST_CASE(PagedKvPoolFixture, reserved_acquire_is_atomic_and_isolated) { + PagedKvPool pool(/*physical_block_count=*/4, + /*max_sequences=*/3, /*block_size=*/16); + PagedKvSequenceHandle first; + CHECK(pool.acquire_reserved(1, /*token_capacity=*/48, first) == + PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 1); + + PagedKvSequenceHandle unchanged{77, 88}; + CHECK(pool.acquire_reserved(2, /*token_capacity=*/32, unchanged) == + PagedKvStatus::BlocksExhausted); + CHECK(unchanged.slot == 77); + CHECK(unchanged.generation == 88); + CHECK(pool.active_sequence_count() == 1); + CHECK(pool.free_block_count() == 1); + + // Unreserved work cannot steal the first sequence's promised pages. + const auto second = acquire(pool, 2); + CHECK(pool.append(second, 32).status == PagedKvStatus::BlocksExhausted); + CHECK(pool.append(first, 48).status == PagedKvStatus::Ok); + CHECK(equals(sequence(pool, first).block_table, {0, 1, 2})); + + CHECK(pool.release(first) == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 4); + CHECK(pool.release(second) == PagedKvStatus::Ok); + + // Retirement also returns reservations that prefill never consumed. + PagedKvSequenceHandle unused; + CHECK(pool.acquire_reserved(3, /*token_capacity=*/32, unused) == + PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 2); + CHECK(sequence(pool, unused).reserved_block_count == 2); + CHECK(pool.release(unused) == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 4); +} + +TEST_CASE(PagedKvPoolFixture, capacity_top_up_is_atomic_and_private) { + PagedKvPool pool(/*physical_block_count=*/3, + /*max_sequences=*/2, /*block_size=*/16); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pool.append(first, 16).status == PagedKvStatus::Ok); + CHECK(pool.append(second, 16).status == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 1); + + // Asking for two more blocks fails atomically: neither the logical length + // nor either ownership view changes. + CHECK(state_unchanged(pool, first, [&] { + CHECK(pool.reserve_capacity(first, /*token_capacity=*/48) == + PagedKvStatus::BlocksExhausted); + })); + + // A one-page top-up succeeds without advancing logical state or exposing + // a block-table entry. + CHECK(pool.reserve_capacity(first, /*token_capacity=*/32) == + PagedKvStatus::Ok); + auto first_state = sequence(pool, first); + CHECK(first_state.kv_seq_len == 16); + CHECK(equals(first_state.block_table, {0})); + CHECK(first_state.reserved_block_count == 1); + CHECK(pool.free_block_count() == 0); + + // The other sequence cannot steal the promised page; its own state stays + // unchanged while the owner can consume the reservation. + CHECK(state_unchanged(pool, second, [&] { + CHECK(pool.append(second, 1).status == + PagedKvStatus::BlocksExhausted); + })); + CHECK(pool.append(first, 1).status == PagedKvStatus::Ok); + first_state = sequence(pool, first); + CHECK(equals(first_state.block_table, {0, 2})); + CHECK(first_state.reserved_block_count == 0); + + // Top-ups are monotonic no-ops when the sequence already owns enough. + CHECK(state_unchanged(pool, first, [&] { + CHECK(pool.reserve_capacity(first, /*token_capacity=*/17) == + PagedKvStatus::Ok); + })); +} + TEST_CASE(PagedKvPoolFixture, zero_token_append_is_a_no_op) { PagedKvPool pool(4, 2, 16); const auto handle = acquire(pool, 7); @@ -273,3 +400,35 @@ TEST_CASE(PagedKvPoolFixture, invalid_arguments) { CHECK(overflow.write_slots.empty()); })); } + +TEST_CASE(PagedKvPoolFixture, auto_pool_sizing) { + PagedKvAutoBudget budget; + budget.free_bytes = 10'000; + budget.fixed_cache_bytes = 1'000; + budget.reserve_bytes = 1'000; + budget.bytes_per_token = 10; + // Memory could hold 800 tokens, but four 128-token logical contexts cap + // the useful physical pool at 512. + CHECK(paged_kv_auto_pool_tokens(128, 4, budget) == 512); + + budget.free_bytes = 5'000; + // 300 raw tokens round down to 288 (18 whole blocks). + CHECK(paged_kv_auto_pool_tokens(128, 4, budget) == 288); + budget.free_bytes = 1'999; + CHECK(paged_kv_auto_pool_tokens(128, 4, budget) == 0); + budget.free_bytes = 5'000; + budget.bytes_per_token = 0; + CHECK(paged_kv_auto_pool_tokens(128, 4, budget) == 0); + + // A representative 24 GiB-card post-weight budget can keep one 32K + // context plus oversubscription headroom without allocating 16 x 32K. + budget.free_bytes = 10LL * 1024 * 1024 * 1024; + budget.fixed_cache_bytes = 4LL * 1024 * 1024 * 1024; + budget.reserve_bytes = 1536LL * 1024 * 1024; + budget.bytes_per_token = 64 * 1024; + const int64_t headline = + paged_kv_auto_pool_tokens(32768, 16, budget); + CHECK(headline == 73728); + CHECK(headline >= 32768); + CHECK(headline < 16LL * 32768); +} diff --git a/server/test/test_seq_batch_plan.cpp b/server/test/test_seq_batch_plan.cpp new file mode 100644 index 000000000..0948c39ff --- /dev/null +++ b/server/test/test_seq_batch_plan.cpp @@ -0,0 +1,172 @@ +#include "common/concurrency/seq_engine.h" +#include "host_check.h" + +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +int main() { + const std::vector pending{ + {7, 30}, {2, 10}, {5, 20}, {1, 40}, + }; + + StepPlanLimits idle_limits{/*max sequences=*/2, + /*per sequence=*/2048, + /*total=*/4096, + /*allocation quantum=*/512}; + StepPlanLimits mixed_limits{/*max sequences=*/2, + /*per sequence=*/512, + /*total=*/1024, + /*allocation quantum=*/512}; + + auto idle = plan_prefill_slices( + pending, idle_limits); + CHECK(idle.size() == 2); + CHECK(idle[0].slot == 2); + CHECK(idle[1].slot == 5); + CHECK(idle[0].max_tokens == 2048); + CHECK(idle[1].max_tokens == 2048); + + auto mixed = plan_prefill_slices( + pending, mixed_limits); + CHECK(mixed.size() == 2); + CHECK(mixed[0].max_tokens == 512); + CHECK(mixed[1].max_tokens == 512); + + mixed_limits.max_prefill_tokens_per_sequence = 256; + mixed_limits.max_prefill_tokens_total = 512; + mixed = plan_prefill_slices(pending, mixed_limits); + CHECK(mixed.size() == 2); + CHECK(mixed[0].max_tokens == 256); + CHECK(mixed[1].max_tokens == 256); + + mixed_limits.max_prefill_tokens_total = 301; + mixed = plan_prefill_slices(pending, mixed_limits); + CHECK(mixed.size() == 2); + CHECK(mixed[0].slot == 2); + CHECK(mixed[0].max_tokens == 256); + CHECK(mixed[1].slot == 5); + CHECK(mixed[1].max_tokens == 45); + + auto rotated = plan_prefill_slices( + pending, mixed_limits, /*round_robin_start=*/1); + CHECK(rotated.size() == 2); + CHECK(rotated[0].slot == 2 && rotated[0].max_tokens == 45); + CHECK(rotated[1].slot == 5 && rotated[1].max_tokens == 256); + + // A budget smaller than one quantum advances one lane; rotation prevents + // the oldest lane from winning every step. + mixed_limits.max_prefill_tokens_total = 200; + auto clamped0 = plan_prefill_slices( + pending, mixed_limits, 0); + auto clamped1 = plan_prefill_slices( + pending, mixed_limits, 1); + CHECK(clamped0.size() == 1 && clamped0[0].slot == 2 && + clamped0[0].max_tokens == 200); + CHECK(clamped1.size() == 1 && clamped1[0].slot == 5 && + clamped1[0].max_tokens == 200); + + // The packed Qwen policy fills all eight idle lanes with one 512-token + // segment, while a mixed step rotates four such segments fairly. + const std::vector packed{ + {0, 0}, {1, 1}, {2, 2}, {3, 3}, + {4, 4}, {5, 5}, {6, 6}, {7, 7}, + }; + StepPlanLimits packed_mixed{/*max sequences=*/8, + /*per sequence=*/512, + /*total=*/2048, + /*allocation quantum=*/512}; + auto packed0 = plan_prefill_slices(packed, packed_mixed, 0); + CHECK(packed0.size() == 4); + for (int i = 0; i < 4; ++i) { + CHECK(packed0[(size_t)i].slot == i); + CHECK(packed0[(size_t)i].max_tokens == 512); + } + auto packed4 = plan_prefill_slices(packed, packed_mixed, 4); + CHECK(packed4.size() == 4); + for (int i = 0; i < 4; ++i) { + CHECK(packed4[(size_t)i].slot == i + 4); + CHECK(packed4[(size_t)i].max_tokens == 512); + } + + packed_mixed.prefill_allocation_quantum = 0; + CHECK(plan_prefill_slices(packed, packed_mixed).empty()); + + mixed_limits.max_prefill_tokens_per_sequence = 0; + CHECK(plan_prefill_slices( + pending, mixed_limits).empty()); + + idle_limits.max_prefill_sequences = 0; + CHECK(plan_prefill_slices( + pending, idle_limits).empty()); + CHECK(plan_prefill_slices({}, idle_limits).empty()); + + // The same model-neutral layer validates engine row ownership before the + // scheduler mutates socket/request state. + SeqEngine::StepPlan work; + work.decode = {{0, 7}}; + work.prefills = {{1, 4}}; + + SeqEngine::StepResult good; + good.decode.push_back({0, 11, false, {}}); + good.prefills.push_back({ + 1, SeqEngine::PrefillOutput::Status::advanced, -1, {}}); + CHECK(validate_step_result(work, good, 2).empty()); + + SeqEngine::StepResult complete = good; + complete.prefills[0] = { + 1, SeqEngine::PrefillOutput::Status::completed, 12, {}}; + CHECK(validate_step_result(work, complete, 2).empty()); + + SeqEngine::StepResult missing_decode = good; + missing_decode.decode.clear(); + CHECK(!validate_step_result(work, missing_decode, 2).empty()); + + SeqEngine::StepResult duplicate_decode = good; + duplicate_decode.decode.push_back({0, 12, false, {}}); + CHECK(!validate_step_result(work, duplicate_decode, 2).empty()); + + SeqEngine::StepResult missing_prefill = good; + missing_prefill.prefills.clear(); + CHECK(!validate_step_result(work, missing_prefill, 2).empty()); + + // 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"}; + CHECK(validate_step_result(work, prefill_failure, 2).empty()); + + SeqEngine::StepResult bad_row_failure = prefill_failure; + bad_row_failure.prefills.back().error.clear(); + CHECK(!validate_step_result(work, bad_row_failure, 2).empty()); + + SeqEngine::StepResult unknown_prefill = good; + unknown_prefill.prefills[0].status = + static_cast(-1); + CHECK(!validate_step_result(work, unknown_prefill, 2).empty()); + + SeqEngine::StepResult bad_decode = good; + bad_decode.decode[0] = {0, -1, true, {}}; + CHECK(!validate_step_result(work, bad_decode, 2).empty()); + + SeqEngine::StepResult failed; + failed.error = "device compute failed"; + CHECK(validate_step_result(work, failed, 2).empty()); + failed.decode.push_back({0, -1, true, "partial"}); + CHECK(!validate_step_result(work, failed, 2).empty()); + + SeqEngine::StepResult idle_result; + CHECK(validate_step_result({}, idle_result, 2).empty()); + CHECK(!validate_step_result(work, idle_result, 2).empty()); + + SeqEngine::StepPlan duplicate_plan = work; + duplicate_plan.decode.push_back({0, 8}); + CHECK(!validate_step_result(duplicate_plan, good, 2).empty()); + + std::printf("test_seq_batch_plan: %d checks passed\n", g_checks); + return 0; +} diff --git a/server/test/test_seq_engine_contract.cpp b/server/test/test_seq_engine_contract.cpp new file mode 100644 index 000000000..be3f031b4 --- /dev/null +++ b/server/test/test_seq_engine_contract.cpp @@ -0,0 +1,323 @@ +// Host-only mutation test for the model-neutral SeqEngine contract. + +#include "seq_engine_contract.h" +#include "host_check.h" + +#include +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +struct Faults { + bool hard_error_when_full = false; + bool reuse_live_slot = false; + bool drop_decode_output = false; + bool serialize_prefills = false; + bool accept_omitted_decode = false; + bool accept_invalid_prefill = false; + bool lose_other_pending = false; + bool overconsume_prefill = false; + bool drop_second_completion = false; + bool retire_leaks = false; +}; + +struct FakeCapabilities { + StepPlanLimits idle{2, 2, 4}; + StepPlanLimits mixed{2, 1, 2}; +}; + +class FakeSeqEngine final : public SeqEngine { +public: + explicit FakeSeqEngine(int count, Faults faults = {}, + FakeCapabilities capabilities = {}) + : slots_((size_t)count), faults_(faults), + capabilities_(capabilities) {} + + int slot_count() const override { return (int)slots_.size(); } + int max_context() const override { return 128; } + StepPlanLimits step_plan_limits(int decode_rows) const override { + return decode_rows > 0 ? capabilities_.mixed : capabilities_.idle; + } + bool token_is_eos(int32_t token) const override { return token == 2; } + + AdmitResult admit(uint64_t, + const std::vector & prompt, + const SamplerCfg &) override { + AdmitResult result; + if (prompt.empty() || prompt.size() > (size_t)max_context()) { + result.error = "invalid prompt"; + return result; + } + + int chosen = -1; + if (faults_.reuse_live_slot) { + chosen = 0; + } else { + for (size_t i = 0; i < slots_.size(); ++i) { + if (!slots_[i].active) { + chosen = (int)i; + break; + } + } + } + if (chosen < 0) { + result.status = faults_.hard_error_when_full + ? AdmitResult::Status::failed + : AdmitResult::Status::busy; + result.error = "all slots are live"; + return result; + } + + Slot & slot = slots_[(size_t)chosen]; + slot.active = true; + slot.prefilling = true; + slot.remaining = (int)prompt.size(); + slot.fed.clear(); + result.status = AdmitResult::Status::admitted; + result.slot = chosen; + return result; + } + + StepResult step(const StepPlan & plan) override { + StepResult result; + std::string validation_error; + if (!valid_decode(plan, validation_error) && + !faults_.accept_omitted_decode) { + result.error = validation_error; + return result; + } + if (!valid_prefills(plan, validation_error) && + !faults_.accept_invalid_prefill) { + result.error = validation_error; + return result; + } + + size_t decode_count = plan.decode.size(); + if (faults_.drop_decode_output && plan.prefills.empty() && + decode_count == 2) { + --decode_count; + } + for (size_t i = 0; i < decode_count; ++i) { + const StepInput & input = plan.decode[i]; + if (input.slot < 0 || input.slot >= slot_count()) continue; + Slot & slot = slots_[(size_t)input.slot]; + slot.fed.push_back(input.token); + result.decode.push_back({ + input.slot, + 100 + input.slot + (int32_t)slot.fed.size(), + false, {}, + }); + } + + std::vector completed_this_step; + for (size_t i = 0; i < plan.prefills.size(); ++i) { + if (faults_.serialize_prefills && i > 0) continue; + const PrefillSlice & slice = plan.prefills[i]; + if (slice.slot < 0 || slice.slot >= slot_count()) continue; + Slot & slot = slots_[(size_t)slice.slot]; + if (!slot.active || !slot.prefilling || slot.remaining <= 0) { + continue; + } + int consumed = std::min(slice.max_tokens, slot.remaining); + if (faults_.overconsume_prefill) consumed = slice.max_tokens + 1; + slot.remaining -= consumed; + if (slot.remaining <= 0) { + slot.prefilling = false; + completed_this_step.push_back(slice.slot); + const bool omit = faults_.drop_second_completion && i > 0; + if (!omit) { + result.prefills.push_back({ + slice.slot, PrefillOutput::Status::completed, + 100 + slice.slot, {}, + }); + } + } else { + result.prefills.push_back({ + slice.slot, PrefillOutput::Status::advanced, -1, {}, + }); + } + } + + if (faults_.lose_other_pending && !completed_this_step.empty()) { + for (Slot & slot : slots_) { + if (slot.active && slot.prefilling) slot.prefilling = false; + } + } + + return result; + } + + void retire(int slot) override { + if (slot < 0 || slot >= slot_count() || faults_.retire_leaks) return; + slots_[(size_t)slot] = Slot{}; + } + +private: + struct Slot { + bool active = false; + bool prefilling = false; + int remaining = 0; + std::vector fed; + }; + + int decoding_count() const { + int count = 0; + for (const Slot & slot : slots_) { + if (slot.active && !slot.prefilling) ++count; + } + return count; + } + + bool valid_decode(const StepPlan & plan, std::string & error) const { + std::vector seen(slots_.size(), false); + if ((int)plan.decode.size() != decoding_count()) { + error = "plan omits a decoding slot"; + return false; + } + for (const StepInput & input : plan.decode) { + if (input.slot < 0 || input.slot >= slot_count() || + seen[(size_t)input.slot] || input.token < 0 || + !slots_[(size_t)input.slot].active || + slots_[(size_t)input.slot].prefilling) { + error = "invalid decode input"; + return false; + } + seen[(size_t)input.slot] = true; + } + for (int slot = 0; slot < slot_count(); ++slot) { + if (slots_[(size_t)slot].active && + !slots_[(size_t)slot].prefilling && + !seen[(size_t)slot]) { + error = "plan omits a decoding slot"; + return false; + } + } + return true; + } + + bool valid_prefills(const StepPlan & plan, std::string & error) const { + const StepPlanLimits limits = + step_plan_limits((int)plan.decode.size()); + const int token_limit = limits.max_prefill_tokens_per_sequence; + if ((int)plan.prefills.size() > limits.max_prefill_sequences) { + error = "too many prefills"; + return false; + } + std::vector seen(slots_.size(), false); + int total_tokens = 0; + for (const PrefillSlice & slice : plan.prefills) { + if (slice.slot < 0 || slice.slot >= slot_count() || + seen[(size_t)slice.slot] || slice.max_tokens <= 0 || + slice.max_tokens > token_limit || + !slots_[(size_t)slice.slot].active || + !slots_[(size_t)slice.slot].prefilling) { + error = "invalid prefill slice"; + return false; + } + seen[(size_t)slice.slot] = true; + total_tokens += slice.max_tokens; + if (total_tokens > limits.max_prefill_tokens_total) { + error = "too many total prefill tokens"; + return false; + } + } + return true; + } + + std::vector slots_; + Faults faults_; + FakeCapabilities capabilities_; +}; + +static void print_violations(const char * label, + const std::vector & violations) { + for (const std::string & violation : violations) { + std::fprintf(stderr, " [%s] %s\n", label, violation.c_str()); + } +} + +static bool mentions(const std::vector & violations, + const char * needle) { + return std::any_of( + violations.begin(), violations.end(), + [&](const std::string & violation) { + return violation.find(needle) != std::string::npos; + }); +} + +int main() { + for (const int slots : {2, 4}) { + FakeSeqEngine engine(slots); + const auto violations = check_seq_engine_contract(engine); + if (!violations.empty()) print_violations("conforming", violations); + CHECK(violations.empty()); + } + + { + FakeCapabilities capabilities; + capabilities.idle = {1, 2, 2}; + capabilities.mixed = {1, 1, 1}; + FakeSeqEngine engine(2, {}, capabilities); + const auto violations = check_seq_engine_contract(engine); + if (!violations.empty()) print_violations("conforming-k1", violations); + CHECK(violations.empty()); + } + + { + FakeCapabilities capabilities; + capabilities.idle = {2, 2, 2}; + capabilities.mixed = {1, 1, 1}; + FakeSeqEngine engine(2, {}, capabilities); + const auto violations = check_seq_engine_contract(engine); + if (!violations.empty()) { + print_violations("conforming-width-dependent", violations); + } + CHECK(violations.empty()); + } + + struct Case { + const char * label; + bool Faults::*fault; + const char * expected; + }; + const Case cases[] = { + {"hard-error-full", &Faults::hard_error_when_full, "full engine"}, + {"reuse-live", &Faults::reuse_live_slot, "reused a live slot"}, + {"drop-decode", &Faults::drop_decode_output, "omitted an output"}, + {"serialize-prefill", &Faults::serialize_prefills, + "omitted an output"}, + {"accept-omitted", &Faults::accept_omitted_decode, + "omits a decoding slot"}, + {"accept-invalid-prefill", &Faults::accept_invalid_prefill, + "prefill work for a decoding slot"}, + {"lose-pending", &Faults::lose_other_pending, + "valid planned work must succeed"}, + {"overconsume", &Faults::overconsume_prefill, + "completion before its final token"}, + {"scalar-completion", &Faults::drop_second_completion, + "omitted an output"}, + {"retire-leak", &Faults::retire_leaks, + "succeed while a slot is free"}, + }; + + for (const Case & test : cases) { + Faults faults; + faults.*test.fault = true; + FakeSeqEngine engine(2, faults); + const auto violations = check_seq_engine_contract(engine); + if (!mentions(violations, test.expected)) { + std::fprintf(stderr, + "FAIL %s: expected violation containing '%s'\n", + test.label, test.expected); + print_violations(test.label, violations); + } + CHECK(mentions(violations, test.expected)); + } + + std::printf("test_seq_engine_contract: %d checks passed\n", g_checks); + return 0; +} diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp new file mode 100644 index 000000000..465f74556 --- /dev/null +++ b/server/test/test_seq_slot_manager.cpp @@ -0,0 +1,426 @@ +// Host-side unit test for SeqSlotManager (concurrent decode slots). +// +// Mirrors test_paged_kv_pool: no model, ggml, or GPU required. Covers the +// admission checks, atomic prompt-plus-headroom reservation, rolling decode +// protection, block-table deltas, exhaustion atomicity, and pool-handle +// lifecycle across retire. + +#include "common/concurrency/seq_slot_manager.h" +#include "host_check.h" + +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +static SamplerCfg greedy_sampler() { + return SamplerCfg{}; +} + +static std::vector prompt_tokens(int count) { + return std::vector((size_t)count, 1); +} + +static SeqEngine::AdmitResult admit( + SeqSlotManager & manager, uint64_t request_id, + const std::vector & prompt, const SamplerCfg & sampler) { + return manager.admit(request_id, prompt, sampler); +} + +static bool is_admitted(const SeqEngine::AdmitResult & result) { + return result.status == SeqEngine::AdmitResult::Status::admitted; +} + +static bool is_busy(const SeqEngine::AdmitResult & result) { + return result.status == SeqEngine::AdmitResult::Status::busy; +} + +int main() { + // 8 blocks x 16 tokens = 128 pool tokens, 2 slots, per-seq max_ctx 64. + { + PagedKvPool pool(/*physical_block_count=*/8, /*max_sequences=*/2, + /*block_size=*/16); + SeqSlotManager mgr(pool, /*max_ctx=*/64); + CHECK(mgr.slot_count() == 2); + CHECK(mgr.max_context() == 64); + CHECK(!mgr.is_active(0) && !mgr.is_active(1)); + + // Invalid asks are hard errors, not busy. + CHECK(!is_admitted(admit(mgr, 1, prompt_tokens(0), greedy_sampler()))); + CHECK(!is_admitted(admit(mgr, 1, prompt_tokens(65), greedy_sampler()))); + CHECK(!mgr.is_active(0) && !mgr.is_active(1)); + CHECK(pool.active_sequence_count() == 0); + + // Admission reserves the prompt's two blocks plus its next decode page + // without advancing logical length; chunked append consumes only the + // prompt portion of that private reservation. + const std::vector admitted_prompt = prompt_tokens(20); + auto a = admit(mgr, 1, admitted_prompt, greedy_sampler()); + CHECK(is_admitted(a) && !is_busy(a)); + CHECK(a.slot == 0); + CHECK(pool.free_block_count() == 5); + CHECK(mgr.is_active(0)); + CHECK(mgr.slot(0).sample_history == admitted_prompt); + + // Prompt allocation follows the chunks actually scheduled. The first + // ten rows open block 0; the next ten consume its tail and open block + // 1, returning only that block-table delta. + auto p0 = mgr.append_prefill(a.slot, 10); + CHECK(p0.ok && !p0.busy); + CHECK(p0.rows.size() == 10 && p0.rows.front() == 0 && + p0.rows.back() == 9); + CHECK(p0.first_new_block == 0 && p0.new_blocks.size() == 1 && + p0.new_blocks[0] == 0); + CHECK(pool.free_block_count() == 5); + auto p1 = mgr.append_prefill(a.slot, 10); + CHECK(p1.ok && !p1.busy); + CHECK(p1.rows.size() == 10 && p1.rows.front() == 10 && + p1.rows.back() == 19); + CHECK(p1.first_new_block == 1 && p1.new_blocks.size() == 1 && + p1.new_blocks[0] == 1); + CHECK(pool.free_block_count() == 5); + mgr.commit_prefill(0); + CHECK(mgr.slot(0).cur_pos == 20); + + // Decode appends: row allocation + sample_history; cur_pos advances + // separately after the step's compute. + auto st = mgr.append_token(0, /*fed_token=*/42); + CHECK(st.ok); + CHECK(st.position == 20); + CHECK(st.physical_row == 20); // tail of the prompt's last block + CHECK(st.new_block < 0 && st.new_block_index < 0); + CHECK(mgr.slot(0).cur_pos == 20); + CHECK(mgr.slot(0).sample_history.size() == 21 && + mgr.slot(0).sample_history.back() == 42); + mgr.commit_step(0); + CHECK(mgr.slot(0).cur_pos == 21); + + // Second admission lands in slot 1 with non-identity rows. + auto b = admit(mgr, 2, prompt_tokens(20), greedy_sampler()); + CHECK(is_admitted(b) && b.slot == 1); + auto pb = mgr.append_prefill(b.slot, 20); + CHECK(pb.ok && pb.rows.front() != 0); + mgr.commit_prefill(1); + + // Third admission: no free slot -> busy. + auto c = admit(mgr, 3, prompt_tokens(8), greedy_sampler()); + CHECK(!is_admitted(c) && is_busy(c)); + + // Retire frees the slot AND the pool blocks. + const uint32_t free_before = pool.free_block_count(); + mgr.retire(0); + CHECK(!mgr.is_active(0)); + CHECK(pool.free_block_count() > free_before); + CHECK(mgr.is_active(1)); + + // The freed slot admits again. + auto d = admit(mgr, 3, prompt_tokens(8), greedy_sampler()); + CHECK(is_admitted(d) && d.slot == 0); + + // Inactive-slot calls are safe no-ops. + mgr.retire(0); + mgr.retire(0); + CHECK(!mgr.append_prefill(0, 1).ok); + CHECK(!mgr.append_token(0, 1).ok); + mgr.commit_step(0); + mgr.retire(-1); + mgr.retire(99); + } + + // A physically impossible headroom page does not make an otherwise useful + // prompt impossible: this one-block pool falls back to prompt-only. + { + PagedKvPool pool(/*physical_block_count=*/1, + /*max_sequences=*/2, /*block_size=*/16); + SeqSlotManager mgr(pool, /*max_ctx=*/128); + auto a = admit(mgr, 1, prompt_tokens(8), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(pool.free_block_count() == 0); + CHECK(mgr.append_prefill(a.slot, 8).ok); + CHECK(pool.free_block_count() == 0); + } + + // Busy-vs-never-fits classification against a small pool. + { + // 4 blocks x 16 = 64 pool tokens, 2 slots, max_ctx 64. + PagedKvPool pool(4, 2, /*block_size=*/16); + SeqSlotManager mgr(pool, 64); + + // A near-limit prompt needs no additional physical page before max_ctx. + auto exact = admit(mgr, 1, prompt_tokens(60), greedy_sampler()); + CHECK(is_admitted(exact) && mgr.max_context() == 64); + CHECK(pool.free_block_count() == 0); + CHECK(mgr.append_prefill(exact.slot, 60).ok); + mgr.retire(exact.slot); + + // Occupy most of the pool, then a second request that WOULD fit an + // empty pool reports busy (blocks held by a live sequence). + auto big = admit(mgr, 2, prompt_tokens(48), greedy_sampler()); // 3 prompt + 1 headroom + CHECK(is_admitted(big)); + CHECK(mgr.append_prefill(big.slot, 48).ok); + auto blocked = admit(mgr, 3, prompt_tokens(32), greedy_sampler()); // 2 blocks + CHECK(!is_admitted(blocked) && is_busy(blocked)); + mgr.retire(big.slot); + auto now_fits = admit(mgr, 3, prompt_tokens(32), greedy_sampler()); + CHECK(is_admitted(now_fits)); + } + + // Never-fits: a prompt beyond the WHOLE pool is a hard error, not + // busy — waiting for a drain could never help. + { + // 4 blocks x 16 = 64 pool tokens, but max_ctx allows asking for more. + PagedKvPool pool(4, 2, /*block_size=*/16); + SeqSlotManager mgr(pool, /*max_ctx=*/128); + auto never = admit(mgr, 1, prompt_tokens(100), greedy_sampler()); + CHECK(!is_admitted(never) && !is_busy(never)); // prompt 100 > pool 64 + CHECK(pool.active_sequence_count() == 0); + + // Impossibility wins over temporary slot pressure: do not queue an + // oversized prompt merely because every sequence slot is occupied. + auto live = admit(mgr, 2, prompt_tokens(16), greedy_sampler()); + CHECK(is_admitted(live)); + auto live2 = admit(mgr, 3, prompt_tokens(16), greedy_sampler()); + CHECK(is_admitted(live2)); + auto still_never = admit(mgr, 4, prompt_tokens(100), greedy_sampler()); + CHECK(!is_admitted(still_never) && !is_busy(still_never)); + } + + // Aggregate prompt reservations prevent two partial prefills from + // consuming the same capacity and reaching a no-progress deadlock. + { + PagedKvPool pool(/*physical_block_count=*/4, + /*max_sequences=*/3, /*block_size=*/16); + SeqSlotManager mgr(pool, /*max_ctx=*/64); + auto a = admit(mgr, 1, prompt_tokens(48), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(pool.free_block_count() == 0); + + // The second prompt fits the whole pool, but cannot reserve its prompt + // plus headroom while the first request owns all four pages. It never + // becomes a second partially-filled live slot. + auto later = admit(mgr, 2, prompt_tokens(32), greedy_sampler()); + CHECK(!is_admitted(later) && is_busy(later)); + CHECK(pool.active_sequence_count() == 1); + + auto first_chunk = mgr.append_prefill(a.slot, 16); + CHECK(first_chunk.ok && !first_chunk.busy); + CHECK(pool.free_block_count() == 0); + // Consuming a reserved page does not make it available to another + // admission, and appending past the admitted prompt is a hard error. + later = admit(mgr, 2, prompt_tokens(32), greedy_sampler()); + CHECK(!is_admitted(later) && is_busy(later)); + auto too_much = mgr.append_prefill(a.slot, 33); + CHECK(!too_much.ok && !too_much.busy); + auto final_chunk = mgr.append_prefill(a.slot, 32); + CHECK(final_chunk.ok && !final_chunk.busy); + mgr.commit_prefill(a.slot); + + mgr.retire(a.slot); + CHECK(pool.free_block_count() == 4); + later = admit(mgr, 2, prompt_tokens(32), greedy_sampler()); + CHECK(is_admitted(later)); + mgr.retire(later.slot); + } + + // When the entire physical pool cannot hold one prompt page plus one decode + // page, admission falls back to prompt-only and exhaustion stays retryable. + { + PagedKvPool pool(/*physical_block_count=*/1, + /*max_sequences=*/1, /*block_size=*/16); + SeqSlotManager mgr(pool, /*max_ctx=*/32); + auto a = admit(mgr, 1, prompt_tokens(16), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(mgr.append_prefill(a.slot, 16).ok); + mgr.commit_prefill(a.slot); + auto decode_blocked = mgr.append_token(a.slot, 77); + CHECK(!decode_blocked.ok && decode_blocked.busy); + CHECK(mgr.slot(a.slot).sample_history == prompt_tokens(16)); + CHECK(mgr.slot(a.slot).cur_pos == 16); + } + + // Rolling headroom has priority over younger admission. Even when a + // decoder consumes its reserved page while the pool is full, blocks freed + // later are topped up for that decoder before a new prompt can claim them. + { + PagedKvPool pool(/*physical_block_count=*/4, + /*max_sequences=*/2, /*block_size=*/16); + SeqSlotManager mgr(pool, /*max_ctx=*/48); + auto older = admit(mgr, 1, prompt_tokens(16), greedy_sampler()); + auto peer = admit(mgr, 2, prompt_tokens(16), greedy_sampler()); + CHECK(is_admitted(older) && is_admitted(peer)); + CHECK(pool.free_block_count() == 0); + CHECK(mgr.append_prefill(older.slot, 16).ok); + CHECK(mgr.append_prefill(peer.slot, 16).ok); + mgr.commit_prefill(older.slot); + mgr.commit_prefill(peer.slot); + + // Enter the initially reserved next page. No block is free yet to + // replenish a third page, but this decode token still succeeds. + auto enter_second = mgr.append_token(older.slot, 70); + CHECK(enter_second.ok && enter_second.new_block >= 0); + mgr.commit_step(older.slot); + PagedKvSequenceSnapshot snapshot; + CHECK(pool.sequence(mgr.slot(older.slot).handle, snapshot) == + PagedKvStatus::Ok); + CHECK(snapshot.block_table.size() == 2); + CHECK(snapshot.reserved_block_count == 0); + + // Retirement frees two pages. A younger admission first restores the + // older decoder's third-page reserve, then reports busy because only + // one page remains for its own two-page admission. + mgr.retire(peer.slot); + CHECK(pool.free_block_count() == 2); + auto younger = admit(mgr, 3, prompt_tokens(16), greedy_sampler()); + CHECK(!is_admitted(younger) && is_busy(younger)); + CHECK(pool.sequence(mgr.slot(older.slot).handle, snapshot) == + PagedKvStatus::Ok); + CHECK(snapshot.reserved_block_count == 1); + CHECK(pool.free_block_count() == 1); + + // The protected third page remains consumable at the next boundary. + for (int i = 0; i < 15; ++i) { + CHECK(mgr.append_token(older.slot, 71 + i).ok); + mgr.commit_step(older.slot); + } + auto enter_third = mgr.append_token(older.slot, 99); + CHECK(enter_third.ok && enter_third.new_block >= 0); + mgr.commit_step(older.slot); + + mgr.retire(older.slot); + younger = admit(mgr, 3, prompt_tokens(16), greedy_sampler()); + CHECK(is_admitted(younger)); + } + + // A row whose current append fits its existing page must not consume the + // last free block as speculative future headroom before a later row that + // needs that block for its current append. Rolling top-up is therefore an + // admission-wide operation, never a side effect of append_token(). + { + PagedKvPool pool(/*physical_block_count=*/5, + /*max_sequences=*/3, /*block_size=*/4); + SeqSlotManager mgr(pool, /*max_ctx=*/16); + auto first = admit(mgr, 1, prompt_tokens(4), greedy_sampler()); + auto boundary = admit(mgr, 2, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(first) && is_admitted(boundary)); + + PagedKvSequenceHandle blocker; + CHECK(pool.acquire_reserved(99, /*token_capacity=*/4, blocker) == + PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 0); + CHECK(mgr.append_prefill(first.slot, 4).ok); + CHECK(mgr.append_prefill(boundary.slot, 4).ok); + mgr.commit_prefill(first.slot); + mgr.commit_prefill(boundary.slot); + + auto first_second_page = mgr.append_token(first.slot, 10); + CHECK(first_second_page.ok && first_second_page.new_block_index == 1); + mgr.commit_step(first.slot); // position 5, inside its second page + for (int i = 0; i < 4; ++i) { + CHECK(mgr.append_token(boundary.slot, 20 + i).ok); + mgr.commit_step(boundary.slot); + } + CHECK(mgr.slot(boundary.slot).cur_pos == 8); // needs page 3 next + + CHECK(pool.release(blocker) == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 1); + CHECK(mgr.append_token(first.slot, 30).ok); + CHECK(pool.free_block_count() == 1); + auto boundary_append = mgr.append_token(boundary.slot, 31); + CHECK(boundary_append.ok && boundary_append.new_block_index == 2); + CHECK(pool.free_block_count() == 0); + } + + // Context exhaustion: append_token refuses past max_ctx. + { + PagedKvPool pool(4, 1, /*block_size=*/16); + SeqSlotManager mgr(pool, /*max_ctx=*/17); + auto a = admit(mgr, 1, prompt_tokens(16), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(mgr.append_prefill(a.slot, 16).ok); + mgr.commit_prefill(0); + auto s1 = mgr.append_token(0, 7); // position 16 (== max_ctx-1) + CHECK(s1.ok && s1.position == 16); + CHECK(s1.new_block == 1 && s1.new_block_index == 1); + mgr.commit_step(0); + auto s2 = mgr.append_token(0, 8); // cur_pos == max_ctx -> refuse + CHECK(!s2.ok); + } + + // Prefilling lifecycle: an admitted slot stays out of the decode batch + // until commit_prefill() makes its first sampled token available. + { + PagedKvPool pool(8, 2, /*block_size=*/16); + SeqSlotManager mgr(pool, 64); + auto a = admit(mgr, 1, prompt_tokens(20), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(mgr.is_prefilling(a.slot)); + CHECK(mgr.is_active(a.slot)); + CHECK(mgr.decoding_count() == 0); + CHECK(!mgr.append_token(a.slot, 42).ok); + + CHECK(mgr.append_prefill(a.slot, 20).ok); + mgr.commit_prefill(a.slot); + CHECK(!mgr.is_prefilling(a.slot)); + CHECK(mgr.decoding_count() == 1); + CHECK(mgr.append_token(a.slot, 42).ok); + CHECK(!mgr.append_prefill(a.slot, 1).ok); + + // A second admission can prefill while the first slot decodes. + auto b = admit(mgr, 2, prompt_tokens(20), greedy_sampler()); + CHECK(is_admitted(b)); + CHECK(mgr.is_active(a.slot) && mgr.is_active(b.slot)); + CHECK(mgr.decoding_count() == 1); + + // Retiring during prefill clears the state before slot reuse. + mgr.retire(b.slot); + CHECK(!mgr.is_prefilling(b.slot)); + auto c = admit(mgr, 3, prompt_tokens(8), greedy_sampler()); + CHECK(is_admitted(c) && c.slot == b.slot); + CHECK(mgr.is_prefilling(c.slot)); + } + + // Seeded sampling RNG is deterministic per admission. The sampler alone + // decides: a seed is honoured exactly when needs_logit_processing() says + // the slot actually draws, so there is no way to ask for seeded sampling + // and be given argmax (or the reverse). + { + PagedKvPool pool(4, 1, /*block_size=*/16); + SeqSlotManager mgr(pool, 64); + SamplerCfg cfg = greedy_sampler(); + cfg.temp = 0.7f; // needs_logit_processing() -> true + cfg.seed = 1234; + CHECK(cfg.needs_logit_processing()); + auto a = admit(mgr, 1, prompt_tokens(4), cfg); + CHECK(is_admitted(a)); + const uint64_t first = mgr.slot(0).rng(); + mgr.retire(0); + auto b = admit(mgr, 2, prompt_tokens(4), cfg); + CHECK(is_admitted(b)); + CHECK(mgr.slot(0).rng() == first); + + // A different seed is a different stream. + mgr.retire(0); + cfg.seed = 5678; + auto c = admit(mgr, 3, prompt_tokens(4), cfg); + CHECK(is_admitted(c)); + CHECK(mgr.slot(0).rng() != first); + } + + // A greedy sampler never draws, so its seed is irrelevant and admission + // must not depend on one being present. + { + PagedKvPool pool(4, 1, /*block_size=*/16); + SeqSlotManager mgr(pool, 64); + SamplerCfg cfg = greedy_sampler(); + cfg.seed = 1234; + CHECK(!cfg.needs_logit_processing()); + auto a = admit(mgr, 1, prompt_tokens(4), cfg); + CHECK(is_admitted(a)); + CHECK(!mgr.slot(0).sampler.needs_logit_processing()); + } + + std::printf("OK test_seq_slot_manager (%d checks)\n", g_checks); + return 0; +} diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 55ae264fe..90a87dc52 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -97,6 +97,13 @@ struct ServerUnitFixture {}; } \ } while (0) +TEST_CASE(ServerUnitFixture, test_api_format_names_are_total) { + CHECK(std::string(api_format_name(ApiFormat::OPENAI_CHAT)) == "chat"); + CHECK(std::string(api_format_name(ApiFormat::ANTHROPIC)) == "anthropic"); + CHECK(std::string(api_format_name(ApiFormat::RESPONSES)) == "responses"); + CHECK(std::string(api_format_name(ApiFormat::COMPLETIONS)) == "completions"); +} + TEST_CASE(ServerUnitFixture, test_daemon_io_external_cancellation_latches) { bool cancel = false; DaemonIO io; @@ -2168,6 +2175,28 @@ TEST_CASE(ServerUnitFixture, test_pflash_config_defaults) { TEST_ASSERT(cfg.draft_residency == DraftResidencyPolicy::Auto); } +TEST_CASE(ServerUnitFixture, test_concurrent_status_is_aggregate_only) { + ServerStatus status; + ServerStatus::RequestInfo info; + info.model = "classic-model"; + status.set_running("classic prompt", 12, true, info); + json snapshot = status.to_json(); + TEST_ASSERT(snapshot["active_requests"] == 0); + TEST_ASSERT(snapshot["current"]["model"] == "classic-model"); + + status.set_concurrent_requests(2); + snapshot = status.to_json(); + TEST_ASSERT(snapshot["phase"] == "decode"); + TEST_ASSERT(snapshot["active_requests"] == 2); + TEST_ASSERT(snapshot["current"].is_null()); + + status.set_idle(); + snapshot = status.to_json(); + TEST_ASSERT(snapshot["phase"] == "idle"); + TEST_ASSERT(snapshot["active_requests"] == 0); + TEST_ASSERT(snapshot["current"].is_null()); +} + TEST_CASE(ServerUnitFixture, test_pflash_config_modes) { ServerConfig cfg; cfg.pflash_mode = ServerConfig::PflashMode::AUTO; @@ -4527,7 +4556,6 @@ TEST_CASE(ServerUnitFixture, test_props_model_card_wholesale_sidecar) { PrefixCache pc(0, tok); ToolMemory tm; json body = build_props_body(cfg, pc, tm); - TEST_ASSERT(body.contains("model_card")); TEST_ASSERT(!body["model_card"].is_null()); // `source` is the upstream URL, NOT the filepath. The filepath label @@ -4640,6 +4668,7 @@ TEST_CASE(ServerUnitFixture, test_props_runtime_shape) { cfg.chunk = 512; cfg.target_device = "auto:0"; cfg.draft_device = "auto:0"; + TEST_ASSERT(cfg.admission_coalesce_ms == 20); Tokenizer tok; PrefixCache pc(0, tok); @@ -4658,12 +4687,17 @@ TEST_CASE(ServerUnitFixture, test_props_runtime_shape) { TEST_ASSERT(rt["chunk"].get() == 512); TEST_ASSERT(rt["target_device"].get() == "auto:0"); TEST_ASSERT(rt["draft_device"].get() == "auto:0"); + TEST_ASSERT(rt["continuous_batching"]["admission_coalesce_ms"] + .get() == 20); TEST_ASSERT(body["pflash"]["draft_residency"].get() == "persistent"); // draft_device is null when no draft model is loaded. cfg.draft_device.clear(); + cfg.admission_coalesce_ms = 7; body = build_props_body(cfg, pc, tm); TEST_ASSERT(body["runtime"]["draft_device"].is_null()); + TEST_ASSERT(body["runtime"]["continuous_batching"] + ["admission_coalesce_ms"].get() == 7); } // ═══════════════════════════════════════════════════════════════════════ diff --git a/server/tests/test_server_parallel.py b/server/tests/test_server_parallel.py new file mode 100644 index 000000000..83a9dc9ad --- /dev/null +++ b/server/tests/test_server_parallel.py @@ -0,0 +1,589 @@ +#!/usr/bin/env python3 +"""End-to-end integration tests for concurrent serving (--max-concurrency N). + +Exercises the qwen35 paged-attention slot engine: true decode overlap across +streams, per-request state isolation, sequential-vs-concurrent consistency, +SSE interleaving, over-subscription (queueing beyond slot count), concurrent +non-streaming completions, and non-pausing admission. + +Usage: + # Start server first with concurrent serving enabled: + ./server/build/dflash_server --port 9099 \ + --paged-attention --max-concurrency 3 + + # Then run tests: + python3 server/tests/test_server_parallel.py \ + --base-url http://127.0.0.1:9099 --max-concurrency 3 + +The non-power-of-two default is intentional: three physical slots produce a +four-row compact decode bucket, with the final row mapped to padding. This +guards the distinction between graph bucket width and physical slot count. + +Note: batched decode may legally differ from single-request decode at the +token level (GEMM reduction order can flip near-tie tokens), so all answer +checks are content-level (the expected number appears), never exact-match. +""" + +import argparse +import json +import re +import sys +import threading +import time +import urllib.request +import urllib.error + + +def make_math_prompts(count: int): + """Deterministic arithmetic prompts with distinct 3-digit answers. + + Sums start at 111 and step by 3 so that no answer is within +/-2 of + another (greedy reasoning that decomposes a sum won't casually emit a + neighboring stream's answer), and for small counts the operands stay + 2-digit while every answer is 3-digit, so an answer never collides with + an operand echoed from another prompt. + """ + prompts = [] + for i in range(count): + s = 111 + 3 * i + a = s // 2 + b = s - a + prompts.append((f"What is {a}+{b}? Answer with just the number.", str(s))) + return prompts + + +def make_long_prompts(count: int): + """Prompts that force long, multi-chunk generations, so overlap and + interleave are observable. Each stream counts a distinct range; the + expected marker is the range start (always emitted early).""" + prompts = [] + for i in range(count): + start = 100 + 50 * i + end = start + 40 + prompts.append(( + f"Count from {start} to {end}, separated by commas, " + f"no other text.", str(start))) + return prompts + + +def contains_number(text: str, number: str) -> bool: + """True if `number` appears in `text` as a standalone number + (not embedded in a longer digit run).""" + return re.search(rf"(? dict: + return { + "model": "dflash", + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, + "temperature": 0.0, + "stream": stream, + } + + def _stream_worker(self, idx: int, prompt: str, max_tokens: int, + results: list, barrier: threading.Barrier | None = None, + timeline: list | None = None, + timeline_lock=None, + timeout: float = 600.0): + """Run one streaming chat completion; record content and chunk timing.""" + r = {"ok": False, "error": None, "content": "", "reasoning": "", + "first_chunk_t": None, "finish_t": None, "n_chunks": 0} + results[idx] = r + try: + if barrier is not None: + barrier.wait(timeout=60) + resp = self._req("POST", "/v1/chat/completions", + self._chat_body(prompt, max_tokens, stream=True), + stream=True, timeout=timeout) + for line in resp: + line = line.decode().strip() + if not line: + continue + if line == "data: [DONE]": + break + if not line.startswith("data: "): + continue + chunk = json.loads(line[6:]) + choices = chunk.get("choices") or [{}] + delta = choices[0].get("delta", {}) + got_piece = False + if delta.get("reasoning_content"): + r["reasoning"] += delta["reasoning_content"] + got_piece = True + if delta.get("content"): + r["content"] += delta["content"] + got_piece = True + if got_piece: + now = time.monotonic() + r["n_chunks"] += 1 + if r["first_chunk_t"] is None: + r["first_chunk_t"] = now + if timeline is not None: + with timeline_lock: + timeline.append((now, idx)) + r["finish_t"] = time.monotonic() + resp.close() + r["ok"] = True + except urllib.error.HTTPError as e: + body = e.read().decode(errors="replace") + r["error"] = f"HTTP {e.code}: {body[:300]}" + except Exception as e: + r["error"] = f"{type(e).__name__}: {e}" + + def _nonstream_worker(self, idx: int, prompt: str, max_tokens: int, + results: list, barrier: threading.Barrier | None = None, + start_delay: float = 0.0, + timeout: float = 600.0): + """Run one non-streaming chat completion.""" + r = {"ok": False, "error": None, "content": "", "reasoning": "", + "usage": {}} + results[idx] = r + try: + if barrier is not None: + barrier.wait(timeout=60) + if start_delay > 0: + time.sleep(start_delay) + resp = self._req("POST", "/v1/chat/completions", + self._chat_body(prompt, max_tokens, stream=False), + timeout=timeout) + msg = resp["choices"][0]["message"] + r["content"] = msg.get("content") or "" + r["reasoning"] = msg.get("reasoning_content") or "" + r["usage"] = resp.get("usage", {}) + r["ok"] = True + except urllib.error.HTTPError as e: + body = e.read().decode(errors="replace") + r["error"] = f"HTTP {e.code}: {body[:300]}" + except Exception as e: + r["error"] = f"{type(e).__name__}: {e}" + + def _run_workers(self, worker, count: int, per_request_kwargs: list, + join_timeout: float = 900.0): + """Launch `count` worker threads through a start barrier; return results.""" + results: list = [None] * count + barrier = threading.Barrier(count) + threads = [] + for i in range(count): + t = threading.Thread(target=worker, + args=(i,), kwargs={**per_request_kwargs[i], + "results": results, + "barrier": barrier}, + daemon=True) + threads.append(t) + t.start() + for t in threads: + t.join(timeout=join_timeout) + return results + + def _launch_streams(self, prompts, max_tokens: int, + timeline: list | None = None, + join_timeout: float = 900.0): + """Fire len(prompts) streaming requests simultaneously.""" + timeline_lock = threading.Lock() if timeline is not None else None + kwargs = [{"prompt": p, "max_tokens": max_tokens, + "timeline": timeline, "timeline_lock": timeline_lock} + for p, _ in prompts] + return self._run_workers(self._stream_worker, len(prompts), kwargs, + join_timeout=join_timeout) + + def _launch_nonstream(self, prompts, max_tokens: int, + join_timeout: float = 900.0): + """Fire len(prompts) non-streaming requests simultaneously.""" + kwargs = [{"prompt": p, "max_tokens": max_tokens} for p, _ in prompts] + return self._run_workers(self._nonstream_worker, len(prompts), kwargs, + join_timeout=join_timeout) + + def _all_completed(self, results, label: str) -> bool: + """Check every worker result; report per-request failures.""" + all_ok = True + for i, r in enumerate(results): + if r is None: + self._check(f"{label} request {i+1} completes", False, + "worker did not finish (timeout)") + all_ok = False + elif not r["ok"]: + self._check(f"{label} request {i+1} completes", False, r["error"]) + all_ok = False + if all_ok: + self._check(f"all {len(results)} {label} requests complete " + "with 200", True) + return all_ok + + @staticmethod + def _combined(r) -> str: + return r["reasoning"] + "\n" + r["content"] + + # ── Tests ──────────────────────────────────────────────────────────── + + def test_parallel_streaming(self): + """N simultaneous streams must overlap and their chunks interleave.""" + n = self.parallel + print(f"\n[PAR-1] Streaming overlap + interleave — " + f"{n} simultaneous requests") + if n == 1: + self._skip("streaming overlap + interleave", + "--max-concurrency 1: concurrency is not expected; " + "run with --max-concurrency > 1") + return + prompts = make_long_prompts(n) + timeline: list = [] + t0 = time.monotonic() + results = self._launch_streams(prompts, max_tokens=192, + timeline=timeline) + elapsed = time.monotonic() - t0 + if not self._all_completed(results, "streaming"): + return + + missing = [i + 1 for i, r in enumerate(results) + if r["first_chunk_t"] is None] + self._check("every stream produced at least one chunk", + not missing, f"streams with no chunks: {missing}") + if missing: + return + + latest_first = max(r["first_chunk_t"] for r in results) + earliest_finish = min(r["finish_t"] for r in results) + rel = [(f"s{i+1}: first={r['first_chunk_t']-t0:.2f}s " + f"finish={r['finish_t']-t0:.2f}s") + for i, r in enumerate(results)] + self._check("all first chunks arrive before earliest finish " + "(true overlap, not serialization)", + latest_first < earliest_finish, + f"latest first chunk at {latest_first-t0:.2f}s, earliest " + f"finish at {earliest_finish-t0:.2f}s; {'; '.join(rel)}") + print(f" → {n} streams completed in {elapsed:.1f}s; " + f"latest first chunk {latest_first-t0:.2f}s, " + f"earliest finish {earliest_finish-t0:.2f}s") + + seq = [idx for _, idx in sorted(timeline)] + distinct = len(set(seq)) + runs = 1 + sum(1 for a, b in zip(seq, seq[1:]) if a != b) if seq else 0 + self._check("chunks arrived from at least two streams", distinct >= 2, + f"streams seen: {sorted(set(seq))}") + # If every stream's chunks formed one contiguous block, runs would + # equal distinct; interleaving means some stream owns >= 2 runs. + self._check("chunks from different streams interleave", + runs > distinct, + f"{len(seq)} chunks arrived in {runs} contiguous runs " + f"across {distinct} streams (fully serialized order)") + print(f" → {len(seq)} chunks, {distinct} streams, {runs} runs") + + def test_parallel_isolation(self): + """Concurrent streams with distinct prompts: each answer belongs to + its own prompt, none leaks into another stream (state isolation).""" + n = self.parallel + print(f"\n[PAR-2] Isolation — {n} concurrent distinct prompts") + prompts = make_math_prompts(n) + results = self._launch_streams(prompts, max_tokens=512) + if not self._all_completed(results, "streaming"): + return + + answers = [ans for _, ans in prompts] + for i, r in enumerate(results): + # Positive: own answer somewhere in reasoning+content. + self._check(f"stream {i+1} contains its own answer {answers[i]}", + contains_number(self._combined(r), answers[i]), + f"content={r['content']!r} " + f"reasoning={r['reasoning'][:200]!r}") + # Negative: no other stream's answer in the final content + # (content only — reasoning legitimately echoes this prompt's + # operands and intermediate sums, content is just the number). + leaked = [answers[j] for j in range(n) + if j != i and contains_number(r["content"], answers[j])] + self._check(f"stream {i+1} contains no other stream's answer", + not leaked, + f"leaked answers {leaked} in content={r['content']!r}") + + def test_parallel_nonstream(self): + """A prompt answered correctly alone must still be answered correctly + when decoded inside an N-way concurrent non-streaming batch. Also + validates every concurrent answer and its completion-token usage. + Content-level match only — batched GEMM reduction order may legally + flip near-tie tokens, so exact token equality is NOT required.""" + n = self.parallel + print(f"\n[PAR-3] Sequential + concurrent non-streaming consistency") + prompts = make_math_prompts(n) + probe_prompt, probe_answer = prompts[0] + + # Solo run. + solo = [None] + self._nonstream_worker(0, probe_prompt, 512, solo) + if not solo[0]["ok"]: + self._check("solo request completes", False, solo[0]["error"]) + return + self._check("solo request completes", True) + solo_ok = contains_number(self._combined(solo[0]), probe_answer) + self._check(f"solo answer contains {probe_answer}", solo_ok, + f"content={solo[0]['content']!r} " + f"reasoning={solo[0]['reasoning'][:200]!r}") + print(f" → solo content: {solo[0]['content']!r}") + + # Same prompt inside an N-way concurrent batch. + results = self._launch_nonstream(prompts, max_tokens=512) + if not self._all_completed(results, "non-streaming"): + return + conc = results[0] + print(f" → concurrent content: {conc['content']!r}") + if solo[0]["content"] != conc["content"]: + print(" → note: solo/concurrent contents differ at token level " + "(allowed — batched reduction order)") + for i, r in enumerate(results): + ans = prompts[i][1] + self._check(f"request {i+1} contains its answer {ans}", + contains_number(self._combined(r), ans), + f"content={r['content']!r} " + f"reasoning={r['reasoning'][:200]!r}") + completion_tokens = r["usage"].get("completion_tokens", 0) + self._check(f"request {i+1} usage.completion_tokens > 0", + completion_tokens > 0, + f"usage={r['usage']}") + prompt_tokens = r["usage"].get("prompt_tokens", 0) + timings = r["usage"].get("timings", {}) + self._check( + f"request {i+1} timings account for the full prompt", + prompt_tokens > 0 + and timings.get("prefilled_tokens") == prompt_tokens + and timings.get("effective_prompt_tokens") == prompt_tokens, + f"usage={r['usage']}") + + def test_parallel_more_than_slots(self): + """2*N requests at once: extras must queue behind the N slots and + all must finish with correct answers.""" + n = self.parallel + count = 2 * n + print(f"\n[PAR-4] Over-subscription — {count} requests on {n} slots") + prompts = make_math_prompts(count) + t0 = time.monotonic() + results = self._launch_nonstream(prompts, max_tokens=512, + join_timeout=1800.0) + elapsed = time.monotonic() - t0 + if not self._all_completed(results, "non-streaming"): + return + for i, r in enumerate(results): + ans = prompts[i][1] + self._check(f"request {i+1} contains its answer {ans}", + contains_number(self._combined(r), ans), + f"content={r['content']!r} " + f"reasoning={r['reasoning'][:200]!r}") + print(f" → {count} requests completed in {elapsed:.1f}s") + + def test_unequal_prefill_staging_leases(self): + """A multi-chunk prefill must keep its staging identity after an + earlier short prefill commits and leaves the FIFO head.""" + print("\n[PAR-5] Unequal prefills retain isolated staging state") + if self.parallel < 2: + self._skip("unequal prefill staging lease", + "--max-concurrency 1: multiple staging sets unavailable") + return + + short_prompt = "What is 55+56? Answer with just the number." + filler = "\n".join( + f"record {i}: alpha beta gamma delta epsilon zeta eta theta" + for i in range(360)) + long_prompt = ( + f"Read and then ignore these padding records:\n{filler}\n" + "What is 4500+4501? Answer with just the number.") + + # Launch both requests close together, with the short request first in + # FIFO order. It should finish first while the long request retains + # staging set 1. + kwargs = [ + {"prompt": short_prompt, "max_tokens": 512, + "start_delay": 0.0}, + {"prompt": long_prompt, "max_tokens": 512, + "start_delay": 0.001}, + ] + results = self._run_workers( + self._nonstream_worker, 2, kwargs, join_timeout=1800.0) + if not self._all_completed(results, "unequal-prefill"): + return + + self._check( + "short prefill answers 111", + contains_number(self._combined(results[0]), "111"), + f"content={results[0]['content']!r}") + self._check( + "long prefill answers 9001 after the short prefill commits", + contains_number(self._combined(results[1]), "9001"), + f"content={results[1]['content']!r} " + f"reasoning={results[1]['reasoning'][:200]!r}") + long_prompt_tokens = results[1]["usage"].get("prompt_tokens", 0) + self._check( + "long prompt crosses the 2048-token initial prefill chunk", + long_prompt_tokens > 2048, + f"usage={results[1]['usage']}") + + def test_parallel_prefill_no_pause(self): + """A long admission must not pause a stream that is already decoding.""" + n = self.parallel + print("\n[PAR-7] Prefill/decode fusion — decode continues during " + "a long prefill") + if n == 1: + self._skip("prefill no-pause", + "--max-concurrency 1: fusion not applicable, skipping") + return + + # Stream A starts alone and reaches steady decode before B arrives. + a_prompt, _ = make_long_prompts(1)[0] + timeline: list = [] + timeline_lock = threading.Lock() + a_res: list = [None] + a_thread = threading.Thread( + target=self._stream_worker, args=(0, a_prompt, 320, a_res), + kwargs={"timeline": timeline, "timeline_lock": timeline_lock}, + daemon=True) + a_thread.start() + deadline = time.monotonic() + 180 + while time.monotonic() < deadline: + ra = a_res[0] + if ra is not None and (ra["n_chunks"] >= 3 or ra["error"]): + break + time.sleep(0.05) + ra = a_res[0] + if ra is None or ra["error"] or ra["n_chunks"] < 3: + self._check("stream A reaches steady decode", False, + "no chunks" if ra is None else + (ra["error"] or f"only {ra['n_chunks']} chunks")) + return + self._check("stream A reaches steady decode", True) + + # Stream B has several 512-token prefill chunks. Answer 167 cannot + # collide with the filler item indices (0..139). + filler = "\n".join( + f"item {i}: the quick brown fox jumps over the lazy dog" + for i in range(140)) + b_prompt = (f"Here is a list:\n{filler}\n" + "Ignore the list entirely. What is 83+84? " + "Answer with just the number.") + b_res: list = [None] + b_started = time.monotonic() + b_thread = threading.Thread( + target=self._stream_worker, args=(0, b_prompt, 512, b_res), + daemon=True) + b_thread.start() + b_thread.join(timeout=900) + a_thread.join(timeout=900) + + rb = b_res[0] + if rb is None or not rb["ok"] or rb["first_chunk_t"] is None: + self._check("stream B completes", False, + "no result" if rb is None else str(rb["error"])) + return + self._check("stream B completes", True) + self._check("stream B answers 167", + contains_number(self._combined(rb), "167"), + f"content={rb['content']!r}") + + # Blocking admission leaves the prefill-dominated first 70% of B's + # TTFT window empty. Fused steps keep producing A outputs there. + window = rb["first_chunk_t"] - b_started + early_end = b_started + 0.7 * window + with timeline_lock: + a_early = [t for t, idx in timeline + if idx == 0 and b_started <= t <= early_end] + ra = a_res[0] + if ra["finish_t"] is not None and ra["finish_t"] < early_end: + self._skip("A kept emitting during B's prefill", + "stream A finished before B's window closed") + return + self._check("stream A kept emitting during B's prefill window", + len(a_early) >= 2, + f"A emitted {len(a_early)} chunks in the first " + f"{0.7 * window:.2f}s of B's {window:.2f}s " + "prefill window") + + # ── Run all ────────────────────────────────────────────────────────── + + def run_all(self): + print("=" * 60) + print("Parallel Serving Tests") + print(f"Target: {self.base} (parallel={self.parallel})") + print("=" * 60) + + self.test_parallel_streaming() + self.test_parallel_isolation() + self.test_parallel_nonstream() + self.test_parallel_more_than_slots() + self.test_unequal_prefill_staging_leases() + self.test_parallel_prefill_no_pause() + + print("\n" + "=" * 60) + total = self.passed + self.failed + self.skipped + print(f"Results: {self.passed}/{total} passed, {self.failed} failed" + + (f", {self.skipped} skipped" if self.skipped else "")) + return 0 if self.failed == 0 else 1 + + +# ─── Main ──────────────────────────────────────────────────────────────── + +def main(): + parser = argparse.ArgumentParser( + description="Concurrent serving tests for dflash_server " + "(run against a server started with " + "--paged-attention --max-concurrency N)") + parser.add_argument("--base-url", default="http://127.0.0.1:9099", + help="Server base URL") + parser.add_argument("--max-concurrency", type=int, default=3, + help="N: --max-concurrency value used to start the server") + args = parser.parse_args() + + if not (1 <= args.max_concurrency <= 64): + print("ERROR: --max-concurrency must be in [1, 64], " + f"got {args.max_concurrency}") + sys.exit(2) + + # Preflight: server must already be running. + base = args.base_url.rstrip("/") + try: + r = urllib.request.urlopen(base + "/health", timeout=10) + health = json.loads(r.read().decode()) + if health.get("status") != "ok": + print(f"ERROR: server unhealthy: {health}") + sys.exit(2) + except Exception as e: + print(f"ERROR: server not reachable at {base}: {e}") + print("Start it first, e.g.:") + print(f" ./server/build/dflash_server --port 9099 " + f"--paged-attention --max-concurrency {args.max_concurrency}") + sys.exit(2) + + suite = ParallelTestSuite(base, args.max_concurrency) + sys.exit(suite.run_all()) + + +if __name__ == "__main__": + main() From f5e8a26b1ce97ad1f68671206657bc296d6c70a9 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 12 Aug 2026 07:52:46 +0000 Subject: [PATCH 2/2] perf(ds4): batch gathered concurrency work --- .../src/deepseek4/deepseek4_fused_verify.inc | 171 ++++++-- server/src/deepseek4/deepseek4_graph.cpp | 373 +++++++++++++----- .../src/deepseek4/deepseek4_paged_cache.cpp | 5 +- server/test/test_deepseek4_paged_cache.cpp | 1 + 4 files changed, 404 insertions(+), 146 deletions(-) diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 637bc6668..e4e126752 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -350,9 +350,27 @@ static bool ds4_build_fused_verify_graph( DeepSeek4PagedCache * paged_cache = nullptr, const std::vector> * paged_rows = nullptr) { const bool paged_mode = paged_cache && paged_rows; - if (fg.sched) { + if (paged_mode) { + if (q < 1 || q > 16 || + paged_rows->size() != (size_t) w.n_layer) { + return false; + } + for (const auto & layer_rows : *paged_rows) { + if (layer_rows.size() != (size_t) q) return false; + } + } + const size_t graph_capacity = q > 8 ? 131072u : 65536u; + std::array sched_backends{}; + if (hybrid) { + sched_backends = {backend, hybrid->cold_backend, hybrid->cpu_backend}; + } + const bool reuse_sched = + hybrid && fg.sched_reusable(sched_backends, graph_capacity); + if (fg.sched && !reuse_sched) { ggml_backend_sched_free(fg.sched); fg.sched = nullptr; + fg.sched_capacity = 0; + fg.sched_backends = {}; } step_graph_free(fg.sg); fg.reset_nodes(); @@ -381,7 +399,6 @@ static bool ds4_build_fused_verify_graph( // sequence. Above eight lanes the resulting whole-model graph exceeds // the verifier-era 64K scheduler hash table even though the metadata // arena still has ample room. - const size_t graph_capacity = q > 8 ? 131072u : 65536u; fg.sg.gf = ggml_new_graph_custom(ctx, graph_capacity, false); ggml_cgraph * gf = fg.sg.gf; @@ -426,6 +443,38 @@ static bool ds4_build_fused_verify_graph( ggml_set_input(fg.mask_bundle); int64_t mask_off = 0; + // Back every gathered lane's scalar and index inputs with three shared + // tensors. Their sizes are fixed by the prepared history lengths encoded + // in this graph's shape key, so the step can upload once per dtype. + int64_t paged_gather_total = 0; + if (paged_mode) { + for (int il = 0; il < w.n_layer; ++il) { + const int layer_ratio = (int) w.compress_ratios[il]; + for (int t = 0; t < q; ++t) { + const auto & rows = (*paged_rows)[(size_t) il][(size_t) t]; + paged_gather_total += + (int64_t) std::max(rows.raw_history.size(), 1); + if (layer_ratio > 0) { + paged_gather_total += (int64_t) std::max( + rows.compressed_history.size(), 1); + } + } + } + ex.paged_i32_n = (int64_t) 5 * w.n_layer * q; + ex.paged_i64_n = (int64_t) 3 * w.n_layer * q; + ex.paged_gather_n = std::max(paged_gather_total, 1); + ex.paged_i32 = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, ex.paged_i32_n); + ggml_set_input(ex.paged_i32); + ex.paged_i64 = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, ex.paged_i64_n); + ggml_set_input(ex.paged_i64); + ex.paged_gather = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, ex.paged_gather_n); + ggml_set_input(ex.paged_gather); + } + int64_t paged_gather_off = 0; + // per-token HC streams std::vector hc_cur(q); for (int t = 0; t < q; ++t) { @@ -471,44 +520,71 @@ static bool ds4_build_fused_verify_graph( // Gather each lane's immutable chronological history and run the // established MLA lane core at q=1; all surrounding HC/MoE/output // machinery remains q-wide and unchanged. - if (q < 1 || q > 16 || paged_rows->size() != (size_t) w.n_layer || - (*paged_rows)[(size_t) il].size() != (size_t) q) return false; DeepSeek4PagedLayerCache & plc = paged_cache->layers[(size_t) il]; ggml_tensor * raw_flat = ggml_reshape_2d( ctx, plc.raw_kv, w.head_dim, (int64_t) DS4_PAGE_TOKENS * paged_cache->plan.slots); + ggml_tensor * attn_normed = + build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); + DeepSeek4PreparedProjectedLane batched_proj = + build_mla_qkv_projection(ctx, attn_normed, w, L, q); + build_mla_qkv_rope( + ctx, batched_proj, w, ds4_rope_params(w, ratio), q, + ex.pos_q, /*fuse_q_rope=*/false); for (int t = 0; t < q; ++t) { const auto & rows = (*paged_rows)[(size_t) il][(size_t) t]; auto & px = ex.paged.emplace_back(); - px.pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(px.pos); - px.neg_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(px.neg_pos); - px.raw_gather = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, - std::max(rows.raw_history.size(), 1)); ggml_set_input(px.raw_gather); + const int64_t lane_idx = (int64_t) il * q + t; + px.i32_base = lane_idx * 5; + px.i64_base = lane_idx * 3; + auto scalar_i32 = [&](int slot) { + return ggml_view_1d( + ctx, ex.paged_i32, 1, + (size_t) (px.i32_base + slot) * sizeof(int32_t)); + }; + auto scalar_i64 = [&](int slot) { + return ggml_view_1d( + ctx, ex.paged_i64, 1, + (size_t) (px.i64_base + slot) * sizeof(int64_t)); + }; + px.pos = scalar_i32(0); + px.neg_pos = scalar_i32(1); + px.raw_n = + (int64_t) std::max(rows.raw_history.size(), 1); + px.raw_off = paged_gather_off; + px.raw_gather = ggml_view_1d( + ctx, ex.paged_gather, px.raw_n, + (size_t) px.raw_off * sizeof(int32_t)); + paged_gather_off += px.raw_n; ggml_tensor * raw_history = rows.raw_history.empty() ? nullptr : ggml_get_rows(ctx, raw_flat, px.raw_gather); ggml_tensor * comp_history = nullptr; ggml_tensor * index_history = nullptr; if (ratio > 0) { - px.comp_gather = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, - std::max(rows.compressed_history.size(), 1)); - ggml_set_input(px.comp_gather); + px.comp_n = (int64_t) std::max( + rows.compressed_history.size(), 1); + px.comp_off = paged_gather_off; + px.comp_gather = ggml_view_1d( + ctx, ex.paged_gather, px.comp_n, + (size_t) px.comp_off * sizeof(int32_t)); + paged_gather_off += px.comp_n; if (!rows.compressed_history.empty()) comp_history = ggml_get_rows(ctx, plc.comp_kv, px.comp_gather); if (ratio == 4) { - px.index_gather = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, - std::max(rows.compressed_history.size(), 1)); - ggml_set_input(px.index_gather); - if (!rows.compressed_history.empty()) - index_history = ggml_get_rows(ctx, plc.index_comp_kv, px.index_gather); + // The indexer and attention compressors address the + // same chronological compressed rows. Explicit mode + // has no indexer consumer, so share the gather and do + // not emit the otherwise-dead index-history read. + px.index_gather = px.comp_gather; } } - px.raw_write = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); ggml_set_input(px.raw_write); + px.raw_write = scalar_i64(0); if (ratio > 0) { - px.comp_write = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); ggml_set_input(px.comp_write); - px.comp_read = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(px.comp_read); - px.ape = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(px.ape); - px.state_row = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); ggml_set_input(px.state_row); - px.comp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(px.comp_pos); + px.comp_write = scalar_i64(1); + px.comp_read = scalar_i32(2); + px.ape = scalar_i32(3); + px.state_row = scalar_i64(2); + px.comp_pos = scalar_i32(4); } DeepSeek4CompressorState attn_state{}, index_state{}; if (ratio > 0 && rows.slot >= 0) { @@ -553,13 +629,19 @@ static bool ds4_build_fused_verify_graph( std::vector ib; std::vector iab; std::vector lab; - ggml_tensor * col = ggml_view_2d(ctx, attn_in, n_embd, 1, - attn_in->nb[1], (size_t) t * attn_in->nb[1]); + ggml_tensor * col = ggml_view_2d( + ctx, attn_normed, n_embd, 1, attn_normed->nb[1], + (size_t) t * attn_normed->nb[1]); + const DeepSeek4PreparedProjectedLane lane_proj = + ds4_slice_projected_lane(ctx, batched_proj, w, t, px.pos); + ggml_tensor * lane_context = nullptr; ggml_tensor * one = build_mla_attention_lane_core( - ctx, gf, build_rms_norm(ctx, col, L.attn_norm, w.rms_eps), - w, L, lane, il, (int) rows.position, 1, &ci, ib, iab, lab, - nullptr, DeepSeek4AttentionImpl::Explicit); - if (!one || !ib.empty() || !iab.empty() || !lab.empty()) { + ctx, gf, col, w, L, lane, il, (int) rows.position, 1, + &ci, ib, iab, lab, nullptr, + DeepSeek4AttentionImpl::Explicit, &lane_proj, + &lane_context); + if (!one || !lane_context || !ib.empty() || !iab.empty() || + !lab.empty()) { std::fprintf(stderr, "[deepseek4-paged] layer %d lane %d attention build " "failed (graph=%d i32=%zu arrays=%zu i64=%zu, " @@ -569,8 +651,12 @@ static bool ds4_build_fused_verify_graph( ? INT32_MIN : iab[0].values[0]); return false; } - attn_out = attn_out ? ggml_concat(ctx, attn_out, one, 1) : one; + attn_out = attn_out + ? ggml_concat(ctx, attn_out, lane_context, 1) + : lane_context; } + attn_out = build_mla_output_projection( + ctx, attn_out, w, L, q, /*allow_grouped=*/false); } else { // ── Batched speculative attention ── DeepSeek4AttentionGraphInputs ain{}; @@ -976,14 +1062,22 @@ static bool ds4_build_fused_verify_graph( warned_cross_vendor_join = true; } } - ggml_backend_t backends[3] = { - backend, peer, hybrid->cpu_backend}; - fg.sched = ggml_backend_sched_new( - backends, nullptr, 3, graph_capacity, false, true); - if (!fg.sched) { - std::fprintf(stderr, - "[ds4-fused-verify] scheduler creation failed\n"); - return false; + if (reuse_sched) { + // Drop prior tensor assignments and deferred events while keeping + // the scheduler-owned pinned staging allocations alive. + ggml_backend_sched_reset(fg.sched); + } else { + ggml_backend_t backends[3] = { + backend, peer, hybrid->cpu_backend}; + fg.sched = ggml_backend_sched_new( + backends, nullptr, 3, graph_capacity, false, true); + if (!fg.sched) { + std::fprintf(stderr, + "[ds4-fused-verify] scheduler creation failed\n"); + return false; + } + fg.sched_capacity = graph_capacity; + fg.sched_backends = sched_backends; } const bool late_join_split = mixed_policy.late_join_split; const MoeHybridGraphPolicy & moe_policy = @@ -1024,6 +1118,9 @@ static bool ds4_build_fused_verify_graph( pin_main(fg.i32_bundle); pin_main(fg.i64_bundle); pin_main(fg.mask_bundle); + pin_main(ex.paged_i32); + pin_main(ex.paged_i64); + pin_main(ex.paged_gather); for (const auto & px : ex.paged) { pin_main(px.pos); pin_main(px.neg_pos); pin_main(px.raw_gather); pin_main(px.comp_gather); pin_main(px.index_gather); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index f196dc7ef..aba56fde6 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -1572,6 +1572,130 @@ struct DeepSeek4PreparedProjectedLane { ggml_tensor * rope_pos = nullptr; }; +// Per-layer RoPE parameters. Compressed layers use YaRN scaling, and +// attn_factor cancels the magnitude scaling rope_yarn applies. +struct Ds4RopeParams { + float freq = 0.0f; + float scale = 1.0f; + float ext = 0.0f; + float attn = 1.0f; + int n_ctx_orig = 0; +}; + +static Ds4RopeParams ds4_rope_params(const DeepSeek4Weights & w, int ratio) { + const bool compressed = ratio > 0; + Ds4RopeParams p; + p.freq = compressed ? w.compress_rope_freq_base : w.rope_freq_base; + p.scale = compressed ? (1.0f / w.rope_scale_factor) : 1.0f; + p.ext = compressed ? 1.0f : 0.0f; + if (p.ext != 0.0f && p.scale > 0.0f) { + p.attn /= (1.0f + 0.1f * logf(1.0f / p.scale)); + } + p.n_ctx_orig = (int) w.rope_orig_ctx; + return p; +} + +// Q/KV projections and their tail RoPE are independent per token. A paged +// caller can evaluate them once at width q and hand each lane a column view, +// avoiding one reread of all three projection weights per active lane. +static DeepSeek4PreparedProjectedLane build_mla_qkv_projection( + ggml_context * ctx, + ggml_tensor * cur, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int n_tokens) { + DeepSeek4PreparedProjectedLane out; + ggml_tensor * qr = ggml_mul_mat(ctx, L.attn_q_a, cur); + qr = build_rms_norm(ctx, qr, L.attn_q_a_norm, w.rms_eps); + ggml_tensor * q = ggml_mul_mat(ctx, L.attn_q_b, qr); + q = ggml_reshape_3d(ctx, q, w.head_dim, w.n_head, n_tokens); + q = ggml_rms_norm(ctx, q, w.rms_eps); + + ggml_tensor * kv = ggml_mul_mat(ctx, L.attn_kv, cur); + kv = build_rms_norm(ctx, kv, L.attn_kv_a_norm, w.rms_eps); + + out.normalized_q_lora = qr; + out.q = q; + out.kv = kv; + return out; +} + +static void build_mla_qkv_rope( + ggml_context * ctx, + DeepSeek4PreparedProjectedLane & p, + const DeepSeek4Weights & w, + const Ds4RopeParams & rope, + int n_tokens, + ggml_tensor * rope_pos, + bool fuse_q_rope) { + if (!fuse_q_rope) { + p.q = build_tail_rope_3d(ctx, p.q, rope_pos, w.n_rot, w.head_dim, + w.n_head, n_tokens, rope.freq, rope.scale, + rope.ext, rope.attn, w.rope_yarn_beta_fast, + w.rope_yarn_beta_slow, rope.n_ctx_orig); + } + p.kv = build_tail_rope_2d(ctx, p.kv, rope_pos, w.n_rot, w.head_dim, + n_tokens, rope.freq, rope.scale, rope.ext, + rope.attn, w.rope_yarn_beta_fast, + w.rope_yarn_beta_slow, rope.n_ctx_orig); + p.rope_pos = rope_pos; +} + +// Grouped low-rank output projection. Several independent gathered lanes can +// concatenate their pre-projection contexts and share one q-wide evaluation. +static ggml_tensor * build_mla_output_projection( + ggml_context * ctx, + ggml_tensor * attn_out, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int n_tokens, + bool allow_grouped) { + const int group_dim = w.head_dim * (w.n_head / w.n_out_group); + attn_out = ggml_reshape_3d( + ctx, attn_out, group_dim, w.n_out_group, n_tokens); + attn_out = ggml_permute(ctx, attn_out, 0, 2, 1, 3); + if (n_tokens == 1) { + attn_out = ggml_cont(ctx, attn_out); + } + ggml_tensor * out_a_3d = ggml_reshape_3d( + ctx, L.attn_output_a, group_dim, w.n_lora_o, w.n_out_group); + ggml_tensor * attn_low = ggml_mul_mat(ctx, out_a_3d, attn_out); + + const bool grouped_output_projection = + allow_grouped && n_tokens > 1 && + !ds4_env_flag("DFLASH_DS4_DISABLE_GROUPED_OUTPUT_PROJECTION"); + if (grouped_output_projection) { + return ggml_mul_mat_grouped_src(ctx, L.attn_output_b, attn_low); + } + attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); + attn_low = ggml_reshape_2d( + ctx, attn_low, (int64_t) w.n_lora_o * w.n_out_group, n_tokens); + return ggml_mul_mat(ctx, L.attn_output_b, attn_low); +} + +// One lane's column of a batched prologue. These are views only. +static DeepSeek4PreparedProjectedLane ds4_slice_projected_lane( + ggml_context * ctx, + const DeepSeek4PreparedProjectedLane & batched, + const DeepSeek4Weights & w, + int lane, + ggml_tensor * lane_rope_pos) { + DeepSeek4PreparedProjectedLane out; + out.normalized_q_lora = ggml_view_2d( + ctx, batched.normalized_q_lora, batched.normalized_q_lora->ne[0], 1, + batched.normalized_q_lora->nb[1], + (size_t) lane * batched.normalized_q_lora->nb[1]); + out.q = ggml_view_3d( + ctx, batched.q, w.head_dim, w.n_head, 1, + batched.q->nb[1], batched.q->nb[2], + (size_t) lane * batched.q->nb[2]); + out.kv = ggml_view_2d( + ctx, batched.kv, batched.kv->ne[0], 1, batched.kv->nb[1], + (size_t) lane * batched.kv->nb[1]); + out.rope_pos = lane_rope_pos; + return out; +} + static DeepSeek4MlaLaneBindings deepseek4_contiguous_lane_bindings( DeepSeek4LayerCache & lc, int ratio, @@ -1606,49 +1730,36 @@ static ggml_tensor * build_mla_attention_lane_core( std::vector & i32_array_inputs, std::vector & i64_array_inputs, std::vector * f32_array_inputs = nullptr, - DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit) { + DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit, + const DeepSeek4PreparedProjectedLane * prepared = nullptr, + ggml_tensor ** out_attn_context = nullptr) { const int n_embd = w.n_embd; const int head_dim = w.head_dim; const int n_head = w.n_head; const int n_rot = w.n_rot; - const int n_out_group = w.n_out_group; - const int n_lora_o = w.n_lora_o; const int ratio = w.compress_ratios[layer_idx]; const bool gathered_history = lane.history_mode == DeepSeek4MlaLaneBindings::HistoryMode::ChronologicalGathered; - // ── Q path: cur → q_a → norm → q_b → per-head norm ───────────── - // q_a: [n_embd, n_tokens] → [n_lora_q, n_tokens] - ggml_tensor * qr = ggml_mul_mat(ctx, L.attn_q_a, cur); - // qr_norm is reused by the ratio-4 indexer before the main q_b projection. - qr = build_rms_norm(ctx, qr, L.attn_q_a_norm, w.rms_eps); - // q_b: [n_lora_q, n_tokens] → [n_head * head_dim, n_tokens] - ggml_tensor * q = ggml_mul_mat(ctx, L.attn_q_b, qr); - // Reshape to [head_dim, n_head, n_tokens] for per-head ops - q = ggml_reshape_3d(ctx, q, head_dim, n_head, n_tokens); - // Reference DS4 applies unweighted RMSNorm independently to every Q head. - q = ggml_rms_norm(ctx, q, w.rms_eps); - - // ── KV path: cur → kv → norm ─────────────────────────────────── - // kv: [n_embd, n_tokens] → [head_dim, n_tokens] - ggml_tensor * kv = ggml_mul_mat(ctx, L.attn_kv, cur); - kv = build_rms_norm(ctx, kv, L.attn_kv_a_norm, w.rms_eps); + // Existing callers leave prepared null and emit the original prologue in + // place. Only gathered paged concurrency supplies a q-wide projection. + DeepSeek4PreparedProjectedLane projected; + if (!prepared) { + projected = build_mla_qkv_projection(ctx, cur, w, L, n_tokens); + } // ── RoPE on Q and KV (tail rotation on last n_rot dims) ──────── - // DS4 uses per-layer RoPE params: compressed layers get YaRN scaling. - const bool compressed = (ratio > 0); - const float rope_freq = compressed ? w.compress_rope_freq_base : w.rope_freq_base; - const float rope_scale = compressed ? (1.0f / w.rope_scale_factor) : 1.0f; - const float rope_ext = compressed ? 1.0f : 0.0f; - // For YaRN: attn_factor cancels the magnitude scaling in rope_yarn - float rope_attn = 1.0f; - if (rope_ext != 0.0f && rope_scale > 0.0f) { - rope_attn /= (1.0f + 0.1f * logf(1.0f / rope_scale)); - } + const Ds4RopeParams rope = ds4_rope_params(w, ratio); + const float rope_freq = rope.freq; + const float rope_scale = rope.scale; + const float rope_ext = rope.ext; + const float rope_attn = rope.attn; + const int rope_n_ctx_orig = rope.n_ctx_orig; // Position tensor for this token batch - ggml_tensor * rope_pos = cached_inputs ? cached_inputs->rope_pos : nullptr; + ggml_tensor * rope_pos = prepared ? prepared->rope_pos + : (cached_inputs ? cached_inputs->rope_pos : nullptr); if (!rope_pos) { rope_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); ggml_set_input(rope_pos); @@ -1657,27 +1768,20 @@ static ggml_tensor * build_mla_attention_lane_core( i32_array_inputs.push_back({rope_pos, std::move(pos_vals)}); } - // n_ctx_orig is critical for YaRN correction on compressed layers - const int rope_n_ctx_orig = (int)w.rope_orig_ctx; // 65536 - // D=512 flash prefill can rotate Q's 64-d tail inside the exact attention // kernel. This avoids materializing cont(nope), cont(tail), rope(tail), // and concat(nope, tail) while retaining the same F32 rounding boundary. const bool fuse_q_rope = attention_impl != DeepSeek4AttentionImpl::Explicit && n_tokens > 1 && head_dim == 512 && n_rot == 64; - if (!fuse_q_rope) { - q = build_tail_rope_3d(ctx, q, rope_pos, n_rot, head_dim, n_head, n_tokens, - rope_freq, rope_scale, rope_ext, rope_attn, - w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + if (prepared) { + projected = *prepared; + } else { + build_mla_qkv_rope( + ctx, projected, w, rope, n_tokens, rope_pos, fuse_q_rope); } - kv = build_tail_rope_2d(ctx, kv, rope_pos, n_rot, head_dim, n_tokens, - rope_freq, rope_scale, rope_ext, rope_attn, - w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); - - const DeepSeek4PreparedProjectedLane projected = {qr, q, kv, rope_pos}; - // Keep the established local names below to make the no-topology-change - // property obvious; the bundle is the handoff seam for a future adapter. - (void) projected; + ggml_tensor * qr = projected.normalized_q_lora; + ggml_tensor * q = projected.q; + ggml_tensor * kv = projected.kv; // ── Causal batched step (exact multi-token target semantics) ─── // The target model is causal: token i must not attend to batch tokens @@ -1825,7 +1929,13 @@ static ggml_tensor * build_mla_attention_lane_core( } ggml_tensor * index_comp_kv_source = lane.index_comp_kv; - if (lane.write_enabled && ratio == 4 && L.indexer_compressor_kv) { + // Gathered paged concurrency always uses Explicit attention, whose + // build_indexer_topk path is disabled. In that mode the indexer compressor + // only writes state that no graph node reads, so omit the dead subgraph. + const bool indexer_compressor_is_dead = + gathered_history && attention_impl != DeepSeek4AttentionImpl::SparseFlash; + if (lane.write_enabled && ratio == 4 && L.indexer_compressor_kv && + !indexer_compressor_is_dead) { build_indexer_compressor_step(ctx, gf, cur_last, w, L, *lane.indexer_compressor, lane.index_comp_kv, token_pos, cached_inputs ? cached_inputs->index_ape_row : nullptr, @@ -2302,46 +2412,13 @@ static ggml_tensor * build_mla_attention_lane_core( // Flatten to [head_dim*n_head, n_tokens] for output projection ggml_tensor * attn_out = ggml_reshape_2d(ctx, context, head_dim * n_head, n_tokens); - // ── Grouped output projection ────────────────────────────────── - // DS4 output uses grouped low-rank projection: - // attn_out: [head_dim*n_head, n_tokens] → reshape [group_dim, n_tokens, n_groups] - // out_a: [group_dim, n_groups*n_lora_o] → reshape [group_dim, n_lora_o, n_groups] - // batched matmul over n_groups: → [n_lora_o, n_tokens, n_groups] - // → reshape [n_lora_o*n_groups, n_tokens] - // out_b: [n_lora_o*n_groups, n_embd] → final: [n_embd, n_tokens] - const int group_dim = head_dim * (n_head / n_out_group); // 512 * 8 = 4096 - // Reshape attn_out: [32768, n_tokens] → [4096, 8, n_tokens] → permute to [4096, n_tokens, 8] - attn_out = ggml_reshape_3d(ctx, attn_out, group_dim, n_out_group, n_tokens); - attn_out = ggml_permute(ctx, attn_out, 0, 2, 1, 3); - if (n_tokens == 1) { - attn_out = ggml_cont(ctx, attn_out); - } - // attn_out is now [group_dim, n_tokens, n_out_group] - ggml_tensor * out_a_3d = ggml_reshape_3d(ctx, L.attn_output_a, group_dim, n_lora_o, n_out_group); - // out_a_3d: [group_dim, n_lora_o, n_out_group] — ne[2] matches - ggml_tensor * attn_low = ggml_mul_mat(ctx, out_a_3d, attn_out); - // attn_low: [n_lora_o, n_tokens, n_out_group] - ggml_tensor * out = nullptr; - const bool grouped_output_projection = - n_tokens > 1 && - !ds4_env_flag("DFLASH_DS4_DISABLE_GROUPED_OUTPUT_PROJECTION"); - if (grouped_output_projection) { - // Batched ROCmFPX MMQ consumes src1's channel stride directly. This - // avoids materializing both permutations (~256 MiB/layer at 2K). - out = ggml_mul_mat_grouped_src(ctx, L.attn_output_b, attn_low); - } else { - // Preserve the established single-token graph and provide an exact - // fallback for heterogeneous runtimes that cannot retain grouped-view - // metadata across a scheduler copy. At verifier widths (q <= 4), this - // materializes at most 128 KiB per layer rather than the long-prefill - // volume avoided by the grouped path. - attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); - attn_low = ggml_reshape_2d( - ctx, attn_low, n_lora_o * n_out_group, n_tokens); - out = ggml_mul_mat(ctx, L.attn_output_b, attn_low); + if (out_attn_context) { + *out_attn_context = attn_out; + return attn_out; } - return out; + return build_mla_output_projection(ctx, attn_out, w, L, n_tokens, + /*allow_grouped=*/true); } // Legacy contiguous-cache adapter. Both decode and the consecutive q>1 @@ -4678,6 +4755,16 @@ struct DeepSeek4FusedDecodeGraph { std::vector authoritative_routes; ggml_tensor * logits = nullptr; ggml_backend_sched_t sched = nullptr; + // The scheduler owns large pinned cross-backend staging buffers. Retain it + // across gathered-paged shape rebuilds when its backend set and capacity + // still match. + size_t sched_capacity = 0; + std::array sched_backends{}; + + bool sched_reusable(const std::array & backends, + size_t capacity) const { + return sched && sched_capacity >= capacity && sched_backends == backends; + } void reset_nodes() { inp_embed = nullptr; @@ -4713,6 +4800,22 @@ struct DeepSeek4FusedDecodeGraph { } } + // Retain shape-independent resources, but first retire native graph + // executables whose keys point into this metadata arena. The allocator and + // scheduler remain alive; the builder resets their per-graph state. + void release_for_rebuild(ggml_backend_t main_backend, + ggml_backend_t peer_backend = nullptr) { + invalidate_native_graphs(main_backend, peer_backend); + // Clear scheduler registrations while their tensor metadata is still + // valid. The builder may reset again after deciding to reuse it; that + // second reset is a no-op but keeps the builder self-contained. + if (sched) { + ggml_backend_sched_reset(sched); + } + step_graph_free(sg); + reset_nodes(); + } + void destroy(ggml_backend_t main_backend, ggml_backend_t peer_backend = nullptr) { // Native graph executables outlive ggml graph metadata in the backend @@ -4723,6 +4826,8 @@ struct DeepSeek4FusedDecodeGraph { ggml_backend_sched_free(sched); sched = nullptr; } + sched_capacity = 0; + sched_backends = {}; step_graph_destroy(sg); reset_nodes(); } @@ -4784,6 +4889,15 @@ struct Ds4FusedVerifyCache { ggml_tensor * ape = nullptr; ggml_tensor * state_row = nullptr; ggml_tensor * comp_pos = nullptr; + // Element offsets into the shared per-dtype upload bundles. + int64_t i32_base = -1; // pos, neg_pos, comp_read, ape, comp_pos + int64_t i64_base = -1; // raw_write, comp_write, state_row + int64_t raw_off = -1; + int64_t raw_n = 0; + int64_t comp_off = -1; + int64_t comp_n = 0; + int64_t index_off = -1; + int64_t index_n = 0; }; ggml_tensor * pos_q = nullptr; // i32 [q] ggml_tensor * neg_q = nullptr; // i32 [q] @@ -4798,6 +4912,12 @@ struct Ds4FusedVerifyCache { // Keeping it per slot removes one allocation from every verify step. std::vector mask_values; std::vector paged; // [layer*q], paged mode only + ggml_tensor * paged_i32 = nullptr; + ggml_tensor * paged_i64 = nullptr; + ggml_tensor * paged_gather = nullptr; + int64_t paged_i32_n = 0; + int64_t paged_i64_n = 0; + int64_t paged_gather_n = 0; int q = 0; void reset() { *this = Extra{}; } @@ -6897,7 +7017,8 @@ bool deepseek4_paged_gathered_step( if (vc.slots[i].last_use < vc.slots[pick].last_use) pick = i; } fg = &vc.slots[pick]; ex = &vc.extra[pick]; - fg->destroy(vc.backend, vc.peer_backend); ex->reset(); + fg->release_for_rebuild(vc.backend, vc.peer_backend); + ex->reset(); if (!ds4_build_fused_verify_graph( mc, *fg, *ex, backend, w, cache.prefill_staging, rt->model.hc_layer_weights, rt->model.hc_output_weights, @@ -6907,47 +7028,83 @@ bool deepseek4_paged_gathered_step( std::fprintf(stderr, "[deepseek4-paged] failed to build gathered graph " "(lanes=%u)\n", lanes); + fg->destroy(vc.backend, vc.peer_backend); + ex->reset(); return false; } } fg->last_use = vc.counter; ds4_fv_set(fg->inp_embed, embeddings, sizeof(float) * (size_t) w.n_embd * lanes); + // The shared Q/KV prologue rotates all gathered lanes at once. Padding + // lanes use position zero, matching their passive prepared row record. + { + std::vector pos_batch(lanes, 0); + std::vector neg_batch(lanes, 0); + for (uint32_t lane = 0; lane < lanes; ++lane) { + if (slots[lane] < 0) continue; + pos_batch[lane] = (int32_t) positions[lane]; + neg_batch[lane] = -(int32_t) positions[lane]; + } + ds4_fv_set(ex->pos_q, pos_batch.data(), sizeof(int32_t) * lanes); + ds4_fv_set(ex->neg_q, neg_batch.data(), sizeof(int32_t) * lanes); + } + + std::vector bundle_i32( + (size_t) std::max(ex->paged_i32_n, 0), 0); + std::vector bundle_i64( + (size_t) std::max(ex->paged_i64_n, 0), 0); + std::vector bundle_gather( + (size_t) std::max(ex->paged_gather_n, 0), 0); size_t pi = 0; for (int il = 0; il < w.n_layer; ++il) { const int ratio = (int) cache.layers[(size_t) il].ratio; for (uint32_t lane = 0; lane < lanes; ++lane, ++pi) { const auto & row = prepared[(size_t) il][lane]; const auto & px = ex->paged[pi]; + if (px.i32_base < 0 || px.i64_base < 0) return false; const int32_t pos = (int32_t) row.position; - const int32_t neg_pos = -pos; - ds4_fv_set(px.pos, &pos, sizeof(pos)); - ds4_fv_set(px.neg_pos, &neg_pos, sizeof(neg_pos)); - std::vector idx(std::max(row.raw_history.size(), 1), 0); - for (size_t i = 0; i < row.raw_history.size(); ++i) idx[i] = (int32_t) row.raw_history[i]; - ds4_fv_set(px.raw_gather, idx.data(), idx.size() * sizeof(int32_t)); - const int64_t raw_write = std::max(row.raw_scatter, 0); - ds4_fv_set(px.raw_write, &raw_write, sizeof(raw_write)); - if (ratio > 0) { - idx.assign(std::max(row.compressed_history.size(), 1), 0); - for (size_t i = 0; i < row.compressed_history.size(); ++i) - idx[i] = (int32_t) row.compressed_history[i]; - ds4_fv_set(px.comp_gather, idx.data(), idx.size() * sizeof(int32_t)); - if (px.index_gather) - ds4_fv_set(px.index_gather, idx.data(), idx.size() * sizeof(int32_t)); + bundle_i32[(size_t) px.i32_base + 0] = pos; + bundle_i32[(size_t) px.i32_base + 1] = -pos; + if (px.raw_off < 0 || + px.raw_off + px.raw_n > ex->paged_gather_n || + (int64_t) row.raw_history.size() > px.raw_n) return false; + for (size_t i = 0; i < row.raw_history.size(); ++i) { + bundle_gather[(size_t) px.raw_off + i] = + (int32_t) row.raw_history[i]; + } + bundle_i64[(size_t) px.i64_base + 0] = + std::max(row.raw_scatter, 0); + if (ratio > 0 && px.comp_off >= 0) { + if (px.comp_off + px.comp_n > ex->paged_gather_n || + (int64_t) row.compressed_history.size() > px.comp_n) { + return false; + } + for (size_t i = 0; i < row.compressed_history.size(); ++i) { + const int32_t value = + (int32_t) row.compressed_history[i]; + bundle_gather[(size_t) px.comp_off + i] = value; + if (px.index_off >= 0) { + bundle_gather[(size_t) px.index_off + i] = value; + } + } const int64_t cw = std::max(row.compressed_scatter, 0); - const int32_t cr = (int32_t) cw; const int32_t ape = pos % ratio; - const int64_t state = ratio == 4 ? 4 + ape : ape; - const int32_t comp_pos = pos + 1 - ratio; - ds4_fv_set(px.comp_write, &cw, sizeof(cw)); - ds4_fv_set(px.comp_read, &cr, sizeof(cr)); - ds4_fv_set(px.ape, &ape, sizeof(ape)); - ds4_fv_set(px.state_row, &state, sizeof(state)); - ds4_fv_set(px.comp_pos, &comp_pos, sizeof(comp_pos)); + bundle_i64[(size_t) px.i64_base + 1] = cw; + bundle_i64[(size_t) px.i64_base + 2] = + ratio == 4 ? 4 + ape : ape; + bundle_i32[(size_t) px.i32_base + 2] = (int32_t) cw; + bundle_i32[(size_t) px.i32_base + 3] = ape; + bundle_i32[(size_t) px.i32_base + 4] = pos + 1 - ratio; } } } + ds4_fv_set(ex->paged_i32, bundle_i32.data(), + bundle_i32.size() * sizeof(int32_t)); + ds4_fv_set(ex->paged_i64, bundle_i64.data(), + bundle_i64.size() * sizeof(int64_t)); + ds4_fv_set(ex->paged_gather, bundle_gather.data(), + bundle_gather.size() * sizeof(int32_t)); if (token_ids) { for (int il = 0; il < w.n_layer; ++il) { ggml_tensor * ids = fg->hash_ids[(size_t) il]; if (!ids) continue; diff --git a/server/src/deepseek4/deepseek4_paged_cache.cpp b/server/src/deepseek4/deepseek4_paged_cache.cpp index cdb395dbb..caca04324 100644 --- a/server/src/deepseek4/deepseek4_paged_cache.cpp +++ b/server/src/deepseek4/deepseek4_paged_cache.cpp @@ -34,8 +34,11 @@ bool prepare_deepseek4_gathered_lane_rows( for (uint32_t lane = 0; lane < lanes; ++lane) { auto & rows = prepared[lane]; rows.slot = slots[lane]; - rows.position = positions[lane]; if (rows.slot < 0) continue; // Padding must remain entirely passive. + // Padding lanes do not have a validated position. Keep the default + // zero so inverse RoPE and compressor bookkeeping cannot inherit an + // arbitrary positions[] value once callers use constant-width steps. + rows.position = positions[lane]; if (rows.position < 0) return false; const uint64_t pos = static_cast(rows.position); // The current row is appended in-graph, so retain at most the 127 diff --git a/server/test/test_deepseek4_paged_cache.cpp b/server/test/test_deepseek4_paged_cache.cpp index dd9681bf0..11df40ef3 100644 --- a/server/test/test_deepseek4_paged_cache.cpp +++ b/server/test/test_deepseek4_paged_cache.cpp @@ -47,6 +47,7 @@ int main() { CHECK(rows[1].compressed_emitted && rows[1].compressed_scatter == 7 * 32); CHECK(rows[2].raw_history.empty() && rows[2].compressed_history.empty()); CHECK(rows[2].raw_scatter == -1 && rows[2].compressed_scatter == -1); + CHECK(rows[2].position == 0); std::vector sixteen_slots(16); std::vector sixteen_positions(16, 0);