From 8902901e61fcc514244bd5cd2ec0109bc56f6571 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 12 Aug 2026 22:55:17 +0000 Subject: [PATCH 01/18] feat(qwen35): complete concurrent speculative serving --- server/CMakeLists.txt | 37 +- server/deps/llama.cpp/ggml/include/ggml.h | 19 +- .../ggml/src/ggml-cuda/paged-attn.cu | 192 +++- server/deps/llama.cpp/ggml/src/ggml.c | 54 +- .../src/common/concurrency/paged_kv_pool.cpp | 110 ++- server/src/common/concurrency/paged_kv_pool.h | 48 + .../common/concurrency/paged_kv_residency.cpp | 929 ++++++++++++++++++ .../common/concurrency/paged_kv_residency.h | 260 +++++ .../concurrency/qwen_paged_kv_transfer.cpp | 405 ++++++++ .../concurrency/qwen_paged_kv_transfer.h | 87 ++ .../qwen_paged_kv_transfer_layout.cpp | 119 +++ server/src/common/concurrency/seq_engine.h | 65 +- server/src/common/ddtree.cpp | 10 + server/src/common/ddtree.h | 10 + server/src/common/feature_gate.cpp | 56 +- server/src/common/gpu_runtime_compat.h | 1 + server/src/common/step_graph.h | 10 +- server/src/internal.h | 36 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 655 +++++++++++- .../qwen35/concurrency/qwen35_seq_engine.h | 36 +- .../concurrency/qwen35_slot_manager.cpp | 251 ++++- .../qwen35/concurrency/qwen35_slot_manager.h | 53 +- server/src/qwen35/graph_builders.cpp | 156 +++ server/src/qwen35/graph_builders.h | 40 + server/src/qwen35/qwen35_backend.cpp | 99 +- server/src/qwen35/qwen35_backend.h | 6 +- server/src/qwen35/qwen35_target_graph.cpp | 187 +++- server/src/server/http_server.cpp | 27 +- server/src/server/http_server.h | 50 + server/src/server/scheduler.cpp | 124 ++- server/test/bench_paged_attention.cpp | 3 +- server/test/seq_engine_contract.h | 13 + server/test/test_ddtree_path.cpp | 46 + server/test/test_feature_gate.cpp | 63 ++ server/test/test_paged_attention.cpp | 192 +++- server/test/test_paged_kv_pool.cpp | 77 ++ server/test/test_paged_kv_residency.cpp | 559 +++++++++++ .../test_qwen_paged_kv_transfer_layout.cpp | 86 ++ server/test/test_recurrent_snapshot.cpp | 33 + server/test/test_seq_batch_plan.cpp | 37 + server/test/test_seq_engine_contract.cpp | 8 + server/test/test_seq_slot_manager.cpp | 58 +- server/test/test_server_unit.cpp | 71 ++ 43 files changed, 5168 insertions(+), 210 deletions(-) create mode 100644 server/src/common/concurrency/paged_kv_residency.cpp create mode 100644 server/src/common/concurrency/paged_kv_residency.h create mode 100644 server/src/common/concurrency/qwen_paged_kv_transfer.cpp create mode 100644 server/src/common/concurrency/qwen_paged_kv_transfer.h create mode 100644 server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp create mode 100644 server/test/test_ddtree_path.cpp create mode 100644 server/test/test_paged_kv_residency.cpp create mode 100644 server/test/test_qwen_paged_kv_transfer_layout.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 46ecfda81..ad6ef874f 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -450,6 +450,9 @@ add_library(dflash_common STATIC src/common/dflash_draft_kv.cpp src/common/dflash_spec_decode.cpp src/common/concurrency/paged_kv_pool.cpp + src/common/concurrency/paged_kv_residency.cpp + src/common/concurrency/qwen_paged_kv_transfer_layout.cpp + src/common/concurrency/qwen_paged_kv_transfer.cpp src/qwen35/concurrency/qwen35_slot_manager.cpp src/common/layer_split_backend.cpp src/common/layer_split_runtime.cpp @@ -1391,12 +1394,33 @@ 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_paged_kv_residency.cpp") + # Pure host-side multi-sequence residency policy + mock DMA test. + add_executable(test_paged_kv_residency + test/test_paged_kv_residency.cpp + src/common/concurrency/paged_kv_pool.cpp + src/common/concurrency/paged_kv_residency.cpp) + target_include_directories(test_paged_kv_residency PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_paged_kv_residency) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_qwen_paged_kv_transfer_layout.cpp") + # Pure host-side validation of packed K/V block byte/stride math. + add_executable(test_qwen_paged_kv_transfer_layout + test/test_qwen_paged_kv_transfer_layout.cpp + src/common/concurrency/qwen_paged_kv_transfer_layout.cpp) + target_include_directories(test_qwen_paged_kv_transfer_layout PRIVATE + ${DFLASH27B_SRC_INCLUDE_DIRS} + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + list(APPEND _raw_unit_test_targets test_qwen_paged_kv_transfer_layout) + 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/qwen35/concurrency/qwen35_slot_manager.cpp - src/common/concurrency/paged_kv_pool.cpp) + src/common/concurrency/paged_kv_pool.cpp + src/common/concurrency/paged_kv_residency.cpp) target_include_directories(test_seq_slot_manager PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) list(APPEND _raw_unit_test_targets test_seq_slot_manager) @@ -1409,6 +1433,15 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/test) list(APPEND _raw_unit_test_targets test_seq_engine_contract) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_ddtree_path.cpp") + # Pure host-side accepted-path/pending-token contract tests. + add_executable(test_ddtree_path + test/test_ddtree_path.cpp + src/common/ddtree.cpp) + target_include_directories(test_ddtree_path PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src) + list(APPEND _raw_unit_test_targets test_ddtree_path) + 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) @@ -1791,6 +1824,8 @@ if(DFLASH27B_TESTS) set(_unit_ctest_name recurrent_snapshot) elseif(_unit_target STREQUAL "test_paged_kv_pool") set(_unit_ctest_name paged_kv_pool) + elseif(_unit_target STREQUAL "test_paged_kv_residency") + set(_unit_ctest_name paged_kv_residency) endif() add_test(NAME "${_unit_ctest_name}" COMMAND ${_unit_target}) set_tests_properties("${_unit_ctest_name}" PROPERTIES SKIP_RETURN_CODE 77) diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index acac40c1f..beb6fc817 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2502,6 +2502,18 @@ extern "C" { // prefill chunks can attend the paged pool causally. A negative position // marks a padding row. NULL keeps the decode semantics (full cached // length per row). + // + // parent_ids/tree_sizes optionally enable packed tree verification. + // Queries are flattened sequence-major: tree sequence s occupies rows + // [s*tree_width, (s+1)*tree_width). parent_ids is contiguous I32 + // [tree_width, n_tree_seq] (root parent -1), and tree_sizes is contiguous + // I32 [n_tree_seq]. active_slot_ids is required and remains per query row; + // it selects the physical block-table column and scratch slab. Each live + // query attends its complete committed prefix from the block table plus + // its own candidate node and ancestors from physical K/V rows + // tree_scratch_base + slot*tree_scratch_stride + node. Siblings and rows + // at or beyond tree_sizes[s] are excluded. query_positions must be NULL + // in tree mode. Pass NULL/NULL/0/0/0 to retain standard paged attention. GGML_API struct ggml_tensor * ggml_paged_attn_ext( struct ggml_context * ctx, struct ggml_tensor * q, @@ -2513,7 +2525,12 @@ extern "C" { struct ggml_tensor * query_positions, float scale, int block_size, - int max_kv_seq_len); + int max_kv_seq_len, + struct ggml_tensor * parent_ids, + struct ggml_tensor * tree_sizes, + int tree_width, + int tree_scratch_base, + int tree_scratch_stride); // TurboQuant FWHT rotation. direction: 0 = forward, 1 = inverse. // Applies signs1 -> FWHT -> signs2 (forward) or signs2 -> FWHT -> signs1 (inverse). diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu index c76d97b36..a881b93d3 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/paged-attn.cu @@ -29,6 +29,42 @@ static __host__ __device__ __forceinline__ int32_t paged_attn_partitions( return requested < available ? requested : available; } +// parent_ids is a sequence-major [tree_width, n_tree_seq] table. Walk only +// from the current query node toward the root; a candidate is visible iff it +// appears on that chain. The bounded walk also turns malformed cycles or +// out-of-range parents into invisible edges instead of an unsafe read. +static __device__ __forceinline__ bool paged_attn_tree_visible( + const char * __restrict__ parent_ids, + int64_t parent_nb0, + int64_t parent_nb1, + int32_t tree_seq, + int32_t query_node, + int32_t candidate, + int32_t tree_size) { + if (candidate < 0 || candidate >= tree_size || + query_node < 0 || query_node >= tree_size) { + return false; + } + + int32_t current = query_node; + for (int32_t depth = 0; depth < tree_size; ++depth) { + if (current == candidate) { + return true; + } + if (current < 0 || current >= tree_size) { + return false; + } + const int32_t parent = *(const int32_t *) ( + parent_ids + (int64_t) current * parent_nb0 + + (int64_t) tree_seq * parent_nb1); + if (parent == current) { + return false; + } + current = parent; + } + return false; +} + // All scores are computed in the log2 domain: log2(e) is folded into the same // Q prescale that already carries the 1/sqrt(D) attention scale, so every // softmax exponential uses the fast exp2f SFU path. @@ -180,6 +216,8 @@ static __global__ void paged_attn_decode( const char * __restrict__ kv_seq_lens, const char * __restrict__ active_slot_ids, const char * __restrict__ query_positions, + const char * __restrict__ parent_ids, + const char * __restrict__ tree_sizes, char * __restrict__ dst, half * __restrict__ partial_acc, float2 * __restrict__ partial_meta, @@ -189,6 +227,7 @@ static __global__ void paged_attn_decode( int64_t bt_nb0, int64_t bt_nb1, int64_t ksl_nb0, int64_t asi_nb0, int64_t qpos_nb0, + int64_t parent_nb0, int64_t parent_nb1, int64_t tree_size_nb0, int64_t dst_nb1, int64_t dst_nb2, int32_t n_table_seq, int32_t n_head, @@ -197,6 +236,9 @@ static __global__ void paged_attn_decode( int32_t max_blocks, int32_t block_size, int32_t min_partitions, + int32_t tree_width, + int32_t tree_scratch_base, + int32_t tree_scratch_stride, float scale) { constexpr int nthreads = WARP_SIZE; constexpr int values_per_load = 4; @@ -222,6 +264,15 @@ static __global__ void paged_attn_decode( const int n_seq = gridDim.y; const int n_partitions = gridDim.z; + const bool tree_mode = parent_ids != nullptr; + const int32_t tree_seq = tree_mode ? seq / tree_width : 0; + const int32_t query_node = + tree_mode ? seq - tree_seq * tree_width : -1; + const int32_t tree_size = tree_mode + ? *(const int32_t *) ( + tree_sizes + (int64_t) tree_seq * tree_size_nb0) + : 0; + const int32_t physical_seq_raw = active_slot_ids ? *(const int32_t *) (active_slot_ids + (int64_t) seq * asi_nb0) : seq; @@ -229,21 +280,22 @@ static __global__ void paged_attn_decode( ? *(const int32_t *) (query_positions + (int64_t) seq * qpos_nb0) : -1; // A row is live when its slot id selects a real block-table column and, - // for ragged batches, its causal position is non-negative. Dead rows are - // pinned to column 0 with kv_seq_len forced to 0, which routes every - // partition through the existing zero-output early path; the block table - // is then never read for them. + // for ragged batches, its causal position is non-negative. Tree padding + // rows are validated by tree_sizes. Dead rows are pinned to column 0 with + // an empty virtual context, so the block table and scratch are never read. const bool valid_query = physical_seq_raw >= 0 && physical_seq_raw < n_table_seq && - (!query_positions || query_pos >= 0); + (!query_positions || query_pos >= 0) && + (!tree_mode || + (tree_size >= 0 && tree_size <= tree_width && + query_node < tree_size)); const int32_t physical_seq = valid_query ? physical_seq_raw : 0; int32_t kv_seq_len_raw = valid_query ? *(const int32_t *) (kv_seq_lens + (int64_t) physical_seq * ksl_nb0) : 0; - // The inclusive clamp IS the causal mask: this row attends tokens - // [0, pos] only, and every downstream bound (partition count, token loop - // extents) already derives from kv_seq_len. + // The inclusive clamp IS the causal mask for non-tree ragged rows. Tree + // rows always read the whole committed prefix carried by kv_seq_lens. if (query_positions && query_pos < kv_seq_len_raw) { kv_seq_len_raw = query_pos + 1; } @@ -254,8 +306,14 @@ static __global__ void paged_attn_decode( : (kv_seq_len_raw < table_capacity ? kv_seq_len_raw : (int32_t) table_capacity); + // Treat the candidate slab as a virtual tail of tree_width tokens. The + // normal partition split then covers prefix and tree candidates in one + // stable softmax; invisible siblings/padding resolve to no physical row. + const int32_t virtual_tokens = valid_query + ? kv_seq_len + (tree_mode ? tree_width : 0) + : 0; const int32_t n_logical_blocks = - (kv_seq_len + block_size - 1) / block_size; + (virtual_tokens + block_size - 1) / block_size; const int32_t active_partitions = paged_attn_partitions(n_logical_blocks, min_partitions, n_partitions); @@ -289,7 +347,7 @@ static __global__ void paged_attn_decode( const int32_t token_begin = logical_block_begin * block_size; const int32_t token_end_blocks = logical_block_end * block_size; const int32_t token_end = - kv_seq_len < token_end_blocks ? kv_seq_len : token_end_blocks; + virtual_tokens < token_end_blocks ? virtual_tokens : token_end_blocks; constexpr bool quantize_q = type_K != GGML_TYPE_F16; constexpr int q_registers = (D / 2) / nthreads; @@ -363,7 +421,12 @@ static __global__ void paged_attn_decode( qk_sum[h] = 0.0f; } - const int32_t n_physical_blocks = pool_tokens / block_size; + // In tree mode the committed block table may address only the prefix + // pool before tree_scratch_base. Candidate rows are addressed directly + // below, keeping uncommitted nodes out of every sequence block table. + const int32_t prefix_pool_tokens = + tree_mode ? tree_scratch_base : pool_tokens; + const int32_t n_physical_blocks = prefix_pool_tokens / block_size; for (int32_t tile_begin = token_begin; tile_begin < token_end; @@ -380,7 +443,7 @@ static __global__ void paged_attn_decode( // read; their tokens contribute nothing, mirroring the CPU reference. int32_t phys_mine = -1; const int32_t my_token = tile_begin + lane; - if (my_token < token_end) { + if (my_token < token_end && my_token < kv_seq_len) { const int32_t logical_block = my_token / block_size; const int32_t physical_block = *(const int32_t *) (block_table + @@ -390,6 +453,19 @@ static __global__ void paged_attn_decode( phys_mine = physical_block * block_size + my_token % block_size; } + } else if (tree_mode && my_token < token_end) { + const int32_t candidate = my_token - kv_seq_len; + if (paged_attn_tree_visible( + parent_ids, parent_nb0, parent_nb1, + tree_seq, query_node, candidate, tree_size)) { + const int64_t physical = + (int64_t) tree_scratch_base + + (int64_t) physical_seq * tree_scratch_stride + + candidate; + if (physical >= 0 && physical < pool_tokens) { + phys_mine = (int32_t) physical; + } + } } float score_mine[n_batch_heads]; @@ -619,6 +695,8 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { const ggml_tensor * kv_seq_lens = dst->src[4]; const ggml_tensor * active_slot_ids = dst->src[5]; const ggml_tensor * query_positions = dst->src[6]; + const ggml_tensor * parent_ids = dst->src[7]; + const ggml_tensor * tree_sizes = dst->src[8]; if (!q || !k || !v || !block_table || !kv_seq_lens) { return false; @@ -628,6 +706,11 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { if (query_positions && !active_slot_ids) { return false; } + const bool tree_mode = parent_ids || tree_sizes; + if ((parent_ids == nullptr) != (tree_sizes == nullptr) || + (tree_mode && (!active_slot_ids || query_positions))) { + return false; + } if (dst->type != GGML_TYPE_F32 || q->type != GGML_TYPE_F32 || @@ -636,7 +719,9 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { block_table->type != GGML_TYPE_I32 || kv_seq_lens->type != GGML_TYPE_I32 || (active_slot_ids && active_slot_ids->type != GGML_TYPE_I32) || - (query_positions && query_positions->type != GGML_TYPE_I32)) { + (query_positions && query_positions->type != GGML_TYPE_I32) || + (parent_ids && parent_ids->type != GGML_TYPE_I32) || + (tree_sizes && tree_sizes->type != GGML_TYPE_I32)) { return false; } @@ -647,6 +732,8 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { kv_seq_lens->nb[0] != sizeof(int32_t) || (active_slot_ids && active_slot_ids->nb[0] != sizeof(int32_t)) || (query_positions && query_positions->nb[0] != sizeof(int32_t)) || + (parent_ids && parent_ids->nb[0] != sizeof(int32_t)) || + (tree_sizes && tree_sizes->nb[0] != sizeof(int32_t)) || dst->nb[0] != sizeof(float)) { return false; } @@ -701,10 +788,47 @@ bool ggml_cuda_paged_attn_supported(const ggml_tensor * dst) { const int32_t block_size = ggml_get_op_params_i32(dst, 1); const int32_t max_kv_seq_len = ggml_get_op_params_i32(dst, 2); - return block_size > 0 && - max_kv_seq_len > 0 && - max_kv_seq_len <= k->ne[1] && - k->ne[1] % block_size == 0; + const int32_t tree_width = ggml_get_op_params_i32(dst, 3); + const int32_t tree_scratch_base = ggml_get_op_params_i32(dst, 4); + const int32_t tree_scratch_stride = ggml_get_op_params_i32(dst, 5); + if (block_size <= 0 || + max_kv_seq_len <= 0 || + (int64_t) max_kv_seq_len + tree_width > INT32_MAX || + k->ne[1] % block_size != 0) { + return false; + } + + if (!tree_mode) { + return tree_width == 0 && + tree_scratch_base == 0 && + tree_scratch_stride == 0; + } + + if (tree_width <= 0 || + tree_scratch_base <= 0 || + tree_scratch_base % block_size != 0 || + tree_scratch_stride < tree_width || + !ggml_is_contiguous(parent_ids) || + !ggml_is_contiguous(tree_sizes) || + parent_ids->ne[0] != tree_width || + parent_ids->ne[1] <= 0 || + parent_ids->ne[1] != tree_sizes->ne[0] || + parent_ids->ne[2] != 1 || + parent_ids->ne[3] != 1 || + tree_sizes->ne[1] != 1 || + tree_sizes->ne[2] != 1 || + tree_sizes->ne[3] != 1 || + parent_ids->ne[1] > INT64_MAX / tree_width || + q->ne[1] != parent_ids->ne[1] * tree_width || + (int64_t) max_kv_seq_len + tree_width > INT32_MAX) { + return false; + } + + const int64_t scratch_end = + (int64_t) tree_scratch_base + + (block_table->ne[1] - 1) * (int64_t) tree_scratch_stride + + tree_width; + return scratch_end <= k->ne[1]; } // Cached max resident blocks/SM for this instantiation at the given block @@ -764,6 +888,12 @@ static bool try_launch_paged_attn( const ggml_tensor * kv_seq_lens = dst->src[4]; const ggml_tensor * active_slot_ids = dst->src[5]; const ggml_tensor * query_positions = dst->src[6]; + const ggml_tensor * parent_ids = dst->src[7]; + const ggml_tensor * tree_sizes = dst->src[8]; + + const int32_t tree_width = ggml_get_op_params_i32(dst, 3); + const int32_t tree_scratch_base = ggml_get_op_params_i32(dst, 4); + const int32_t tree_scratch_stride = ggml_get_op_params_i32(dst, 5); const int32_t n_head = (int32_t) q->ne[2]; const int32_t n_head_kv = (int32_t) k->ne[2]; @@ -837,15 +967,21 @@ static bool try_launch_paged_attn( if (min_partitions > partition_limit) { min_partitions = partition_limit; } - if (min_partitions > block_table->ne[0]) { - min_partitions = (int32_t) block_table->ne[0]; + const int32_t tree_blocks = + (tree_width + block_size - 1) / block_size; + const int64_t partitionable_blocks = + block_table->ne[0] + (parent_ids ? tree_blocks : 0); + if (min_partitions > partitionable_blocks) { + min_partitions = (int32_t) partitionable_blocks; } - // Size the launch from the live maximum sequence length carried in the - // graph op, not the block-table capacity. Ragged sequences still clamp - // their own active partition count from kv_seq_lens on device. + // Size the launch from the live maximum committed prefix plus the virtual + // tree tail. Ragged/tree rows still clamp their own active partition count + // from device metadata. + const int32_t live_tokens = + max_kv_seq_len + (parent_ids ? tree_width : 0); const int32_t live_blocks = - (max_kv_seq_len + block_size - 1) / block_size; + (live_tokens + block_size - 1) / block_size; int32_t n_partitions = paged_attn_partitions( live_blocks, min_partitions, PAGED_ATTN_MAX_PARTITIONS); @@ -861,7 +997,7 @@ static bool try_launch_paged_attn( }(); if (forced_partitions >= 1 && forced_partitions <= PAGED_ATTN_MAX_PARTITIONS && - forced_partitions <= block_table->ne[0]) { + forced_partitions <= partitionable_blocks) { min_partitions = forced_partitions; n_partitions = forced_partitions; } @@ -914,6 +1050,8 @@ static bool try_launch_paged_attn( (const char *) kv_seq_lens->data, active_slot_ids ? (const char *) active_slot_ids->data : nullptr, query_positions ? (const char *) query_positions->data : nullptr, + parent_ids ? (const char *) parent_ids->data : nullptr, + tree_sizes ? (const char *) tree_sizes->data : nullptr, (char *) dst->data, partial_acc, partial_meta, @@ -924,6 +1062,9 @@ static bool try_launch_paged_attn( kv_seq_lens->nb[0], active_slot_ids ? active_slot_ids->nb[0] : 0, query_positions ? query_positions->nb[0] : 0, + parent_ids ? parent_ids->nb[0] : 0, + parent_ids ? parent_ids->nb[1] : 0, + tree_sizes ? tree_sizes->nb[0] : 0, dst->nb[1], dst->nb[2], (int32_t) block_table->ne[1], n_head, @@ -932,6 +1073,9 @@ static bool try_launch_paged_attn( (int32_t) block_table->ne[0], block_size, min_partitions, + tree_width, + tree_scratch_base, + tree_scratch_stride, scale); if (n_partitions > 1) { diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 86e4520d6..77a590cca 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5708,7 +5708,12 @@ struct ggml_tensor * ggml_paged_attn_ext( struct ggml_tensor * query_positions, float scale, int block_size, - int max_kv_seq_len) { + int max_kv_seq_len, + struct ggml_tensor * parent_ids, + struct ggml_tensor * tree_sizes, + int tree_width, + int tree_scratch_base, + int tree_scratch_stride) { GGML_ASSERT(q->type == GGML_TYPE_F32); GGML_ASSERT(k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_Q4_0 || k->type == GGML_TYPE_Q8_0); GGML_ASSERT(v->type == GGML_TYPE_F16 || v->type == GGML_TYPE_Q4_0 || v->type == GGML_TYPE_Q8_0); @@ -5720,6 +5725,13 @@ struct ggml_tensor * ggml_paged_attn_ext( GGML_ASSERT(query_positions == NULL || active_slot_ids != NULL); GGML_ASSERT(query_positions == NULL || query_positions->type == GGML_TYPE_I32); + const bool tree_mode = parent_ids != NULL || tree_sizes != NULL; + GGML_ASSERT((parent_ids == NULL) == (tree_sizes == NULL)); + GGML_ASSERT(!tree_mode || active_slot_ids != NULL); + GGML_ASSERT(!tree_mode || query_positions == NULL); + GGML_ASSERT(!tree_mode || parent_ids->type == GGML_TYPE_I32); + GGML_ASSERT(!tree_mode || tree_sizes->type == GGML_TYPE_I32); + GGML_ASSERT(q->ne[0] == k->ne[0] && q->ne[0] == v->ne[0]); GGML_ASSERT(k->ne[1] == v->ne[1]); GGML_ASSERT(k->ne[2] > 0); @@ -5749,13 +5761,49 @@ struct ggml_tensor * ggml_paged_attn_ext( GGML_ASSERT(block_size > 0); GGML_ASSERT(k->ne[1] % block_size == 0); GGML_ASSERT(max_kv_seq_len > 0); - GGML_ASSERT(max_kv_seq_len <= k->ne[1]); + // This is a padded logical launch bound, not a physical-cache extent. + // Each row clamps its actual sequence length to the block-table capacity + // and validates every resolved physical block before dereferencing K/V. + GGML_ASSERT((int64_t) max_kv_seq_len + tree_width <= INT32_MAX); + + if (tree_mode) { + GGML_ASSERT(tree_width > 0); + GGML_ASSERT(tree_scratch_base > 0); + GGML_ASSERT(tree_scratch_base % block_size == 0); + GGML_ASSERT(tree_scratch_stride >= tree_width); + GGML_ASSERT(ggml_is_contiguous(parent_ids)); + GGML_ASSERT(ggml_is_contiguous(tree_sizes)); + GGML_ASSERT(parent_ids->ne[0] == tree_width); + GGML_ASSERT(parent_ids->ne[1] == tree_sizes->ne[0]); + GGML_ASSERT(parent_ids->ne[2] == 1 && parent_ids->ne[3] == 1); + GGML_ASSERT(tree_sizes->ne[1] == 1 && tree_sizes->ne[2] == 1 && tree_sizes->ne[3] == 1); + GGML_ASSERT(parent_ids->ne[1] > 0); + GGML_ASSERT(parent_ids->ne[1] <= INT64_MAX / tree_width); + GGML_ASSERT(q->ne[1] == parent_ids->ne[1] * tree_width); + + // Every physical sequence slot owns one non-overlapping scratch slab. + // Bound the largest address with int64 arithmetic before the GPU sees + // the int32 op parameters. + const int64_t scratch_end = + (int64_t) tree_scratch_base + + (block_table->ne[1] - 1) * (int64_t) tree_scratch_stride + + tree_width; + GGML_ASSERT(scratch_end <= k->ne[1]); + GGML_ASSERT((int64_t) max_kv_seq_len + tree_width <= INT32_MAX); + } else { + GGML_ASSERT(tree_width == 0); + GGML_ASSERT(tree_scratch_base == 0); + GGML_ASSERT(tree_scratch_stride == 0); + } struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, q->ne); ggml_set_op_params_f32(result, 0, scale); ggml_set_op_params_i32(result, 1, block_size); ggml_set_op_params_i32(result, 2, max_kv_seq_len); + ggml_set_op_params_i32(result, 3, tree_width); + ggml_set_op_params_i32(result, 4, tree_scratch_base); + ggml_set_op_params_i32(result, 5, tree_scratch_stride); result->op = GGML_OP_PAGED_ATTN; result->src[0] = q; @@ -5765,6 +5813,8 @@ struct ggml_tensor * ggml_paged_attn_ext( result->src[4] = kv_seq_lens; result->src[5] = active_slot_ids; result->src[6] = query_positions; + result->src[7] = parent_ids; + result->src[8] = tree_sizes; return result; } diff --git a/server/src/common/concurrency/paged_kv_pool.cpp b/server/src/common/concurrency/paged_kv_pool.cpp index 6bdf0f3fe..741ace600 100644 --- a/server/src/common/concurrency/paged_kv_pool.cpp +++ b/server/src/common/concurrency/paged_kv_pool.cpp @@ -20,6 +20,12 @@ const char * paged_kv_status_string(PagedKvStatus status) { return "physical blocks exhausted"; case PagedKvStatus::StaleHandle: return "stale sequence handle"; + case PagedKvStatus::LogicalBlockOutOfRange: + return "logical block out of range"; + case PagedKvStatus::BlockNotResident: + return "logical block is not resident"; + case PagedKvStatus::BlockAlreadyResident: + return "logical block is already resident"; } return "unknown paged KV status"; } @@ -138,8 +144,35 @@ PagedKvAppendResult PagedKvPool::append(PagedKvSequenceHandle handle, const uint32_t old_kv_seq_len = sequence.kv_seq_len; const uint32_t new_kv_seq_len = old_kv_seq_len + token_count; - result.status = - extend_block_table(sequence, blocks_for_tokens(new_kv_seq_len)); + const uint32_t required_blocks = blocks_for_tokens(new_kv_seq_len); + const uint32_t additional_blocks = + required_blocks - static_cast(sequence.block_table.size()); + + // Appending into a partially-filled cold head needs one physical remap in + // addition to any newly opened logical blocks. Preflight the aggregate so + // append remains all-or-nothing on BlocksExhausted. + const bool remap_head = + old_kv_seq_len % block_size_ != 0 && + sequence.block_table[old_kv_seq_len / block_size_] == + PAGED_KV_COLD_BLOCK; + const uint32_t allocations = additional_blocks + (remap_head ? 1u : 0u); + const uint64_t available = + static_cast(sequence.reserved_blocks.size()) + + free_blocks_.size(); + if (allocations > available) { + result.status = PagedKvStatus::BlocksExhausted; + return result; + } + + sequence.block_table.reserve(required_blocks); + if (remap_head) { + const uint32_t logical_block = old_kv_seq_len / block_size_; + const uint32_t physical_block = take_append_block(sequence); + sequence.block_table[logical_block] = physical_block; + result.remapped_cold_blocks.push_back( + {logical_block, physical_block}); + } + result.status = extend_block_table(sequence, required_blocks); if (result.status != PagedKvStatus::Ok) return result; const auto make_slot = [&](uint32_t logical_position) { @@ -176,7 +209,9 @@ PagedKvStatus PagedKvPool::release(PagedKvSequenceHandle handle) { SequenceState & sequence = sequences_[handle.slot]; request_to_slot_.erase(sequence.request_id); for (uint32_t block : sequence.block_table) { - give_back(free_blocks_, block); + if (block != PAGED_KV_COLD_BLOCK) { + give_back(free_blocks_, block); + } } for (uint32_t block : sequence.reserved_blocks) { give_back(free_blocks_, block); @@ -226,11 +261,70 @@ PagedKvStatus PagedKvPool::owned_block_count( if (status != PagedKvStatus::Ok) return status; const SequenceState & sequence = sequences_[handle.slot]; + // Cold logical blocks remain appended sequence capacity even though they + // no longer own a physical page. Preserve the pre-residency API contract: + // appended logical blocks plus blocks reserved for future append. out_count = static_cast( sequence.block_table.size() + sequence.reserved_blocks.size()); return PagedKvStatus::Ok; } +PagedKvStatus PagedKvPool::page_out_block( + PagedKvSequenceHandle handle, uint32_t logical_block, + uint32_t & out_physical_block) { + const PagedKvStatus status = validate(handle); + if (status != PagedKvStatus::Ok) return status; + + SequenceState & sequence = sequences_[handle.slot]; + if (logical_block >= sequence.block_table.size()) { + return PagedKvStatus::LogicalBlockOutOfRange; + } + const uint32_t physical_block = sequence.block_table[logical_block]; + if (physical_block == PAGED_KV_COLD_BLOCK) { + return PagedKvStatus::BlockNotResident; + } + + sequence.block_table[logical_block] = PAGED_KV_COLD_BLOCK; + give_back(free_blocks_, physical_block); + out_physical_block = physical_block; + return PagedKvStatus::Ok; +} + +PagedKvStatus PagedKvPool::page_in_block( + PagedKvSequenceHandle handle, uint32_t logical_block, + uint32_t & out_physical_block) { + const PagedKvStatus status = validate(handle); + if (status != PagedKvStatus::Ok) return status; + + SequenceState & sequence = sequences_[handle.slot]; + if (logical_block >= sequence.block_table.size()) { + return PagedKvStatus::LogicalBlockOutOfRange; + } + if (sequence.block_table[logical_block] != PAGED_KV_COLD_BLOCK) { + return PagedKvStatus::BlockAlreadyResident; + } + if (free_blocks_.empty()) { + return PagedKvStatus::BlocksExhausted; + } + + const uint32_t physical_block = take_lowest(free_blocks_); + sequence.block_table[logical_block] = physical_block; + out_physical_block = physical_block; + return PagedKvStatus::Ok; +} + +PagedKvStatus PagedKvPool::resident_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(std::count_if( + sequence.block_table.begin(), sequence.block_table.end(), + [](uint32_t block) { return block != PAGED_KV_COLD_BLOCK; })); + 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_; @@ -262,13 +356,17 @@ PagedKvStatus PagedKvPool::extend_block_table(SequenceState & sequence, sequence.block_table.reserve(required_blocks); for (uint32_t i = 0; i < additional_blocks; ++i) { - std::vector & source = sequence.reserved_blocks.empty() - ? free_blocks_ : sequence.reserved_blocks; - sequence.block_table.push_back(take_lowest(source)); + sequence.block_table.push_back(take_append_block(sequence)); } return PagedKvStatus::Ok; } +uint32_t PagedKvPool::take_append_block(SequenceState & sequence) { + std::vector & source = sequence.reserved_blocks.empty() + ? free_blocks_ : sequence.reserved_blocks; + return take_lowest(source); +} + void PagedKvPool::take_reserved_blocks(SequenceState & sequence, uint32_t additional_blocks) { sequence.reserved_blocks.reserve( diff --git a/server/src/common/concurrency/paged_kv_pool.h b/server/src/common/concurrency/paged_kv_pool.h index 9d745739f..3cf90964a 100644 --- a/server/src/common/concurrency/paged_kv_pool.h +++ b/server/src/common/concurrency/paged_kv_pool.h @@ -33,6 +33,9 @@ enum class PagedKvStatus : uint8_t { SequenceSlotsExhausted, // all max_sequences slots are in use BlocksExhausted, // not enough free physical blocks for the growth StaleHandle, // handle refers to a released or reused slot + LogicalBlockOutOfRange, // logical block is not materialized by the sequence + BlockNotResident, // page_out_block() targeted an already-cold block + BlockAlreadyResident, // page_in_block() targeted a resident block }; // Human-readable status name for logs and error messages. @@ -57,6 +60,23 @@ struct PagedKvWriteSlot { uint64_t physical_token_index = 0; }; +// Sentinel in PagedKvSequenceSnapshot::block_table for a logical block whose +// full-attention K/V bytes are host-backed by PagedKvResidencyManager. It is +// never returned as a PagedKvWriteSlot::physical_block. +inline constexpr uint32_t PAGED_KV_COLD_BLOCK = + std::numeric_limits::max(); + +// A formerly-cold append-head block that append() had to remap. The caller +// must restore the block's host-backed bytes before consuming or overwriting +// any row in the returned physical block. PagedKvResidencyManager::append() +// and prepare_append() do that before returning to the engine; this record +// primarily makes a direct pool append fail-obvious instead of losing the +// remap information. +struct PagedKvBlockRemap { + uint32_t logical_block = 0; + uint32_t physical_block = 0; +}; + // Outcome of append(). On success, `token_count` is the number of appended // tokens. By default, `write_slots` holds one entry per token in logical // order. With `only_first_last_slots`, it is empty and `first` and `last` @@ -68,6 +88,7 @@ struct PagedKvAppendResult { std::vector write_slots; PagedKvWriteSlot first; PagedKvWriteSlot last; + std::vector remapped_cold_blocks; explicit operator bool() const { return status == PagedKvStatus::Ok; } }; @@ -75,6 +96,8 @@ struct PagedKvAppendResult { // Copy of one sequence's bookkeeping state, as returned by sequence(). struct PagedKvSequenceSnapshot { uint32_t kv_seq_len = 0; + // Entries are physical block indices or PAGED_KV_COLD_BLOCK. Logical + // length and block-table length do not shrink when a block is paged out. 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. @@ -155,6 +178,26 @@ class PagedKvPool { PagedKvStatus owned_block_count(PagedKvSequenceHandle handle, uint32_t & out_count) const; + // Relinquish one materialized physical block while preserving its logical + // block-table position as PAGED_KV_COLD_BLOCK. The caller must first copy + // the complete block to host backing. On success, out_physical_block is + // the returned pool block and may be reused immediately. + PagedKvStatus page_out_block(PagedKvSequenceHandle handle, + uint32_t logical_block, + uint32_t & out_physical_block); + + // Allocate a physical block for one cold logical entry. The caller must + // restore its complete host-backed bytes before attention or append reads + // it. On failure, the cold mapping and output argument are unchanged. + PagedKvStatus page_in_block(PagedKvSequenceHandle handle, + uint32_t logical_block, + uint32_t & out_physical_block); + + // Number of materialized logical blocks that currently own physical + // storage. Reserved append capacity is deliberately excluded. + PagedKvStatus resident_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. @@ -182,6 +225,11 @@ class PagedKvPool { PagedKvStatus extend_block_table(SequenceState & sequence, uint32_t required_blocks); + // Allocate a physical block from a sequence's reservation first, then the + // global free list. Used only by append(), where consuming promised future + // capacity is correct. + uint32_t take_append_block(SequenceState & sequence); + // Move exactly `additional_blocks` globally free blocks into a sequence's // private reservation. Caller must preflight availability. void take_reserved_blocks(SequenceState & sequence, diff --git a/server/src/common/concurrency/paged_kv_residency.cpp b/server/src/common/concurrency/paged_kv_residency.cpp new file mode 100644 index 000000000..43c8443b0 --- /dev/null +++ b/server/src/common/concurrency/paged_kv_residency.cpp @@ -0,0 +1,929 @@ +#include "paged_kv_residency.h" + +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +PagedKvResidencyStatus from_pool_status(PagedKvStatus status) { + switch (status) { + case PagedKvStatus::Ok: + return PagedKvResidencyStatus::Ok; + case PagedKvStatus::StaleHandle: + return PagedKvResidencyStatus::StaleHandle; + case PagedKvStatus::BlocksExhausted: + return PagedKvResidencyStatus::PoolExhausted; + case PagedKvStatus::InvalidArgument: + case PagedKvStatus::LogicalBlockOutOfRange: + case PagedKvStatus::BlockNotResident: + case PagedKvStatus::BlockAlreadyResident: + return PagedKvResidencyStatus::InvalidArgument; + case PagedKvStatus::DuplicateRequest: + case PagedKvStatus::SequenceSlotsExhausted: + return PagedKvResidencyStatus::InconsistentPoolState; + } + return PagedKvResidencyStatus::InconsistentPoolState; +} + +} // namespace + +const char * paged_kv_residency_status_string( + PagedKvResidencyStatus status) { + switch (status) { + case PagedKvResidencyStatus::Ok: + return "ok"; + case PagedKvResidencyStatus::InvalidArgument: + return "invalid argument"; + case PagedKvResidencyStatus::SequenceNotRegistered: + return "sequence not registered"; + case PagedKvResidencyStatus::StaleHandle: + return "stale sequence handle"; + case PagedKvResidencyStatus::PoolExhausted: + return "physical pool exhausted"; + case PagedKvResidencyStatus::NoEvictableBlock: + return "no evictable block"; + case PagedKvResidencyStatus::HostAllocationFailed: + return "pinned host allocation failed"; + case PagedKvResidencyStatus::TransferFailed: + return "K/V transfer failed"; + case PagedKvResidencyStatus::HostCopyMissing: + return "cold block has no valid host copy"; + case PagedKvResidencyStatus::InconsistentPoolState: + return "residency state disagrees with paged pool"; + } + return "unknown paged K/V residency status"; +} + +PagedKvResidencyManager::PagedKvResidencyManager( + PagedKvPool & pool, PagedKvResidencyConfig config, + PagedKvResidencyTransferOps transfers) + : pool_(pool), config_(config), transfers_(std::move(transfers)) { + if (config_.block_bytes == 0 || !transfers_.allocate_pinned || + !transfers_.free_pinned || !transfers_.copy_out_async || + !transfers_.copy_in_async || !transfers_.synchronize) { + throw std::invalid_argument("invalid paged K/V residency callbacks"); + } + if (config_.resident_budget_blocks == 0) { + config_.resident_budget_blocks = pool_.physical_block_count(); + } + if (config_.resident_budget_blocks > pool_.physical_block_count()) { + throw std::invalid_argument("resident budget exceeds paged K/V pool"); + } + sequences_.resize(pool_.max_sequences()); +} + +PagedKvResidencyManager::~PagedKvResidencyManager() { + try { + reset(); + } catch (...) { + // Destructors must not propagate callback failures. Production pinned + // allocators/free functions are non-throwing; this catch protects test + // and plugin callbacks from terminating teardown. + } +} + +PagedKvResidencyStatus PagedKvResidencyManager::validate_registered( + PagedKvSequenceHandle handle) const { + if (transfer_barrier_failed_) { + return PagedKvResidencyStatus::TransferFailed; + } + if (handle.slot >= sequences_.size()) { + return PagedKvResidencyStatus::StaleHandle; + } + const SequenceState & state = sequences_[handle.slot]; + if (!state.active) { + return PagedKvResidencyStatus::SequenceNotRegistered; + } + if (state.generation != handle.generation) { + return PagedKvResidencyStatus::StaleHandle; + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::register_sequence( + PagedKvSequenceHandle handle) { + if (transfer_barrier_failed_) { + return PagedKvResidencyStatus::TransferFailed; + } + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) { + return from_pool_status(pool_status); + } + if (handle.slot >= sequences_.size()) { + return PagedKvResidencyStatus::StaleHandle; + } + if (std::find(snapshot.block_table.begin(), snapshot.block_table.end(), + PAGED_KV_COLD_BLOCK) != snapshot.block_table.end()) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + + SequenceState & state = sequences_[handle.slot]; + if (state.active && state.generation == handle.generation) { + return PagedKvResidencyStatus::Ok; + } + if (state.active) { + const auto synced = synchronize_before_read(); + if (synced != PagedKvResidencyStatus::Ok) return synced; + free_sequence_buffers(state); + } + state.active = true; + state.generation = handle.generation; + state.blocks.resize(snapshot.block_table.size()); + for (BlockState & block : state.blocks) block.last_use = ++clock_; + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::forget_sequence( + PagedKvSequenceHandle handle) { + // Teardown is also the recovery path after a failed copy-stream barrier. + // validate_registered() deliberately rejects ordinary operations while a + // transfer is quarantined, so validate the generation directly here and + // allow synchronize_before_read() to retry the barrier. + if (handle.slot >= sequences_.size()) { + return PagedKvResidencyStatus::StaleHandle; + } + const SequenceState & state = sequences_[handle.slot]; + if (!state.active) { + return PagedKvResidencyStatus::SequenceNotRegistered; + } + if (state.generation != handle.generation) { + return PagedKvResidencyStatus::StaleHandle; + } + const auto synced = synchronize_before_read(); + if (synced != PagedKvResidencyStatus::Ok) return synced; + free_sequence_buffers(sequences_[handle.slot]); + return PagedKvResidencyStatus::Ok; +} + +void PagedKvResidencyManager::reset() { + // A failed barrier should not let teardown free memory still referenced by + // an async transfer. Leak those buffers rather than creating a use-after- + // free; the process/backend is already unhealthy in this case. + if (synchronize_before_read() != PagedKvResidencyStatus::Ok) return; + for (SequenceState & state : sequences_) free_sequence_buffers(state); + stats_ = {}; + clock_ = 0; +} + +void PagedKvResidencyManager::free_sequence_buffers( + SequenceState & state) noexcept { + for (BlockState & block : state.blocks) { + if (!block.host) continue; + try { + transfers_.free_pinned(block.host); + } catch (...) { + // The callback contract is non-throwing. Continue freeing other + // pages if a third-party implementation violates it. + } + block.host = nullptr; + if (stats_.host_bytes >= config_.block_bytes) { + stats_.host_bytes -= config_.block_bytes; + } + } + state = {}; +} + +PagedKvResidencyStatus PagedKvResidencyManager::refresh_sequence( + PagedKvSequenceHandle handle, bool appended_rows, + uint32_t append_first, uint32_t append_count) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + + SequenceState & state = sequences_[handle.slot]; + if (snapshot.block_table.size() < state.blocks.size()) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + state.blocks.resize(snapshot.block_table.size()); + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (snapshot.block_table[logical] == PAGED_KV_COLD_BLOCK && + !state.blocks[logical].host_valid) { + return PagedKvResidencyStatus::HostCopyMissing; + } + } + + if (appended_rows && append_count > 0) { + const uint64_t last = static_cast(append_first) + + append_count - 1; + const uint32_t first_block = append_first / pool_.block_size(); + const uint32_t last_block = + static_cast(last / pool_.block_size()); + if (last_block >= state.blocks.size()) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + for (uint32_t logical = first_block; logical <= last_block; ++logical) { + // The target graph has not written the returned rows yet. Retain a + // restored partial page's old host image until commit, but make the + // physical page ineligible for eviction in every policy path. + state.blocks[logical].write_pending = true; + state.blocks[logical].last_use = ++clock_; + } + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::prepare_append( + PagedKvSequenceHandle handle, uint32_t token_count) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (token_count > std::numeric_limits::max() - + snapshot.kv_seq_len) { + return PagedKvResidencyStatus::InvalidArgument; + } + if (token_count == 0) return finish_transfers(PagedKvResidencyStatus::Ok); + + bool restore_partial_head = false; + uint32_t partial_head = 0; + if (snapshot.kv_seq_len % pool_.block_size() != 0) { + partial_head = snapshot.kv_seq_len / pool_.block_size(); + restore_partial_head = + snapshot.block_table[partial_head] == PAGED_KV_COLD_BLOCK; + } + + const uint32_t new_length = snapshot.kv_seq_len + token_count; + const uint32_t required_blocks = + 1 + (new_length - 1) / pool_.block_size(); + const uint32_t additional_blocks = required_blocks - + static_cast(snapshot.block_table.size()); + const uint32_t globally_needed = additional_blocks > + snapshot.reserved_block_count + ? additional_blocks - snapshot.reserved_block_count : 0; + // Reserve the partial-head restoration and every new physical page as one + // transaction. Restoring first and making room later can otherwise select + // the just-restored append head as the next eviction victim. + const auto room = make_room( + handle, globally_needed + (restore_partial_head ? 1u : 0u), + additional_blocks + (restore_partial_head ? 1u : 0u)); + if (room != PagedKvResidencyStatus::Ok) return finish_transfers(room); + if (restore_partial_head) { + const auto restored = restore_block_async( + handle, partial_head); + if (restored != PagedKvResidencyStatus::Ok) { + return finish_transfers(restored); + } + } + return finish_transfers(PagedKvResidencyStatus::Ok); +} + +PagedKvResidentAppendResult PagedKvResidencyManager::append( + PagedKvSequenceHandle handle, uint32_t token_count, + bool only_first_last_slots) { + PagedKvResidentAppendResult result; + result.status = prepare_append(handle, token_count); + if (result.status != PagedKvResidencyStatus::Ok) return result; + + result.pool_result = + pool_.append(handle, token_count, only_first_last_slots); + if (result.pool_result.status != PagedKvStatus::Ok) { + result.status = from_pool_status(result.pool_result.status); + return result; + } + result.status = observe_append(handle, result.pool_result); + return result; +} + +PagedKvResidencyStatus PagedKvResidencyManager::observe_append( + PagedKvSequenceHandle handle, const PagedKvAppendResult & result) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + if (result.status != PagedKvStatus::Ok) return from_pool_status(result.status); + + SequenceState & state = sequences_[handle.slot]; + for (const PagedKvBlockRemap & remap : result.remapped_cold_blocks) { + if (remap.logical_block >= state.blocks.size() || + !state.blocks[remap.logical_block].host_valid) { + return finish_transfers(PagedKvResidencyStatus::HostCopyMissing); + } + bool queued = false; + try { + queued = transfers_.copy_in_async( + handle, remap.logical_block, remap.physical_block, + state.blocks[remap.logical_block].host, config_.block_bytes); + } catch (...) { + queued = false; + } + if (!queued) return finish_transfers(PagedKvResidencyStatus::TransferFailed); + transfers_pending_ = true; + stats_.page_ins++; + stats_.moved_bytes += config_.block_bytes; + } + const auto synced = finish_transfers(PagedKvResidencyStatus::Ok); + if (synced != PagedKvResidencyStatus::Ok) return synced; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (result.token_count > snapshot.kv_seq_len) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + return refresh_sequence( + handle, /*appended_rows=*/true, + snapshot.kv_seq_len - result.token_count, result.token_count); +} + +PagedKvResidencyStatus PagedKvResidencyManager::commit_pending_writes( + PagedKvSequenceHandle handle) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + + // The caller supplies the target-compute -> pager dependency by + // synchronizing the target backend before this call. Any old host image is + // stale only now, after the device rows have actually been overwritten. + SequenceState & state = sequences_[handle.slot]; + for (BlockState & block : state.blocks) { + if (!block.write_pending) continue; + block.write_pending = false; + block.host_valid = false; + block.last_use = ++clock_; + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::finish_transfers( + PagedKvResidencyStatus status) { + const auto synced = synchronize_before_read(); + return synced == PagedKvResidencyStatus::Ok ? status : synced; +} + +PagedKvResidencyStatus PagedKvResidencyManager::synchronize_before_read() { + if (!transfers_pending_) { + return transfer_barrier_failed_ + ? PagedKvResidencyStatus::TransferFailed + : PagedKvResidencyStatus::Ok; + } + bool ok = false; + try { + ok = transfers_.synchronize(); + } catch (...) { + ok = false; + } + if (!ok) { + // A failed barrier does not prove that the stream stopped using its + // source and destination blocks. Keep every transfer pending and every + // H2D destination mapped (therefore quarantined from the free list), + // and reject all manager operations until an explicit retry confirms + // that the stream has drained. + transfer_barrier_failed_ = true; + return PagedKvResidencyStatus::TransferFailed; + } + + transfers_pending_ = false; + transfer_barrier_failed_ = false; + + PagedKvResidencyStatus result = PagedKvResidencyStatus::Ok; + for (const PendingTransfer & transfer : pending_page_outs_) { + if (validate_registered(transfer.handle) != + PagedKvResidencyStatus::Ok) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + BlockState & block = + sequences_[transfer.handle.slot].blocks[transfer.logical_block]; + block.page_out_pending = false; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(transfer.handle, snapshot) != PagedKvStatus::Ok || + transfer.logical_block >= snapshot.block_table.size() || + snapshot.block_table[transfer.logical_block] != + transfer.physical_block) { + block.host_valid = false; + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + uint32_t released = PAGED_KV_COLD_BLOCK; + if (pool_.page_out_block( + transfer.handle, transfer.logical_block, released) != + PagedKvStatus::Ok || released != transfer.physical_block) { + block.host_valid = false; + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + block.host_valid = true; + stats_.page_outs++; + stats_.moved_bytes += config_.block_bytes; + } + for (const PendingTransfer & transfer : pending_page_ins_) { + if (validate_registered(transfer.handle) != + PagedKvResidencyStatus::Ok) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + BlockState & block = + sequences_[transfer.handle.slot].blocks[transfer.logical_block]; + block.page_in_pending = false; + stats_.page_ins++; + stats_.moved_bytes += config_.block_bytes; + } + pending_page_outs_.clear(); + pending_page_ins_.clear(); + if (result != PagedKvResidencyStatus::Ok) return result; + return PagedKvResidencyStatus::Ok; +} + +bool PagedKvResidencyManager::is_protected( + PagedKvSequenceHandle handle, uint32_t logical_block) const { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(handle, snapshot) != PagedKvStatus::Ok || + logical_block >= snapshot.block_table.size()) { + return true; + } + if (logical_block < config_.sink_blocks) return true; + const uint32_t tail_begin = snapshot.block_table.size() > config_.tail_blocks + ? static_cast(snapshot.block_table.size()) - + config_.tail_blocks + : 0; + return logical_block >= tail_begin; +} + +uint32_t PagedKvResidencyManager::sequence_resident_count( + PagedKvSequenceHandle handle) const { + uint32_t count = 0; + return pool_.resident_block_count(handle, count) == PagedKvStatus::Ok + ? count : 0; +} + +uint32_t PagedKvResidencyManager::sequence_pending_page_out_count( + PagedKvSequenceHandle handle) const { + if (handle.slot >= sequences_.size()) return 0; + const SequenceState & state = sequences_[handle.slot]; + if (!state.active || state.generation != handle.generation) return 0; + return static_cast(std::count_if( + state.blocks.begin(), state.blocks.end(), + [](const BlockState & block) { return block.page_out_pending; })); +} + +uint32_t PagedKvResidencyManager::sequence_protected_count( + PagedKvSequenceHandle handle) const { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(handle, snapshot) != PagedKvStatus::Ok) return 0; + uint32_t count = 0; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (snapshot.block_table[logical] != PAGED_KV_COLD_BLOCK && + is_protected(handle, logical)) { + ++count; + } + } + return count; +} + +uint32_t PagedKvResidencyManager::total_resident_count() const { + uint32_t total = 0; + for (uint32_t slot = 0; slot < sequences_.size(); ++slot) { + const SequenceState & state = sequences_[slot]; + if (!state.active) continue; + total += sequence_resident_count({slot, state.generation}); + } + return total; +} + +uint32_t PagedKvResidencyManager::quota_for_slot(uint32_t slot) const { + uint32_t active = 0; + uint32_t rank = 0; + for (uint32_t i = 0; i < sequences_.size(); ++i) { + if (!sequences_[i].active) continue; + if (i < slot) ++rank; + ++active; + } + if (active == 0 || slot >= sequences_.size() || + !sequences_[slot].active) { + return 0; + } + const uint32_t base = config_.resident_budget_blocks / active; + const uint32_t remainder = config_.resident_budget_blocks % active; + const uint32_t fair = base + (rank < remainder ? 1u : 0u); + const PagedKvSequenceHandle handle{slot, sequences_[slot].generation}; + return std::max(fair, sequence_protected_count(handle)); +} + +uint32_t PagedKvResidencyManager::fair_quota( + PagedKvSequenceHandle handle) const { + return validate_registered(handle) == PagedKvResidencyStatus::Ok + ? quota_for_slot(handle.slot) : 0; +} + +PagedKvResidencyManager::Victim PagedKvResidencyManager::choose_victim( + PagedKvSequenceHandle requester) const { + Victim best; + for (uint32_t slot = 0; slot < sequences_.size(); ++slot) { + const SequenceState & state = sequences_[slot]; + if (!state.active) continue; + const PagedKvSequenceHandle owner{slot, state.generation}; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(owner, snapshot) != PagedKvStatus::Ok) continue; + const uint32_t resident = sequence_resident_count(owner); + const uint32_t pending_page_outs = + sequence_pending_page_out_count(owner); + const uint32_t resident_after_pending = + resident > pending_page_outs ? resident - pending_page_outs : 0; + const uint32_t quota = quota_for_slot(slot); + + int class_rank = 2; + if (resident_after_pending > quota) class_rank = 0; + else if (slot == requester.slot && + state.generation == requester.generation) class_rank = 1; + + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (snapshot.block_table[logical] == PAGED_KV_COLD_BLOCK || + state.blocks[logical].write_pending || + state.blocks[logical].page_out_pending || + state.blocks[logical].page_in_pending || + state.blocks[logical].reservation_pending || + is_protected(owner, logical)) { + continue; + } + const BlockState & block = state.blocks[logical]; + // An unscored page is the first eviction candidate once relevance + // scoring is active for only part of a sequence. + const float score = block.score_valid + ? block.score : -std::numeric_limits::infinity(); + const bool better = !best.found || + class_rank < best.class_rank || + (class_rank == best.class_rank && score < best.score) || + (class_rank == best.class_rank && score == best.score && + block.last_use < best.last_use) || + (class_rank == best.class_rank && score == best.score && + block.last_use == best.last_use && + (slot < best.handle.slot || + (slot == best.handle.slot && + logical < best.logical_block))); + if (better) { + best.found = true; + best.handle = owner; + best.logical_block = logical; + best.class_rank = class_rank; + best.score = score; + best.last_use = block.last_use; + } + } + } + return best; +} + +PagedKvResidencyStatus PagedKvResidencyManager::make_room( + PagedKvSequenceHandle requester, uint32_t pool_blocks_needed, + uint32_t future_resident_blocks) { + const uint32_t free = pool_.free_block_count(); + const uint32_t need_for_pool = pool_blocks_needed > free + ? pool_blocks_needed - free : 0; + const uint32_t resident = total_resident_count(); + const uint64_t projected = + static_cast(resident) + future_resident_blocks; + const uint32_t need_for_budget = + projected > config_.resident_budget_blocks + ? static_cast(projected - config_.resident_budget_blocks) + : 0; + const uint32_t evictions = std::max(need_for_pool, need_for_budget); + for (uint32_t i = 0; i < evictions; ++i) { + const Victim victim = choose_victim(requester); + if (!victim.found) { + return finish_transfers(PagedKvResidencyStatus::NoEvictableBlock); + } + const auto status = evict_block_async( + victim.handle, victim.logical_block, + /*allow_protected=*/false); + if (status != PagedKvResidencyStatus::Ok) return finish_transfers(status); + } + // A page is recyclable only after the complete D2H batch succeeds. + return finish_transfers(PagedKvResidencyStatus::Ok); +} + +PagedKvResidencyStatus PagedKvResidencyManager::evict_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block, + bool allow_protected) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (logical_block >= snapshot.block_table.size()) { + return PagedKvResidencyStatus::InvalidArgument; + } + if (snapshot.block_table[logical_block] == PAGED_KV_COLD_BLOCK) { + return PagedKvResidencyStatus::InvalidArgument; + } + + BlockState & block = sequences_[handle.slot].blocks[logical_block]; + if (block.write_pending || block.page_out_pending || + block.page_in_pending || + (!allow_protected && is_protected(handle, logical_block))) { + return PagedKvResidencyStatus::NoEvictableBlock; + } + if (!block.host) { + try { + block.host = transfers_.allocate_pinned(config_.block_bytes); + } catch (...) { + block.host = nullptr; + } + if (!block.host) return PagedKvResidencyStatus::HostAllocationFailed; + stats_.host_bytes += config_.block_bytes; + } + + const uint32_t physical = snapshot.block_table[logical_block]; + bool queued = false; + try { + queued = transfers_.copy_out_async( + handle, logical_block, physical, block.host, + config_.block_bytes); + } catch (...) { + queued = false; + } + if (!queued) return PagedKvResidencyStatus::TransferFailed; + transfers_pending_ = true; + block.host_valid = false; + block.page_out_pending = true; + pending_page_outs_.push_back({handle, logical_block, physical}); + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::restore_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + if (logical_block >= snapshot.block_table.size()) { + return PagedKvResidencyStatus::InvalidArgument; + } + if (snapshot.block_table[logical_block] != PAGED_KV_COLD_BLOCK) { + sequences_[handle.slot].blocks[logical_block].last_use = ++clock_; + return PagedKvResidencyStatus::Ok; + } + BlockState & block = sequences_[handle.slot].blocks[logical_block]; + if (block.page_out_pending || block.page_in_pending) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + if (!block.host || !block.host_valid) { + return PagedKvResidencyStatus::HostCopyMissing; + } + + uint32_t physical = PAGED_KV_COLD_BLOCK; + const PagedKvStatus page_status = + pool_.page_in_block(handle, logical_block, physical); + if (page_status != PagedKvStatus::Ok) return from_pool_status(page_status); + + bool queued = false; + try { + queued = transfers_.copy_in_async( + handle, logical_block, physical, block.host, + config_.block_bytes); + } catch (...) { + queued = false; + } + if (!queued) { + uint32_t released = PAGED_KV_COLD_BLOCK; + if (pool_.page_out_block(handle, logical_block, released) != + PagedKvStatus::Ok) { + return PagedKvResidencyStatus::InconsistentPoolState; + } + return PagedKvResidencyStatus::TransferFailed; + } + transfers_pending_ = true; + block.page_in_pending = true; + pending_page_ins_.push_back({handle, logical_block, physical}); + block.last_use = ++clock_; + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::ensure_resident( + PagedKvSequenceHandle handle, + const std::vector & logical_blocks) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + std::vector cold; + std::vector seen(snapshot.block_table.size(), 0); + // Validate the complete request before reserving any member. Otherwise a + // later bad index leaves earlier blocks permanently ineligible as victims. + if (std::any_of(logical_blocks.begin(), logical_blocks.end(), + [&](uint32_t logical) { + return logical >= snapshot.block_table.size(); + })) { + return PagedKvResidencyStatus::InvalidArgument; + } + for (uint32_t logical : logical_blocks) { + if (seen[logical]) continue; + seen[logical] = 1; + sequences_[handle.slot].blocks[logical].reservation_pending = true; + if (snapshot.block_table[logical] == PAGED_KV_COLD_BLOCK) { + cold.push_back(logical); + } else { + sequences_[handle.slot].blocks[logical].last_use = ++clock_; + } + } + const auto clear_reservations = [&] { + for (uint32_t logical = 0; logical < seen.size(); ++logical) { + if (seen[logical]) { + sequences_[handle.slot].blocks[logical].reservation_pending = false; + } + } + }; + const auto room = make_room( + handle, static_cast(cold.size()), + static_cast(cold.size())); + if (room != PagedKvResidencyStatus::Ok) { + clear_reservations(); + return room; + } + + // Restore the requested set from one reserved capacity pool. No member can + // become the victim of a later member in this same operation. + for (uint32_t logical : cold) { + const auto status = restore_block_async( + handle, logical); + if (status != PagedKvResidencyStatus::Ok) { + const auto finished = finish_transfers(status); + clear_reservations(); + return finished; + } + } + const auto finished = finish_transfers(PagedKvResidencyStatus::Ok); + clear_reservations(); + return finished; +} + +PagedKvResidencyStatus PagedKvResidencyManager::evict_block( + PagedKvSequenceHandle handle, uint32_t logical_block, + bool allow_protected) { + return finish_transfers( + evict_block_async(handle, logical_block, allow_protected)); +} + +PagedKvResidencyStatus PagedKvResidencyManager::touch( + PagedKvSequenceHandle handle, uint32_t logical_block) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + if (logical_block >= sequences_[handle.slot].blocks.size()) { + return PagedKvResidencyStatus::InvalidArgument; + } + sequences_[handle.slot].blocks[logical_block].last_use = ++clock_; + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::set_scores( + PagedKvSequenceHandle handle, const std::vector & scores) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + SequenceState & state = sequences_[handle.slot]; + if (scores.empty()) { + for (BlockState & block : state.blocks) block.score_valid = false; + return PagedKvResidencyStatus::Ok; + } + if (scores.size() != state.blocks.size() || + std::any_of(scores.begin(), scores.end(), + [](float score) { return !std::isfinite(score); })) { + return PagedKvResidencyStatus::InvalidArgument; + } + for (uint32_t i = 0; i < scores.size(); ++i) { + state.blocks[i].score = scores[i]; + state.blocks[i].score_valid = true; + } + return PagedKvResidencyStatus::Ok; +} + +PagedKvResidencyStatus PagedKvResidencyManager::reselect( + PagedKvSequenceHandle handle) { + const auto registered = validate_registered(handle); + if (registered != PagedKvResidencyStatus::Ok) return registered; + stats_.reselects++; + + PagedKvSequenceSnapshot snapshot; + const PagedKvStatus pool_status = pool_.sequence(handle, snapshot); + if (pool_status != PagedKvStatus::Ok) return from_pool_status(pool_status); + SequenceState & state = sequences_[handle.slot]; + + std::vector wanted(snapshot.block_table.size(), 0); + uint32_t wanted_count = 0; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (state.blocks[logical].write_pending || + is_protected(handle, logical)) { + wanted[logical] = 1; + ++wanted_count; + } + } + const uint32_t target = std::max( + wanted_count, + std::min(quota_for_slot(handle.slot), + snapshot.block_table.size())); + + std::vector candidates; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (!wanted[logical]) candidates.push_back(logical); + } + std::sort(candidates.begin(), candidates.end(), + [&](uint32_t a, uint32_t b) { + const BlockState & lhs = state.blocks[a]; + const BlockState & rhs = state.blocks[b]; + if (lhs.score_valid != rhs.score_valid) return lhs.score_valid; + if (lhs.score_valid && lhs.score != rhs.score) return lhs.score > rhs.score; + if (lhs.last_use != rhs.last_use) return lhs.last_use > rhs.last_use; + return a < b; + }); + for (uint32_t logical : candidates) { + if (wanted_count >= target) break; + wanted[logical] = 1; + ++wanted_count; + } + + for (uint32_t logical = 0; logical < wanted.size(); ++logical) { + if (wanted[logical]) state.blocks[logical].reservation_pending = true; + } + const auto clear_reservations = [&] { + for (uint32_t logical = 0; logical < wanted.size(); ++logical) { + if (wanted[logical]) state.blocks[logical].reservation_pending = false; + } + }; + + // Out first, making one batch of free physical pages before recalls. + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (!wanted[logical] && + snapshot.block_table[logical] != PAGED_KV_COLD_BLOCK) { + const auto status = evict_block_async( + handle, logical, /*allow_protected=*/false); + if (status != PagedKvResidencyStatus::Ok) { + const auto finished = finish_transfers(status); + clear_reservations(); + return finished; + } + } + } + auto finished = finish_transfers(PagedKvResidencyStatus::Ok); + if (finished != PagedKvResidencyStatus::Ok) { + clear_reservations(); + return finished; + } + + std::vector cold_wanted; + for (uint32_t logical = 0; logical < snapshot.block_table.size(); ++logical) { + if (wanted[logical] && !is_resident(handle, logical)) { + cold_wanted.push_back(logical); + } + } + const auto room = make_room( + handle, static_cast(cold_wanted.size()), + static_cast(cold_wanted.size())); + if (room != PagedKvResidencyStatus::Ok) { + clear_reservations(); + return room; + } + for (uint32_t logical : cold_wanted) { + const auto status = restore_block_async( + handle, logical); + if (status != PagedKvResidencyStatus::Ok) { + finished = finish_transfers(status); + clear_reservations(); + return finished; + } + } + finished = finish_transfers(PagedKvResidencyStatus::Ok); + clear_reservations(); + return finished; +} + +PagedKvResidencyStatus PagedKvResidencyManager::rebalance() { + while (true) { + // The dummy requester prevents any real sequence from getting the + // requester-swap class. A class-0 victim is necessarily over quota. + const PagedKvSequenceHandle none{ + std::numeric_limits::max(), 0}; + const Victim victim = choose_victim(none); + if (!victim.found || victim.class_rank != 0) break; + const auto status = evict_block_async( + victim.handle, victim.logical_block, + /*allow_protected=*/false); + if (status != PagedKvResidencyStatus::Ok) { + return finish_transfers(status); + } + } + return finish_transfers(PagedKvResidencyStatus::Ok); +} + +bool PagedKvResidencyManager::is_resident( + PagedKvSequenceHandle handle, uint32_t logical_block) const { + if (validate_registered(handle) != PagedKvResidencyStatus::Ok) return false; + PagedKvSequenceSnapshot snapshot; + return pool_.sequence(handle, snapshot) == PagedKvStatus::Ok && + logical_block < snapshot.block_table.size() && + snapshot.block_table[logical_block] != PAGED_KV_COLD_BLOCK; +} + +PagedKvResidencyStats PagedKvResidencyManager::stats() const { + PagedKvResidencyStats out = stats_; + out.resident_blocks = total_resident_count(); + return out; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/paged_kv_residency.h b/server/src/common/concurrency/paged_kv_residency.h new file mode 100644 index 000000000..dc6f67cf5 --- /dev/null +++ b/server/src/common/concurrency/paged_kv_residency.h @@ -0,0 +1,260 @@ +// Multi-sequence host residency for PagedKvPool full-attention K/V blocks. +// +// The policy and bookkeeping are backend-neutral. The Qwen engine supplies a +// pinned allocator plus async full-block D2H/H2D callbacks that capture its +// dedicated copy stream and know how to concatenate every full-attention K/V +// tensor into one host block. The manager batches transfers and synchronizes +// once before returning any operation that permits device K/V to be read or a +// recycled physical block to be written. + +#pragma once + +#include "paged_kv_pool.h" + +#include +#include +#include +#include + +namespace dflash::common { + +enum class PagedKvResidencyStatus : uint8_t { + Ok = 0, + InvalidArgument, + SequenceNotRegistered, + StaleHandle, + PoolExhausted, + NoEvictableBlock, + HostAllocationFailed, + TransferFailed, + HostCopyMissing, + InconsistentPoolState, +}; + +const char * paged_kv_residency_status_string(PagedKvResidencyStatus status); + +struct PagedKvResidencyConfig { + // Bytes copied for one physical pool block across all full-attention K/V + // tensors. Recurrent/DeltaNet state is intentionally outside this pager. + size_t block_bytes = 0; + + // Hard bound for materialized logical blocks. Zero uses the complete + // physical PagedKvPool. Reserved-but-unappended pool blocks still consume + // physical capacity and may temporarily make the effective bound smaller. + uint32_t resident_budget_blocks = 0; + + // Attention sinks and the trailing local window are never automatic + // eviction victims. Explicit evict_block(..., allow_protected=true) is + // available for teardown/tests only. + uint32_t sink_blocks = 1; + uint32_t tail_blocks = 4; +}; + +struct PagedKvResidencyTransferOps { + // Production implementations should use cudaMallocHost/hipHostMalloc (or + // the equivalent backend pinned allocator). Every allocation is exactly + // config.block_bytes and lives until forget_sequence()/reset(). + std::function allocate_pinned; + std::function free_pinned; + + // Queue a complete physical-block copy on one dedicated copy stream. + // handle/logical_block are supplied for diagnostics only; physical_block + // identifies the source/destination in the Qwen paged K/V tensors. + std::function copy_out_async; + std::function copy_in_async; + + // Synchronize that copy stream. The manager calls it once after a batch + // and before attention reads or append writes can observe remapped blocks. + std::function synchronize; +}; + +struct PagedKvResidencyStats { + uint64_t page_ins = 0; + uint64_t page_outs = 0; + uint64_t resident_blocks = 0; + uint64_t reselects = 0; + uint64_t host_bytes = 0; + uint64_t moved_bytes = 0; +}; + +struct PagedKvResidentAppendResult { + PagedKvResidencyStatus status = PagedKvResidencyStatus::Ok; + PagedKvAppendResult pool_result; + + explicit operator bool() const { + return status == PagedKvResidencyStatus::Ok && + pool_result.status == PagedKvStatus::Ok; + } +}; + +class PagedKvResidencyManager { +public: + // Throws std::invalid_argument for an empty transfer callback, zero block + // bytes, or a resident budget larger than the physical pool. + PagedKvResidencyManager(PagedKvPool & pool, + PagedKvResidencyConfig config, + PagedKvResidencyTransferOps transfers); + ~PagedKvResidencyManager(); + + PagedKvResidencyManager(const PagedKvResidencyManager &) = delete; + PagedKvResidencyManager & operator=(const PagedKvResidencyManager &) = delete; + + // Call immediately after PagedKvPool::acquire[_reserved](). Registration + // rejects pre-existing cold entries because their host backing is unknown. + PagedKvResidencyStatus register_sequence(PagedKvSequenceHandle handle); + + // Free this sequence's pinned backing. Call before pool.release, and do + // not release the pool handle unless this returns Ok: a failed transfer + // barrier keeps its physical pages and pinned buffers quarantined until + // forget_sequence is retried. The handle must still match the manager's + // registered generation. + PagedKvResidencyStatus forget_sequence(PagedKvSequenceHandle handle); + void reset(); + + // Recommended append integration: restores a cold partial append head, + // fairly evicts enough blocks, synchronizes the copy stream, appends in + // PagedKvPool, and marks every touched block pending until the engine has + // finished the target graph that writes the returned physical rows. + PagedKvResidentAppendResult append(PagedKvSequenceHandle handle, + uint32_t token_count, + bool only_first_last_slots = false); + + // Split integration for callers whose slot manager owns pool.append(): + // prepare_append(handle, n); pool.append(handle, n); observe_append(...) + // prepare_append always synchronizes before returning Ok. observe_append + // must be called before the next residency operation. + PagedKvResidencyStatus prepare_append(PagedKvSequenceHandle handle, + uint32_t token_count); + PagedKvResidencyStatus observe_append( + PagedKvSequenceHandle handle, const PagedKvAppendResult & result); + + // Clear all blocks staged by append()/observe_append() for this sequence. + // Call only after synchronizing the target compute that wrote every returned + // physical row. Pending blocks cannot be evicted or deselected, preventing a + // later slot staged in the same packed step from recycling an unwritten page. + PagedKvResidencyStatus commit_pending_writes( + PagedKvSequenceHandle handle); + + // Restore the requested host-backed logical blocks as one transfer batch. + // The operation synchronizes before returning Ok, so attention may read + // the returned pool block table immediately. + PagedKvResidencyStatus ensure_resident( + PagedKvSequenceHandle handle, + const std::vector & logical_blocks); + + // Explicit single-block eviction. Automatic policy never evicts sink/tail + // blocks; protected eviction requires an explicit opt-in. Pending target + // writes are never evictable, including with allow_protected=true. + PagedKvResidencyStatus evict_block(PagedKvSequenceHandle handle, + uint32_t logical_block, + bool allow_protected = false); + + // Update recency after a block participates in attention. + PagedKvResidencyStatus touch(PagedKvSequenceHandle handle, + uint32_t logical_block); + + // Optional relevance array, one score per materialized logical block. + // Higher values are retained. Without scores, reselection uses LRU. + PagedKvResidencyStatus set_scores(PagedKvSequenceHandle handle, + const std::vector & scores); + PagedKvResidencyStatus reselect(PagedKvSequenceHandle handle); + + // Evict unprotected blocks above each active sequence's deterministic + // fair share. Spare capacity remains borrowable; it is reclaimed first + // from borrowers when another sequence needs a block. + PagedKvResidencyStatus rebalance(); + + // Explicit barrier for engine paths that directly issue operations and + // then read the paged K/V tensors. Normally append/ensure/reselect/evict + // already provide this barrier. + PagedKvResidencyStatus synchronize_before_read(); + + uint32_t fair_quota(PagedKvSequenceHandle handle) const; + bool is_resident(PagedKvSequenceHandle handle, + uint32_t logical_block) const; + PagedKvResidencyStats stats() const; + +private: + struct BlockState { + void * host = nullptr; + bool host_valid = false; + bool write_pending = false; + bool page_out_pending = false; + bool page_in_pending = false; + bool reservation_pending = false; + bool score_valid = false; + float score = 0.0f; + uint64_t last_use = 0; + }; + + struct SequenceState { + bool active = false; + uint64_t generation = 0; + std::vector blocks; + }; + + struct Victim { + bool found = false; + PagedKvSequenceHandle handle; + uint32_t logical_block = 0; + int class_rank = 0; + float score = 0.0f; + uint64_t last_use = 0; + }; + + struct PendingTransfer { + PagedKvSequenceHandle handle; + uint32_t logical_block = 0; + uint32_t physical_block = PAGED_KV_COLD_BLOCK; + }; + + PagedKvResidencyStatus validate_registered( + PagedKvSequenceHandle handle) const; + PagedKvResidencyStatus refresh_sequence( + PagedKvSequenceHandle handle, bool appended_rows = false, + uint32_t append_first = 0, uint32_t append_count = 0); + PagedKvResidencyStatus finish_transfers( + PagedKvResidencyStatus status); + + bool is_protected(PagedKvSequenceHandle handle, + uint32_t logical_block) const; + uint32_t sequence_resident_count(PagedKvSequenceHandle handle) const; + uint32_t sequence_pending_page_out_count( + PagedKvSequenceHandle handle) const; + uint32_t sequence_protected_count(PagedKvSequenceHandle handle) const; + uint32_t total_resident_count() const; + uint32_t quota_for_slot(uint32_t slot) const; + + Victim choose_victim(PagedKvSequenceHandle requester) const; + PagedKvResidencyStatus make_room(PagedKvSequenceHandle requester, + uint32_t pool_blocks_needed, + uint32_t future_resident_blocks); + PagedKvResidencyStatus evict_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block, + bool allow_protected); + PagedKvResidencyStatus restore_block_async( + PagedKvSequenceHandle handle, uint32_t logical_block); + + void free_sequence_buffers(SequenceState & state) noexcept; + + PagedKvPool & pool_; + PagedKvResidencyConfig config_; + PagedKvResidencyTransferOps transfers_; + std::vector sequences_; + PagedKvResidencyStats stats_; + uint64_t clock_ = 0; + bool transfers_pending_ = false; + bool transfer_barrier_failed_ = false; + std::vector pending_page_outs_; + std::vector pending_page_ins_; +}; + +} // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer.cpp b/server/src/common/concurrency/qwen_paged_kv_transfer.cpp new file mode 100644 index 000000000..111e6a36f --- /dev/null +++ b/server/src/common/concurrency/qwen_paged_kv_transfer.cpp @@ -0,0 +1,405 @@ +#include "qwen_paged_kv_transfer.h" + +#include "common/gpu_runtime_compat.h" +#include "internal.h" + +#include "ggml.h" + +#include +#include +#include + +namespace dflash::common { +namespace { + +void set_error(std::string * error, const std::string & message) { + if (error) *error = message; +} + +bool runtime_pointer_is_device(const void * pointer, int device) { + cudaPointerAttributes attributes{}; + const cudaError_t status = cudaPointerGetAttributes(&attributes, pointer); + if (status != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + return attributes.type == hipMemoryTypeDevice && + attributes.device == device; +#else + return attributes.type == cudaMemoryTypeDevice && + attributes.device == device; +#endif +} + +bool tensor_is_device_backed(const ggml_tensor * tensor, int device, + std::string * error) { + if (!tensor || !tensor->buffer || !tensor->data) { + set_error(error, "paged K/V tensor is null or unallocated"); + return false; + } + const ggml_backend_buffer_type_t buft = + ggml_backend_buffer_get_type(tensor->buffer); + if (!buft || ggml_backend_buft_is_meta(buft)) { + set_error(error, + "meta/tensor-parallel paged K/V buffers are unsupported"); + return false; + } + if (ggml_backend_buft_is_host(buft)) { + set_error(error, "host-backed paged K/V buffers are unsupported"); + return false; + } + const ggml_backend_dev_t tensor_device = + ggml_backend_buft_get_device(buft); + if (!tensor_device) { + set_error(error, "paged K/V buffer has no backend device"); + return false; + } + const enum ggml_backend_dev_type tensor_device_type = + ggml_backend_dev_type(tensor_device); + if (tensor_device_type != GGML_BACKEND_DEVICE_TYPE_GPU && + tensor_device_type != GGML_BACKEND_DEVICE_TYPE_IGPU) { + set_error(error, "paged K/V buffer is not GPU device memory"); + return false; + } + if (!runtime_pointer_is_device(tensor->data, device)) { + set_error(error, + "paged K/V tensor is not on the requested HIP/CUDA device"); + return false; + } + return true; +} + +} // namespace + +struct QwenPagedKvResidencyTransfer::State { + struct TensorCopy { + uint8_t * device_data = nullptr; + QwenPagedKvTensorLayout layout; + size_t host_offset = 0; + size_t host_head_bytes = 0; + }; + + int device = -1; + uint32_t block_size = 0; + size_t block_bytes = 0; + uint32_t physical_block_count = 0; + cudaStream_t stream = nullptr; + bool stream_may_reference_host = false; + std::vector tensors; + std::vector pinned_allocations; + + ~State() { + if (device >= 0 && cudaSetDevice(device) == cudaSuccess) { + if (stream && cudaStreamSynchronize(stream) != cudaSuccess) { + // Match PagedKvResidencyManager's fail-safe teardown rule: + // leak stream/backing rather than free host memory that a + // failed runtime may still reference. + pinned_allocations.clear(); + stream = nullptr; + return; + } + for (void * pointer : pinned_allocations) { + if (pointer) (void)cudaFreeHost(pointer); + } + pinned_allocations.clear(); + if (stream) (void)cudaStreamDestroy(stream); + } + stream = nullptr; + } + + bool select_device() const { + return device >= 0 && cudaSetDevice(device) == cudaSuccess; + } + + void * allocate_pinned(size_t bytes) { + if (bytes != block_bytes || !select_device()) return nullptr; + void * pointer = nullptr; + if (cudaMallocHost(&pointer, bytes) != cudaSuccess || !pointer) { + (void)cudaGetLastError(); + return nullptr; + } + try { + pinned_allocations.push_back(pointer); + } catch (...) { + (void)cudaFreeHost(pointer); + return nullptr; + } + return pointer; + } + + void free_pinned(void * pointer) { + if (!pointer) return; + const auto found = std::find( + pinned_allocations.begin(), pinned_allocations.end(), pointer); + if (found == pinned_allocations.end()) return; + if (stream_may_reference_host) { + if (!stream || !select_device() || + cudaStreamSynchronize(stream) != cudaSuccess) { + return; + } + stream_may_reference_host = false; + } + if (select_device() && cudaFreeHost(pointer) == cudaSuccess) { + pinned_allocations.erase(found); + } + } + + bool synchronize() { + if (!stream || !select_device()) return false; + if (cudaStreamSynchronize(stream) != cudaSuccess) { + (void)cudaGetLastError(); + stream_may_reference_host = true; + return false; + } + stream_may_reference_host = false; + return true; + } + + bool queue_copy(uint32_t physical_block, void * host, size_t bytes, + bool copy_out) { + if (!host || bytes != block_bytes || !stream || + physical_block >= physical_block_count || !select_device()) { + return false; + } + + const size_t first_row = + static_cast(physical_block) * block_size; + uint8_t * const host_bytes = static_cast(host); + for (const TensorCopy & tensor : tensors) { + uint8_t * const device_block = tensor.device_data + + first_row * tensor.layout.row_stride; + uint8_t * const host_tensor = + host_bytes + tensor.host_offset; + + cudaError_t status = cudaSuccess; + if (tensor.layout.row_stride == tensor.layout.row_bytes) { + // Common Qwen cache layout: block rows are contiguous inside + // each head plane, so one 2D copy covers every head. + status = copy_out + ? cudaMemcpy2DAsync( + host_tensor, tensor.host_head_bytes, + device_block, tensor.layout.head_stride, + tensor.host_head_bytes, tensor.layout.heads, + cudaMemcpyDeviceToHost, stream) + : cudaMemcpy2DAsync( + device_block, tensor.layout.head_stride, + host_tensor, tensor.host_head_bytes, + tensor.host_head_bytes, tensor.layout.heads, + cudaMemcpyHostToDevice, stream); + } else { + // Preserve uncommon row padding with one pitched async copy + // per head. Host images remain tightly packed payload bytes. + for (size_t head = 0; + head < static_cast(tensor.layout.heads); + ++head) { + uint8_t * const device_head = + device_block + head * tensor.layout.head_stride; + uint8_t * const host_head = + host_tensor + head * tensor.host_head_bytes; + status = copy_out + ? cudaMemcpy2DAsync( + host_head, tensor.layout.row_bytes, + device_head, tensor.layout.row_stride, + tensor.layout.row_bytes, block_size, + cudaMemcpyDeviceToHost, stream) + : cudaMemcpy2DAsync( + device_head, tensor.layout.row_stride, + host_head, tensor.layout.row_bytes, + tensor.layout.row_bytes, block_size, + cudaMemcpyHostToDevice, stream); + if (status != cudaSuccess) break; + } + } + if (status != cudaSuccess) { + // A callback that returns false is not marked pending by the + // residency manager. Drain any prefix already queued here so + // its pinned buffer can still be released safely. + (void)cudaGetLastError(); + if (cudaStreamSynchronize(stream) != cudaSuccess) { + stream_may_reference_host = true; + } + return false; + } + } + return true; + } +}; + +QwenPagedKvResidencyTransfer::QwenPagedKvResidencyTransfer( + std::shared_ptr state) + : state_(std::move(state)) {} + +QwenPagedKvResidencyTransfer::~QwenPagedKvResidencyTransfer() = default; + +std::unique_ptr +QwenPagedKvResidencyTransfer::create( + const TargetCache & cache, + ggml_backend_t backend, + int device, + uint32_t block_size, + std::string * error) { + if (error) error->clear(); + if (!backend || device < 0 || block_size == 0) { + set_error(error, "invalid paged K/V transfer backend/device/block size"); + return nullptr; + } + const ggml_backend_dev_t backend_device = ggml_backend_get_device(backend); + if (!backend_device) { + set_error(error, "paged K/V transfer backend has no device"); + return nullptr; + } + const enum ggml_backend_dev_type backend_type = + ggml_backend_dev_type(backend_device); + if (backend_type == GGML_BACKEND_DEVICE_TYPE_META || + ggml_backend_buft_is_meta( + ggml_backend_get_default_buffer_type(backend))) { + set_error(error, + "meta/tensor-parallel paged K/V transfer is unsupported"); + return nullptr; + } + if (backend_type != GGML_BACKEND_DEVICE_TYPE_GPU && + backend_type != GGML_BACKEND_DEVICE_TYPE_IGPU) { + set_error(error, "paged K/V transfer requires a GPU backend"); + return nullptr; + } + if (cache.backend && cache.backend != backend) { + set_error(error, "paged K/V cache belongs to a different backend"); + return nullptr; + } + if (cache.attn_k.empty() || + cache.attn_k.size() != cache.attn_v.size()) { + set_error(error, "paged K/V cache has no complete attention layers"); + return nullptr; + } + + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || + device >= device_count || cudaSetDevice(device) != cudaSuccess) { + (void)cudaGetLastError(); + set_error(error, "invalid HIP/CUDA device for paged K/V transfer"); + return nullptr; + } + + std::vector cache_tensors; + std::vector layouts; + try { + cache_tensors.reserve(cache.attn_k.size() * 2); + layouts.reserve(cache.attn_k.size() * 2); + } catch (...) { + set_error(error, "paged K/V transfer layout allocation failed"); + return nullptr; + } + + for (size_t layer = 0; layer < cache.attn_k.size(); ++layer) { + const ggml_tensor * pair[] = { + cache.attn_k[layer], cache.attn_v[layer], + }; + if (!pair[0] || !pair[1]) { + set_error(error, + "partial/tensor-parallel paged K/V cache is unsupported"); + return nullptr; + } + if (pair[0]->ne[1] != pair[1]->ne[1] || + pair[0]->ne[2] != pair[1]->ne[2]) { + set_error(error, "K/V cache pair dimensions do not match"); + return nullptr; + } + for (const ggml_tensor * tensor : pair) { + if (!tensor_is_device_backed(tensor, device, error)) return nullptr; + if (tensor->ne[0] <= 0 || tensor->ne[1] <= 0 || + tensor->ne[2] <= 0 || tensor->ne[3] != 1 || + tensor->ne[0] % ggml_blck_size(tensor->type) != 0) { + set_error(error, "unsupported paged K/V tensor dimensions"); + return nullptr; + } + const size_t row_bytes = + ggml_row_size(tensor->type, tensor->ne[0]); + cache_tensors.push_back(tensor); + layouts.push_back({ + row_bytes, + tensor->nb[1], + tensor->nb[2], + ggml_nbytes(tensor), + static_cast(tensor->ne[1]), + static_cast(tensor->ne[2]), + }); + } + } + + QwenPagedKvBlockLayout plan; + if (!plan_qwen_paged_kv_block_layout( + layouts, block_size, plan, error)) { + return nullptr; + } + + std::shared_ptr state; + try { + state = std::make_shared(); + state->device = device; + state->block_size = block_size; + state->block_bytes = plan.block_bytes; + state->physical_block_count = plan.physical_block_count; + state->tensors.reserve(cache_tensors.size()); + for (size_t i = 0; i < cache_tensors.size(); ++i) { + state->tensors.push_back({ + static_cast(cache_tensors[i]->data), + layouts[i], + plan.tensor_offsets[i], + plan.tensor_head_bytes[i], + }); + } + } catch (...) { + set_error(error, "paged K/V transfer state allocation failed"); + return nullptr; + } + + if (cudaStreamCreateWithFlags( + &state->stream, cudaStreamNonBlocking) != cudaSuccess) { + (void)cudaGetLastError(); + set_error(error, "failed to create paged K/V transfer stream"); + return nullptr; + } + + try { + return std::unique_ptr( + new QwenPagedKvResidencyTransfer(std::move(state))); + } catch (...) { + set_error(error, "paged K/V transfer wrapper allocation failed"); + return nullptr; + } +} + +size_t QwenPagedKvResidencyTransfer::block_bytes() const noexcept { + return state_ ? state_->block_bytes : 0; +} + +PagedKvResidencyTransferOps +QwenPagedKvResidencyTransfer::callbacks() const { + PagedKvResidencyTransferOps result; + const std::shared_ptr state = state_; + if (!state) return result; + result.allocate_pinned = [state](size_t bytes) { + return state->allocate_pinned(bytes); + }; + result.free_pinned = [state](void * pointer) { + state->free_pinned(pointer); + }; + result.copy_out_async = + [state](PagedKvSequenceHandle, uint32_t, uint32_t physical_block, + void * host, size_t bytes) { + return state->queue_copy( + physical_block, host, bytes, /*copy_out=*/true); + }; + result.copy_in_async = + [state](PagedKvSequenceHandle, uint32_t, uint32_t physical_block, + const void * host, size_t bytes) { + return state->queue_copy( + physical_block, const_cast(host), bytes, + /*copy_out=*/false); + }; + result.synchronize = [state] { return state->synchronize(); }; + return result; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer.h b/server/src/common/concurrency/qwen_paged_kv_transfer.h new file mode 100644 index 000000000..2aee297a4 --- /dev/null +++ b/server/src/common/concurrency/qwen_paged_kv_transfer.h @@ -0,0 +1,87 @@ +// Production HIP/CUDA block transfers for Qwen paged K/V residency. +// +// A host image packs the payload bytes for one physical block from every +// full-attention K and V tensor. Device row/head strides are retained in the +// copy plan, so quantized K/V types and padded tensor layouts are copied +// without reinterpretation. + +#pragma once + +#include "paged_kv_residency.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct TargetCache; + +// GPU-independent input/output for block-layout validation. storage_bytes is +// the accessible span starting at the tensor data pointer (ggml_nbytes for a +// normal cache tensor). +struct QwenPagedKvTensorLayout { + size_t row_bytes = 0; + size_t row_stride = 0; + size_t head_stride = 0; + size_t storage_bytes = 0; + uint64_t physical_rows = 0; + uint64_t heads = 0; +}; + +struct QwenPagedKvBlockLayout { + size_t block_bytes = 0; + uint32_t physical_block_count = 0; + // One entry per input tensor. Tensor payloads are packed consecutively; + // each head occupies tensor_head_bytes[i] bytes in the host image. + std::vector tensor_offsets; + std::vector tensor_head_bytes; +}; + +// Validates all dimensions/strides and computes the packed host block image. +// Every tensor must cover the same physical row count and head count. Row +// padding is supported; padding bytes are not copied into the host image. +bool plan_qwen_paged_kv_block_layout( + const std::vector & tensors, + uint32_t block_size, + QwenPagedKvBlockLayout & out, + std::string * error = nullptr); + +// Owns the dedicated nonblocking transfer stream. callbacks() retains shared +// ownership of the stream/layout state, so this wrapper may be destroyed as +// soon as callbacks are handed to PagedKvResidencyManager. TargetCache and +// its K/V buffers must outlive those callbacks. Before asking the residency +// manager to evict, the engine must have completed the compute work that last +// wrote the source block; the manager's synchronize callback supplies the +// opposite copy-stream -> later-attention/append barrier. +class QwenPagedKvResidencyTransfer { +public: + static std::unique_ptr create( + const TargetCache & cache, + ggml_backend_t backend, + int device, + uint32_t block_size, + std::string * error = nullptr); + + ~QwenPagedKvResidencyTransfer(); + + QwenPagedKvResidencyTransfer( + const QwenPagedKvResidencyTransfer &) = delete; + QwenPagedKvResidencyTransfer & operator=( + const QwenPagedKvResidencyTransfer &) = delete; + + size_t block_bytes() const noexcept; + PagedKvResidencyTransferOps callbacks() const; + +private: + struct State; + explicit QwenPagedKvResidencyTransfer(std::shared_ptr state); + + std::shared_ptr state_; +}; + +} // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp b/server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp new file mode 100644 index 000000000..d75a3f864 --- /dev/null +++ b/server/src/common/concurrency/qwen_paged_kv_transfer_layout.cpp @@ -0,0 +1,119 @@ +#include "qwen_paged_kv_transfer.h" + +#include + +namespace dflash::common { +namespace { + +bool fail(QwenPagedKvBlockLayout & out, std::string * error, + const char * message) { + out = {}; + if (error) *error = message; + return false; +} + +bool checked_mul(size_t lhs, size_t rhs, size_t & out) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return false; + } + out = lhs * rhs; + return true; +} + +bool checked_add(size_t lhs, size_t rhs, size_t & out) { + if (rhs > std::numeric_limits::max() - lhs) return false; + out = lhs + rhs; + return true; +} + +} // namespace + +bool plan_qwen_paged_kv_block_layout( + const std::vector & tensors, + uint32_t block_size, + QwenPagedKvBlockLayout & out, + std::string * error) { + out = {}; + if (tensors.empty() || block_size == 0) { + return fail(out, error, "empty paged K/V layout or zero block size"); + } + + const uint64_t physical_rows = tensors.front().physical_rows; + const uint64_t heads = tensors.front().heads; + if (physical_rows == 0 || heads == 0 || + physical_rows % block_size != 0 || + physical_rows / block_size > std::numeric_limits::max()) { + return fail(out, error, "invalid paged K/V physical dimensions"); + } + out.physical_block_count = + static_cast(physical_rows / block_size); + + try { + out.tensor_offsets.reserve(tensors.size()); + out.tensor_head_bytes.reserve(tensors.size()); + } catch (...) { + return fail(out, error, "paged K/V layout allocation failed"); + } + + size_t host_offset = 0; + for (const QwenPagedKvTensorLayout & tensor : tensors) { + if (tensor.row_bytes == 0 || tensor.row_stride < tensor.row_bytes || + tensor.head_stride == 0 || tensor.storage_bytes == 0 || + tensor.physical_rows != physical_rows || tensor.heads != heads) { + return fail(out, error, "inconsistent paged K/V tensor layout"); + } + if (physical_rows > std::numeric_limits::max() || + heads > std::numeric_limits::max()) { + return fail(out, error, "paged K/V dimensions exceed host size_t"); + } + + size_t last_row_offset = 0; + size_t last_head_offset = 0; + size_t used_end = 0; + if (!checked_mul(static_cast(physical_rows - 1), + tensor.row_stride, last_row_offset) || + !checked_mul(static_cast(heads - 1), + tensor.head_stride, last_head_offset) || + !checked_add(last_head_offset, last_row_offset, used_end) || + !checked_add(used_end, tensor.row_bytes, used_end) || + used_end > tensor.storage_bytes) { + return fail(out, error, "paged K/V tensor strides exceed storage"); + } + + // Heads must not overlap. Larger head strides (alignment/padding) are + // fine because device copies retain the original pitch. + size_t plane_span = 0; + if (!checked_mul(static_cast(physical_rows - 1), + tensor.row_stride, plane_span) || + !checked_add(plane_span, tensor.row_bytes, plane_span) || + tensor.head_stride < plane_span) { + return fail(out, error, "overlapping paged K/V head planes"); + } + + size_t head_bytes = 0; + size_t tensor_bytes = 0; + size_t next_offset = 0; + if (!checked_mul(tensor.row_bytes, block_size, head_bytes) || + !checked_mul(head_bytes, static_cast(heads), + tensor_bytes) || + !checked_add(host_offset, tensor_bytes, next_offset)) { + return fail(out, error, "paged K/V host block size overflow"); + } + try { + out.tensor_offsets.push_back(host_offset); + out.tensor_head_bytes.push_back(head_bytes); + } catch (...) { + return fail(out, error, "paged K/V layout allocation failed"); + } + host_offset = next_offset; + } + + if (host_offset == 0) { + return fail(out, error, "zero-byte paged K/V host block"); + } + out.block_bytes = host_offset; + if (error) error->clear(); + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index fa4dbba35..fa2909da3 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -8,9 +8,10 @@ // 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. +// scheduler feeds the final sampled token back as the next step's input. An +// engine may also return speculative children it already committed before that +// final pending token. The scheduler can disable that burst path per slot when +// it must retain authority to substitute a thinking-budget close token. // // The split of duties is deliberate and is the reason this interface exists // apart from ModelBackend: @@ -181,14 +182,31 @@ class SeqEngine { struct StepInput { int slot = -1; int32_t token = -1; // token to commit at this slot's next position + // False when scheduler-side policy may replace the sampled token + // before it is committed (currently the thinking-budget close hook). + bool allow_speculation = true; }; struct DecodeOutput { int slot = -1; - int32_t token = -1; // newly sampled token (pending until next step) + // Final newly sampled token, pending until the scheduler feeds it into + // the next step. `committed_tokens`, when non-empty, precede this token + // and are already present in backend state. + int32_t token = -1; 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; + std::vector committed_tokens; + + // Per-slot deltas for this engine step. The scheduler aggregates them + // until retirement and emits one machine-readable proof record. + uint64_t ddtree_steps = 0; + uint64_t ddtree_accepted_tokens = 0; + uint64_t target_forwards = 0; + uint64_t kvflash_page_ins = 0; + uint64_t kvflash_page_outs = 0; + uint64_t kvflash_resident_blocks = 0; + uint64_t kvflash_reselects = 0; }; struct PrefillOutput { @@ -250,6 +268,19 @@ class SeqEngine { virtual bool token_is_eos(int32_t token) const = 0; }; +// Deliver a successful decode result in wire order. The visitor returns false +// after a stop/EOS/output-cap decision; in that case later committed children +// and the final pending token are intentionally hidden and the slot is retired. +template +inline bool consume_decode_output_tokens( + const SeqEngine::DecodeOutput & output, Advance advance) { + if (output.failed) return false; + for (int32_t token : output.committed_tokens) { + if (!advance(token)) return false; + } + return advance(output.token); +} + // 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. @@ -265,6 +296,7 @@ inline std::string validate_step_result( } std::vector decode_planned((size_t)slot_count, 0); + std::vector speculation_allowed((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) @@ -272,6 +304,8 @@ inline std::string validate_step_result( if (decode_planned[(size_t)input.slot]) return "decode plan contains a duplicate slot"; decode_planned[(size_t)input.slot] = 1; + speculation_allowed[(size_t)input.slot] = + input.allow_speculation ? 1 : 0; } for (const PrefillSlice & slice : plan.prefills) { if (slice.slot < 0 || slice.slot >= slot_count || @@ -290,10 +324,25 @@ inline std::string validate_step_result( return "decode output names an unplanned slot"; if (decode_seen[(size_t)output.slot]) return "step returned duplicate decode outputs"; - if (output.failed && (output.token >= 0 || output.error.empty())) - return "failed decode has invalid payload"; - if (!output.failed && (output.token < 0 || !output.error.empty())) - return "successful decode has invalid payload"; + if (output.failed) { + if (output.error.empty()) + return "failed decode has no diagnostic"; + if (output.token >= 0 || !output.committed_tokens.empty()) + return "failed decode exposes token payload"; + } else { + if (output.token < 0) + return "successful decode has no pending token"; + if (!output.error.empty()) + return "successful decode carries an error diagnostic"; + if (!speculation_allowed[(size_t)output.slot] && + !output.committed_tokens.empty()) + return "decode output burst violates disabled speculation"; + if (std::any_of( + output.committed_tokens.begin(), + output.committed_tokens.end(), + [](int32_t token) { return token < 0; })) + return "decode output burst contains an invalid token"; + } decode_seen[(size_t)output.slot] = 1; } diff --git a/server/src/common/ddtree.cpp b/server/src/common/ddtree.cpp index 08ca33464..54402c2cc 100644 --- a/server/src/common/ddtree.cpp +++ b/server/src/common/ddtree.cpp @@ -223,4 +223,14 @@ std::vector follow_verified_tree(const DDTree & tree, return accepted; } +bool truncate_verified_path(std::vector & accepted, + std::size_t max_committed, + const int32_t * posterior, + int & out_next_token) { + if (accepted.size() <= max_committed) return false; + accepted.resize(max_committed); + out_next_token = accepted.empty() ? -1 : posterior[accepted.back()]; + return true; +} + } // namespace dflash::common diff --git a/server/src/common/ddtree.h b/server/src/common/ddtree.h index afe22f226..718c05b57 100644 --- a/server/src/common/ddtree.h +++ b/server/src/common/ddtree.h @@ -8,6 +8,7 @@ #pragma once +#include #include #include #include @@ -61,4 +62,13 @@ std::vector follow_verified_tree(const DDTree & tree, int & out_next_token, int * out_node_idx = nullptr); +// Bound a verified path to the number of tokens that can actually be +// committed. When truncation removes the old tip, the pending token must be +// recomputed from the posterior at the new tip; otherwise it describes model +// state that was never committed. +bool truncate_verified_path(std::vector & accepted, + std::size_t max_committed, + const int32_t * posterior, + int & out_next_token); + } // namespace dflash::common diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index f655f5d4f..63099832b 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -163,6 +163,26 @@ std::string check_feature_compatibility( "' does not support PFlash compression"; } + const bool concurrent_paged_qwen = + arch == "qwen35" && args.paged_attention && + args.max_concurrency > 1 && + !args.device.is_layer_split() && + !args.remote_target_shard.enabled() && + args.fa_window == 0; + const bool concurrent_local_paged_qwen = + concurrent_paged_qwen && !args.remote_draft.enabled() && + target_backend == draft_backend && + args.device.gpu == args.draft_device.gpu && + !args.device.is_tensor_parallel(); + const bool concurrent_local_ddtree = + concurrent_local_paged_qwen && args.draft_path != nullptr && + args.ddtree_mode; + + if (args.ddtree_mode && + (args.ddtree_budget < 1 || args.ddtree_budget > 255)) { + return "--ddtree-budget must be in [1, 255]"; + } + // ── --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 @@ -180,20 +200,24 @@ std::string check_feature_compatibility( args.remote_target_shard.enabled()) { return "--paged-attention requires one local target device"; } - if (args.draft_path != nullptr || args.remote_draft.enabled() || - args.ddtree_mode) { - return "--paged-attention requires autoregressive decode without a " - "draft or DDTree"; + if (args.remote_draft.enabled()) { + return "concurrent paged DDTree requires a local draft on the target device"; + } + if ((args.draft_path != nullptr || args.ddtree_mode) && + !concurrent_local_ddtree) { + return "paged draft decode is supported only as concurrent local DDTree " + "on one target/draft device"; } if (args.fa_window != 0) { return "--paged-attention requires full attention (--fa-window 0)"; } - if (features.pflash_enabled) { - return "--paged-attention cannot be combined with PFlash prefill " - "compression"; + if (features.pflash_enabled && !concurrent_local_paged_qwen) { + return "paged PFlash prefill compression requires concurrent local " + "Qwen3.5/Qwen3.6 serving on one target/draft device"; } - if (features.kvflash_enabled) { - return "--paged-attention cannot be combined with KVFlash"; + if (features.kvflash_enabled && !concurrent_paged_qwen) { + return "paged KVFlash requires concurrent local Qwen3.5/Qwen3.6 " + "serving with full attention"; } // 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. @@ -232,10 +256,16 @@ std::string check_feature_compatibility( if (args.max_concurrency <= 1) { return "--kv-pool-tokens requires --max-concurrency greater than 1"; } - // The cache appends one scratch block after the physical pool, and - // the requested pool itself is rounded up to a whole block. Cap the - // request at the largest aligned pool that leaves room for scratch. - const int64_t max_pool_tokens = paged_kv_address_cap(); + // Reserve the dead-row block plus one rounded DDTree candidate slab + // per slot, exactly matching Qwen35Backend's cache allocation. + const int64_t tree_scratch = concurrent_local_ddtree + ? (int64_t)args.max_concurrency * + paged_token_capacity(args.ddtree_budget + 1) + : 0; + const int64_t scratch_tokens = PAGED_BLOCK_SIZE + tree_scratch; + const int64_t max_pool_tokens = + ((int64_t)INT32_MAX - scratch_tokens) / PAGED_BLOCK_SIZE * + PAGED_BLOCK_SIZE; if (args.kv_pool_tokens < PAGED_BLOCK_SIZE || args.kv_pool_tokens > max_pool_tokens) { return "--kv-pool-tokens must be in [" + diff --git a/server/src/common/gpu_runtime_compat.h b/server/src/common/gpu_runtime_compat.h index dba8eaa7c..8a21fe025 100644 --- a/server/src/common/gpu_runtime_compat.h +++ b/server/src/common/gpu_runtime_compat.h @@ -60,6 +60,7 @@ #define cudaPointerAttributes hipPointerAttribute_t #define cudaPointerGetAttributes hipPointerGetAttributes #define cudaStreamCreate hipStreamCreate +#define cudaStreamCreateWithFlags hipStreamCreateWithFlags #define cudaStreamDefault hipStreamDefault #define cudaStreamDestroy hipStreamDestroy #define cudaStreamNonBlocking hipStreamNonBlocking diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index cdd4d0210..5e1814a72 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -35,7 +35,8 @@ struct StepGraph { ggml_tensor * inp_embed = nullptr; ggml_tensor * positions = nullptr; ggml_tensor * attn_mask = nullptr; // may be null - ggml_tensor * parent_ids = nullptr; // DDTree tree-mode; null for chain mode + ggml_tensor * parent_ids = nullptr; // DDTree [tree_width,n_tree_seqs] + ggml_tensor * tree_sizes = nullptr; // DDTree [n_tree_seqs], 0 = padding ggml_tensor * target_hidden_cat = nullptr; // draft only ggml_tensor * positions_k = nullptr; // draft only ggml_tensor * pad_mask_full = nullptr; // draft only; padded-ctx mask @@ -55,11 +56,16 @@ struct StepGraph { // state_slot_ids has the same shape but maps padding to a safe readable // slot for graph-level conv-state gathers. ggml_tensor * active_slot_ids = nullptr; + // Recurrent gather rows. Unlike active/paged IDs, padding must name a + // valid harmless slot (normally 0): ggml_get_rows does not mask -1. ggml_tensor * state_slot_ids = nullptr; // Ragged paged read (concurrent prefill): per-row block-table column and // inclusive causal position, [n_tokens] i32 each. Padding rows carry -1. ggml_tensor * paged_query_seq_ids = nullptr; ggml_tensor * paged_query_positions = nullptr; + // DFlash target-feature destination rows. Multi-slot replay maps each + // token to its slot-local ring; padding maps to the cache's dead row. + ggml_tensor * target_feat_rows = nullptr; // Multi-prompt steps: i32 row indices gathered from the final norm // before the LM head (committing rows + decode rows). ggml_tensor * logits_row_indices = nullptr; @@ -92,11 +98,13 @@ inline void step_graph_free(StepGraph & sg) { sg.built_view = false; sg.hidden_input = nullptr; sg.parent_ids = nullptr; + sg.tree_sizes = nullptr; sg.kv_write_rows = nullptr; sg.active_slot_ids = nullptr; sg.state_slot_ids = nullptr; sg.paged_query_seq_ids = nullptr; sg.paged_query_positions = nullptr; + sg.target_feat_rows = nullptr; sg.logits_row_indices = nullptr; sg.logits = nullptr; sg.hidden_states = nullptr; diff --git a/server/src/internal.h b/server/src/internal.h index fffda41ac..9cb2a03b6 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -414,13 +414,12 @@ struct TargetCache { std::vector conv_input_cache; // size = n_delta (48) // Rolling target layer features captured during target forward passes. - // Shape [5 * hidden, target_feat_cap] bf16. target_feat_cap is typically - // << max_ctx (e.g. 4096) so the buffer stays small at 128K context. The - // graph writes to slot `(kv_start + i) % target_feat_cap` so positions - // beyond the cap wrap and overwrite older entries. Readers (draft) only - // need the last DRAFT_CTX_MAX positions, so wrap is invisible in - // practice. Fed into the draft graph's fc projection after a bf16→f32 - // cast (ggml_get_to_fp32_cuda). + // Single-sequence shape: [5 * hidden, target_feat_cap] bf16. A multi-slot + // cache owns one ring per physical sequence slot and one final dead row: + // [5 * hidden, target_feat_cap * n_seq_slots + 1]. Live row P in slot S + // maps to S*target_feat_cap + P%target_feat_cap; bucket padding maps to the + // dead final row because ggml_set_rows does not accept a negative index. + // target_feat_cap remains the per-sequence ring width. ggml_tensor * target_feat = nullptr; int target_feat_cap = 0; @@ -559,6 +558,10 @@ bool restore_target_cache_chain(const PrefixSnapshot * thick, // decode is AR-only). With // n_seq_slots > 1 the attention K/V tensors are sized by ctx_alloc (the shared // pool capacity plus one scratch block) rather than one sequence's max_ctx. +// `concurrent_tree` declares that a paged multi-slot caller will build packed +// DDTree verification graphs. Those graphs are deliberately side-effect-free +// for recurrent state and commit accepted paths through a later replay, so no +// rollback snapshots/intermediates are allocated. bool create_target_cache(const TargetWeights & w, int max_ctx, int max_verify_tokens, @@ -567,7 +570,8 @@ bool create_target_cache(const TargetWeights & w, bool prefill_only = false, int ctx_alloc = 0, bool paged_attention = false, - int n_seq_slots = 1); + int n_seq_slots = 1, + bool concurrent_tree = false); // `f32_ssm_intermediates` enables exact per-token checkpoints for the opt-in // layer-split fast rollback path. The default preserves the established Q8_0 @@ -584,7 +588,8 @@ bool create_target_cache_partial(const TargetWeights & w, int ctx_alloc = 0, bool f32_ssm_intermediates = false, bool paged_attention = false, - int n_seq_slots = 1); + int n_seq_slots = 1, + bool concurrent_tree = false); void free_target_cache(TargetCache & c); @@ -657,7 +662,8 @@ struct QwenGraphInputs { bool capture_moe_router = false; // if true, expose selected expert ids for MoE layers int fa_window = 0; // sliding window for FA layers: 0 = full attention int logits_tail_rows = 0; // compute logits only for last n rows; 0 = all - ggml_tensor * parent_ids = nullptr; // [n_tokens] i32; tree mode when non-null + ggml_tensor * parent_ids = nullptr; // tree: [tree_width,n_tree_seqs] i32 + ggml_tensor * tree_sizes = nullptr; // tree: [n_tree_seqs] i32; 0 = padding tree // [n_tokens,n_head_kv] i64 physical destination rows for the // ggml_set_rows KV write; step-invariant. ggml_tensor * kv_write_rows = nullptr; @@ -683,6 +689,10 @@ struct QwenGraphInputs { // last row plus the decode rows), which a tail view cannot express. // Non-null overrides logits_tail_rows. ggml_tensor * logits_row_indices = nullptr; + // Optional replay-stable DFlash capture destinations. When present, all + // captured layers are concatenated once and written with ggml_set_rows. + // Multi-slot callers provide per-slot ring rows (padding uses dead row). + ggml_tensor * target_feat_rows = nullptr; // [n_tokens] i32 // Prefill segments on the leading token axis (see QwenPrefillSegment). // n_prefill_tokens is their total row count. seq_slot is ignored when // segments are present. @@ -717,6 +727,12 @@ struct QwenGraphInputs { int seq_slot = 0; int paged_max_kv_len = 0; int n_prefill_tokens = 0; + // Packed paged-tree metadata. Tokens are flattened sequence-major: + // row = sequence*tree_width + node. tree_scratch_* describe the physical + // KV scratch slab owned by each physical sequence slot. + int tree_width = 0; + int tree_scratch_base = 0; + int tree_scratch_stride = 0; // Capture the LAST token's post-RoPE/post-rotation Q per full-attention // layer into cache.q_cap (KVFlash target-QK scorer). Step-invariant: // node properties depend only on n_tokens and the layer index. diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 0a54761f9..545e7d9bf 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -12,6 +12,8 @@ #include "attn_masks.h" #include "prefill_helpers.h" #include "common/sampler.h" +#include "common/ddtree.h" +#include "common/geometric_draft_topk_cuda.h" #include "internal.h" #include @@ -35,6 +37,506 @@ int decode_bucket_width(int live_count) { } // namespace +Qwen35SeqEngine::Qwen35SeqEngine( + Qwen35Backend & backend, PagedKvPool & pool, int max_ctx, + int64_t scratch_row, int tree_width, int tree_scratch_base, + int tree_scratch_stride, int max_prefills, + int mixed_prefill_tokens, int long_mixed_prefill_tokens, + int long_prefill_threshold, int idle_prefill_tokens, + int prefill_quantum) + : max_prefills_(std::max(1, max_prefills)), + mixed_prefill_tokens_(std::max(1, mixed_prefill_tokens)), + long_mixed_prefill_tokens_(std::max(1, long_mixed_prefill_tokens)), + long_prefill_threshold_(std::max(1, long_prefill_threshold)), + idle_prefill_tokens_(std::max(1, idle_prefill_tokens)), + prefill_quantum_(std::max(1, prefill_quantum)), b_(backend), + slots_(pool, max_ctx, std::max(1, tree_width), + backend.paged_kv_residency_.get()), + scratch_row_(scratch_row), tree_width_(tree_width), + tree_scratch_base_(tree_scratch_base), + tree_scratch_stride_(tree_scratch_stride) { + const int n_slots = slots_.slot_count(); + slot_draft_kv_.resize((size_t)n_slots); + + // The concurrent DDTree stack is gated to a local same-device drafter. + // Build metadata-only BF16 views over each slot's disjoint target feature + // ring; draft_kv_begin_step converts only newly committed rows to its F32 + // append input instead of syncing the entire 200 MiB ring per round. + capture_features_ = tree_width_ > 0 && b_.cache_.target_feat && + b_.cache_.target_feat_cap > 0 && b_.cfg_.draft_path && + !b_.cfg_.remote_draft.enabled() && !b_.split_gpus_ && + b_.cfg_.draft_gpu == b_.cfg_.device.gpu; + if (!capture_features_) return; + + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead() * (size_t)(n_slots + 1); + ip.no_alloc = true; + feature_view_ctx_ = ggml_init(ip); + if (!feature_view_ctx_) { + capture_features_ = false; + return; + } + + const int cap = b_.cache_.target_feat_cap; + const int64_t fc_in = + (int64_t)b_.w_.n_capture_layers * b_.w_.n_embd; + slot_feature_mirrors_.resize((size_t)n_slots); + for (int slot = 0; slot < n_slots; ++slot) { + DraftFeatureMirror & mirror = slot_feature_mirrors_[(size_t)slot]; + mirror.target_feat = ggml_view_2d( + feature_view_ctx_, b_.cache_.target_feat, fc_in, cap, + b_.cache_.target_feat->nb[1], + (size_t)slot * (size_t)cap * b_.cache_.target_feat->nb[1]); + mirror.device = b_.cfg_.draft_gpu; + mirror.target_device = b_.cfg_.device.gpu; + mirror.cap = cap; + mirror.n_target_layers = b_.w_.n_capture_layers; + mirror.hidden_size = b_.w_.n_embd; + mirror.storage_type = b_.cache_.target_feat->type; + } +} + +Qwen35SeqEngine::~Qwen35SeqEngine() { + for (std::unique_ptr & state : slot_draft_kv_) { + if (state) draft_kv_free(*state); + } + slot_draft_kv_.clear(); + for (DraftFeatureMirror & mirror : slot_feature_mirrors_) { + draft_feature_mirror_free(mirror); + } + slot_feature_mirrors_.clear(); + if (feature_view_ctx_) { + ggml_free(feature_view_ctx_); + feature_view_ctx_ = nullptr; + } +} + +DraftFeatureMirror * Qwen35SeqEngine::slot_feature_mirror(int slot) { + if (!capture_features_ || slot < 0 || + slot >= (int)slot_feature_mirrors_.size()) { + return nullptr; + } + return &slot_feature_mirrors_[(size_t)slot]; +} + +DraftKvState * Qwen35SeqEngine::ensure_slot_draft_kv(int slot) { + DraftFeatureMirror * mirror = slot_feature_mirror(slot); + if (!mirror || slot < 0 || slot >= (int)slot_draft_kv_.size()) { + return nullptr; + } + std::unique_ptr & state = slot_draft_kv_[(size_t)slot]; + if (state && state->gf && state->built_for == (const void *)&b_.dw_) { + return state.get(); + } + if (state) draft_kv_free(*state); + state = std::make_unique(); + const int cap = std::min( + mirror->cap, std::max(1, b_.cfg_.draft_ctx_max)); + if (!draft_kv_init(*state, b_.dw_, b_.draft_backend_, cap, nullptr)) { + draft_kv_free(*state); + state.reset(); + return nullptr; + } + return state.get(); +} + +bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { + if (tree_width_ <= 1 || !capture_features_ || !plan.prefills.empty() || + plan.decode.empty() || b_.dw_.block_size <= 1 || + b_.cfg_.ddtree_budget + 1 != tree_width_) { + return false; + } + const int min_floor = []() { + const char * value = std::getenv("DFLASH_MIN_TOKENS"); + return value ? std::max(0, std::atoi(value)) : 0; + }(); + for (const StepInput & in : plan.decode) { + if (!in.allow_speculation || in.slot < 0 || + in.slot >= slots_.slot_count() || + !slots_.slot(in.slot).decoding() || + slots_.slot(in.slot).sampler.needs_logit_processing() || + slots_.slot(in.slot).cur_pos < 1 || + slots_.slot(in.slot).cur_pos >= slots_.max_context()) { + return false; + } + const Qwen35Slot & seq = slots_.slot(in.slot); + const int generated = seq.generated_tokens(); + if (generated < min_floor) return false; + } + return true; +} + +std::optional Qwen35SeqEngine::step_ddtree( + const StepPlan & plan) { + StepResult result; + const int active = (int)plan.decode.size(); + const int bucket = decode_bucket_width(active); + const int T = tree_width_; + const int q_len = b_.dw_.block_size; + const int hidden = b_.w_.n_embd; + const int n_head_kv = b_.w_.n_head_kv; + const int n_slots = slots_.slot_count(); + const int K = b_.cfg_.ddtree_budget > q_len - 1 ? 8 : 1; + + struct Proposal { + int slot = -1; + int32_t root = -1; + DDTree tree; + std::vector flat; + std::vector accepted; + int32_t bonus = -1; + }; + std::vector proposals; + proposals.reserve((size_t)active); + std::vector noise((size_t)q_len, b_.w_.mask_token_id); + std::vector noise_embed((size_t)hidden * q_len); + std::vector logits((size_t)b_.w_.n_vocab * q_len); + std::vector top_lp((size_t)q_len * K); + std::vector top_ids((size_t)q_len * K); + + auto proposal_fallback = [&]() -> std::optional { + // begin_step updates persistent drafter bookkeeping before compute. + // A failed proposal graph may therefore leave only part of that + // cache valid. Reset all participating draft rings so a later + // speculative round rebuilds them from committed target features. + for (const StepInput & in : plan.decode) { + if (in.slot >= 0 && in.slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)in.slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)in.slot]); + } + } + return std::nullopt; + }; + + if (!build_lm_head_projection_step( + b_.proj_sg_, b_.w_, b_.target_backend_, q_len)) { + return proposal_fallback(); + } + + // Proposal is sequential by slot: immutable draft weights are shared, + // while each slot owns an independent persistent context-KV ring. + for (const StepInput & in : plan.decode) { + DraftKvState * draft = ensure_slot_draft_kv(in.slot); + DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); + if (!draft || !mirror) return proposal_fallback(); + noise[0] = in.token; + std::fill(noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed( + noise.data(), q_len, noise_embed.data()) || + !draft_kv_begin_step(*draft, b_.dw_, b_.draft_backend_, + *mirror, slots_.slot(in.slot).cur_pos)) { + return proposal_fallback(); + } + ggml_backend_tensor_set( + draft->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + if (ggml_backend_graph_compute(b_.draft_backend_, draft->gf) != + GGML_STATUS_SUCCESS) { + return proposal_fallback(); + } + ggml_backend_tensor_copy( + draft->hidden_states, b_.proj_sg_.hidden_input); + if (ggml_backend_graph_compute( + b_.target_backend_, b_.proj_sg_.gf) != GGML_STATUS_SUCCESS) { + return proposal_fallback(); + } + bool topk_ready = false; +#ifdef DFLASH27B_HAVE_DRAFT_TOPK + topk_ready = geometric_extract_draft_topk_cuda( + b_.proj_sg_.logits->data, q_len, b_.w_.n_vocab, K, + top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); +#endif + if (!topk_ready) { + ggml_backend_tensor_get( + b_.proj_sg_.logits, logits.data(), 0, + sizeof(float) * logits.size()); + extract_draft_topk( + logits.data(), q_len, b_.w_.n_vocab, K, + top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); + } + + Proposal p; + p.slot = in.slot; + p.root = in.token; + p.tree = build_ddtree( + top_lp.data() + K, top_ids.data() + K, + q_len - 1, K, b_.cfg_.ddtree_budget, + b_.cfg_.ddtree_chain_seed); + p.flat.assign((size_t)T, 0); + p.flat[0] = in.token; + for (int node = 0; node < p.tree.n_nodes; ++node) { + p.flat[(size_t)node + 1] = p.tree.token_ids[(size_t)node]; + } + proposals.push_back(std::move(p)); + } + + StepGraph & tree_sg = b_.sg_; + int max_prefix = 1; + for (const Proposal & p : proposals) { + max_prefix = std::max(max_prefix, slots_.slot(p.slot).cur_pos); + } + if (!build_target_step_paged_tree( + tree_sg, b_.w_, b_.cache_, b_.target_backend_, T, bucket, + max_prefix, tree_scratch_base_, tree_scratch_stride_, + b_.cfg_.kq_stride_pad)) { + result.error = "packed DDTree verify graph build failed"; + return result; + } + + const int total_tree = T * bucket; + std::vector flat_tokens((size_t)total_tree, 0); + std::vector parents((size_t)total_tree, -1); + std::vector sizes((size_t)bucket, 0); + // Negative active IDs identify bucket padding. Recurrent gathers cannot + // index a negative slab, so padded trees use slot 0 only for their + // read-only base-state gather; tree_size=0/query_slot=-1 keeps all of + // their attention/output rows inactive and tree mode never persists it. + std::vector tree_slots((size_t)bucket, -1); + std::vector tree_state_slots((size_t)bucket, 0); + std::vector query_slots((size_t)total_tree, -1); + std::vector tree_rows( + (size_t)total_tree * n_head_kv, scratch_row_); + std::vector tree_pos((size_t)4 * total_tree, 0); + std::vector tree_embed((size_t)hidden * total_tree, 0.0f); + seq_lens_.assign((size_t)n_slots, 0); + + for (int s = 0; s < active; ++s) { + const Proposal & p = proposals[(size_t)s]; + const int base = s * T; + sizes[(size_t)s] = p.tree.n_nodes + 1; + tree_slots[(size_t)s] = p.slot; + tree_state_slots[(size_t)s] = p.slot; + seq_lens_[(size_t)p.slot] = slots_.slot(p.slot).cur_pos; + for (int node = 0; node < sizes[(size_t)s]; ++node) { + const int row = base + node; + flat_tokens[(size_t)row] = p.flat[(size_t)node]; + parents[(size_t)row] = node == 0 ? -1 : + p.tree.parents[(size_t)node]; + query_slots[(size_t)row] = p.slot; + const int depth = node == 0 ? 0 : + p.tree.depths[(size_t)node - 1]; + const int pos = slots_.slot(p.slot).cur_pos + depth; + tree_pos[(size_t)0 * total_tree + row] = pos; + tree_pos[(size_t)1 * total_tree + row] = pos; + tree_pos[(size_t)2 * total_tree + row] = pos; + for (int h = 0; h < n_head_kv; ++h) { + tree_rows[(size_t)h * total_tree + row] = + (int64_t)tree_scratch_base_ + + (int64_t)p.slot * tree_scratch_stride_ + node; + } + } + } + if (!b_.w_.embedder.embed( + flat_tokens.data(), total_tree, tree_embed.data())) { + result.error = "packed DDTree embedding failed"; + return result; + } + ggml_backend_tensor_set(tree_sg.inp_embed, tree_embed.data(), 0, + sizeof(float) * tree_embed.size()); + ggml_backend_tensor_set(tree_sg.positions, tree_pos.data(), 0, + sizeof(int32_t) * tree_pos.size()); + ggml_backend_tensor_set(tree_sg.parent_ids, parents.data(), 0, + sizeof(int32_t) * parents.size()); + ggml_backend_tensor_set(tree_sg.tree_sizes, sizes.data(), 0, + sizeof(int32_t) * sizes.size()); + ggml_backend_tensor_set(tree_sg.active_slot_ids, tree_slots.data(), 0, + sizeof(int32_t) * tree_slots.size()); + ggml_backend_tensor_set(tree_sg.state_slot_ids, tree_state_slots.data(), 0, + sizeof(int32_t) * tree_state_slots.size()); + ggml_backend_tensor_set(tree_sg.paged_query_seq_ids, query_slots.data(), 0, + sizeof(int32_t) * query_slots.size()); + ggml_backend_tensor_set(tree_sg.kv_write_rows, tree_rows.data(), 0, + sizeof(int64_t) * tree_rows.size()); + ggml_backend_tensor_set(b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + if (ggml_backend_graph_compute(b_.target_backend_, tree_sg.gf) != + GGML_STATUS_SUCCESS) { + result.error = "packed DDTree verify compute failed"; + return result; + } + std::vector posterior((size_t)total_tree, -1); + ggml_backend_tensor_get(tree_sg.argmax_tokens, posterior.data(), 0, + sizeof(int32_t) * posterior.size()); + + int replay_total = 0; + for (int s = 0; s < active; ++s) { + Proposal & p = proposals[(size_t)s]; + p.accepted = follow_verified_tree( + p.tree, posterior.data() + (size_t)s * T, p.bonus); + const int room = slots_.max_context() - slots_.slot(p.slot).cur_pos; + truncate_verified_path( + p.accepted, (size_t)std::max(0, room), + posterior.data() + (size_t)s * T, p.bonus); + if (p.accepted.empty()) { + result.error = "DDTree accepted path has no context headroom"; + return result; + } + replay_total += (int)p.accepted.size(); + } + + std::vector replay_segments; + std::vector replay_tokens; + std::vector replay_slots; + std::vector replay_positions; + std::vector replay_rows; + std::vector replay_logits_rows; + replay_segments.reserve((size_t)active); + replay_tokens.reserve((size_t)replay_total); + replay_slots.reserve((size_t)replay_total); + replay_positions.reserve((size_t)replay_total); + replay_rows.assign((size_t)replay_total * n_head_kv, scratch_row_); + replay_logits_rows.reserve((size_t)active); + seq_lens_.assign((size_t)n_slots, 0); + + int replay_offset = 0; + for (Proposal & p : proposals) { + std::vector path; + path.reserve(p.accepted.size()); + for (int dfs : p.accepted) { + path.push_back(dfs == 0 ? p.root : + p.tree.token_ids[(size_t)dfs - 1]); + } + const Qwen35SlotManager::StepAppend app = slots_.append_tokens( + p.slot, path.data(), (int)path.size()); + const bool table_ok = slots_.residency_active() || + upload_block_table_delta(p.slot, app.first_new_block, + app.new_blocks.data(), app.new_blocks.size()); + if (!app.ok || app.physical_rows.size() != path.size() || + !table_ok) { + result.error = app.busy + ? "paged KV pool exhausted during DDTree replay" + : "DDTree replay K/V append failed"; + return result; + } + replay_segments.push_back( + {replay_offset, (int)path.size(), p.slot}); + for (size_t i = 0; i < path.size(); ++i) { + replay_tokens.push_back(path[i]); + replay_slots.push_back(p.slot); + replay_positions.push_back(app.position + (int)i); + for (int h = 0; h < n_head_kv; ++h) { + replay_rows[(size_t)h * replay_total + replay_offset + i] = + app.physical_rows[i]; + } + } + replay_offset += (int)path.size(); + replay_logits_rows.push_back(replay_offset - 1); + seq_lens_[(size_t)p.slot] = app.position + (int)path.size(); + } + + if (!upload_all_active_block_tables()) { + result.error = "DDTree replay block-table refresh failed"; + return result; + } + + StepGraph & replay_sg = b_.sg_; + if (!build_target_step( + replay_sg, b_.w_, b_.cache_, b_.target_backend_, + 0, replay_total, false, true, false, 0, 0, + b_.cfg_.kq_stride_pad, false, false, false, true, + 1, 0, *std::max_element(seq_lens_.begin(), seq_lens_.end()), + replay_total, replay_segments.data(), + (int)replay_segments.size(), active, false) || + !replay_sg.target_feat_rows || !replay_sg.paged_query_seq_ids || + !replay_sg.paged_query_positions || !replay_sg.logits_row_indices || + !replay_sg.argmax_tokens) { + result.error = "DDTree accepted-path replay graph build failed"; + return result; + } + embed_buf_.resize((size_t)hidden * replay_total); + if (!b_.w_.embedder.embed( + replay_tokens.data(), replay_total, embed_buf_.data())) { + result.error = "DDTree replay embedding failed"; + return result; + } + pos_buf_.assign((size_t)4 * replay_total, 0); + feature_rows_.resize((size_t)replay_total); + const int cap = b_.cache_.target_feat_cap; + for (int row = 0; row < replay_total; ++row) { + const int pos = replay_positions[(size_t)row]; + pos_buf_[(size_t)0 * replay_total + row] = pos; + pos_buf_[(size_t)1 * replay_total + row] = pos; + pos_buf_[(size_t)2 * replay_total + row] = pos; + feature_rows_[(size_t)row] = + replay_slots[(size_t)row] * cap + pos % cap; + } + ggml_backend_tensor_set(replay_sg.inp_embed, embed_buf_.data(), 0, + sizeof(float) * embed_buf_.size()); + ggml_backend_tensor_set(replay_sg.positions, pos_buf_.data(), 0, + sizeof(int32_t) * pos_buf_.size()); + ggml_backend_tensor_set(replay_sg.kv_write_rows, replay_rows.data(), 0, + sizeof(int64_t) * replay_rows.size()); + ggml_backend_tensor_set(replay_sg.paged_query_seq_ids, + replay_slots.data(), 0, + sizeof(int32_t) * replay_slots.size()); + ggml_backend_tensor_set(replay_sg.paged_query_positions, + replay_positions.data(), 0, + sizeof(int32_t) * replay_positions.size()); + ggml_backend_tensor_set(replay_sg.logits_row_indices, + replay_logits_rows.data(), 0, + sizeof(int32_t) * replay_logits_rows.size()); + ggml_backend_tensor_set(replay_sg.target_feat_rows, + feature_rows_.data(), 0, + sizeof(int32_t) * feature_rows_.size()); + ggml_backend_tensor_set(b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, + sizeof(int32_t) * seq_lens_.size()); + if (ggml_backend_graph_compute(b_.target_backend_, replay_sg.gf) != + GGML_STATUS_SUCCESS) { + result.error = "DDTree accepted-path replay compute failed"; + return result; + } + std::vector replay_write_slots; + replay_write_slots.reserve(proposals.size()); + for (const Proposal & p : proposals) replay_write_slots.push_back(p.slot); + if (!commit_residency_writes(replay_write_slots)) { + result.error = "DDTree replay KV write commit failed"; + return result; + } + + // The replay is the durable target forward: its recurrent/KV/feature + // state is what the next step consumes. Use its posterior rather than + // the tree-verify posterior so the pending scalar remains exact even if + // the two graph shapes differ numerically. + std::vector replay_next((size_t)active, -1); + ggml_backend_tensor_get( + replay_sg.argmax_tokens, replay_next.data(), 0, + sizeof(int32_t) * replay_next.size()); + for (int s = 0; s < active; ++s) { + if (replay_next[(size_t)s] < 0) { + result.error = "DDTree replay produced an invalid pending token"; + return result; + } + proposals[(size_t)s].bonus = replay_next[(size_t)s]; + } + + for (Proposal & p : proposals) { + slots_.commit_step(p.slot); + std::string reselect_error; + if (!maybe_reselect_residency(p.slot, reselect_error)) { + result.error = reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error; + return result; + } + } + result.decode.reserve((size_t)active); + for (Proposal & p : proposals) { + DecodeOutput out; + out.slot = p.slot; + out.token = p.bonus; + out.ddtree_steps = 1; + out.ddtree_accepted_tokens = p.accepted.size() - 1; + out.target_forwards = 2; + for (size_t i = 1; i < p.accepted.size(); ++i) { + const int dfs = p.accepted[i]; + out.committed_tokens.push_back( + p.tree.token_ids[(size_t)dfs - 1]); + } + attach_residency_telemetry(out); + result.decode.push_back(std::move(out)); + } + return result; +} + bool Qwen35SeqEngine::token_is_eos(int32_t token) const { return b_.token_is_eos(token); } @@ -46,6 +548,14 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( AdmitResult result = slots_.admit(request_id, prompt, sampler); if (result.status == AdmitResult::Status::admitted) { reset_recurrent_slot(b_.cache_, result.slot); + if (slots_.residency_active()) { + slots_.slot(result.slot).kvflash_last_reselect_generated = + -std::max(1, b_.kvflash_tau_); + } + if (result.slot >= 0 && result.slot < (int)slot_draft_kv_.size() && + slot_draft_kv_[(size_t)result.slot]) { + draft_kv_reset(*slot_draft_kv_[(size_t)result.slot]); + } } return result; } @@ -102,6 +612,77 @@ bool Qwen35SeqEngine::upload_block_table_delta( return true; } +bool Qwen35SeqEngine::upload_all_active_block_tables() { + if (!slots_.residency_active()) return true; + ggml_tensor * table = b_.cache_.paged_block_table; + if (!table) return false; + std::vector column((size_t)table->ne[0], -1); + std::vector snapshot; + for (int slot = 0; slot < slots_.slot_count(); ++slot) { + if (!slots_.is_active(slot)) continue; + std::fill(column.begin(), column.end(), -1); + if (!slots_.block_table_snapshot(slot, snapshot) || + snapshot.size() > column.size()) { + return false; + } + std::copy(snapshot.begin(), snapshot.end(), column.begin()); + ggml_backend_tensor_set( + table, column.data(), (size_t)slot * table->nb[1], + sizeof(int32_t) * column.size()); + } + return true; +} + +bool Qwen35SeqEngine::commit_residency_writes( + const std::vector & slots) { + if (!slots_.residency_active()) return true; + // Graph completion is not a host barrier for every backend. Pending pages + // become evictable only after all target writes are device-complete. + ggml_backend_synchronize(b_.target_backend_); + for (int slot : slots) { + if (!slots_.commit_residency_writes(slot)) return false; + } + return true; +} + +bool Qwen35SeqEngine::maybe_reselect_residency( + int slot, std::string & error) { + if (!slots_.residency_active()) return true; + Qwen35Slot & seq = slots_.slot(slot); + const int generated = seq.generated_tokens(); + const int tau = std::max( + b_.kvflash_tau_, (int)(seq.sample_history.size() / 45)); + if (generated - seq.kvflash_last_reselect_generated < tau) return true; + + b_.kvflash_ensure_scorer(); + std::vector scores; + const std::vector * score_ptr = nullptr; + if (b_.kvflash_scorer_) { + if (!b_.kvflash_scorer_->score_chunks( + seq.sample_history, PAGED_BLOCK_SIZE, scores)) { + // Short histories and recoverable drafter failures are expected + // scorer outcomes. Preserve service with the pager's explicit + // recency/LRU policy; only residency or transfer errors below are + // fatal to the request. + std::fprintf(stderr, + "[parallel-kvflash] scorer unavailable for slot %d; using LRU\n", + slot); + } else { + const size_t blocks = (seq.sample_history.size() + + PAGED_BLOCK_SIZE - 1) / PAGED_BLOCK_SIZE; + scores.resize(blocks, scores.empty() ? 0.0f : scores.back()); + score_ptr = &scores; + } + } + if (!slots_.reselect_residency(slot, score_ptr, &error)) return false; + seq.kvflash_last_reselect_generated = generated; + return upload_all_active_block_tables(); +} + +void Qwen35SeqEngine::attach_residency_telemetry(DecodeOutput & out) { + slots_.take_residency_telemetry(out.slot, out); +} + void Qwen35SeqEngine::fail_prefill( int slot, std::vector & prefill_outputs, const char * log_message, const char * client_message) { @@ -135,9 +716,10 @@ Qwen35SeqEngine::PrefillStage Qwen35SeqEngine::stage_prefill_chunk( "prefill K/V allocation failed"); return PrefillStage{}; } - if (!upload_block_table_delta( - slot, chunk.first_new_block, chunk.new_blocks.data(), - chunk.new_blocks.size())) { + const bool table_ok = slots_.residency_active() || upload_block_table_delta( + slot, chunk.first_new_block, chunk.new_blocks.data(), + chunk.new_blocks.size()); + if (!table_ok) { fail_prefill( slot, prefill_outputs, "prefill block-table delta exceeds device capacity", "prefill block-table update failed"); @@ -207,6 +789,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (inputs.empty() && plan.prefills.empty()) return result; + if (ddtree_eligible(plan)) { + std::optional speculative = step_ddtree(plan); + if (speculative) return std::move(*speculative); + // Proposal setup failed before target/cache mutation. Preserve service + // by taking the existing packed AR path for this iteration. + } + const TargetWeights & w = b_.w_; StepGraph & sg = b_.sg_; const int hidden = w.n_embd; @@ -242,9 +831,10 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { output_rows_.push_back(compact_row); continue; } - if (app.new_block >= 0 && - !upload_block_table_delta( - in.slot, app.new_block_index, &app.new_block, 1)) { + const bool table_ok = slots_.residency_active() || app.new_block < 0 || + upload_block_table_delta(in.slot, app.new_block_index, + &app.new_block, 1); + if (!table_ok) { out.error = "decode block-table entry exceeds device capacity"; decode_outputs.push_back(std::move(out)); output_rows_.push_back(compact_row); @@ -278,6 +868,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } prefills.push_back(std::move(prefill)); } + if (!upload_all_active_block_tables()) { + return fail_step("active KVFlash block-table refresh failed"); + } const int live_count = (int)live_tokens_.size(); const bool with_decode = live_count > 0; @@ -329,7 +922,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { built = build_target_step( sg, w, b_.cache_, b_.target_backend_, /*kv_start=*/0, /*n_tokens=*/n_total, - /*with_mask=*/false, /*capture=*/false, + /*with_mask=*/false, /*capture=*/capture_features_, /*capture_delta_intermediate=*/false, /*fa_window=*/0, /*logits_tail_rows=*/0, b_.cfg_.kq_stride_pad, @@ -347,7 +940,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { built = build_target_step( sg, w, b_.cache_, b_.target_backend_, /*kv_start=*/0, /*n_tokens=*/decode_bucket, - /*with_mask=*/false, /*capture=*/false, + /*with_mask=*/false, /*capture=*/capture_features_, /*capture_delta_intermediate=*/false, /*fa_window=*/0, /*logits_tail_rows=*/0, b_.cfg_.kq_stride_pad, @@ -365,6 +958,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { /*compact_slots=*/true); } if (!built || !sg.kv_write_rows || + (capture_features_ && !sg.target_feat_rows) || (with_prefill && (!sg.paged_query_seq_ids || !sg.paged_query_positions || !sg.logits_row_indices))) { @@ -428,6 +1022,30 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { b_.target_backend_, sg.kv_write_rows, rows_buf_.data(), 0, sizeof(int64_t) * rows_buf_.size()); + if (capture_features_) { + const int cap = b_.cache_.target_feat_cap; + const int dead_row = cap * n_slots; + feature_rows_.assign((size_t)n_total, dead_row); + int feature_offset = 0; + for (size_t i = 0; i < prefills.size(); ++i) { + const PrefillStage & prefill = prefills[i]; + const int slot = plan.prefills[i].slot; + for (int row = 0; row < prefill.chunk; ++row) { + feature_rows_[(size_t)(feature_offset + row)] = + slot * cap + (prefill.kv_pos + row) % cap; + } + feature_offset += prefill.chunk; + } + for (int row = 0; row < live_count; ++row) { + feature_rows_[(size_t)(n_prefill + row)] = + live_slot_ids_[(size_t)row] * cap + + live_positions_[(size_t)row] % cap; + } + ggml_backend_tensor_set_async( + b_.target_backend_, sg.target_feat_rows, feature_rows_.data(), 0, + sizeof(int32_t) * feature_rows_.size()); + } + if (with_prefill) { query_slot_ids_.assign((size_t)n_total, -1); query_positions_.assign((size_t)n_total, -1); @@ -507,6 +1125,16 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { "qwen35.argmax_readback", roctx_metadata); ggml_backend_synchronize(b_.target_backend_); } + std::vector write_slots; + write_slots.reserve(live_slot_ids_.size() + prefills.size()); + write_slots.insert(write_slots.end(), + live_slot_ids_.begin(), live_slot_ids_.end()); + for (size_t i = 0; i < prefills.size(); ++i) { + write_slots.push_back(plan.prefills[i].slot); + } + if (!commit_residency_writes(write_slots)) { + return fail_step("KVFlash pending write commit failed"); + } for (size_t oi = 0; oi < inputs.size(); ++oi) { DecodeOutput & out = decode_outputs[oi]; @@ -517,6 +1145,17 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { out.slot, row, &argmax_buf_[(size_t)row], &logits_buf_); } + for (DecodeOutput & out : decode_outputs) { + if (out.failed) continue; + std::string reselect_error; + if (!maybe_reselect_residency(out.slot, reselect_error)) { + return fail_step(reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error); + } + out.target_forwards = 1; + attach_residency_telemetry(out); + } + int commit_row = 0; for (size_t i = 0; i < prefills.size(); ++i) { const int slot = plan.prefills[i].slot; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index d9391784c..75275df39 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -22,10 +22,15 @@ #pragma once #include "common/concurrency/seq_engine.h" +#include "common/dflash_draft_kv.h" +#include "common/dflash_feature_ring.h" +#include "common/ddtree.h" #include "qwen35_slot_manager.h" #include #include +#include +#include #include namespace dflash::common { @@ -40,19 +45,14 @@ class Qwen35SeqEngine final : public SeqEngine { // `max_prefills` bounds scheduler-selected prompt slices per traversal. Qwen35SeqEngine(Qwen35Backend & backend, PagedKvPool & pool, int max_ctx, int64_t scratch_row, - int max_prefills = 8, + int tree_width = 0, int tree_scratch_base = 0, + int tree_scratch_stride = 0, int max_prefills = 8, int mixed_prefill_tokens = 2048, int long_mixed_prefill_tokens = 4096, int long_prefill_threshold = 768, int idle_prefill_tokens = 4096, - int prefill_quantum = 512) - : max_prefills_(std::max(1, max_prefills)), - mixed_prefill_tokens_(std::max(1, mixed_prefill_tokens)), - long_mixed_prefill_tokens_(std::max(1, long_mixed_prefill_tokens)), - long_prefill_threshold_(std::max(1, long_prefill_threshold)), - idle_prefill_tokens_(std::max(1, idle_prefill_tokens)), - prefill_quantum_(std::max(1, prefill_quantum)), b_(backend), - slots_(pool, max_ctx), scratch_row_(scratch_row) {} + int prefill_quantum = 512); + ~Qwen35SeqEngine() override; int slot_count() const override { return slots_.slot_count(); } int max_context() const override { return slots_.max_context(); } @@ -104,6 +104,10 @@ class Qwen35SeqEngine final : public SeqEngine { bool upload_block_table_delta(int slot, int first_block, const int32_t * blocks, size_t count); + bool upload_all_active_block_tables(); + bool commit_residency_writes(const std::vector & slots); + bool maybe_reselect_residency(int slot, std::string & error); + void attach_residency_telemetry(DecodeOutput & out); void fail_prefill(int slot, std::vector & outputs, const char * log_message, const char * client_message); @@ -112,10 +116,23 @@ class Qwen35SeqEngine final : public SeqEngine { int32_t sample_graph_row(int slot, int logits_row, const int32_t * cached_argmax = nullptr, std::vector * logits_scratch = nullptr); + DraftFeatureMirror * slot_feature_mirror(int slot); + DraftKvState * ensure_slot_draft_kv(int slot); + bool ddtree_eligible(const StepPlan & plan) const; + // nullopt means proposal setup failed before target/cache mutation and the + // caller may safely use the ordinary packed AR path for this iteration. + std::optional step_ddtree(const StepPlan & plan); Qwen35Backend & b_; Qwen35SlotManager slots_; int64_t scratch_row_ = 0; + int tree_width_ = 0; + int tree_scratch_base_ = 0; + int tree_scratch_stride_ = 0; + bool capture_features_ = false; + ggml_context * feature_view_ctx_ = nullptr; + std::vector slot_feature_mirrors_; + std::vector> slot_draft_kv_; // Hoisted per-step buffers (reused across step() calls). std::vector output_rows_; @@ -131,6 +148,7 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector query_slot_ids_; std::vector query_positions_; std::vector logits_rows_; + std::vector feature_rows_; std::vector embed_buf_; std::vector pos_buf_; std::vector rows_buf_; diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index b8dcc36ec..f740a283b 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -5,8 +5,12 @@ namespace dflash::common { -Qwen35SlotManager::Qwen35SlotManager(PagedKvPool & pool, int max_ctx) - : pool_(pool), max_ctx_(max_ctx) { +Qwen35SlotManager::Qwen35SlotManager( + PagedKvPool & pool, int max_ctx, int speculative_headroom, + PagedKvResidencyManager * residency) + : pool_(pool), max_ctx_(max_ctx), + headroom_tokens_(std::max(pool.block_size(), speculative_headroom)), + residency_(residency) { slots_.assign(pool.max_sequences(), Qwen35Slot{}); } @@ -21,7 +25,7 @@ int Qwen35SlotManager::decoding_count() const { uint32_t Qwen35SlotManager::decode_headroom_capacity(int logical_tokens) const { const uint64_t extended = static_cast(std::max(0, logical_tokens)) + - pool_.block_size(); + static_cast(headroom_tokens_); return static_cast(std::min( static_cast(max_ctx_), extended)); } @@ -83,6 +87,20 @@ bool Qwen35SlotManager::is_active(int slot) const { bool Qwen35SlotManager::is_prefilling(int slot) const { return is_active(slot) && slots_[(size_t)slot].prefilling(); } +void Qwen35SlotManager::accumulate_residency_delta( + Qwen35Slot & slot, const PagedKvResidencyStats & before) { + if (!residency_) return; + const PagedKvResidencyStats after = residency_->stats(); + if (after.page_ins >= before.page_ins) { + slot.kvflash_page_ins_pending += after.page_ins - before.page_ins; + } + if (after.page_outs >= before.page_outs) { + slot.kvflash_page_outs_pending += after.page_outs - before.page_outs; + } + if (after.reselects >= before.reselects) { + slot.kvflash_reselects_pending += after.reselects - before.reselects; + } +} bool Qwen35SlotManager::has_prefill_prompt_at_least(int tokens) const { if (tokens <= 0) return true; @@ -112,7 +130,7 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( // 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) { + if (!residency_ && (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) + @@ -132,7 +150,8 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( // 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(); + const PagedKvStatus headroom_status = residency_ + ? PagedKvStatus::Ok : protect_decode_headroom(); if (headroom_status != PagedKvStatus::Ok) { r.status = headroom_status == PagedKvStatus::BlocksExhausted ? AdmitStatus::busy : AdmitStatus::failed; @@ -143,7 +162,7 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( } PagedKvSequenceHandle handle; - uint32_t reservation_capacity = + uint32_t reservation_capacity = residency_ ? 0 : decode_headroom_capacity(prompt_len); if (!capacity_fits_pool(reservation_capacity)) { // The prompt itself fits, but this physical pool can never hold its @@ -151,8 +170,9 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( // 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); + const PagedKvStatus status = residency_ + ? pool_.acquire(request_id, handle) + : pool_.acquire_reserved(request_id, reservation_capacity, handle); if (status != PagedKvStatus::Ok) { r.status = status == PagedKvStatus::SequenceSlotsExhausted || status == PagedKvStatus::BlocksExhausted @@ -163,6 +183,17 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( return r; } + if (residency_) { + const PagedKvResidencyStatus resident_status = + residency_->register_sequence(handle); + if (resident_status != PagedKvResidencyStatus::Ok) { + (void)pool_.release(handle); + r.error = std::string("KV residency registration failed: ") + + paged_kv_residency_status_string(resident_status); + return r; + } + } + Qwen35Slot & s = slots_[(size_t)slot]; s.phase = Qwen35SlotPhase::prefill; s.handle = handle; @@ -194,7 +225,22 @@ Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( return out; } - PagedKvAppendResult app = pool_.append(s.handle, (uint32_t)n_tokens); + PagedKvAppendResult app; + if (residency_) { + const PagedKvResidencyStats before = residency_->stats(); + PagedKvResidentAppendResult resident = residency_->append( + s.handle, (uint32_t)n_tokens); + accumulate_residency_delta(s, before); + if (!resident) { + std::fprintf(stderr, + "[parallel-kvflash] prefill append failed for slot %d: %s\n", + slot, paged_kv_residency_status_string(resident.status)); + return out; + } + app = std::move(resident.pool_result); + } else { + 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 @@ -218,6 +264,17 @@ Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( out.new_blocks.push_back((int32_t)write.physical_block); } } + if (residency_) { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(s.handle, snapshot) != PagedKvStatus::Ok) { + return out; + } + out.full_block_table.reserve(snapshot.block_table.size()); + for (uint32_t block : snapshot.block_table) { + out.full_block_table.push_back( + block == PAGED_KV_COLD_BLOCK ? -1 : (int32_t)block); + } + } s.cur_pos += n_tokens; out.ok = true; return out; @@ -230,43 +287,189 @@ void Qwen35SlotManager::commit_prefill(int slot) { s.phase = Qwen35SlotPhase::decode; } -Qwen35SlotManager::StepAppend Qwen35SlotManager::append_token(int slot, - int32_t fed_token) { +Qwen35SlotManager::StepAppend Qwen35SlotManager::append_tokens( + int slot, const int32_t * fed_tokens, int n_tokens) { StepAppend out; - if (!is_active(slot) || !slots_[(size_t)slot].decoding()) return out; + if (!is_active(slot) || !slots_[(size_t)slot].decoding() || + !fed_tokens || n_tokens < 1) { + return out; + } Qwen35Slot & s = slots_[(size_t)slot]; - if (s.cur_pos >= max_ctx_) { - // No context left; the scheduler should have stopped this slot. + if (!s.staged_tokens.empty() || s.cur_pos > max_ctx_ || + n_tokens > max_ctx_ - s.cur_pos) { 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) { + + PagedKvAppendResult app; + if (residency_) { + const PagedKvResidencyStats before = residency_->stats(); + PagedKvResidentAppendResult resident = residency_->append( + s.handle, static_cast(n_tokens)); + accumulate_residency_delta(s, before); + if (!resident) { + out.busy = resident.status == + PagedKvResidencyStatus::PoolExhausted || + resident.status == PagedKvResidencyStatus::NoEvictableBlock; + return out; + } + app = std::move(resident.pool_result); + } else { + app = pool_.append(s.handle, static_cast(n_tokens)); + } + if (!app || app.token_count != static_cast(n_tokens)) { out.busy = app.status == PagedKvStatus::BlocksExhausted; return out; } - s.sample_history.push_back(fed_token); + if (app.write_slots.size() != static_cast(n_tokens) || + app.write_slots.front().logical_position != + static_cast(s.cur_pos) || + app.write_slots.back().logical_position != + static_cast(s.cur_pos + n_tokens - 1)) { + // Pool success guarantees this shape. Retain staged ownership so a + // fatal caller retires the sequence rather than double-appending. + s.staged_tokens.assign(fed_tokens, fed_tokens + n_tokens); + return out; + } + + out.physical_rows.reserve(app.write_slots.size()); + for (const PagedKvWriteSlot & write : app.write_slots) { + out.physical_rows.push_back( + static_cast(write.physical_token_index)); + if (write.block_offset == 0) { + if (out.first_new_block < 0) { + out.first_new_block = static_cast( + write.logical_position / pool_.block_size()); + } + out.new_blocks.push_back( + static_cast(write.physical_block)); + } + } + s.staged_tokens.assign(fed_tokens, fed_tokens + n_tokens); + if (residency_) { + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(s.handle, snapshot) != PagedKvStatus::Ok) { + return out; + } + out.full_block_table.reserve(snapshot.block_table.size()); + for (uint32_t block : snapshot.block_table) { + out.full_block_table.push_back( + block == PAGED_KV_COLD_BLOCK ? -1 : (int32_t)block); + } + } out.ok = true; - out.physical_row = (int64_t)app.last.physical_token_index; + out.count = n_tokens; 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(); + if (n_tokens == 1) { + out.physical_row = out.physical_rows.front(); + if (!out.new_blocks.empty()) { + out.new_block = out.new_blocks.front(); + out.new_block_index = out.first_new_block; + } } return out; } +Qwen35SlotManager::StepAppend Qwen35SlotManager::append_token( + int slot, int32_t fed_token) { + return append_tokens(slot, &fed_token, 1); +} + void Qwen35SlotManager::commit_step(int slot) { if (!is_active(slot)) return; - slots_[(size_t)slot].cur_pos += 1; + Qwen35Slot & s = slots_[(size_t)slot]; + if (s.staged_tokens.empty()) return; + s.sample_history.insert( + s.sample_history.end(), s.staged_tokens.begin(), + s.staged_tokens.end()); + s.cur_pos += static_cast(s.staged_tokens.size()); + s.staged_tokens.clear(); +} + +bool Qwen35SlotManager::block_table_snapshot( + int slot, std::vector & out) const { + out.clear(); + if (!is_active(slot)) return false; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(slots_[(size_t)slot].handle, snapshot) != + PagedKvStatus::Ok) { + return false; + } + out.reserve(snapshot.block_table.size()); + for (uint32_t block : snapshot.block_table) { + out.push_back(block == PAGED_KV_COLD_BLOCK ? -1 : (int32_t)block); + } + return true; +} + +bool Qwen35SlotManager::commit_residency_writes(int slot) { + if (!residency_) return true; + if (!is_active(slot)) return false; + return residency_->commit_pending_writes(slots_[(size_t)slot].handle) == + PagedKvResidencyStatus::Ok; +} + +bool Qwen35SlotManager::reselect_residency( + int slot, const std::vector * scores, std::string * error) { + if (!residency_) return true; + if (!is_active(slot)) { + if (error) *error = "inactive KVFlash slot"; + return false; + } + Qwen35Slot & s = slots_[(size_t)slot]; + const PagedKvResidencyStats before = residency_->stats(); + const std::vector no_scores; + PagedKvResidencyStatus status = residency_->set_scores( + s.handle, scores ? *scores : no_scores); + if (status == PagedKvResidencyStatus::Ok) { + status = residency_->reselect(s.handle); + } + accumulate_residency_delta(s, before); + if (status != PagedKvResidencyStatus::Ok) { + if (error) { + *error = std::string("KVFlash reselect failed: ") + + paged_kv_residency_status_string(status); + } + return false; + } + return true; +} + +void Qwen35SlotManager::take_residency_telemetry( + int slot, SeqEngine::DecodeOutput & out) { + if (!residency_ || !is_active(slot)) return; + Qwen35Slot & s = slots_[(size_t)slot]; + out.kvflash_page_ins = s.kvflash_page_ins_pending; + out.kvflash_page_outs = s.kvflash_page_outs_pending; + out.kvflash_reselects = s.kvflash_reselects_pending; + uint32_t resident = 0; + if (pool_.resident_block_count(s.handle, resident) == PagedKvStatus::Ok) { + out.kvflash_resident_blocks = resident; + } + s.kvflash_page_ins_pending = 0; + s.kvflash_page_outs_pending = 0; + s.kvflash_reselects_pending = 0; } void Qwen35SlotManager::retire(int slot) { if (slot < 0 || slot >= (int)slots_.size()) return; Qwen35Slot & s = slots_[(size_t)slot]; if (!s.active()) return; + if (residency_) { + const PagedKvResidencyStatus resident_status = + residency_->forget_sequence(s.handle); + if (resident_status != PagedKvResidencyStatus::Ok && + resident_status != PagedKvResidencyStatus::StaleHandle && + resident_status != PagedKvResidencyStatus::SequenceNotRegistered) { + std::fprintf(stderr, + "[parallel-kvflash] slot %d residency release failed: %s\n", + slot, paged_kv_residency_status_string(resident_status)); + // A failed copy-stream barrier leaves physical pages in flight. + // Keep the slot and its pool handle intact so a later retirement + // can retry forget_sequence without recycling those pages. + 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", diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index 1da009f69..dd077496c 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -17,6 +17,7 @@ #pragma once #include "common/concurrency/paged_kv_pool.h" +#include "common/concurrency/paged_kv_residency.h" #include "common/sampler.h" #include "common/concurrency/seq_engine.h" @@ -45,6 +46,17 @@ struct Qwen35Slot { // Penalty history is recorded as fed rather than sampled: the scheduler // may override a sample before the model consumes it. std::vector sample_history; + // Decode rows allocated from the pool but not yet made durable by a + // successful target forward. A step stages one contiguous token range + // per slot, then commit_step() publishes all of it atomically. + std::vector staged_tokens; + + // Residency operation deltas are attributed to the requesting slot and + // held across prefill until a DecodeOutput can carry them upstream. + uint64_t kvflash_page_ins_pending = 0; + uint64_t kvflash_page_outs_pending = 0; + uint64_t kvflash_reselects_pending = 0; + int kvflash_last_reselect_generated = 0; int generated_tokens() const { return sample_history.size() > (size_t)prompt_len @@ -61,7 +73,13 @@ class Qwen35SlotManager { public: // `max_ctx` is the per-sequence logical bound; slot count comes from the // pool's max_sequences. The pool must outlive the manager. - Qwen35SlotManager(PagedKvPool & pool, int max_ctx); + Qwen35SlotManager(PagedKvPool & pool, int max_ctx, + int speculative_headroom = 1, + PagedKvResidencyManager * residency = nullptr); + Qwen35SlotManager(const Qwen35SlotManager &) = delete; + Qwen35SlotManager & operator=(const Qwen35SlotManager &) = delete; + Qwen35SlotManager(Qwen35SlotManager &&) = delete; + Qwen35SlotManager & operator=(Qwen35SlotManager &&) = delete; // 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 @@ -80,6 +98,7 @@ class Qwen35SlotManager { // Delta to patch into the slot's device block-table column. std::vector new_blocks; int first_new_block = -1; + std::vector full_block_table; }; // Append `n_tokens` more prompt rows for a prefilling slot. Physical block @@ -93,17 +112,37 @@ class Qwen35SlotManager { struct StepAppend { bool ok = false; bool busy = false; // no physical block available right now + std::vector physical_rows; + std::vector new_blocks; + int first_new_block = -1; + int count = 0; + // Compatibility fields for the common one-token append. int64_t physical_row = -1; - int position = -1; // logical position the fed token is written at + int position = -1; // logical position of the first staged token int32_t new_block = -1; int new_block_index = -1; + std::vector full_block_table; }; - // 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(). + // Atomically allocate and stage a contiguous accepted path. The pool + // append is all-or-nothing; sample_history and cur_pos remain unchanged + // until commit_step(). A slot may have only one staged range at a time. + StepAppend append_tokens(int slot, const int32_t * fed_tokens, + int n_tokens); + + // Complete residency snapshots encode cold logical pages as -1. The + // engine pads each device column before upload. + bool residency_active() const { return residency_ != nullptr; } + bool block_table_snapshot(int slot, std::vector & out) const; + bool commit_residency_writes(int slot); + bool reselect_residency(int slot, const std::vector * scores, + std::string * error = nullptr); + void take_residency_telemetry(int slot, SeqEngine::DecodeOutput & out); + + // One-token compatibility wrapper used by ordinary autoregressive decode. StepAppend append_token(int slot, int32_t fed_token); - // The batched step's compute succeeded: cur_pos++. + // Publish every staged token after a successful target forward. void commit_step(int slot); // Release the slot's blocks and clear its state. Safe on inactive slots @@ -111,6 +150,8 @@ class Qwen35SlotManager { void retire(int slot); int slot_count() const { return (int)slots_.size(); } + void accumulate_residency_delta(Qwen35Slot & slot, + const PagedKvResidencyStats & before); int max_context() const { return max_ctx_; } int decoding_count() const; bool is_active(int slot) const; @@ -131,6 +172,8 @@ class Qwen35SlotManager { PagedKvPool & pool_; int max_ctx_ = 0; + int headroom_tokens_; + PagedKvResidencyManager * residency_ = nullptr; std::vector slots_; }; diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index ad3d56b58..0ca65ecdb 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -9,6 +9,44 @@ namespace dflash::common { +bool detail::validate_target_paged_tree_layout( + const TargetCache & cache, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride) { + static constexpr int tree_buckets[] = { + 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, + }; + if (tree_width < 1 || + std::find(std::begin(tree_buckets), std::end(tree_buckets), + n_tree_seqs) == std::end(tree_buckets) || + cache.n_seq_slots <= 1 || !cache.paged_block_table || + !cache.paged_kv_seq_lens || paged_max_kv_len < 1 || + tree_scratch_base <= 0 || + tree_scratch_base % PAGED_BLOCK_SIZE != 0 || + tree_scratch_stride < tree_width) { + return false; + } + + int physical_kv_rows = 0; + for (ggml_tensor * tensor : cache.attn_k) { + if (tensor) { + physical_kv_rows = (int)tensor->ne[1]; + break; + } + } + if (physical_kv_rows < 1) return false; + + const int64_t scratch_end = + (int64_t)tree_scratch_base + + (int64_t)(cache.n_seq_slots - 1) * tree_scratch_stride + + tree_width; + return scratch_end <= physical_kv_rows && + (int64_t)tree_width * n_tree_seqs <= INT32_MAX; +} + // ── build_layer_step ──────────────────────────────────────────── bool build_layer_step( @@ -490,6 +528,12 @@ bool build_target_step( ggml_set_name(sg.logits_row_indices, "logits_row_indices"); ggml_set_input(sg.logits_row_indices); } + if (capture && paged_attention && cache.target_feat) { + sg.target_feat_rows = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + ggml_set_name(sg.target_feat_rows, "target_feat_rows"); + ggml_set_input(sg.target_feat_rows); + } sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); @@ -539,6 +583,7 @@ bool build_target_step( gi.paged_query_seq_ids = sg.paged_query_seq_ids; gi.paged_query_positions = sg.paged_query_positions; gi.logits_row_indices = sg.logits_row_indices; + gi.target_feat_rows = sg.target_feat_rows; gi.prefill_segments = prefill_segments; gi.n_prefill_segments = n_prefill_segments; @@ -630,6 +675,117 @@ bool build_target_step_tree( return ggml_gallocr_alloc_graph(sg.alloc, sg.gf); } +// ── build_target_step_paged_tree ──────────────────────────────── + +bool build_target_step_paged_tree( + StepGraph & sg, + const TargetWeights & w, + TargetCache & cache, + ggml_backend_t backend, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride, + int kq_stride_pad) { + (void)kq_stride_pad; + step_graph_free(sg); + + if (!detail::validate_target_paged_tree_layout( + cache, tree_width, n_tree_seqs, paged_max_kv_len, + tree_scratch_base, tree_scratch_stride)) { + return false; + } + const int n_tokens = tree_width * n_tree_seqs; + + ggml_init_params ip{}; + ip.mem_size = 512 * 1024 * 1024; + static thread_local std::vector g_tree_arena; + if (g_tree_arena.size() < ip.mem_size) g_tree_arena.resize(ip.mem_size); + ip.mem_buffer = g_tree_arena.data(); + ip.no_alloc = true; + sg.ctx = ggml_init(ip); + if (!sg.ctx) return false; + + // Salt graph addresses by the stable bucket shape so captured graphs for + // different T/S buckets never alias in ggml-cuda's topology cache. + for (int i = 0; i < tree_width + n_tree_seqs; ++i) { + (void)ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 1); + } + + sg.inp_embed = ggml_new_tensor_3d( + sg.ctx, GGML_TYPE_F32, w.n_embd, n_tokens, 1); + sg.positions = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 4 * n_tokens); + sg.parent_ids = ggml_new_tensor_2d( + sg.ctx, GGML_TYPE_I32, tree_width, n_tree_seqs); + sg.tree_sizes = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tree_seqs); + sg.active_slot_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tree_seqs); + sg.state_slot_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tree_seqs); + sg.paged_query_seq_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + sg.kv_write_rows = ggml_new_tensor_2d( + sg.ctx, GGML_TYPE_I64, n_tokens, w.n_head_kv); + + const struct NamedInput { + ggml_tensor * tensor; + const char * name; + } inputs[] = { + {sg.inp_embed, "inp_embed"}, + {sg.positions, "positions"}, + {sg.parent_ids, "parent_ids"}, + {sg.tree_sizes, "tree_sizes"}, + {sg.active_slot_ids, "active_slot_ids"}, + {sg.state_slot_ids, "state_slot_ids"}, + {sg.paged_query_seq_ids, "paged_query_seq_ids"}, + {sg.kv_write_rows, "kv_write_rows"}, + }; + for (const NamedInput & input : inputs) { + ggml_set_name(input.tensor, input.name); + ggml_set_input(input.tensor); + } + + sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); + QwenGraphInputs gi{}; + gi.inp_embed = sg.inp_embed; + gi.positions = sg.positions; + gi.n_tokens = n_tokens; + gi.kv_start = 0; + gi.capture_layers = false; + gi.capture_delta_intermediate = false; + gi.parent_ids = sg.parent_ids; + gi.tree_sizes = sg.tree_sizes; + gi.kv_write_rows = sg.kv_write_rows; + gi.paged_block_table = cache.paged_block_table; + gi.paged_kv_seq_lens = cache.paged_kv_seq_lens; + gi.active_slot_ids = sg.active_slot_ids; + gi.state_slot_ids = sg.state_slot_ids; + gi.paged_query_seq_ids = sg.paged_query_seq_ids; + gi.n_seqs = n_tree_seqs; + gi.paged_max_kv_len = paged_max_kv_len; + gi.tree_width = tree_width; + gi.tree_scratch_base = tree_scratch_base; + gi.tree_scratch_stride = tree_scratch_stride; + + QwenGraphOutputs go = build_qwen35_graph(sg.ctx, sg.gf, w, cache, gi); + if (!go.logits) return false; + sg.logits = go.logits; + ggml_set_output(sg.logits); + sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); + ggml_set_name(sg.argmax_tokens, "paged_tree_verify_argmax"); + ggml_set_output(sg.argmax_tokens); + ggml_build_forward_expand(sg.gf, sg.argmax_tokens); + + if (!sg.alloc) { + sg.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + } + return ggml_gallocr_alloc_graph(sg.alloc, sg.gf); +} + // ── build_lm_head_projection_step ─────────────────────────────── diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index cdbaf75ed..fac406863 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -23,6 +23,22 @@ namespace dflash::common { +namespace detail { + +// Model-free validation shared by the packed-tree builder and its shape +// tests. paged_max_kv_len is a logical launch bound and may exceed the +// bounded physical K/V pool; only the per-slot scratch slabs must fit in the +// physical tensor rows. +bool validate_target_paged_tree_layout( + const TargetCache & cache, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride); + +} // namespace detail + // Layer-segmented prefill: process one target layer for chunk_start..chunk_start+n_tokens. bool build_layer_step( StepGraph & sg, @@ -109,6 +125,10 @@ bool build_hybrid_full_layer_step( // overrides logits_tail_rows. Multi-prompt steps need it because // committing rows are scattered. 0 keeps the tail-view behavior. // `logits_tail_rows` — logits/argmax only for the last n rows (0 = all). +// When `capture && paged_attention`, sg.target_feat_rows is an I32 graph +// input mapping every token to its slot-local feature-ring destination. This +// keeps accepted-path replay graph-stable and leaves legacy offset capture +// unchanged for callers that do not use paged serving. bool build_target_step( StepGraph & sg, const TargetWeights & w, @@ -146,6 +166,26 @@ bool build_target_step_tree( int fa_window = 0, int kq_stride_pad = KQ_MASK_PAD); +// Packed concurrent DDTree verify over a paged multi-slot cache. Tokens are +// flattened sequence-major as [tree_width*n_tree_seqs]. n_tree_seqs is a +// stable graph-bucket width; inactive trees use tree_size=0 and dead/safe row +// mappings. In particular state_slot_ids padding must map to a valid harmless +// slot (normally 0), while active/paged sequence IDs may use -1. The graph +// writes candidate K/V into per-slot scratch slabs but +// does not mutate persistent recurrent state or target features. Accepted +// paths are committed by a later row-indexed replay through build_target_step. +bool build_target_step_paged_tree( + StepGraph & sg, + const TargetWeights & w, + TargetCache & cache, + ggml_backend_t backend, + int tree_width, + int n_tree_seqs, + int paged_max_kv_len, + int tree_scratch_base, + int tree_scratch_stride, + int kq_stride_pad = KQ_MASK_PAD); + // LM-head projection: project draft hidden states through the target output matrix. bool build_lm_head_projection_step( StepGraph & sg, diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 490f86d3b..c640d778f 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -173,7 +173,7 @@ static FILE * open_dflash_floor_log() { // staging K/V or staging recurrent slab to reserve. static int64_t concurrent_fixed_cache_bytes( const TargetWeights & w, int max_ctx, int n_slots, - int64_t kv_bytes_per_token) { + int64_t kv_bytes_per_token, int64_t scratch_tokens) { const int64_t n_full_attn = w.n_layer / w.full_attention_interval; const int64_t n_delta = w.n_layer - n_full_attn; @@ -188,7 +188,8 @@ static int64_t concurrent_fixed_cache_bytes( state_per_layer * n_delta * (int64_t)n_slots; const int64_t target_feat = (int64_t)w.n_capture_layers * w.n_embd * - std::min(max_ctx, 4096) * (int64_t)sizeof(uint16_t); + ((int64_t)std::min(max_ctx, 4096) * n_slots + 1) * + (int64_t)sizeof(uint16_t); const int64_t q_capture = (int64_t)w.n_embd_head_k * w.n_head * n_full_attn * (int64_t)sizeof(float); @@ -196,7 +197,7 @@ static int64_t concurrent_fixed_cache_bytes( ((int64_t)paged_block_count(max_ctx) * n_slots + n_slots) * (int64_t)sizeof(int32_t); const int64_t scratch = - kv_bytes_per_token * PAGED_BLOCK_SIZE; + kv_bytes_per_token * scratch_tokens; return recurrent + target_feat + q_capture + paged_metadata + scratch; } @@ -331,10 +332,20 @@ bool Qwen35Backend::init() { } } + // Feature-gate validation normally catches this before construction, but + // keep backend arithmetic safe for direct/test callers too. + if (cfg_.ddtree_mode && + (cfg_.ddtree_budget < 1 || cfg_.ddtree_budget > 255)) { + set_last_error("--ddtree-budget must be in [1, 255]"); + return false; + } + // Create KV cache const int max_verify_tokens = cfg_.ddtree_mode ? std::max(dw_.block_size, cfg_.ddtree_budget + 1) : dw_.block_size; + const int n_slots = concurrent_slots(); + // kvflash (bounded residency): pool size from the env, rounded/floored/ // clamped by the shared reader (256-stride keeps FA vec-kernel // eligibility; the floor keeps eviction from deadlocking). @@ -366,16 +377,16 @@ bool Qwen35Backend::init() { if (!post_kvflash_init_gate()) return false; // KVFlash is resolved from env at init; this is the authoritative // paged×KVFlash compatibility check. - if (cfg_.paged_attention && kvflash_active()) { + if (cfg_.paged_attention && kvflash_active() && n_slots <= 1) { std::fprintf(stderr, - "[paged-attention] cannot be combined with KVFlash " - "(resident pool %d tokens)\n", kvflash_tokens_); - set_last_error("paged attention cannot be combined with KVFlash"); + "[paged-attention] single-sequence paged decode cannot be combined " + "with KVFlash (resident pool %d tokens)\n", kvflash_tokens_); + set_last_error( + "paged KVFlash requires --max-concurrency greater than 1"); return false; } // Paged mode sizes the KV cache to whole blocks; otherwise KVFlash // decides the allocation (0 = full max_ctx). - const int n_slots = concurrent_slots(); const int max_concurrent_prefills = n_slots > 1 ? std::clamp( env_int_or_default("DFLASH_MAX_CONCURRENT_PREFILLS", 8), @@ -395,6 +406,17 @@ bool Qwen35Backend::init() { set_last_error("--max-concurrency requires --paged-attention"); return false; } + const bool concurrent_local_ddtree = + n_slots > 1 && cfg_.ddtree_mode && cfg_.draft_path && + !use_remote_draft && !tensor_parallel && !split_gpus_ && + target_backend_ == draft_backend_; + const int tree_width = concurrent_local_ddtree + ? cfg_.ddtree_budget + 1 : 0; + const int tree_stride = concurrent_local_ddtree + ? paged_token_capacity(tree_width) : 0; + const int64_t concurrent_scratch_tokens = + (int64_t)n_slots * tree_stride + PAGED_BLOCK_SIZE; + // Concurrent slots share one physical pool. An explicit // --kv-pool-tokens is rounded up to a whole block; otherwise capacity is // derived from device-free memory after subtracting fixed concurrent cache @@ -404,10 +426,17 @@ bool Qwen35Backend::init() { // pool's index space) as the write target of dead decode-batch rows. int64_t pool_tokens = 0; if (n_slots > 1) { - if (cfg_.kv_pool_tokens > 0) { + if (kvflash_active()) { + pool_tokens = paged_token_capacity(kvflash_tokens_); + std::fprintf(stderr, + "[parallel-kvflash] physical resident pool %lld tokens; " + "logical per-slot cap %d across %d slots " + "(--kv-pool-tokens does not expand resident VRAM)\n", + (long long)pool_tokens, cfg_.device.max_ctx, n_slots); + } else if (cfg_.kv_pool_tokens > 0) { pool_tokens = (int64_t)paged_token_capacity( (int)std::min( - cfg_.kv_pool_tokens, INT32_MAX - PAGED_BLOCK_SIZE)); + cfg_.kv_pool_tokens, INT32_MAX - concurrent_scratch_tokens)); } else { PagedKvAutoBudget budget; // TODO: Size tensor-parallel pools from each device's free memory @@ -416,7 +445,8 @@ bool Qwen35Backend::init() { budget.bytes_per_token = kvf_budget.bytes_per_token; budget.reserve_bytes = kvf_budget.reserve_bytes; budget.fixed_cache_bytes = concurrent_fixed_cache_bytes( - w_, cfg_.device.max_ctx, n_slots, budget.bytes_per_token); + w_, cfg_.device.max_ctx, n_slots, budget.bytes_per_token, + concurrent_scratch_tokens); pool_tokens = paged_kv_auto_pool_tokens( cfg_.device.max_ctx, n_slots, budget); const int64_t one_context = @@ -438,19 +468,19 @@ bool Qwen35Backend::init() { return false; } } - if (pool_tokens + PAGED_BLOCK_SIZE > INT32_MAX) { + if (pool_tokens + concurrent_scratch_tokens > INT32_MAX) { set_last_error("paged KV pool exceeds INT32_MAX tokens"); return false; } } const int ctx_alloc = n_slots > 1 - ? (int)(pool_tokens + PAGED_BLOCK_SIZE) + ? (int)(pool_tokens + concurrent_scratch_tokens) : (cfg_.paged_attention ? paged_token_capacity(cfg_.device.max_ctx) : kvflash_tokens_); if (!create_target_cache(w_, cfg_.device.max_ctx, max_verify_tokens, target_backend_, cache_, /*prefill_only=*/true, ctx_alloc, - cfg_.paged_attention, n_slots)) { + cfg_.paged_attention, n_slots, concurrent_local_ddtree)) { std::fprintf(stderr, "cache: %s\n", dflash27b_last_error()); return false; } @@ -481,10 +511,40 @@ bool Qwen35Backend::init() { e.what()); return false; } + if (n_slots > 1 && kvflash_active()) { + std::string transfer_error; + paged_kv_transfer_ = QwenPagedKvResidencyTransfer::create( + cache_, target_backend_, cfg_.device.gpu, PAGED_BLOCK_SIZE, + &transfer_error); + if (!paged_kv_transfer_) { + set_last_error("concurrent KVFlash transfer init failed: " + + transfer_error); + return false; + } + PagedKvResidencyConfig residency_cfg; + residency_cfg.block_bytes = paged_kv_transfer_->block_bytes(); + residency_cfg.resident_budget_blocks = + paged_kv_pool_->physical_block_count(); + residency_cfg.sink_blocks = 1; + residency_cfg.tail_blocks = 4; + try { + paged_kv_residency_ = + std::make_unique( + *paged_kv_pool_, residency_cfg, + paged_kv_transfer_->callbacks()); + } catch (const std::exception & e) { + set_last_error(std::string( + "concurrent KVFlash residency init failed: ") + e.what()); + return false; + } + } if (n_slots > 1) { + const int tree_scratch_base = (int)pool_tokens; + const int64_t dead_scratch_row = + pool_tokens + (int64_t)n_slots * tree_stride; seq_engine_ = std::make_unique( *this, *paged_kv_pool_, cfg_.device.max_ctx, - /*scratch_row=*/pool_tokens, + dead_scratch_row, tree_width, tree_scratch_base, tree_stride, max_concurrent_prefills, mixed_prefill_tokens, long_mixed_prefill_tokens, long_prefill_threshold, idle_prefill_tokens, prefill_quantum); @@ -509,7 +569,7 @@ bool Qwen35Backend::init() { cfg_.device.max_ctx); std::fflush(stdout); } - if (kvflash_active()) { + if (kvflash_active() && !(cfg_.paged_attention && n_slots > 1)) { KvFlashConfig pc; pc.pool_tokens = kvflash_tokens_; if (!kvflash_pager_.attach(pc, cache_.attn_k, cache_.attn_v)) { @@ -540,7 +600,8 @@ bool Qwen35Backend::init() { // Init feature mirror when draft model is available (needed for spec decode). // On single-GPU, this is an F32 conversion buffer; on split-GPU, a cross-device mirror. - if (cfg_.draft_path && !use_remote_draft) { + if (cfg_.draft_path && !use_remote_draft && + !(n_slots > 1 && concurrent_local_ddtree)) { const int mirror_cap = std::min({cfg_.draft_ctx_max, cfg_.device.max_ctx, cache_.target_feat_cap > 0 ? cache_.target_feat_cap : cfg_.device.max_ctx}); if (!draft_feature_mirror_init(feature_mirror_, draft_backend_, @@ -1184,7 +1245,11 @@ DFlashTarget * Qwen35Backend::dflash_target() { void Qwen35Backend::shutdown() { const bool use_remote_draft = cfg_.remote_draft.enabled(); + seq_engine_.reset(); end_paged_sequence(); + paged_kv_residency_.reset(); + paged_kv_transfer_.reset(); + paged_kv_pool_.reset(); free_drafter(); step_graph_destroy(sg_); step_graph_destroy(draft_sg_); diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index a6a508f73..62fee8524 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -24,7 +24,9 @@ #include "concurrency/qwen35_seq_engine.h" #include "internal.h" // TargetWeights, TargetCache, DraftWeights, PrefixSnapshot #include "qwen3/qwen3_drafter.h" // DrafterContext, load_drafter, free_drafter, drafter_score_and_compress -#include "kvflash_pager.h" // bounded KV residency pool +#include "kvflash_pager.h" +#include "common/concurrency/paged_kv_residency.h" +#include "common/concurrency/qwen_paged_kv_transfer.h" #include "kvflash_scorer.h" // chunk-relevance policy interface #include "kvflash_qk.h" // target-QK scorer (pooled keys + query) @@ -295,6 +297,8 @@ class Qwen35Backend : public ModelBackend { // Page size comes from PAGED_BLOCK_SIZE (paged_attention_config.h), // shared with the graph builder and the cache's block-aligned sizing. std::unique_ptr paged_kv_pool_; + std::unique_ptr paged_kv_transfer_; + std::unique_ptr paged_kv_residency_; std::optional paged_sequence_; PagedKvRequestId paged_request_id_ = 0; diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 1082df6f5..5f10dcdc9 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -80,12 +80,14 @@ bool create_target_cache(const TargetWeights & w, bool prefill_only, int ctx_alloc, bool paged_attention, - int n_seq_slots) { + int n_seq_slots, + bool concurrent_tree) { return create_target_cache_partial(w, max_ctx, max_verify_tokens, backend, out, prefill_only, 0, w.n_layer, true, ctx_alloc, /*f32_ssm_intermediates=*/false, - paged_attention, n_seq_slots); + paged_attention, n_seq_slots, + concurrent_tree); } // concurrent_fixed_cache_bytes() in qwen35_backend.cpp mirrors this @@ -103,7 +105,8 @@ bool create_target_cache_partial(const TargetWeights & w, int ctx_alloc, bool f32_ssm_intermediates, bool paged_attention, - int n_seq_slots) { + int n_seq_slots, + bool concurrent_tree) { if (layer_begin < 0) layer_begin = 0; if (layer_end < 0 || layer_end > w.n_layer) layer_end = w.n_layer; if (layer_begin > layer_end) { @@ -115,6 +118,11 @@ bool create_target_cache_partial(const TargetWeights & w, set_last_error("multi-slot target cache requires paged attention"); return false; } + if (concurrent_tree && (!paged_attention || n_seq_slots <= 1)) { + set_last_error( + "concurrent tree cache requires paged multi-slot serving"); + return false; + } out.backend = backend; out.max_ctx = max_ctx; out.cur_pos = 0; @@ -227,7 +235,14 @@ bool create_target_cache_partial(const TargetWeights & w, out.target_feat_cap = std::min(max_ctx, TARGET_FEAT_CAP_DEFAULT); if (allocate_target_feat) { const int fc_in = w.n_capture_layers * w.n_embd; - out.target_feat = ggml_new_tensor_2d(out.base_ctx, GGML_TYPE_BF16, fc_in, out.target_feat_cap); + // Concurrent slots own disjoint feature rings. The final row is + // dead scratch for padded bucket rows because set_rows does not + // accept negative destination indices. + const int feat_rows = multi_slot + ? out.target_feat_cap * n_seq_slots + 1 + : out.target_feat_cap; + out.target_feat = ggml_new_tensor_2d( + out.base_ctx, GGML_TYPE_BF16, fc_in, feat_rows); ggml_set_name(out.target_feat, "target_feat"); } else { out.target_feat = nullptr; @@ -266,9 +281,10 @@ bool create_target_cache_partial(const TargetWeights & w, } // ── Rollback context: snapshots + intermediates ─────────────────── - // Multi-slot caches skip these entirely: concurrent serving is paged and - // therefore AR-only (no spec-decode rollback), and the tensors are the - // single largest optional allocation (~0.8 GB at 48 delta layers). + // Multi-slot caches skip these entirely. Packed tree verification gathers + // the selected slots' base state without mutating it, then a bounded + // replay commits accepted paths. T*S recurrent captures would be tens of + // GiB at useful Strix concurrency and are intentionally not allocated. if (!prefill_only && !multi_slot) { const int rb_tensors = 4 * n_delta; ggml_init_params ip{}; @@ -722,7 +738,16 @@ static ggml_tensor * build_full_attn_block( int paged_max_kv_len = 0, // Compact decode row -> physical block-table column. Negative ids are // graph-bucket padding rows. - ggml_tensor * active_slot_ids = nullptr + ggml_tensor * active_slot_ids = nullptr, + // Packed paged-tree verification. Query rows are flattened + // sequence-major; row mappings are supplied through + // paged_query_seq_ids, while parent/tree metadata describes each tree. + ggml_tensor * paged_tree_parent_ids = nullptr, + ggml_tensor * paged_tree_sizes = nullptr, + int tree_width = 0, + int tree_scratch_base = 0, + int tree_scratch_stride = 0, + int paged_logical_max_ctx = 0 ) { const int head_dim = w.n_embd_head_k; const int n_head = w.n_head; @@ -810,9 +835,14 @@ static ggml_tensor * build_full_attn_block( Kcur_T = ggml_turbo_wht(ctx, Kcur_T, 0); } + const bool paged_tree = paged_tree_parent_ids || paged_tree_sizes; + GGML_ASSERT((paged_tree_parent_ids == nullptr) == + (paged_tree_sizes == nullptr)); const bool ragged = paged_query_seq_ids != nullptr; - GGML_ASSERT(!ragged || (paged_block_table && paged_query_positions && - kv_write_rows)); + GGML_ASSERT(!ragged || (paged_block_table && kv_write_rows)); + GGML_ASSERT(!ragged || paged_tree || paged_query_positions); + GGML_ASSERT(!paged_tree || + (ragged && !paged_query_positions && tree_width > 0)); if (kv_write_rows) { // Step-invariant: the destination tensor stays fixed while the input // indices carry contiguous, KVFlash, or paged physical rows. @@ -878,12 +908,31 @@ static ggml_tensor * build_full_attn_block( ggml_tensor * row_seq_ids, ggml_tensor * row_positions, bool dense_token_layout) { - const int padded = ((std::max(1, launch_kv_len) + 255) / 256) * 256; - const int launch_len = std::min(padded, (int)cache_k->ne[1]); + // max_kv_seq_len sizes the logical partition grid. In bounded + // KVFlash mode the block table maps that logical range onto a much + // smaller physical K/V pool, so cache_k->ne[1] is not a valid clamp. + // Bound against both sources of logical capacity instead, doing the + // 256-window rounding in i64 to avoid signed overflow at large + // configured contexts. Per-row kv_seq_lens remains the exact runtime + // bound, and the paged kernel bounds every resolved physical row. + GGML_ASSERT(paged_block_table && cache_k && cache_v); + const int64_t table_capacity = + (int64_t)paged_block_table->ne[0] * PAGED_BLOCK_SIZE; + const int64_t logical_capacity = + std::min(paged_logical_max_ctx, table_capacity); + GGML_ASSERT(logical_capacity > 0 && logical_capacity <= INT32_MAX); + const int64_t requested = + std::min(std::max(1, launch_kv_len), + logical_capacity); + const int64_t padded = ((requested + 255) / 256) * 256; + const int launch_len = + (int)std::min(padded, logical_capacity); ggml_tensor * out = ggml_paged_attn_ext( ctx, q, cache_k, cache_v, paged_block_table, paged_kv_seq_lens, row_seq_ids, row_positions, kq_scale, - PAGED_BLOCK_SIZE, launch_len); + PAGED_BLOCK_SIZE, launch_len, + paged_tree_parent_ids, paged_tree_sizes, + tree_width, tree_scratch_base, tree_scratch_stride); if (dense_token_layout) { out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3)); } @@ -891,7 +940,19 @@ static ggml_tensor * build_full_attn_block( }; ggml_tensor * attn = nullptr; - if (ragged) { + if (paged_tree) { + // ── Packed concurrent tree verify. Every query row selects its + // physical sequence/scratch slab. The paged kernel combines the + // committed block-table prefix with only this node's ancestor chain; + // query_positions is intentionally absent in tree mode. + ggml_tensor * Qfa = q_segment(0, n_tokens); + if (q_fa_out) *q_fa_out = Qfa; + const int launch_kv_len = paged_max_kv_len > 0 + ? paged_max_kv_len : kv_start + n_tokens; + attn = paged_read(Qfa, launch_kv_len, + paged_query_seq_ids, /*row_positions=*/nullptr, + /*dense_token_layout=*/n_tokens > 1); + } else if (ragged) { // ── Ragged concurrent step: prefill chunk rows and decode rows all // read the pool through one call, each row clamped to its own // inclusive position. This step's chunk rows are visible to their @@ -899,6 +960,7 @@ static ggml_tensor * build_full_attn_block( // attention in the graph; cross-sequence isolation is structural // (each row's seq id selects its own block-table column). ggml_tensor * Qfa = q_segment(0, n_tokens); + if (q_fa_out) *q_fa_out = Qfa; const int launch_kv_len = paged_max_kv_len > 0 ? paged_max_kv_len : kv_start + n_tokens; attn = paged_read(Qfa, launch_kv_len, @@ -920,8 +982,8 @@ static ggml_tensor * build_full_attn_block( // bound only over-sizes the partition grid, and partitions past the // real length exit with a zero-weight sentinel. // Batched decode: kv_len (kv_start + n_tokens) describes one sequence; - // the launch bound must cover the longest live slot instead. Clamped - // because ggml_paged_attn asserts max_kv_seq_len <= k->ne[1]. + // the launch bound must cover the longest live slot instead. Bounded + // paged pools may be physically smaller than this logical span. const int launch_kv_len = paged_max_kv_len > 0 ? paged_max_kv_len : kv_len; attn = paged_read( Qfa, launch_kv_len, active_slot_ids, /*row_positions=*/nullptr, @@ -1032,8 +1094,13 @@ static ggml_tensor * build_delta_net_block( prefill_total += prefill_segments[i].n_tokens; } GGML_ASSERT((active_slot_ids == nullptr) == (state_slot_ids == nullptr)); + const bool mapped_tree = active_slot_ids && parent_ids; + GGML_ASSERT(!active_slot_ids || !cap); GGML_ASSERT(!active_slot_ids || - (!cap && !parent_ids && prefill_total + n_seqs == n_tokens)); + (mapped_tree + ? (!ragged && prefill_total == 0 && + n_tokens % n_seqs == 0) + : (prefill_total + n_seqs == n_tokens))); if (!active_slot_ids) { GGML_ASSERT(n_seqs == 1); GGML_ASSERT(prefill_total == 0 || prefill_total == n_tokens); @@ -1065,6 +1132,7 @@ static ggml_tensor * build_delta_net_block( int T; // timesteps per sequence int S; // sequences bool active; // compact decode segment (slot-mapped) + bool tree; // mapped tree: gather-only, no persistence ggml_tensor * conv_st; ggml_tensor * ssm_st; }; @@ -1082,14 +1150,17 @@ static ggml_tensor * build_delta_net_block( ssm_state->ne[0], ssm_state->ne[1], ssm_state->ne[2], 1, ssm_state->nb[1], ssm_state->nb[2], ssm_state->nb[3], (size_t)pf.seq_slot * ssm_state->nb[3]); - segs.push_back({pf.token_offset, pf.n_tokens, 1, false, c, s}); + segs.push_back({pf.token_offset, pf.n_tokens, 1, + false, false, c, s}); } if (active_slot_ids) { - segs.push_back({prefill_total, 1, n_seqs, true, - conv_state, ssm_state}); + const int tree_tokens = mapped_tree ? n_tokens / n_seqs : 1; + segs.push_back({prefill_total, tree_tokens, n_seqs, true, + mapped_tree, conv_state, ssm_state}); } else if (segs.empty()) { // No general [timesteps x sequences] mode: one multi-token sequence. - segs.push_back({0, n_tokens, n_seqs, false, conv_state, ssm_state}); + segs.push_back({0, n_tokens, n_seqs, false, false, + conv_state, ssm_state}); } const int n_segs = (int)segs.size(); @@ -1109,13 +1180,14 @@ static ggml_tensor * build_delta_net_block( const int seg_seqs = seg.S; const int seg_tokens = seg.T * seg.S; const bool seg_active = seg.active; + const bool seg_tree = seg.tree; // Plain one-token decode has no in-graph consumer of the updated state: // the next graph evaluation is the first read. Write the final state // directly into its persistent slab and avoid materializing/copying a // second S_v x S_v x H_v state. The active-aware path also updates each // mapped physical slab directly; only its negative bucket-padding rows // use the result tensor's retained scratch state region. - const bool inplace_state = seg_active || + const bool inplace_state = (seg_active && !seg_tree) || (allow_inplace_state && can_skip_gdn_intermediate && !ragged && n_seq_tokens == 1); @@ -1183,7 +1255,7 @@ static ggml_tensor * build_delta_net_block( w.ssm_d_conv - 1, conv_channels, seg_seqs, conv_input->nb[1], conv_input->nb[2], (conv_input->ne[0] - (w.ssm_d_conv - 1)) * ggml_element_size(conv_input)); - if (seg_active) { + if (seg_active && !seg_tree) { const int64_t slab = (int64_t)(w.ssm_d_conv - 1) * conv_channels; ggml_tensor * compact_last = ggml_reshape_2d( @@ -1193,7 +1265,7 @@ static ggml_tensor * build_delta_net_block( ggml_build_forward_expand( gf, ggml_set_rows_masked( ctx, all_conv, compact_last, active_slot_ids)); - } else { + } else if (!seg_tree) { ggml_build_forward_expand(gf, ggml_cpy(ctx, last_conv, seg.conv_st)); } @@ -1246,10 +1318,25 @@ static ggml_tensor * build_delta_net_block( } // ── SSM state (recurrent): reshape to [S_v, S_v, H_v, n_seqs] - ggml_tensor * s = seg_active - ? seg.ssm_st - : ggml_reshape_4d(ctx, seg.ssm_st, + ggml_tensor * s = nullptr; + if (seg_tree) { + // Packed tree verification starts each tree from the owning slot's + // base state. Gather compact slabs, then leave the persistent tensor + // untouched; accepted paths are committed by a later replay. + const int64_t slab = + (int64_t)head_v_dim * head_v_dim * num_v_heads; + ggml_tensor * all_ssm = ggml_reshape_2d( + ctx, seg.ssm_st, slab, seg.ssm_st->ne[3]); + ggml_tensor * gathered = + ggml_get_rows(ctx, all_ssm, state_slot_ids); + s = ggml_reshape_4d(ctx, gathered, head_v_dim, head_v_dim, num_v_heads, seg_seqs); + } else { + s = seg_active + ? seg.ssm_st + : ggml_reshape_4d(ctx, seg.ssm_st, + head_v_dim, head_v_dim, num_v_heads, seg_seqs); + } // ── Fused Gated DeltaNet op — returns packed (output | new_state [| intermediates]). // In tree mode, the kernel uses parent_ids to reload state at DFS @@ -1302,7 +1389,7 @@ static ggml_tensor * build_delta_net_block( ggml_build_forward_expand(gf, ggml_cpy(ctx, r.new_state, s)); } else { ggml_tensor * result; - if (seg_active) { + if (seg_active && !seg_tree) { result = ggml_gated_delta_net_active_inplace( ctx, q_c, k_c, v_c, g_tensor, beta, s, active_slot_ids); } else if (parent_ids) { @@ -1337,7 +1424,7 @@ static ggml_tensor * build_delta_net_block( S_v * H_v * r_elt, S_v * H_v * n_seq_tokens * r_elt, 0); - if (!inplace_state) { + if (!inplace_state && !seg_tree) { ggml_tensor * new_state = ggml_view_4d(ctx, result, S_v, S_v, H_v, seg_seqs, S_v * r_elt, @@ -1345,8 +1432,8 @@ static ggml_tensor * build_delta_net_block( S_v * S_v * H_v * r_elt, S_v * H_v * n_seq_tokens * seg_seqs * r_elt); - // Persist new_state back to cache. Both compact active decode and the - // plain in-place AR path write state from the GDN kernel directly. + // Persist new_state back to cache. Mapped trees deliberately skip + // this branch: their gathered base state is read-only. ggml_build_forward_expand(gf, ggml_cpy(ctx, new_state, seg.ssm_st)); } @@ -1555,6 +1642,12 @@ QwenGraphOutputs build_qwen35_graph( const int hidden = w.n_embd; const float eps = w.rms_eps; + const bool capture_with_rows = + in.capture_layers && cache.target_feat && in.target_feat_rows; + std::vector capture_slices; + if (capture_with_rows) { + capture_slices.assign((size_t)N_CAPTURE, nullptr); + } for (int il = 0; il < w.n_layer; il++) { const TargetLayer & L = w.layers[il]; @@ -1584,7 +1677,13 @@ QwenGraphOutputs build_qwen35_graph( in.paged_query_seq_ids, in.paged_query_positions, in.paged_max_kv_len, - in.active_slot_ids); + in.active_slot_ids, + in.parent_ids, + in.tree_sizes, + in.tree_width, + in.tree_scratch_base, + in.tree_scratch_stride, + cache.max_ctx); if (want_q_cap && q_fa) { // Last token's Q, all heads: src [head_dim, 1, n_head] view of // [head_dim, n_tokens, n_head]; dst = q_cap plane fa_idx @@ -1676,6 +1775,13 @@ QwenGraphOutputs build_qwen35_graph( if (CAPTURE_LAYERS[k] == il) { capture_idx = k; break; } } if (capture_idx >= 0) { + ggml_tensor * cur_2d = + ggml_reshape_2d(ctx, cur, hidden, n_tokens); + if (capture_with_rows) { + capture_slices[(size_t)capture_idx] = cur_2d; + inpL = cur; + continue; + } const size_t elt = ggml_element_size(cache.target_feat); const size_t col_stride = cache.target_feat->nb[1]; const int cap = cache.target_feat_cap; @@ -1683,8 +1789,6 @@ QwenGraphOutputs build_qwen35_graph( const int pre_n = std::min(n_tokens, cap - slot_start); const int post_n = n_tokens - pre_n; - ggml_tensor * cur_2d = ggml_reshape_2d(ctx, cur, hidden, n_tokens); - // First slice: [slot_start..slot_start+pre_n) in the ring. { const size_t offset = @@ -1714,6 +1818,21 @@ QwenGraphOutputs build_qwen35_graph( inpL = cur; } + if (capture_with_rows) { + GGML_ASSERT(!capture_slices.empty()); + ggml_tensor * feat_cat = capture_slices[0]; + GGML_ASSERT(feat_cat); + for (int k = 1; k < (int)capture_slices.size(); ++k) { + GGML_ASSERT(capture_slices[(size_t)k]); + feat_cat = ggml_concat( + ctx, feat_cat, capture_slices[(size_t)k], 0); + } + feat_cat = ggml_cont(ctx, feat_cat); + ggml_build_forward_expand( + gf, ggml_set_rows( + ctx, cache.target_feat, feat_cat, in.target_feat_rows)); + } + // 2. Final norm ggml_tensor * out = rms_norm_mul(ctx, inpL, w.out_norm, w.rms_eps); diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index e92f495bf..52a4a845c 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -1379,11 +1379,27 @@ int HttpServer::run() { std::fprintf(stderr, "[server] listening on http://%s:%d\n", config_.host.c_str(), config_.port); - // 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()) { + // A backend-provided sequence engine replaces the one-request worker. + // Local PFlash stays on this path too: scheduler admission prepares each + // prompt exactly once, then admits the effective tokens. Persistent + // residency is the explicit safety contract that lets compression run + // without parking model state owned by other live slots. + SeqEngine * engine = backend_.seq_engine(); + if (engine && config_.pflash_upstream_base.empty()) { + const ConcurrentPflashPlan pflash_plan = + resolve_concurrent_pflash_plan(config_, drafter_tokenizer_ != nullptr); + if (!pflash_plan.ok()) { + std::fprintf(stderr, "[server] %s\n", pflash_plan.error.c_str()); + socket_close(listen_fd_); + listen_fd_ = kInvalidSocket; + return 2; + } + if (pflash_plan.force_skip_park && !config_.pflash_skip_park) { + config_.pflash_skip_park = true; + std::fprintf(stderr, + "[server] concurrent PFlash: persistent residency enables " + "skip-park for live sequence safety\n"); + } worker_thread_ = std::thread([this, engine]() { scheduler_loop(*engine); }); } else { @@ -4231,6 +4247,7 @@ std::string HttpServer::format_http_response( case 400: reason = "Bad Request"; break; case 404: reason = "Not Found"; break; case 405: reason = "Method Not Allowed"; break; + case 409: reason = "Conflict"; break; case 413: reason = "Payload Too Large"; break; case 500: reason = "Internal Server Error"; break; case 503: reason = "Service Unavailable"; break; diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 90e02412e..0fa0ef718 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -250,6 +250,49 @@ bool should_clamp_flowkv_disk_cache( bool flowkv, const DiskPrefixCachePolicy & policy); } // namespace http_detail +// Resolve the prompt-compression contract for a backend-provided sequence +// engine. Concurrent PFlash must keep both models resident: parking either +// model while another slot owns live device state invalidates that slot. The +// explicit persistent policy is therefore the opt-in that also implies the +// effective skip-park behavior; callers do not need a second, redundant CLI +// flag on large-memory concurrent hosts. +struct ConcurrentPflashPlan { + bool enabled = false; + bool force_skip_park = false; + std::string error; + + bool ok() const { return error.empty(); } +}; + +inline ConcurrentPflashPlan resolve_concurrent_pflash_plan( + const ServerConfig & config, bool drafter_tokenizer_available) { + ConcurrentPflashPlan plan; + if (config.pflash_mode == ServerConfig::PflashMode::OFF || + !config.pflash_upstream_base.empty()) { + return plan; + } + plan.enabled = true; + if (!drafter_tokenizer_available) { + plan.error = + "concurrent PFlash requires a loaded --prefill-drafter tokenizer"; + return plan; + } + if (config.draft_residency != DraftResidencyPolicy::Persistent) { + plan.error = + "concurrent PFlash requires --draft-residency persistent so " + "prompt compression cannot park live target/draft state"; + return plan; + } + if (config.prefix_cache_cap > 0 || config.prefill_cache_cap > 0 || + !config.disk_cache_dir.empty()) { + plan.error = + "concurrent paged PFlash does not support prefix/prefill " + "snapshots; disable the snapshot caches"; + return plan; + } + plan.force_skip_park = true; + return plan; +} // ─── Parsed request ───────────────────────────────────────────────────── @@ -605,6 +648,13 @@ struct ServerJob { // server-side prefill/elapsed telemetry does not erase queueing delay. std::chrono::steady_clock::time_point parallel_started_at{}; std::unique_ptr emitter; + // Prompt preparation (FlowKV/PFlash) is expensive and may load a resident + // drafter. Cache its result on the job so a pool-full retry never runs it + // twice. The original request tokens remain untouched for API accounting; + // this vector is the effective prompt admitted to the sequence engine. + bool parallel_prompt_prepared = false; + bool parallel_prompt_compressed = false; + std::vector parallel_prompt_tokens; }; // ─── Parse session_id from a chat-completion JSON body ────────────────── diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 3c0b8f7c3..c34c71fb6 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -36,6 +36,16 @@ struct SchedSlot { double prefill_s = 0.0; int n_gen_cap = 0; int completion_tokens = 0; + int effective_prompt_tokens = 0; + bool prompt_compressed = false; + uint64_t engine_request_id = 0; + uint64_t ddtree_steps = 0; + uint64_t ddtree_accepted_tokens = 0; + uint64_t target_forwards = 0; + uint64_t kvflash_page_ins = 0; + uint64_t kvflash_page_outs = 0; + uint64_t kvflash_resident_blocks = 0; + uint64_t kvflash_reselects = 0; bool client_disconnected = false; bool failed = false; std::string error; @@ -238,7 +248,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { stop_job_stream(s.job, &s.send_buffer); 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(); + const int prompt_tokens = s.effective_prompt_tokens; GenTimings gen_timings{ s.prefill_s, decode_s, @@ -253,9 +263,10 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { 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; + ? (double)prompt_tokens / s.prefill_s : 0.0; perf.decode_tok_s = decode_s > 0.0 ? (double)s.completion_tokens / decode_s : 0.0; + perf.pflash = s.prompt_compressed; status_.record_perf(perf); } @@ -292,17 +303,38 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { 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 " + "[server] chat DONE %s ok=%s in=%zu effective_in=%d 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, + req.prompt_tokens.size(), prompt_tokens, 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); + const json concurrency_metrics = { + {"request_id", req.response_id}, + {"response_id", req.response_id}, + {"engine_request_id", s.engine_request_id}, + {"raw_prompt_tokens", req.prompt_tokens.size()}, + {"effective_prompt_tokens", prompt_tokens}, + {"output_tokens", out_tokens}, + {"pflash_applied", s.prompt_compressed}, + {"pflash_input_tokens", req.prompt_tokens.size()}, + {"pflash_output_tokens", prompt_tokens}, + {"ddtree_steps", s.ddtree_steps}, + {"ddtree_accepted_tokens", s.ddtree_accepted_tokens}, + {"target_forwards", s.target_forwards}, + {"kvflash_page_ins", s.kvflash_page_ins}, + {"kvflash_page_outs", s.kvflash_page_outs}, + {"kvflash_resident_blocks", s.kvflash_resident_blocks}, + {"kvflash_reselects", s.kvflash_reselects}, + }; + std::fprintf(stderr, "[concurrency-metrics] %s\n", + concurrency_metrics.dump().c_str()); + engine.retire(idx); // A retirement may have released the blocks the head job needs. deferred_retry_at = {}; @@ -431,9 +463,65 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { start_job_stream(job); } + // Apply the same FlowKV/PFlash precedence and overflow checks as the + // classic worker. Keep the prepared prompt on the job because an + // atomically-busy admission is retried at the FIFO head later. + if (!job->parallel_prompt_prepared) { + PreparedPrompt prepared = prepare_prompt(req); + if (prepared.error_status != 0) { + const std::string message = prepared.error.empty() + ? "prompt preparation failed" + : prepared.error; + std::fprintf(stderr, + "[server] concurrent prompt preparation failed: %s\n", + message.c_str()); + if (req.stream && job->sse_started) { + stop_job_stream(job); + for (const std::string & chunk : + sse_error_close_chunks(message)) { + send_job_bytes(job, chunk.data(), chunk.size()); + } + } else { + send_error(job->fd, prepared.error_status, message); + } + finish_job(job); + return AdmissionDisposition::Retired; + } + // Paged sequence engines cannot restore the classic snapshot + // format. Startup normally disables those caches; keep this + // check as a hard guard for embedded/non-CLI callers. + if (prepared.full_cache_hit_slot >= 0 || + prepared.full_cache_served_tokens >= 0) { + const std::string message = + "concurrent paged serving cannot restore a prefix snapshot"; + if (req.stream && job->sse_started) { + stop_job_stream(job); + for (const std::string & chunk : + sse_error_close_chunks(message)) { + send_job_bytes(job, chunk.data(), chunk.size()); + } + } else { + send_error(job->fd, 409, message); + } + finish_job(job); + return AdmissionDisposition::Retired; + } + job->parallel_prompt_tokens = std::move(prepared.tokens); + job->parallel_prompt_compressed = prepared.compressed; + job->parallel_prompt_prepared = true; + std::fprintf(stderr, + "[server] concurrent prompt READY %s raw=%zu effective=%zu " + "pflash=%s\n", + req.response_id.c_str(), req.prompt_tokens.size(), + job->parallel_prompt_tokens.size(), + prepared.compressed ? "true" : "false"); + } + const auto & effective_prompt = job->parallel_prompt_tokens; + // 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, + const uint64_t engine_request_id = next_request_id; + auto ar = engine.admit(engine_request_id, effective_prompt, req.sampler); if (ar.status == SeqEngine::AdmitResult::Status::busy) return AdmissionDisposition::Deferred; @@ -464,9 +552,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.admission_order = next_admission_order++; s.started_at = started_at; s.decode_started_at = started_at; // sane on prefill failure + s.effective_prompt_tokens = (int)effective_prompt.size(); + s.prompt_compressed = job->parallel_prompt_compressed; + s.engine_request_id = engine_request_id; s.n_gen_cap = std::min( n_gen_cap, - engine.max_context() - (int)req.prompt_tokens.size() + 1); + engine.max_context() - s.effective_prompt_tokens + 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() && @@ -606,8 +697,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { 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}); + SeqEngine::StepInput input; + input.slot = i; + input.token = slots[(size_t)i].pending_tok; + input.allow_speculation = + slots[(size_t)i].hook.close_token_ids.empty(); + step_plan.decode.push_back(input); } else if (slots[(size_t)i].job) { prefill_candidates.push_back( {i, slots[(size_t)i].admission_order}); @@ -655,7 +750,18 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { s.finished = true; continue; } - advance_slot(s, out.token); + s.ddtree_steps += out.ddtree_steps; + s.ddtree_accepted_tokens += out.ddtree_accepted_tokens; + s.target_forwards += out.target_forwards; + s.kvflash_page_ins += out.kvflash_page_ins; + s.kvflash_page_outs += out.kvflash_page_outs; + s.kvflash_resident_blocks = std::max( + s.kvflash_resident_blocks, out.kvflash_resident_blocks); + s.kvflash_reselects += out.kvflash_reselects; + consume_decode_output_tokens(out, [&](int32_t token) { + advance_slot(s, token); + return !s.finished; + }); } using PrefillStatus = SeqEngine::PrefillOutput::Status; for (const auto & out : step_result.prefills) { diff --git a/server/test/bench_paged_attention.cpp b/server/test/bench_paged_attention.cpp index 49fc00960..1e5afe9c8 100644 --- a/server/test/bench_paged_attention.cpp +++ b/server/test/bench_paged_attention.cpp @@ -725,7 +725,8 @@ bool run_case( ggml_tensor * paged_output = ggml_paged_attn_ext( ctx.value, q_paged, k_paged, v_paged, table, kv_seq_lens_tensor, nullptr, nullptr, 1.0f / std::sqrt(static_cast(D)), - BLOCK_SIZE, max_context); + BLOCK_SIZE, max_context, + nullptr, nullptr, 0, 0, 0); ggml_tensor * contiguous_output = ggml_flash_attn_ext( ctx.value, q_contiguous, k_contiguous, v_contiguous, padding_mask, 1.0f / std::sqrt(static_cast(D)), diff --git a/server/test/seq_engine_contract.h b/server/test/seq_engine_contract.h index 63ca035ab..0361a893c 100644 --- a/server/test/seq_engine_contract.h +++ b/server/test/seq_engine_contract.h @@ -258,6 +258,19 @@ inline std::vector check_seq_engine_contract(SeqEngine & engine) { return violations; } + // Scheduler policy can retain commit authority for selected slots. A + // conforming engine may still use its ordinary AR implementation, but it + // must not return already-committed children for a disabled input. + SeqEngine::StepPlan no_speculation; + no_speculation.decode = decode_inputs(); + for (SeqEngine::StepInput & input : no_speculation.decode) { + input.allow_speculation = false; + } + if (!execute(no_speculation)) { + 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); diff --git a/server/test/test_ddtree_path.cpp b/server/test/test_ddtree_path.cpp new file mode 100644 index 000000000..e7d2bbf3a --- /dev/null +++ b/server/test/test_ddtree_path.cpp @@ -0,0 +1,46 @@ +#include "common/ddtree.h" + +#include +#include +#include +#include + +using dflash::common::DDTree; +using dflash::common::follow_verified_tree; +using dflash::common::truncate_verified_path; + +int main() { + DDTree tree; + tree.n_nodes = 2; + tree.token_ids = {11, 22}; + tree.depths = {1, 2}; + tree.parents = {-1, 0, 1}; + tree.child_maps.resize(3); + tree.child_maps[0][11] = 1; + tree.child_maps[1][22] = 2; + + const int32_t posterior[] = {11, 22, 33}; + int pending = -1; + std::vector accepted = + follow_verified_tree(tree, posterior, pending); + assert((accepted == std::vector{0, 1, 2})); + assert(pending == 33); + + // Truncating after node 1 means node 2's token becomes pending. Keeping + // the old value (33) would skip token 22 and describe uncommitted state. + assert(truncate_verified_path(accepted, 2, posterior, pending)); + assert((accepted == std::vector{0, 1})); + assert(pending == 22); + + // An unchanged path preserves the already-computed pending token. + assert(!truncate_verified_path(accepted, 2, posterior, pending)); + assert(pending == 22); + + // No headroom is represented explicitly and never dereferences a tip. + assert(truncate_verified_path(accepted, 0, posterior, pending)); + assert(accepted.empty()); + assert(pending == -1); + + std::puts("ddtree path tests passed"); + return 0; +} diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index d01b9da04..fb1a5571a 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -381,6 +381,53 @@ void test_feature_gate_paged_attention_requires_plain_ar_decode() { ddtree.ddtree_mode = true; CHECK(!gate_result(ddtree, "qwen35", PlacementBackend::Cuda).empty()); + BackendArgs concurrent_ddtree = base; + concurrent_ddtree.max_concurrency = 16; + concurrent_ddtree.draft_path = "/nonexistent/draft.gguf"; + concurrent_ddtree.ddtree_mode = true; + concurrent_ddtree.ddtree_budget = 22; + TEST_ASSERT(gate_accepts( + concurrent_ddtree, "qwen35", PlacementBackend::Cuda)); + TEST_ASSERT(gate_accepts( + concurrent_ddtree, "qwen35", PlacementBackend::Hip)); + + BackendArgs tensor_ddtree = concurrent_ddtree; + TEST_ASSERT(parse_placement_device_list( + "cuda:0,cuda:1", tensor_ddtree.device)); + tensor_ddtree.device.split_mode = TargetSplitMode::Tensor; + TEST_ASSERT(!gate_accepts( + tensor_ddtree, "qwen35", PlacementBackend::Cuda)); + + + BackendFeatureConfig concurrent_pflash; + concurrent_pflash.pflash_enabled = true; + concurrent_pflash.pflash_drafter_configured = true; + TEST_ASSERT(gate_accepts(concurrent_ddtree, "qwen35", + PlacementBackend::Hip, concurrent_pflash)); + + BackendArgs concurrent_plain = base; + concurrent_plain.max_concurrency = 16; + TEST_ASSERT(gate_accepts(concurrent_plain, "qwen35", + PlacementBackend::Hip, concurrent_pflash)); + BackendFeatureConfig concurrent_kvflash; + concurrent_kvflash.kvflash_enabled = true; + TEST_ASSERT(gate_accepts(concurrent_plain, "qwen35", + PlacementBackend::Hip, concurrent_kvflash)); + + BackendArgs bad_budget = concurrent_ddtree; + for (int value : {0, -1, 256, INT_MAX}) { + bad_budget.ddtree_budget = value; + TEST_ASSERT(!gate_accepts( + bad_budget, "qwen35", PlacementBackend::Hip)); + } + + BackendArgs remote_ddtree = concurrent_ddtree; + remote_ddtree.remote_draft.ipc_bin = "/usr/bin/draft-ipc"; + remote_ddtree.draft_device.backend = PlacementBackend::Cuda; + remote_ddtree.device.backend = PlacementBackend::Hip; + TEST_ASSERT(!gate_accepts( + remote_ddtree, "qwen35", PlacementBackend::Hip)); + BackendArgs windowed = base; windowed.fa_window = 4096; CHECK(!gate_result( @@ -481,6 +528,22 @@ void test_feature_gate_parallel_and_kv_pool_rules() { pool.kv_pool_tokens = max_pool_tokens; CHECK(gate_result(pool, "qwen35", PlacementBackend::Cuda).empty()); + BackendArgs tree_pool = paged; + tree_pool.max_concurrency = 16; + tree_pool.draft_path = "/nonexistent/draft.gguf"; + tree_pool.ddtree_mode = true; + tree_pool.ddtree_budget = 22; + const long long tree_scratch = + (long long)tree_pool.max_concurrency * + paged_token_capacity(tree_pool.ddtree_budget + 1); + const long long max_tree_pool_tokens = + ((long long)INT_MAX - PAGED_BLOCK_SIZE - tree_scratch) / + PAGED_BLOCK_SIZE * PAGED_BLOCK_SIZE; + tree_pool.kv_pool_tokens = max_tree_pool_tokens; + TEST_ASSERT(gate_accepts(tree_pool, "qwen35", PlacementBackend::Cuda)); + tree_pool.kv_pool_tokens = max_tree_pool_tokens + PAGED_BLOCK_SIZE; + TEST_ASSERT(!gate_accepts(tree_pool, "qwen35", PlacementBackend::Cuda)); + // The automatic pool is memory-derived, so a logical slot/context product // larger than the physical tensor address space is legal. BackendArgs overflow = paged; diff --git a/server/test/test_paged_attention.cpp b/server/test/test_paged_attention.cpp index 78338f717..3e550c1de 100644 --- a/server/test/test_paged_attention.cpp +++ b/server/test/test_paged_attention.cpp @@ -30,6 +30,13 @@ struct TestCase { bool corrupt_blocks; }; +struct TreeMetadata { + int width; + int scratch_stride; + std::vector parent_ids; + std::vector tree_sizes; +}; + int clamped_seq_len(const TestCase & test_case, int seq) { return std::max( 0, std::min( @@ -51,6 +58,36 @@ bool block_is_valid(int32_t block, int physical_blocks) { return block >= 0 && block < physical_blocks; } +bool tree_visible( + const TreeMetadata & tree, + int tree_seq, + int query_node, + int candidate) { + const int tree_size = tree.tree_sizes[tree_seq]; + if (tree_size < 0 || tree_size > tree.width || + query_node < 0 || query_node >= tree_size || + candidate < 0 || candidate >= tree_size) { + return false; + } + + int current = query_node; + for (int depth = 0; depth < tree_size; ++depth) { + if (current == candidate) { + return true; + } + if (current < 0 || current >= tree_size) { + return false; + } + const int parent = + tree.parent_ids[tree_seq * tree.width + current]; + if (parent == current) { + return false; + } + current = parent; + } + return false; +} + std::vector make_block_table( const TestCase & test_case, int physical_blocks) { @@ -127,7 +164,9 @@ std::vector reference_attention( const std::vector & k, const std::vector & v, const std::vector * active_slot_ids = nullptr, - const std::vector * query_positions = nullptr) { + const std::vector * query_positions = nullptr, + const TreeMetadata * tree = nullptr, + int tree_scratch_base = 0) { std::vector output(q.size(), 0.0f); const float scale = 1.0f / std::sqrt(static_cast(D)); const int q_per_kv = N_HEAD / N_HEAD_KV; @@ -147,50 +186,74 @@ std::vector reference_attention( // cached tokens [0, position]. kv_seq_len = (*query_positions)[seq] + 1; } + const int tree_seq = tree ? seq / tree->width : 0; + const int query_node = tree ? seq % tree->width : -1; + const int tree_size = tree ? tree->tree_sizes[tree_seq] : 0; + if (tree && + (tree_size < 0 || tree_size > tree->width || + query_node >= tree_size)) { + continue; + } + + std::vector physical_rows; + physical_rows.reserve(kv_seq_len + (tree ? tree->width : 0)); + for (int token = 0; token < kv_seq_len; ++token) { + const int block = + block_table[ + physical_seq * test_case.max_blocks + + token / BLOCK_SIZE]; + physical_rows.push_back( + block_is_valid(block, physical_blocks) + ? block * BLOCK_SIZE + token % BLOCK_SIZE + : -1); + } + if (tree) { + for (int candidate = 0; candidate < tree->width; ++candidate) { + physical_rows.push_back( + tree_visible(*tree, tree_seq, query_node, candidate) + ? tree_scratch_base + + physical_seq * tree->scratch_stride + candidate + : -1); + } + } for (int head = 0; head < N_HEAD; ++head) { const int kv_head = head / q_per_kv; const float * q_row = q.data() + (static_cast(head) * n_seq + seq) * D; - std::vector scores(kv_seq_len); + std::vector scores(physical_rows.size(), -INFINITY); float max_score = -INFINITY; - for (int token = 0; token < kv_seq_len; ++token) { - const int block = - block_table[ - physical_seq * test_case.max_blocks + token / BLOCK_SIZE]; - if (!block_is_valid(block, physical_blocks)) { - // Mirrors the kernel: invalid blocks contribute nothing. - scores[token] = -INFINITY; - continue; - } - const int physical = block * BLOCK_SIZE + token % BLOCK_SIZE; + for (size_t row = 0; row < physical_rows.size(); ++row) { + const int physical = physical_rows[row]; + if (physical < 0) continue; const float * k_row = k.data() + (static_cast(kv_head) * pool_tokens + physical) * D; float dot = 0.0f; for (int d = 0; d < D; ++d) dot += q_row[d] * k_row[d]; - scores[token] = dot * scale; - max_score = std::max(max_score, scores[token]); + scores[row] = dot * scale; + max_score = std::max(max_score, scores[row]); } float denominator = 0.0f; for (float & score : scores) { + if (!std::isfinite(score)) { + score = 0.0f; + continue; + } score = std::exp(score - max_score); denominator += score; } float * out_row = output.data() + (static_cast(head) * n_seq + seq) * D; - for (int token = 0; token < kv_seq_len; ++token) { - const int block = - block_table[ - physical_seq * test_case.max_blocks + token / BLOCK_SIZE]; - if (!block_is_valid(block, physical_blocks)) continue; - const int physical = block * BLOCK_SIZE + token % BLOCK_SIZE; + for (size_t row = 0; row < physical_rows.size(); ++row) { + const int physical = physical_rows[row]; + if (physical < 0 || denominator == 0.0f) continue; const float * v_row = v.data() + (static_cast(kv_head) * pool_tokens + physical) * D; - const float probability = scores[token] / denominator; + const float probability = scores[row] / denominator; for (int d = 0; d < D; ++d) { out_row[d] += probability * v_row[d]; } @@ -205,7 +268,8 @@ bool run_case(ggml_backend_t backend, ggml_type k_type, ggml_type v_type, const std::vector * active_slot_ids = nullptr, - const std::vector * query_positions = nullptr) { + const std::vector * query_positions = nullptr, + const TreeMetadata * tree = nullptr) { const int physical_n_seq = static_cast(test_case.kv_seq_lens.size()); const int n_seq = active_slot_ids ? static_cast(active_slot_ids->size()) @@ -214,8 +278,23 @@ bool run_case(ggml_backend_t backend, GGML_ASSERT(!query_positions || (active_slot_ids && query_positions->size() == active_slot_ids->size())); + GGML_ASSERT(!tree || (active_slot_ids && !query_positions)); + if (tree) { + GGML_ASSERT(tree->width > 0); + GGML_ASSERT(tree->scratch_stride >= tree->width); + GGML_ASSERT( + tree->parent_ids.size() == + static_cast(tree->width) * tree->tree_sizes.size()); + GGML_ASSERT( + n_seq == + tree->width * static_cast(tree->tree_sizes.size())); + } const int physical_blocks = count_physical_blocks(test_case); - const int pool_tokens = physical_blocks * BLOCK_SIZE; + const int tree_scratch_base = physical_blocks * BLOCK_SIZE; + const int pool_tokens = tree + ? tree_scratch_base + physical_n_seq * tree->scratch_stride + : tree_scratch_base; + GGML_ASSERT(pool_tokens % BLOCK_SIZE == 0); const std::vector block_table = make_block_table(test_case, physical_blocks); @@ -250,13 +329,26 @@ bool run_case(ggml_backend_t backend, positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seq); ggml_set_input(positions); } + ggml_tensor * parents = nullptr; + ggml_tensor * sizes = nullptr; + if (tree) { + parents = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, tree->width, tree->tree_sizes.size()); + sizes = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, tree->tree_sizes.size()); + ggml_set_input(parents); + ggml_set_input(sizes); + } const float scale = 1.0f / std::sqrt(static_cast(D)); const int max_kv_seq_len = *std::max_element( test_case.kv_seq_lens.begin(), test_case.kv_seq_lens.end()); ggml_tensor * output = ggml_paged_attn_ext( ctx, q, k, v, table, kv_seq_lens, active, positions, - scale, BLOCK_SIZE, max_kv_seq_len); + scale, BLOCK_SIZE, max_kv_seq_len, parents, sizes, + tree ? tree->width : 0, + tree ? tree_scratch_base : 0, + tree ? tree->scratch_stride : 0); ggml_set_output(output); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, output); @@ -315,6 +407,14 @@ bool run_case(ggml_backend_t backend, positions, query_positions->data(), 0, query_positions->size() * sizeof((*query_positions)[0])); } + if (tree) { + ggml_backend_tensor_set( + parents, tree->parent_ids.data(), 0, + tree->parent_ids.size() * sizeof(tree->parent_ids[0])); + ggml_backend_tensor_set( + sizes, tree->tree_sizes.data(), 0, + tree->tree_sizes.size() * sizeof(tree->tree_sizes[0])); + } ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; } @@ -327,7 +427,7 @@ bool run_case(ggml_backend_t backend, reference_attention( test_case, block_table, pool_tokens, physical_blocks, q_data, k_reference, v_reference, active_slot_ids, - query_positions); + query_positions, tree, tree ? tree_scratch_base : 0); max_abs_error = 0.0f; for (size_t i = 0; i < actual.size(); ++i) { if (!std::isfinite(actual[i])) { @@ -340,10 +440,11 @@ bool run_case(ggml_backend_t backend, ok = ok && max_abs_error < MAX_ABS_ERROR; } - std::printf("paged attention %-11s K=%-4s V=%-4s active=%s pos=%s max_abs=%.6g %s\n", + std::printf("paged attention %-11s K=%-4s V=%-4s active=%s pos=%s tree=%s max_abs=%.6g %s\n", test_case.name, ggml_type_name(k_type), ggml_type_name(v_type), active_slot_ids ? "yes" : "no", query_positions ? "yes" : "no", + tree ? "yes" : "no", max_abs_error, ok ? "PASS" : "FAIL"); ggml_gallocr_free(allocator); ggml_free(ctx); @@ -373,7 +474,8 @@ bool rejects_unlaunchable_gqa(ggml_backend_t backend) { ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_tensor * output = ggml_paged_attn_ext( ctx, q, k, v, table, kv_seq_lens, nullptr, nullptr, - 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE, 1); + 1.0f / std::sqrt(static_cast(D)), BLOCK_SIZE, 1, + nullptr, nullptr, 0, 0, 0); const bool rejected = !ggml_backend_supports_op(backend, output); std::printf("paged attention unlaunchable GQA support %s\n", @@ -404,6 +506,38 @@ void run_paged_attention_case(const TestCase & test_case) { ggml_backend_free(backend); } +void run_tree_case() { + ggml_backend_t backend = ggml_backend_cuda_init(0); + REQUIRE_NOT_NULL(backend); + const TestCase tree_case{"tree", 65, {1025, 17, 257}, false}; + const TreeMetadata tree_metadata{ + 22, 32, + { + -1, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, + 5, 5, 6, 6, 7, 8, 9, 10, 11, 12, 13, + -1, 0, 0, 1, 1, 2, 2, 3, 4, + -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, + }, + {22, 9}, + }; + const std::vector tree_slot_ids{ + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, + 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, + }; + CHECK(run_case(backend, tree_case, GGML_TYPE_F16, GGML_TYPE_F16, + &tree_slot_ids, nullptr, &tree_metadata)); + CHECK(run_case(backend, tree_case, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, + &tree_slot_ids, nullptr, &tree_metadata)); + CHECK(run_case(backend, tree_case, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, + &tree_slot_ids, nullptr, &tree_metadata)); + ggml_backend_free(backend); +} + void run_active_slot_case(const TestCase & test_case, const std::vector & active_slot_ids) { ggml_backend_t backend = ggml_backend_cuda_init(0); @@ -466,6 +600,10 @@ TEST_CASE(PagedAttention, CompactThreeSlotBucketMatchesReference) { }, {2, 1, 0, -1}); } +TEST_CASE(PagedAttention, PackedTreesMatchReference) { + run_tree_case(); +} + TEST_CASE(PagedAttention, RaggedCausalPositionsMatchReference) { // Interleaved query rows from two sequences attend the paged pool causally // through per-row positions. Sequence 1 spans 65 logical blocks, so its diff --git a/server/test/test_paged_kv_pool.cpp b/server/test/test_paged_kv_pool.cpp index e7fd60c07..a91d5775e 100644 --- a/server/test/test_paged_kv_pool.cpp +++ b/server/test/test_paged_kv_pool.cpp @@ -401,6 +401,83 @@ TEST_CASE(PagedKvPoolFixture, invalid_arguments) { })); } + +TEST_CASE(PagedKvPoolFixture, cold_block_roundtrip_and_release) { + PagedKvPool pool(/*physical_block_count=*/4, + /*max_sequences=*/2, /*block_size=*/16); + const auto first = acquire(pool, 1); + CHECK(pool.append(first, 33).status == PagedKvStatus::Ok); + CHECK(equals(sequence(pool, first).block_table, {0, 1, 2})); + CHECK(pool.free_block_count() == 1); + + uint32_t released = 99; + CHECK(pool.page_out_block(first, 1, released) == PagedKvStatus::Ok); + CHECK(released == 1); + CHECK(pool.free_block_count() == 2); + CHECK(sequence(pool, first).block_table[1] == PAGED_KV_COLD_BLOCK); + uint32_t resident = 0; + uint32_t owned = 0; + CHECK(pool.resident_block_count(first, resident) == PagedKvStatus::Ok); + CHECK(resident == 2); + CHECK(pool.owned_block_count(first, owned) == PagedKvStatus::Ok); + CHECK(owned == 3); // logical appended capacity includes the cold block + + uint32_t unchanged = 123; + CHECK(pool.page_out_block(first, 1, unchanged) == + PagedKvStatus::BlockNotResident); + CHECK(unchanged == 123); + CHECK(pool.page_in_block(first, 99, unchanged) == + PagedKvStatus::LogicalBlockOutOfRange); + CHECK(unchanged == 123); + + uint32_t restored = 99; + CHECK(pool.page_in_block(first, 1, restored) == PagedKvStatus::Ok); + CHECK(restored == 1); + CHECK(equals(sequence(pool, first).block_table, {0, 1, 2})); + CHECK(pool.page_in_block(first, 1, unchanged) == + PagedKvStatus::BlockAlreadyResident); + + CHECK(pool.page_out_block(first, 1, released) == PagedKvStatus::Ok); + // release() skips the sentinel and returns only resident blocks. + CHECK(pool.release(first) == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 4); +} + +TEST_CASE(PagedKvPoolFixture, append_remaps_cold_partial_head_atomically) { + PagedKvPool pool(/*physical_block_count=*/2, + /*max_sequences=*/2, /*block_size=*/16); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pool.append(first, 8).status == PagedKvStatus::Ok); + CHECK(pool.append(second, 16).status == PagedKvStatus::Ok); + CHECK(pool.free_block_count() == 0); + + uint32_t released = 99; + CHECK(pool.page_out_block(first, 0, released) == PagedKvStatus::Ok); + CHECK(released == 0); + auto appended = pool.append(first, 4); + CHECK(appended.status == PagedKvStatus::Ok); + CHECK(appended.remapped_cold_blocks.size() == 1); + CHECK(appended.remapped_cold_blocks[0].logical_block == 0); + CHECK(appended.remapped_cold_blocks[0].physical_block == 0); + CHECK(appended.write_slots.front().physical_block == 0); + CHECK(appended.write_slots.front().logical_position == 8); + + // A cold head plus a newly-opened block needs two physical allocations. + CHECK(pool.page_out_block(first, 0, released) == PagedKvStatus::Ok); + const auto before = sequence(pool, first); + appended = pool.append(first, 8); // positions 12..19 + CHECK(appended.status == PagedKvStatus::BlocksExhausted); + const auto after = sequence(pool, first); + CHECK(after.kv_seq_len == before.kv_seq_len); + CHECK(after.block_table == before.block_table); + CHECK(after.block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(appended.remapped_cold_blocks.empty()); + + pool.reset(); + CHECK(pool.free_block_count() == 2); +} + TEST_CASE(PagedKvPoolFixture, auto_pool_sizing) { PagedKvAutoBudget budget; budget.free_bytes = 10'000; diff --git a/server/test/test_paged_kv_residency.cpp b/server/test/test_paged_kv_residency.cpp new file mode 100644 index 000000000..d39a5a0a8 --- /dev/null +++ b/server/test/test_paged_kv_residency.cpp @@ -0,0 +1,559 @@ +// Pure-host tests for multi-sequence paged K/V residency. Transfers use a +// deterministic mock device array but preserve async queue/barrier semantics. + +#define GENERATE_UNIT_TEST_MAIN +#include "CppUnitTestFramework.hpp" +#include "common/concurrency/paged_kv_residency.h" + +#include +#include +#include +#include +#include +#include +#include + +using namespace dflash::common; + +namespace { +struct PagedKvResidencyFixture {}; + +PagedKvSequenceHandle acquire(PagedKvPool & pool, uint64_t request) { + PagedKvSequenceHandle handle; + if (pool.acquire(request, handle) != PagedKvStatus::Ok) { + throw std::runtime_error("acquire failed"); + } + return handle; +} + +PagedKvSequenceSnapshot snapshot(PagedKvPool & pool, + PagedKvSequenceHandle handle) { + PagedKvSequenceSnapshot out; + if (pool.sequence(handle, out) != PagedKvStatus::Ok) { + throw std::runtime_error("sequence failed"); + } + return out; +} + +struct MockTransfers { + struct Pending { + std::function apply; + }; + + explicit MockTransfers(uint32_t blocks, size_t bytes) + : block_bytes(bytes), device((size_t)blocks * bytes, 0) {} + + PagedKvResidencyTransferOps callbacks() { + return { + [this](size_t bytes) -> void * { + if (fail_alloc || bytes != block_bytes) return nullptr; + allocations++; + return new uint8_t[bytes]; + }, + [this](void * ptr) { + frees++; + delete[] static_cast(ptr); + }, + [this](PagedKvSequenceHandle, uint32_t, uint32_t physical, + void * host, size_t bytes) { + if (fail_copy_out || bytes != block_bytes) return false; + pending.push_back({[this, physical, host, bytes] { + std::memcpy(host, &device[(size_t)physical * block_bytes], bytes); + }}); + return true; + }, + [this](PagedKvSequenceHandle, uint32_t, uint32_t physical, + const void * host, size_t bytes) { + if (fail_copy_in || bytes != block_bytes) return false; + pending.push_back({[this, physical, host, bytes] { + std::memcpy(&device[(size_t)physical * block_bytes], host, bytes); + }}); + return true; + }, + [this] { + syncs++; + if (fail_sync) return false; + for (Pending & op : pending) op.apply(); + pending.clear(); + return true; + }, + }; + } + + void fill(uint32_t physical, uint8_t value) { + std::fill_n(&device[(size_t)physical * block_bytes], block_bytes, value); + } + + bool block_is(uint32_t physical, uint8_t value) const { + const auto begin = device.begin() + (size_t)physical * block_bytes; + return std::all_of(begin, begin + block_bytes, + [value](uint8_t byte) { return byte == value; }); + } + + size_t block_bytes; + std::vector device; + std::vector pending; + int allocations = 0; + int frees = 0; + int syncs = 0; + bool fail_alloc = false; + bool fail_copy_out = false; + bool fail_copy_in = false; + bool fail_sync = false; +}; + +PagedKvResidencyConfig config(size_t block_bytes, uint32_t budget, + uint32_t sink = 0, uint32_t tail = 0) { + return {block_bytes, budget, sink, tail}; +} + +} // namespace + +TEST_CASE(PagedKvResidencyFixture, page_roundtrip_is_bit_exact_and_barriered) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/1, /*block size=*/4); + MockTransfers io(3, 32); + PagedKvResidencyManager pager(pool, config(32, 3), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 12)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + const auto table = snapshot(pool, handle).block_table; + io.fill(table[0], 0x11); + io.fill(table[1], 0x22); + io.fill(table[2], 0x33); + + CHECK(pager.evict_block(handle, 1) == PagedKvResidencyStatus::Ok); + CHECK(io.syncs == 1); + CHECK(snapshot(pool, handle).block_table[1] == PAGED_KV_COLD_BLOCK); + io.fill(table[1], 0xEE); // recycled device bytes must not affect backing + + CHECK(pager.ensure_resident(handle, {1}) == PagedKvResidencyStatus::Ok); + CHECK(io.syncs == 2); + const uint32_t restored = snapshot(pool, handle).block_table[1]; + CHECK(io.block_is(restored, 0x22)); + CHECK(pager.stats().page_outs == 1); + CHECK(pager.stats().page_ins == 1); + CHECK(pager.stats().resident_blocks == 3); + CHECK(pager.stats().host_bytes == 32); + CHECK(pager.stats().moved_bytes == 64); +} + +TEST_CASE(PagedKvResidencyFixture, fair_share_reclaims_borrowed_pages) { + PagedKvPool pool(/*physical blocks=*/6, /*sequences=*/2, /*block size=*/4); + MockTransfers io(6, 16); + PagedKvResidencyManager pager(pool, config(16, 6), io.callbacks()); + const auto first = acquire(pool, 1); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 20)); // borrows five of six pages while alone + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(first) == 6); + + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(first) == 3); + CHECK(pager.fair_quota(second) == 3); + CHECK(pager.append(second, 12)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().page_outs == 2); + CHECK(pager.stats().resident_blocks == 6); + + uint32_t first_resident = 0; + uint32_t second_resident = 0; + CHECK(pool.resident_block_count(first, first_resident) == PagedKvStatus::Ok); + CHECK(pool.resident_block_count(second, second_resident) == PagedKvStatus::Ok); + CHECK(first_resident == 3); + CHECK(second_resident == 3); +} + +TEST_CASE(PagedKvResidencyFixture, sink_and_tail_are_never_auto_evicted) { + PagedKvPool pool(/*physical blocks=*/5, /*sequences=*/2, /*block size=*/4); + MockTransfers io(5, 16); + PagedKvResidencyManager pager( + pool, config(16, 5, /*sink=*/1, /*tail=*/1), io.callbacks()); + const auto first = acquire(pool, 1); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 16)); // logical blocks 0..3 + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 8)); // needs two pages; only one was free + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + const auto first_table = snapshot(pool, first).block_table; + CHECK(first_table[0] != PAGED_KV_COLD_BLOCK); // sink + CHECK(first_table[3] != PAGED_KV_COLD_BLOCK); // tail + CHECK(first_table[1] == PAGED_KV_COLD_BLOCK || + first_table[2] == PAGED_KV_COLD_BLOCK); + CHECK(pager.evict_block(first, 0) == + PagedKvResidencyStatus::NoEvictableBlock); +} + +TEST_CASE(PagedKvResidencyFixture, append_restores_cold_partial_head) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 2)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + const uint32_t head = snapshot(pool, handle).block_table[0]; + io.fill(head, 0xA5); + CHECK(pager.evict_block(handle, 0) == PagedKvResidencyStatus::Ok); + io.fill(head, 0x00); + + const auto append = pager.append(handle, 1); + CHECK(append); + CHECK(append.pool_result.remapped_cold_blocks.empty()); + CHECK(append.pool_result.write_slots[0].logical_position == 2); + CHECK(io.block_is(append.pool_result.write_slots[0].physical_block, 0xA5)); + CHECK(pager.stats().page_ins == 1); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, scores_drive_reselection) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); + MockTransfers io(3, 16); + PagedKvResidencyManager pager(pool, config(16, 3), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 12)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 2) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().resident_blocks == 1); + const auto peer = acquire(pool, 2); + CHECK(pager.register_sequence(peer) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(handle) == 2); + CHECK(pager.set_scores(handle, {1.0f, 2.0f, 9.0f}) == + PagedKvResidencyStatus::Ok); + CHECK(pager.reselect(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().reselects == 1); + CHECK(pager.is_resident(handle, 2)); + CHECK(pager.is_resident(handle, 1)); + CHECK(!pager.is_resident(handle, 0)); +} + +TEST_CASE(PagedKvResidencyFixture, allocation_and_copy_failures_are_explicit) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 8)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + + io.fail_alloc = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::HostAllocationFailed); + CHECK(pager.is_resident(handle, 0)); + io.fail_alloc = false; + io.fail_copy_out = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(pager.is_resident(handle, 0)); +} + +TEST_CASE(PagedKvResidencyFixture, forget_frees_host_backing_and_stale_is_rejected) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto old_handle = acquire(pool, 1); + CHECK(pager.register_sequence(old_handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(old_handle, 8)); + CHECK(pager.commit_pending_writes(old_handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(old_handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(io.allocations == 1); + CHECK(pager.forget_sequence(old_handle) == PagedKvResidencyStatus::Ok); + CHECK(io.frees == 1); + CHECK(pool.release(old_handle) == PagedKvStatus::Ok); + + const auto replacement = acquire(pool, 2); + CHECK(replacement.generation != old_handle.generation); + CHECK(pager.register_sequence(replacement) == PagedKvResidencyStatus::Ok); + CHECK(pager.touch(old_handle, 0) == PagedKvResidencyStatus::StaleHandle); +} + +TEST_CASE(PagedKvResidencyFixture, + forget_retries_failed_barrier_before_releasing_backing) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + + io.fail_sync = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(pool.free_block_count() == 0); + CHECK(io.frees == 0); + CHECK(pager.forget_sequence(handle) == + PagedKvResidencyStatus::TransferFailed); + CHECK(pool.free_block_count() == 0); + CHECK(io.frees == 0); + + // Teardown retries the quarantined barrier instead of being rejected by + // the manager-wide failure latch. + io.fail_sync = false; + CHECK(pager.forget_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(io.frees == 1); + CHECK(pool.release(handle) == PagedKvStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, + invalid_restore_request_does_not_leak_reservations) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/2, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 4)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + + CHECK(pager.ensure_resident(first, {0, 1}) == + PagedKvResidencyStatus::InvalidArgument); + + // The valid prefix remains evictable, allowing the peer to make progress. + CHECK(pager.append(second, 4)); + CHECK(snapshot(pool, first).block_table[0] == PAGED_KV_COLD_BLOCK); +} + +TEST_CASE(PagedKvResidencyFixture, + staged_writes_are_never_recycled_across_slots) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); + MockTransfers io(3, 16); + PagedKvResidencyManager pager(pool, config(16, 3), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + + CHECK(pager.append(first, 8)); + CHECK(pager.append(second, 4)); + const auto first_staged = snapshot(pool, first).block_table; + const auto second_staged = snapshot(pool, second).block_table; + CHECK(first_staged.size() == 2); + CHECK(second_staged.size() == 1); + + const auto blocked = pager.append(first, 4); + CHECK(blocked.status == PagedKvResidencyStatus::NoEvictableBlock); + CHECK(pager.stats().page_outs == 0); + CHECK(pager.evict_block(second, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::NoEvictableBlock); + CHECK(pager.reselect(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.stats().page_outs == 0); + CHECK(snapshot(pool, first).block_table == first_staged); + CHECK(snapshot(pool, second).block_table == second_staged); + + // Simulate one synchronized packed target graph, then make room. The still + // pending peer page cannot be the victim even though another slot grows. + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const auto grown = pager.append(first, 4); + CHECK(grown); + CHECK(pager.stats().page_outs == 1); + CHECK(snapshot(pool, second).block_table == second_staged); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, + cross_slot_eviction_is_visible_in_full_victim_snapshot) { + PagedKvPool pool(/*physical blocks=*/4, /*sequences=*/2, /*block size=*/4); + MockTransfers io(4, 16); + PagedKvResidencyManager pager(pool, config(16, 4), io.callbacks()); + const auto first = acquire(pool, 1); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 12)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const auto before = snapshot(pool, first).block_table; + + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 8)); + const auto victim = snapshot(pool, first).block_table; + const auto requester = snapshot(pool, second).block_table; + CHECK(victim.size() == before.size()); + CHECK(requester.size() == 2); + + size_t cold = victim.size(); + for (size_t logical = 0; logical < victim.size(); ++logical) { + if (victim[logical] == PAGED_KV_COLD_BLOCK) { + CHECK(cold == victim.size()); + cold = logical; + } + } + CHECK(cold < victim.size()); + CHECK(std::find(requester.begin(), requester.end(), before[cold]) != + requester.end()); + CHECK(pager.stats().page_outs == 1); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); +} + +TEST_CASE(PagedKvResidencyFixture, + sixteen_slots_progress_with_adaptive_fair_quotas) { + constexpr uint32_t kSlots = 16; + constexpr uint32_t kPoolBlocks = 96; + PagedKvPool pool(kPoolBlocks, kSlots, /*block size=*/4); + MockTransfers io(kPoolBlocks, 16); + PagedKvResidencyManager pager( + pool, config(16, kPoolBlocks, /*sink=*/1, /*tail=*/4), + io.callbacks()); + + std::vector handles; + handles.reserve(kSlots); + for (uint32_t slot = 0; slot < kSlots; ++slot) { + handles.push_back(acquire(pool, slot + 1)); + CHECK(pager.register_sequence(handles.back()) == + PagedKvResidencyStatus::Ok); + } + for (const auto handle : handles) { + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == + PagedKvResidencyStatus::Ok); + } + for (const auto handle : handles) { + CHECK(pager.append(handle, 28)); // eight logical blocks total + CHECK(pager.commit_pending_writes(handle) == + PagedKvResidencyStatus::Ok); + } + + CHECK(pager.stats().resident_blocks == kPoolBlocks); + CHECK(pager.stats().page_outs == kSlots * 2); + for (const auto handle : handles) { + CHECK(pager.fair_quota(handle) == kPoolBlocks / kSlots); + const auto table = snapshot(pool, handle).block_table; + CHECK(table.size() == 8); + CHECK(table[0] != PAGED_KV_COLD_BLOCK); + for (size_t logical = 4; logical < 8; ++logical) { + CHECK(table[logical] != PAGED_KV_COLD_BLOCK); + } + } +} + +TEST_CASE(PagedKvResidencyFixture, + failed_copy_out_barrier_retains_the_device_mapping) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 8)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + const auto before = snapshot(pool, handle).block_table; + + io.fail_sync = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(snapshot(pool, handle).block_table == before); + CHECK(pool.free_block_count() == 0); + CHECK(pager.stats().page_outs == 0); +} + +TEST_CASE(PagedKvResidencyFixture, + failed_copy_in_barrier_quarantines_mapping_until_stream_drains) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + + io.fail_sync = true; + CHECK(pager.ensure_resident(handle, {0}) == + PagedKvResidencyStatus::TransferFailed); + const uint32_t quarantined = snapshot(pool, handle).block_table[0]; + CHECK(quarantined != PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 0); + CHECK(pager.stats().page_ins == 0); + + // No operation may reuse or mutate the quarantined destination while the + // failed stream is not known to be drained. + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(io.pending.size() == 1); + + io.fail_sync = false; + CHECK(pager.synchronize_before_read() == PagedKvResidencyStatus::Ok); + CHECK(snapshot(pool, handle).block_table[0] == quarantined); + CHECK(pool.free_block_count() == 0); + CHECK(pager.stats().page_ins == 1); +} + +TEST_CASE(PagedKvResidencyFixture, + rebalance_evicts_only_blocks_above_fair_quota) { + PagedKvPool pool(/*physical blocks=*/6, /*sequences=*/2, /*block size=*/4); + MockTransfers io(6, 16); + PagedKvResidencyManager pager(pool, config(16, 6), io.callbacks()); + const auto borrower = acquire(pool, 1); + CHECK(pager.register_sequence(borrower) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(borrower, 20)); + CHECK(pager.commit_pending_writes(borrower) == + PagedKvResidencyStatus::Ok); + + const auto peer = acquire(pool, 2); + CHECK(pager.register_sequence(peer) == PagedKvResidencyStatus::Ok); + CHECK(pager.fair_quota(borrower) == 3); + CHECK(pager.rebalance() == PagedKvResidencyStatus::Ok); + + uint32_t borrower_resident = 0; + CHECK(pool.resident_block_count(borrower, borrower_resident) == + PagedKvStatus::Ok); + CHECK(borrower_resident == 3); + CHECK(pager.stats().page_outs == 2); +} + +TEST_CASE(PagedKvResidencyFixture, + append_reserves_partial_head_and_new_block_together) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/2, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager(pool, config(16, 2), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 2)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(first, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 4)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + const auto grown = pager.append(first, 3); + CHECK(grown); + const auto table = snapshot(pool, first).block_table; + CHECK(table.size() == 2); + CHECK(table[0] != PAGED_KV_COLD_BLOCK); + CHECK(table[1] != PAGED_KV_COLD_BLOCK); + CHECK(snapshot(pool, second).block_table[0] == PAGED_KV_COLD_BLOCK); +} + +TEST_CASE(PagedKvResidencyFixture, + requested_restore_set_is_protected_as_one_batch) { + PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); + MockTransfers io(3, 16); + PagedKvResidencyManager pager(pool, config(16, 3), io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 8)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(first, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(first, 1, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(pager.append(second, 8)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + CHECK(pager.ensure_resident(first, {0, 1}) == + PagedKvResidencyStatus::Ok); + CHECK(pager.is_resident(first, 0)); + CHECK(pager.is_resident(first, 1)); +} diff --git a/server/test/test_qwen_paged_kv_transfer_layout.cpp b/server/test/test_qwen_paged_kv_transfer_layout.cpp new file mode 100644 index 000000000..59c9d0e7a --- /dev/null +++ b/server/test/test_qwen_paged_kv_transfer_layout.cpp @@ -0,0 +1,86 @@ +#define GENERATE_UNIT_TEST_MAIN +#include "CppUnitTestFramework.hpp" +#include "common/concurrency/qwen_paged_kv_transfer.h" + +#include +#include +#include + +using namespace dflash::common; + +namespace { +struct QwenPagedKvTransferLayoutFixture {}; + +QwenPagedKvTensorLayout packed(size_t row_bytes, uint64_t rows, + uint64_t heads) { + return { + row_bytes, + row_bytes, + row_bytes * static_cast(rows), + row_bytes * static_cast(rows) * + static_cast(heads), + rows, + heads, + }; +} +} // namespace + +TEST_CASE(QwenPagedKvTransferLayoutFixture, + packs_mixed_kv_types_and_all_heads) { + const std::vector tensors = { + packed(8, 64, 2), + packed(4, 64, 2), + packed(6, 64, 2), + packed(2, 64, 2), + }; + QwenPagedKvBlockLayout plan; + std::string error; + CHECK(plan_qwen_paged_kv_block_layout( + tensors, /*block_size=*/16, plan, &error)); + CHECK(error.empty()); + CHECK(plan.physical_block_count == 4); + CHECK(plan.tensor_offsets == + std::vector({0, 256, 384, 576})); + CHECK(plan.tensor_head_bytes == + std::vector({128, 64, 96, 32})); + CHECK(plan.block_bytes == 640); +} + +TEST_CASE(QwenPagedKvTransferLayoutFixture, + accepts_row_and_head_padding_without_storing_padding) { + QwenPagedKvTensorLayout tensor; + tensor.row_bytes = 6; + tensor.row_stride = 8; + tensor.head_stride = 520; + tensor.physical_rows = 64; + tensor.heads = 2; + tensor.storage_bytes = 520 + 63 * 8 + 6; + + QwenPagedKvBlockLayout plan; + CHECK(plan_qwen_paged_kv_block_layout( + {tensor}, /*block_size=*/16, plan)); + CHECK(plan.tensor_head_bytes == std::vector({96})); + CHECK(plan.block_bytes == 192); +} + +TEST_CASE(QwenPagedKvTransferLayoutFixture, + rejects_mismatched_rows_short_storage_and_overflow) { + QwenPagedKvBlockLayout plan; + std::string error; + CHECK(!plan_qwen_paged_kv_block_layout( + {packed(8, 64, 2), packed(8, 32, 2)}, 16, plan, &error)); + CHECK(!error.empty()); + + QwenPagedKvTensorLayout short_tensor = packed(8, 64, 2); + short_tensor.storage_bytes--; + CHECK(!plan_qwen_paged_kv_block_layout( + {short_tensor}, 16, plan, &error)); + + QwenPagedKvTensorLayout overflow = packed(8, 64, 2); + overflow.row_bytes = std::numeric_limits::max(); + overflow.row_stride = overflow.row_bytes; + overflow.head_stride = overflow.row_bytes; + overflow.storage_bytes = overflow.row_bytes; + CHECK(!plan_qwen_paged_kv_block_layout( + {overflow}, 16, plan, &error)); +} diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index 6fb5de4d6..05e8a59cb 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -1,5 +1,6 @@ #include "CppUnitTestFramework.hpp" #include "internal.h" +#include "qwen35/graph_builders.h" #include "ggml-backend.h" #include "ggml-cpu.h" @@ -32,6 +33,38 @@ static std::vector get_tensor(const ggml_tensor * tensor) { return values; } +TEST_CASE(RecurrentSnapshotFixture, validates_paged_tree_layout) { + // The packed-tree launch length is logical. KVFlash may keep a much + // smaller physical resident pool, provided every tree scratch slab still + // fits within that pool. + { + ggml_init_params shape_params{}; + shape_params.mem_size = 8 * ggml_tensor_overhead(); + shape_params.no_alloc = true; + ggml_context * shape_ctx = ggml_init(shape_params); + CHECK(shape_ctx != nullptr); + if (shape_ctx) { + TargetCache shape_cache; + shape_cache.n_seq_slots = 2; + shape_cache.paged_block_table = + ggml_new_tensor_2d(shape_ctx, GGML_TYPE_I32, 4, 2); + shape_cache.paged_kv_seq_lens = + ggml_new_tensor_1d(shape_ctx, GGML_TYPE_I32, 2); + shape_cache.attn_k = { + ggml_new_tensor_4d(shape_ctx, GGML_TYPE_F16, 4, 64, 1, 1), + }; + CHECK(dflash::common::detail::validate_target_paged_tree_layout( + shape_cache, 8, 2, 4096, 32, 16)); + CHECK(!dflash::common::detail::validate_target_paged_tree_layout( + shape_cache, 8, 2, 4096, 48, 16)); + CHECK(!dflash::common::detail::validate_target_paged_tree_layout( + shape_cache, 8, 5, 4096, 32, 16)); + ggml_free(shape_ctx); + } + } + +} + TEST_CASE(RecurrentSnapshotFixture, snapshot_and_restore_recurrent_state) { ggml_backend_t backend = ggml_backend_cpu_init(); CHECK(backend != nullptr); diff --git a/server/test/test_seq_batch_plan.cpp b/server/test/test_seq_batch_plan.cpp index 8f3b87178..a771d160b 100644 --- a/server/test/test_seq_batch_plan.cpp +++ b/server/test/test_seq_batch_plan.cpp @@ -119,6 +119,7 @@ int main() { SeqEngine::StepPlan work; work.decode = {{0, 7}}; work.prefills = {{1, 4}}; + CHECK(work.decode[0].allow_speculation); SeqEngine::StepResult good; good.decode.push_back({0, 11, false, {}}); @@ -126,6 +127,35 @@ int main() { 1, SeqEngine::PrefillOutput::Status::advanced, -1, {}}); CHECK(validate_step_result(work, good, 2).empty()); + SeqEngine::StepResult burst = good; + burst.decode[0].committed_tokens = {8, 9, 10}; + burst.decode[0].ddtree_steps = 1; + burst.decode[0].ddtree_accepted_tokens = 3; + burst.decode[0].target_forwards = 1; + CHECK(validate_step_result(work, burst, 2).empty()); + + // A scheduler stop in the committed prefix must hide the remaining burst + // and final pending token. The backend state is discarded at retirement. + std::vector delivered; + const bool delivered_all = consume_decode_output_tokens( + burst.decode[0], [&](int32_t token) { + delivered.push_back(token); + return token != 9; + }); + CHECK(!delivered_all); + CHECK((delivered == std::vector{8, 9})); + + SeqEngine::StepResult malformed_burst = burst; + malformed_burst.decode[0].committed_tokens = {8, -1}; + CHECK(!validate_step_result(work, malformed_burst, 2).empty()); + + SeqEngine::StepPlan speculation_disabled = work; + speculation_disabled.decode[0].allow_speculation = false; + CHECK(!validate_step_result( + speculation_disabled, burst, 2).empty()); + CHECK(validate_step_result( + speculation_disabled, good, 2).empty()); + SeqEngine::StepResult complete = good; complete.prefills[0] = { 1, SeqEngine::PrefillOutput::Status::completed, 12, {}}; @@ -171,6 +201,13 @@ int main() { success_with_error.decode[0].error = "contradictory diagnostic"; CHECK(!validate_step_result(work, success_with_error, 2).empty()); + SeqEngine::StepResult failed_burst = good; + failed_burst.decode[0].failed = true; + failed_burst.decode[0].token = -1; + failed_burst.decode[0].error = "decode failed"; + failed_burst.decode[0].committed_tokens = {8}; + CHECK(!validate_step_result(work, failed_burst, 2).empty()); + SeqEngine::StepResult failed; failed.error = "device compute failed"; CHECK(validate_step_result(work, failed, 2).empty()); diff --git a/server/test/test_seq_engine_contract.cpp b/server/test/test_seq_engine_contract.cpp index be3f031b4..63a62d18a 100644 --- a/server/test/test_seq_engine_contract.cpp +++ b/server/test/test_seq_engine_contract.cpp @@ -23,6 +23,7 @@ struct Faults { bool overconsume_prefill = false; bool drop_second_completion = false; bool retire_leaks = false; + bool burst_when_speculation_disabled = false; }; struct FakeCapabilities { @@ -111,6 +112,10 @@ class FakeSeqEngine final : public SeqEngine { 100 + input.slot + (int32_t)slot.fed.size(), false, {}, }); + if (faults_.burst_when_speculation_disabled && + !input.allow_speculation) { + result.decode.back().committed_tokens.push_back(91); + } } std::vector completed_this_step; @@ -302,6 +307,9 @@ int main() { "omitted an output"}, {"retire-leak", &Faults::retire_leaks, "succeed while a slot is free"}, + {"ignore-speculation-gate", + &Faults::burst_when_speculation_disabled, + "disabled speculation"}, }; for (const Case & test : cases) { diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index a64329e5f..46a5287d7 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -8,6 +8,7 @@ #include "qwen35/concurrency/qwen35_slot_manager.h" #include "host_check.h" +#include #include #include @@ -85,18 +86,22 @@ int main() { 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. + // Decode append stages row allocation and fed-token history; both + // history and cur_pos publish only after the target compute succeeds. 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); + CHECK(mgr.slot(0).sample_history.size() == 20); + CHECK(mgr.slot(0).staged_tokens.size() == 1 && + mgr.slot(0).staged_tokens.back() == 42); mgr.commit_step(0); CHECK(mgr.slot(0).cur_pos == 21); + CHECK(mgr.slot(0).sample_history.size() == 21 && + mgr.slot(0).sample_history.back() == 42); + CHECK(mgr.slot(0).staged_tokens.empty()); // Second admission lands in slot 1 with non-identity rows. auto b = admit(mgr, 2, prompt_tokens(20), greedy_sampler()); @@ -333,6 +338,51 @@ int main() { CHECK(pool.free_block_count() == 0); } + // Accepted-path replay stages multiple rows and publishes history and + // logical position only after the target compute succeeds. + { + PagedKvPool pool(8, 1, /*block_size=*/4); + Qwen35SlotManager mgr(pool, /*max_ctx=*/32, + /*speculative_headroom=*/7); + auto a = admit(mgr, 1, prompt_tokens(3), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(mgr.append_prefill(a.slot, 3).ok); + mgr.commit_prefill(a.slot); + + const int32_t accepted[] = {41, 42, 43, 44, 45, 46}; + auto staged = mgr.append_tokens(a.slot, accepted, 6); + CHECK(staged.ok && !staged.busy && staged.count == 6); + CHECK(staged.position == 3 && staged.physical_rows.size() == 6); + CHECK(staged.first_new_block == 1); + CHECK(staged.new_blocks.size() == 2); + CHECK(mgr.slot(a.slot).cur_pos == 3); + CHECK(mgr.slot(a.slot).sample_history.size() == 3); + CHECK(mgr.slot(a.slot).staged_tokens == + std::vector(accepted, accepted + 6)); + CHECK(!mgr.append_token(a.slot, 99).ok); + + // A failed target step never publishes staged history. Scheduler + // retirement releases both materialized rows and the staged host range. + const uint32_t free_while_staged = pool.free_block_count(); + mgr.retire(a.slot); + CHECK(!mgr.is_active(a.slot)); + CHECK(pool.free_block_count() > free_while_staged); + + a = admit(mgr, 2, prompt_tokens(3), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(mgr.append_prefill(a.slot, 3).ok); + mgr.commit_prefill(a.slot); + staged = mgr.append_tokens(a.slot, accepted, 6); + CHECK(staged.ok); + + mgr.commit_step(a.slot); + CHECK(mgr.slot(a.slot).cur_pos == 9); + CHECK(mgr.slot(a.slot).staged_tokens.empty()); + CHECK(mgr.slot(a.slot).sample_history.size() == 9); + CHECK(std::equal(accepted, accepted + 6, + mgr.slot(a.slot).sample_history.end() - 6)); + } + // Context exhaustion: append_token refuses past max_ctx. { PagedKvPool pool(4, 1, /*block_size=*/16); diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 8373859dc..2ff471a3d 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -21,6 +21,7 @@ #include "server/http_server.h" #include "server/chat_template.h" #include "common/sampler.h" +#include "common/concurrency/seq_engine.h" #include "common/backend_precision.h" #include "common/backend_ipc.h" #include "common/moe_hybrid_ffn_eval.h" @@ -1947,6 +1948,33 @@ TEST_CASE(ServerUnitFixture, test_stop_sequence_holdback_extends) { TEST_ASSERT(em.accumulated_text().find("suffix") == std::string::npos); } +TEST_CASE(ServerUnitFixture, + test_concurrent_scheduler_burst_stops_at_eos) { + SeqEngine::DecodeOutput burst; + burst.slot = 0; + burst.committed_tokens = {101, 2, 103}; + burst.token = 104; + + std::vector emitted; + int completion_tokens = 0; + const bool consumed_all = consume_decode_output_tokens( + burst, [&](int32_t token) { + emitted.push_back(token); + ++completion_tokens; + // Mirrors scheduler.cpp: advance_slot marks the slot finished on + // EOS and its callback immediately stops the rest of the burst. + return token != 2; + }); + + TEST_ASSERT(!consumed_all); + TEST_ASSERT((emitted == std::vector{101, 2})); + TEST_ASSERT(completion_tokens == 2); + TEST_ASSERT(std::find(emitted.begin(), emitted.end(), 103) == + emitted.end()); + TEST_ASSERT(std::find(emitted.begin(), emitted.end(), 104) == + emitted.end()); +} + // ═══════════════════════════════════════════════════════════════════════ // Prefix cache hash tests (model-free) // ═══════════════════════════════════════════════════════════════════════ @@ -2369,6 +2397,49 @@ TEST_CASE(ServerUnitFixture, test_pflash_config_modes) { TEST_ASSERT(cfg.pflash_mode != ServerConfig::PflashMode::AUTO); } +TEST_CASE(ServerUnitFixture, test_concurrent_pflash_persistent_forces_skip_park) { + ServerConfig cfg; + cfg.pflash_mode = ServerConfig::PflashMode::AUTO; + cfg.draft_residency = DraftResidencyPolicy::Persistent; + cfg.prefix_cache_cap = 0; + + const auto plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/true); + TEST_ASSERT(plan.ok()); + TEST_ASSERT(plan.enabled); + TEST_ASSERT(plan.force_skip_park); +} + +TEST_CASE(ServerUnitFixture, test_concurrent_pflash_rejects_unsafe_residency) { + ServerConfig cfg; + cfg.pflash_mode = ServerConfig::PflashMode::ALWAYS; + cfg.draft_residency = DraftResidencyPolicy::Auto; + cfg.prefix_cache_cap = 0; + + auto plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/true); + TEST_ASSERT(!plan.ok()); + TEST_ASSERT(plan.error.find("--draft-residency persistent") != + std::string::npos); + + cfg.draft_residency = DraftResidencyPolicy::Persistent; + plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/false); + TEST_ASSERT(!plan.ok()); + TEST_ASSERT(plan.error.find("--prefill-drafter") != std::string::npos); +} + +TEST_CASE(ServerUnitFixture, test_concurrent_pflash_rejects_snapshot_caches) { + ServerConfig cfg; + cfg.pflash_mode = ServerConfig::PflashMode::AUTO; + cfg.draft_residency = DraftResidencyPolicy::Persistent; + // Default prefix_cache_cap is intentionally non-zero. + const auto plan = resolve_concurrent_pflash_plan( + cfg, /*drafter_tokenizer_available=*/true); + TEST_ASSERT(!plan.ok()); + TEST_ASSERT(plan.error.find("snapshots") != std::string::npos); +} + TEST_CASE(ServerUnitFixture, test_pflash_compress_request_struct) { ModelBackend::CompressRequest req; req.input_ids = {1, 2, 3, 4, 5}; From 3ad0796906fb085156a472e9066a49ed65e5b5d8 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 12 Aug 2026 22:55:43 +0000 Subject: [PATCH 02/18] bench(qwen36): verify concurrent feature matrix --- .../benchmarks/concurrency/FEATURE_MATRIX.md | 113 ++++ harness/benchmarks/concurrency/README.md | 70 +++ .../concurrency/concurrent_benchmark.py | 348 +++++++++++ .../feature_concurrent_benchmark.py | 255 ++++++++ .../concurrency/generate_feature_prompts.py | 57 ++ .../concurrency/generate_ragged_prompts.py | 111 ++++ .../concurrency/record_feature_runtime.py | 119 ++++ .../concurrency/run_qwen36_concurrency.sh | 194 ++++++ .../concurrency/run_qwen36_feature_matrix.sh | 362 ++++++++++++ .../concurrency/summarize_concurrency.py | 219 +++++++ .../concurrency/summarize_feature_matrix.py | 208 +++++++ .../concurrency/test_concurrency_tools.py | 218 +++++++ .../concurrency/test_concurrent_benchmark.py | 131 +++++ .../test_feature_concurrent_benchmark.py | 123 ++++ .../concurrency/test_feature_metadata.py | 153 +++++ .../concurrency/test_feature_tools.py | 556 ++++++++++++++++++ .../concurrency/verify_feature_metrics.py | 322 ++++++++++ .../concurrency/write_feature_metadata.py | 164 ++++++ 18 files changed, 3723 insertions(+) create mode 100644 harness/benchmarks/concurrency/FEATURE_MATRIX.md create mode 100644 harness/benchmarks/concurrency/README.md create mode 100755 harness/benchmarks/concurrency/concurrent_benchmark.py create mode 100755 harness/benchmarks/concurrency/feature_concurrent_benchmark.py create mode 100644 harness/benchmarks/concurrency/generate_feature_prompts.py create mode 100755 harness/benchmarks/concurrency/generate_ragged_prompts.py create mode 100644 harness/benchmarks/concurrency/record_feature_runtime.py create mode 100755 harness/benchmarks/concurrency/run_qwen36_concurrency.sh create mode 100755 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh create mode 100755 harness/benchmarks/concurrency/summarize_concurrency.py create mode 100755 harness/benchmarks/concurrency/summarize_feature_matrix.py create mode 100644 harness/benchmarks/concurrency/test_concurrency_tools.py create mode 100644 harness/benchmarks/concurrency/test_concurrent_benchmark.py create mode 100644 harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py create mode 100644 harness/benchmarks/concurrency/test_feature_metadata.py create mode 100644 harness/benchmarks/concurrency/test_feature_tools.py create mode 100644 harness/benchmarks/concurrency/verify_feature_metrics.py create mode 100644 harness/benchmarks/concurrency/write_feature_metadata.py diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md new file mode 100644 index 000000000..7a1885f18 --- /dev/null +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -0,0 +1,113 @@ +# Qwen3.6 concurrent feature matrix + +`run_qwen36_feature_matrix.sh` extends the PR #596 protocol with feature +ablations for the complete Strix Halo configuration: + +- `ar`: concurrent paged autoregressive control. +- `ddtree`: adds the decode draft, DDTree, and the recorded budget. +- `pflash`: adds auto prefill compression, its drafter, and persistent draft + residency. +- `kvflash`: adds bounded KV residency in auto mode and explicitly supplies + the hashed prefill drafter for relevance-scored page selection; prefill + compression remains off in this ablation. +- `full`: enables DDTree, PFlash, and KVFlash together with both devices on + `hip:0`. +- `llama`: optional external comparison; its binary is required only when this + variant is explicitly requested. + +Run the default bounded C4 screening repeat (seven applicable fresh-server cases): + +```bash +MODEL=/opt/models/Qwen3.6-27B-Q4_K_M.gguf \ +DRAFT_MODEL=/opt/models/draft/dflash-draft-3.6-q4_k_m.gguf \ +PREFILL_DRAFTER=/opt/models/Qwen3-0.6B-BF16.gguf \ +REPEATS=1 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +On a 128 GiB Strix Halo host, budget roughly 45–90 minutes for this smoke run; +the long-context AR controls dominate and actual time depends on the build. +Every row remains independently selectable through `VARIANTS`. For example: + +```bash +WORKLOADS=short CLIENTS=4 VARIANTS=ar,ddtree MAX_TOKENS=256 REPEATS=1 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +Use five fresh-process, paired repeats for published measurements: + +```bash +WORKLOADS=short,compression CLIENTS=1,4,8,16 \ +VARIANTS=ar,ddtree,pflash,kvflash,full MAX_TOKENS=256 REPEATS=5 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +That full matrix can take roughly 15–30 hours on Strix Halo; keep the generated +case directory so interrupted or suspect rows can be diagnosed rather than +quoted. + +## Activation workloads + +Auto features cannot be validated with the original 400–4,000 word prompts. +The extension adds two deterministic, disjoint 29-prompt cohorts: + +- `compression`: 34K–40K words, chosen after observing 38,130–44,856 + tokens with the development Qwen GGUF tokenizer. +- `kv-pressure`: 12K–18K words, which produced 13,463–20,190 tokens with + that tokenizer against the runner's explicitly recorded 8K pool cap. + +The observed counts are a fixture sanity check, not a claim derived from word +count and not publication evidence. Each row records the model hash; the proof +cross-checks logged raw PFlash input against `usage.prompt_tokens`, requires +it to meet the recorded auto threshold, and uses server-reported effective +tokens plus actual page traffic for KVFlash. + +The bounded default uses `short,compression` at C4. It runs PFlash and the +full configuration only on `compression`; KVFlash can also be selected on +`kv-pressure`. Inapplicable pairs are printed as skips and never appear as +successful rows. AR controls use the same prompts, so feature deltas remain +paired. + +## Fail-closed feature proof + +The server must write one JSON object per completed request with this prefix: + +```text +[concurrency-metrics] {"request_id":"...", ...} +``` + +Required fields are `effective_prompt_tokens`, `ddtree_steps`, +`ddtree_accepted_tokens`, `target_forwards`, `kvflash_page_ins`, +`kvflash_page_outs`, `kvflash_resident_blocks`, `kvflash_reselects`, +`pflash_applied`, `pflash_input_tokens`, and `pflash_output_tokens`. + +The proof tool correlates log objects with measured SSE request IDs and also +checks the log's effective token count against +`usage.timings.effective_prompt_tokens`. A requested feature fails the case +unless: + +- DDTree has positive step and target-forward counts. Acceptance may be zero. +- PFlash reports `pflash_applied=true`, a smaller output prompt, and (in auto + mode) an input token count at or above the recorded activation threshold. +- KVFlash always records an explicit hashed scorer drafter, reports its + startup-observed physical pool and enabled metadata, and has a positive + resident-block count. `kvflash`-only and `kv-pressure` rows must also show page-in or + page-out traffic. For a `full` row, traffic is required only when a + server-reported `effective_prompt_tokens` value exceeds that observed pool + token limit; zero traffic is valid when PFlash compression fits in the pool. + +After health succeeds, the runner fail-closed parses the backend's +`[parallel-kvflash] physical resident pool ...` and `[paged-attention] ...` +startup markers into `runtime_observed`. This distinguishes the actual resident +pool from both the requested auto cap and `--kv-pool-tokens`, which concurrent +KVFlash intentionally does not use to expand VRAM. + +Each case retains the exact shell-escaped command, controlled launch +environment, literal client process arguments and client-script hash, +binary/shared-library/target/draft/PFlash-and-KV-scorer hashes, the ordered +`literal_screenshot_flags` array, all feature values, raw request report, +server log, and `feature-proof.json`. The summary refuses to +include a Lucebox row whose proof is missing or invalid. + + + diff --git a/harness/benchmarks/concurrency/README.md b/harness/benchmarks/concurrency/README.md new file mode 100644 index 000000000..b39db289e --- /dev/null +++ b/harness/benchmarks/concurrency/README.md @@ -0,0 +1,70 @@ +# Qwen3.6 concurrency benchmark + +This protocol measures the serving behavior targeted by packed continuous +prefill and concurrent decode. It is intentionally small: one streaming client, +one fresh-process runner, one deterministic prompt generator, and one summary +script. + +Run a quick screening repeat: + +```bash +MODEL=/path/to/Qwen3.6-27B-Q4_K_M.gguf \ +LUCE_SERVER_BIN=server/build-hip/dflash_server \ +LLAMA_SERVER_BIN=/path/to/llama-server \ +harness/benchmarks/concurrency/run_qwen36_concurrency.sh +``` + +Run a decode-heavy comparison with the same harness: + +```bash +MODEL=/path/to/Qwen3.6-27B-Q4_K_M.gguf \ +LUCE_SERVER_BIN=server/build-hip/dflash_server \ +LLAMA_SERVER_BIN=/path/to/llama-server \ +WORKLOADS=short MAX_TOKENS=256 VARIANTS=luce-k8,llama REPEATS=3 \ +harness/benchmarks/concurrency/run_qwen36_concurrency.sh +``` + +The short ragged prompts keep admission realistic while 256 forced output +tokens make generation dominate the measured window. Use `REPEATS=5` for +publication. Every measured case starts a fresh server and first runs a +discarded warmup at the same concurrency. The variants are: + +- `luce-k8`: packed prefill with up to eight concurrent prefills. +- `luce-k1`: the same binary/configuration with packing width limited to one. +- `llama`: llama.cpp continuous batching with fixed `-b 2048 -ub 512`. + +The 29 generated prompts are disjoint cohorts for C1/C4/C8/C16. C4 and above +contain four substantial length strata while holding the mean target length +constant. The default short, medium, and long profiles target mean lengths of +400, 1,000, and 3,000 words per request. Those generator targets are not token +counts; reports retain the exact server-observed token counts for the selected +model and tokenizer. The client refuses to wrap or reuse a prompt. + +The headline metric is aggregate output goodput: exact server-reported +completion tokens divided by level wall time. It includes queueing, prefill, +and decode and must not be called decode throughput. + +`Output-window tok/s` divides exact completion tokens by the interval from the +earliest observed first output to the final request completion. It removes the +initial all-prefill interval and is decode-facing, but it can still contain +staggered prefill while later requests await their first token. +`Request decode tok/s` is the median per-request estimate +`(completion_tokens - 1) / (end - first_output)`; it assumes the first +observed streaming event accounts for one token. Neither metric is pure kernel +decode throughput. + +`Prompt tok/s to first` is the sum of server-reported prompt tokens divided by +the latest first-token arrival; it is a useful prefill-facing metric but still +includes admission, queueing, and transport. Report TTFT median/max alongside +all throughput metrics. + +The K8-vs-K1 comparison is the causal packing ablation. The K8-vs-llama +comparison is the product comparison. Five paired repeats, the exact command +and hashes recorded in each case, zero failures, and a fixed declared output +length are required before using results in a post. The standard prefill-facing +protocol uses 64 output tokens; the decode-heavy protocol above uses 256. +Variant gains are computed as the median of same-repeat ratios, not as a ratio +of independently aggregated medians. The summarizer rejects mismatched repeat +sets. It also marks whether each variant produced the same ordered output +hashes across at least two repeats; a one-repeat screen reports stability as +`n/a`, and an unstable result is a correctness warning, not a performance win. diff --git a/harness/benchmarks/concurrency/concurrent_benchmark.py b/harness/benchmarks/concurrency/concurrent_benchmark.py new file mode 100755 index 000000000..773260393 --- /dev/null +++ b/harness/benchmarks/concurrency/concurrent_benchmark.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 +"""Measure end-to-end output goodput and TTFT under concurrent streaming load.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import sys +import threading +import time +import urllib.request +from pathlib import Path +from typing import Any, Iterable + + +def sha256_text(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def load_prompts(path: Path) -> list[str]: + prompts = [] + for line_no, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + line = raw.strip() + if not line: + continue + if line.startswith("{"): + value = json.loads(line).get("prompt") + if not isinstance(value, str) or not value: + raise ValueError(f"{path}:{line_no}: missing string 'prompt'") + prompts.append(value) + else: + prompts.append(line) + if not prompts: + raise ValueError(f"{path}: no prompts") + return prompts + + +def request_prompts(prompts: list[str], count: int, offset: int) -> list[str]: + if offset < 0: + raise ValueError("--prompt-offset must be >= 0") + if offset + count > len(prompts): + raise ValueError( + f"need prompts [{offset}, {offset + count}), but only " + f"{len(prompts)} were supplied; refusing to reuse prompts" + ) + return prompts[offset:offset + count] + + +def iter_sse_data(lines: Iterable[bytes]) -> Iterable[str]: + data: list[str] = [] + for raw in lines: + line = raw.decode("utf-8", errors="replace").rstrip("\r\n") + if not line: + if data: + yield "\n".join(data) + data.clear() + elif line.startswith("data:"): + data.append(line[5:].lstrip()) + if data: + yield "\n".join(data) + + +def stream_request(args: argparse.Namespace, prompt: str) -> dict[str, Any]: + started = time.perf_counter() + first = None + content: list[str] = [] + reasoning: list[str] = [] + completion_tokens = None + prompt_tokens = None + finish_reason = None + done_received = False + error = None + payload = { + "model": args.model, + "messages": [{"role": "user", "content": prompt}], + "stream": True, + "stream_options": {"include_usage": True}, + "max_tokens": args.max_tokens, + "temperature": args.temperature, + "seed": args.seed, + } + if args.ignore_eos: + payload["ignore_eos"] = True + headers = {"Content-Type": "application/json"} + if args.api_key: + headers["Authorization"] = f"Bearer {args.api_key}" + request = urllib.request.Request( + args.base_url.rstrip("/") + "/chat/completions", + data=json.dumps(payload).encode("utf-8"), headers=headers, method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=args.timeout) as response: + for data in iter_sse_data(response): + if data == "[DONE]": + done_received = True + break + event = json.loads(data) + usage = event.get("usage") or {} + if isinstance(usage.get("completion_tokens"), int): + completion_tokens = usage["completion_tokens"] + if isinstance(usage.get("prompt_tokens"), int): + prompt_tokens = usage["prompt_tokens"] + for choice in event.get("choices") or []: + if choice.get("finish_reason") is not None: + finish_reason = choice["finish_reason"] + delta = choice.get("delta") or {} + piece = delta.get("content") + thought = delta.get("reasoning_content") + if isinstance(piece, str) and piece: + first = first or time.perf_counter() + content.append(piece) + if isinstance(thought, str) and thought: + first = first or time.perf_counter() + reasoning.append(thought) + except Exception as exc: # preserve partial timing/output for diagnosis + error = f"{type(exc).__name__}: {exc}" + if error is None and not done_received: + error = "ProtocolError: stream ended before [DONE]" + elif error is None and finish_reason is None: + error = "ProtocolError: stream ended without a terminal finish_reason" + ended = time.perf_counter() + output = "".join(content) + reasoning_output = "".join(reasoning) + decode_duration = ended - first if first is not None and ended > first else None + request_decode_tok_s = ( + (completion_tokens - 1) / decode_duration + if isinstance(completion_tokens, int) and completion_tokens > 0 + and decode_duration is not None else None + ) + return { + "t_start": started, "t_first": first, "t_end": ended, + "duration_s": ended - started, + "ttft_s": first - started if first is not None else None, + "decode_duration_s": decode_duration, + "completion_tokens": completion_tokens, "prompt_tokens": prompt_tokens, + "finish_reason": finish_reason, "done_received": done_received, "error": error, + "content_sha256": sha256_text(output), + "reasoning_content_sha256": sha256_text(reasoning_output), + "content_chars": len(output), "reasoning_content_chars": len(reasoning_output), + "request_output_tok_s": ( + completion_tokens / (ended - started) + if completion_tokens is not None and ended > started else None + ), + "request_decode_tok_s": request_decode_tok_s, + } + + +def run_level( + clients: int, args: argparse.Namespace, prompts: list[str], offset: int, +) -> dict[str, Any]: + selected = request_prompts(prompts, clients, offset) + barrier = threading.Barrier(clients) + records: list[dict[str, Any] | None] = [None] * clients + + def worker(index: int) -> None: + barrier.wait() + record = stream_request(args, selected[index]) + record["prompt_index"] = offset + index + record["prompt_sha256"] = sha256_text(selected[index]) + records[index] = record + + threads = [threading.Thread(target=worker, args=(i,), daemon=True) for i in range(clients)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(args.timeout + 30) + completed = [record for record in records if record is not None] + hung = sum(thread.is_alive() for thread in threads) + failures = hung + sum(record["error"] is not None for record in completed) + ok = [record for record in completed if record["error"] is None] + starts = [record["t_start"] for record in completed] + ends = [record["t_end"] for record in completed] + level_start = min(starts) if starts else time.perf_counter() + wall = max(ends) - level_start if ends else 0.0 + for record in completed: + record["start_offset_s"] = record["t_start"] - level_start + + completion_counts = [r["completion_tokens"] for r in ok] + prompt_counts = [r["prompt_tokens"] for r in ok] + completion_complete = bool(ok) and all(isinstance(v, int) for v in completion_counts) + prompt_complete = bool(ok) and all(isinstance(v, int) for v in prompt_counts) + ttfts = [r["ttft_s"] for r in ok if r["ttft_s"] is not None] + first_window = ( + max(r["start_offset_s"] + r["ttft_s"] for r in ok) + if len(ttfts) == len(ok) and ok else None + ) + first_times = [r["t_first"] for r in ok if r["t_first"] is not None] + output_window = ( + max(r["t_end"] for r in ok) - min(first_times) + if len(first_times) == len(ok) and ok else None + ) + request_decode_rates = [ + r["request_decode_tok_s"] for r in ok + if r.get("request_decode_tok_s") is not None + ] + fixed_valid = ( + failures == 0 and len(ok) == clients + and completion_complete + and all(v == args.max_tokens for v in completion_counts) + ) if args.ignore_eos else None + prompt_hashes = [r["prompt_sha256"] for r in ok] + output_hashes = [ + [r["content_sha256"], r["reasoning_content_sha256"]] for r in ok + ] + digest = lambda value: sha256_text(json.dumps(value, separators=(",", ":"))) + return { + "clients": clients, "requests": clients, "requests_ok": len(ok), + "failures": failures, "wall_s": wall, + "start_skew_s": max(starts) - min(starts) if starts else None, + "completion_tokens_total": sum(completion_counts) if completion_complete else None, + "token_count_complete": completion_complete, + "fixed_token_workload_valid": fixed_valid, + "aggregate_tok_s": ( + sum(completion_counts) / wall if completion_complete and wall > 0 else None + ), + "aggregate_metric": "completion_tokens_per_level_wall_second", + "output_window_s": output_window, + "output_window_tok_s": ( + sum(completion_counts) / output_window + if completion_complete and output_window is not None and output_window > 0 + else None + ), + "output_window_metric": "completion_tokens_per_first_output_to_final_completion_second", + "request_decode_tok_s_median": ( + statistics.median(request_decode_rates) + if len(request_decode_rates) == len(ok) and ok else None + ), + "prompt_tokens_total": sum(prompt_counts) if prompt_complete else None, + "prompt_tokens_min": min(prompt_counts) if prompt_complete else None, + "prompt_tokens_max": max(prompt_counts) if prompt_complete else None, + "prompt_tokens_distinct": len(set(prompt_counts)) if prompt_complete else None, + "prompt_token_count_complete": prompt_complete, + "prompt_to_first_token_s": first_window, + "prompt_tokens_per_s_to_first_token": ( + sum(prompt_counts) / first_window + if prompt_complete and first_window is not None and first_window > 0 else None + ), + "ttft_median_s": statistics.median(ttfts) if ttfts else None, + "ttft_max_s": max(ttfts) if ttfts else None, + "selected_prompt_set_sha256": digest(prompt_hashes), + "selected_output_set_sha256": digest(output_hashes), + "requests_detail": completed, + } + + +def fmt(value: Any, spec: str = ".2f") -> str: + return format(value, spec) if isinstance(value, (int, float)) else "n/a" + + +def markdown(report: dict[str, Any]) -> str: + lines = [ + f"# Concurrent benchmark — {report['label']}", "", + "| C | Ok | Output goodput tok/s | Output-window tok/s | " + "Request decode tok/s | Prompt tok/s to first | Prompt range | " + "TTFT median s | TTFT max s | Wall s |", + "| ---: | ---: | ---: | ---: | ---: | ---: | :--- | ---: | ---: | ---: |", + ] + for level in report["levels"]: + lines.append( + f"| {level['clients']} | {level['requests_ok']}/{level['requests']} | " + f"{fmt(level['aggregate_tok_s'])} | " + f"{fmt(level['output_window_tok_s'])} | " + f"{fmt(level['request_decode_tok_s_median'])} | " + f"{fmt(level['prompt_tokens_per_s_to_first_token'])} | " + f"{fmt(level['prompt_tokens_min'], '.0f')}–{fmt(level['prompt_tokens_max'], '.0f')} | " + f"{fmt(level['ttft_median_s'], '.3f')} | {fmt(level['ttft_max_s'], '.3f')} | " + f"{fmt(level['wall_s'])} |" + ) + return "\n".join(lines) + "\n" + + +def level_failed(level: dict[str, Any], ignore_eos: bool) -> bool: + return bool( + level["failures"] + or not level["token_count_complete"] + or not level["prompt_token_count_complete"] + or (ignore_eos and level["fixed_token_workload_valid"] is not True) + ) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18080/v1") + parser.add_argument("--api-key", default="") + parser.add_argument("--model", default="luce-dflash") + parser.add_argument("--clients", type=int, action="append", dest="client_levels") + parser.add_argument("--prompt-file", type=Path, required=True) + parser.add_argument("--prompt-offset", type=int, default=0) + parser.add_argument("--require-distinct-prompts", action="store_true", + help="Compatibility flag; this client always refuses reuse") + parser.add_argument("--max-tokens", type=int, default=64) + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--seed", type=int, default=1) + parser.add_argument("--ignore-eos", action="store_true") + parser.add_argument("--timeout", type=float, default=1200.0) + parser.add_argument("--cooldown", type=float, default=0.0) + parser.add_argument("--server-metadata-json", type=Path) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--label", default="") + return parser + + +def run(args: argparse.Namespace) -> int: + levels = args.client_levels or [1, 4, 8, 16] + if any(level < 1 for level in levels): + raise ValueError("--clients must be positive") + if args.prompt_offset < 0 or args.max_tokens < 1 or args.timeout <= 0: + raise ValueError("invalid offset, max-tokens, or timeout") + prompts = load_prompts(args.prompt_file) + results = [] + offset = args.prompt_offset + for index, clients in enumerate(levels): + if index and args.cooldown > 0: + time.sleep(args.cooldown) + print(f"[bench] C={clients} max_tokens={args.max_tokens}", flush=True) + results.append(run_level(clients, args, prompts, offset)) + offset += clients + metadata = ( + json.loads(args.server_metadata_json.read_text(encoding="utf-8")) + if args.server_metadata_json else {} + ) + report = { + "schema_version": 2, "label": args.label, "base_url": args.base_url, + "model": args.model, "max_tokens": args.max_tokens, + "temperature": args.temperature, "seed": args.seed, + "ignore_eos": args.ignore_eos, "prompt_offset": args.prompt_offset, + "prompt_file_sha256": hashlib.sha256(args.prompt_file.read_bytes()).hexdigest(), + "server_metadata": metadata, "levels": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(markdown(report), end="") + bad = any(level_failed(level, args.ignore_eos) for level in results) + return 1 if bad else 0 + + +def main() -> int: + try: + return run(build_parser().parse_args()) + except Exception as exc: + print(f"[bench] error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/feature_concurrent_benchmark.py new file mode 100755 index 000000000..bc509c590 --- /dev/null +++ b/harness/benchmarks/concurrency/feature_concurrent_benchmark.py @@ -0,0 +1,255 @@ +#!/usr/bin/env python3 +"""Concurrency client with request IDs and effective-prompt telemetry.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import sys +import time +import urllib.request +from pathlib import Path +from typing import Any + +import concurrent_benchmark as base + + +CLIENT_SCRIPT = Path(__file__).resolve() + + +def client_provenance(argv: list[str] | None = None) -> dict[str, Any]: + """Return the literal process argv and the exact client source digest.""" + process_argv = list(sys.orig_argv if argv is None else argv) + if not process_argv or not all(isinstance(value, str) for value in process_argv): + raise ValueError("client process argv must be a non-empty string array") + return { + "client_argv": process_argv, + "client_script": str(CLIENT_SCRIPT), + "client_script_sha256": hashlib.sha256(CLIENT_SCRIPT.read_bytes()).hexdigest(), + } + + +def stream_request(args: argparse.Namespace, prompt: str) -> dict[str, Any]: + started = time.perf_counter() + first = None + request_id = None + content: list[str] = [] + reasoning: list[str] = [] + completion_tokens = None + prompt_tokens = None + finish_reason = None + done_received = False + timings: dict[str, Any] = {} + wire_metrics: dict[str, Any] = {} + error = None + payload = { + "model": args.model, + "messages": [{"role": "user", "content": prompt}], + "stream": True, + "stream_options": {"include_usage": True}, + "max_tokens": args.max_tokens, + "temperature": args.temperature, + "seed": args.seed, + } + if args.ignore_eos: + payload["ignore_eos"] = True + headers = {"Content-Type": "application/json"} + if args.api_key: + headers["Authorization"] = f"Bearer {args.api_key}" + request = urllib.request.Request( + args.base_url.rstrip("/") + "/chat/completions", + data=json.dumps(payload).encode("utf-8"), headers=headers, method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=args.timeout) as response: + for data in base.iter_sse_data(response): + if data == "[DONE]": + done_received = True + break + event = json.loads(data) + if isinstance(event.get("id"), str): + request_id = event["id"] + usage = event.get("usage") or {} + if type(usage.get("completion_tokens")) is int: + completion_tokens = usage["completion_tokens"] + if type(usage.get("prompt_tokens")) is int: + prompt_tokens = usage["prompt_tokens"] + if isinstance(usage.get("timings"), dict): + timings = dict(usage["timings"]) + # Log telemetry is authoritative, but retaining a future wire + # copy makes reports forward-compatible without weakening proof. + if isinstance(usage.get("concurrency_metrics"), dict): + wire_metrics = dict(usage["concurrency_metrics"]) + for choice in event.get("choices") or []: + if choice.get("finish_reason") is not None: + finish_reason = choice["finish_reason"] + delta = choice.get("delta") or {} + piece = delta.get("content") + thought = delta.get("reasoning_content") + if isinstance(piece, str) and piece: + first = first or time.perf_counter() + content.append(piece) + if isinstance(thought, str) and thought: + first = first or time.perf_counter() + reasoning.append(thought) + except Exception as exc: # retain partial data for diagnosis + error = f"{type(exc).__name__}: {exc}" + if error is None and not done_received: + error = "ProtocolError: stream ended before [DONE]" + elif error is None and finish_reason is None: + error = "ProtocolError: stream ended without a terminal finish_reason" + ended = time.perf_counter() + output = "".join(content) + reasoning_output = "".join(reasoning) + decode_duration = ended - first if first is not None and ended > first else None + request_decode_tok_s = ( + (completion_tokens - 1) / decode_duration + if type(completion_tokens) is int and completion_tokens > 0 + and decode_duration is not None else None + ) + return { + "request_id": request_id, + "t_start": started, "t_first": first, "t_end": ended, + "duration_s": ended - started, + "ttft_s": first - started if first is not None else None, + "decode_duration_s": decode_duration, + "completion_tokens": completion_tokens, "prompt_tokens": prompt_tokens, + "effective_prompt_tokens": timings.get("effective_prompt_tokens"), + "prefilled_tokens": timings.get("prefilled_tokens"), + "cached_prefix_tokens": timings.get("cached_prefix_tokens"), + "cache_hit": timings.get("cache_hit"), + "server_prefill_ms": timings.get("prefill_ms"), + "server_decode_ms": timings.get("decode_ms"), + "server_decode_tokens_per_sec": timings.get("decode_tokens_per_sec"), + "server_timings": timings, + "wire_concurrency_metrics": wire_metrics, + "finish_reason": finish_reason, "done_received": done_received, "error": error, + "content_sha256": base.sha256_text(output), + "reasoning_content_sha256": base.sha256_text(reasoning_output), + "content_chars": len(output), "reasoning_content_chars": len(reasoning_output), + "request_output_tok_s": ( + completion_tokens / (ended - started) + if completion_tokens is not None and ended > started else None + ), + "request_decode_tok_s": request_decode_tok_s, + } + + +# The base client owns the concurrency/barrier/accounting implementation. Its +# module-global hook is intentional: this process runs one benchmark at a time. +base.stream_request = stream_request + + +def enrich_level(level: dict[str, Any]) -> None: + ok = [r for r in level["requests_detail"] if r.get("error") is None] + effective = [r.get("effective_prompt_tokens") for r in ok] + effective_complete = bool(ok) and all(type(v) is int for v in effective) + request_ids = [r.get("request_id") for r in ok] + request_ids_complete = ( + bool(ok) and all(isinstance(v, str) and v for v in request_ids) + and len(set(request_ids)) == len(request_ids) + ) + level.update({ + "request_ids_complete": request_ids_complete, + "effective_prompt_token_count_complete": effective_complete, + "effective_prompt_tokens_total": sum(effective) if effective_complete else None, + "effective_prompt_tokens_min": min(effective) if effective_complete else None, + "effective_prompt_tokens_max": max(effective) if effective_complete else None, + "effective_to_wire_prompt_ratio": ( + sum(effective) / level["prompt_tokens_total"] + if effective_complete and level.get("prompt_tokens_total") else None + ), + }) + for key in ("server_prefill_ms", "server_decode_ms", "server_decode_tokens_per_sec"): + values = [r.get(key) for r in ok if type(r.get(key)) in (int, float)] + level[f"{key}_median"] = statistics.median(values) if len(values) == len(ok) and ok else None + + +def markdown(report: dict[str, Any]) -> str: + lines = [ + f"# Concurrent feature benchmark — {report['label']}", "", + "| C | Ok | Output goodput tok/s | Output-window tok/s | " + "Request decode tok/s | Wire prompt range | Effective prompt range | " + "Effective/wire | TTFT max s |", + "| ---: | ---: | ---: | ---: | ---: | :--- | :--- | ---: | ---: |", + ] + for level in report["levels"]: + lines.append( + f"| {level['clients']} | {level['requests_ok']}/{level['requests']} | " + f"{base.fmt(level['aggregate_tok_s'])} | " + f"{base.fmt(level['output_window_tok_s'])} | " + f"{base.fmt(level['request_decode_tok_s_median'])} | " + f"{base.fmt(level['prompt_tokens_min'], '.0f')}–{base.fmt(level['prompt_tokens_max'], '.0f')} | " + f"{base.fmt(level['effective_prompt_tokens_min'], '.0f')}–" + f"{base.fmt(level['effective_prompt_tokens_max'], '.0f')} | " + f"{base.fmt(level['effective_to_wire_prompt_ratio'], '.3f')} | " + f"{base.fmt(level['ttft_max_s'], '.3f')} |" + ) + return "\n".join(lines) + "\n" + + +def build_parser() -> argparse.ArgumentParser: + parser = base.build_parser() + parser.description = __doc__ + parser.add_argument( + "--require-effective-prompt-telemetry", action="store_true", + help="fail when usage.timings.effective_prompt_tokens is absent", + ) + return parser + + +def run(args: argparse.Namespace) -> int: + levels = args.client_levels or [1, 4, 8, 16] + if any(level < 1 for level in levels): + raise ValueError("--clients must be positive") + if args.prompt_offset < 0 or args.max_tokens < 1 or args.timeout <= 0: + raise ValueError("invalid offset, max-tokens, or timeout") + prompts = base.load_prompts(args.prompt_file) + results = [] + offset = args.prompt_offset + for index, clients in enumerate(levels): + if index and args.cooldown > 0: + time.sleep(args.cooldown) + print(f"[bench] C={clients} max_tokens={args.max_tokens}", flush=True) + level = base.run_level(clients, args, prompts, offset) + enrich_level(level) + results.append(level) + offset += clients + metadata = ( + json.loads(args.server_metadata_json.read_text(encoding="utf-8")) + if args.server_metadata_json else {} + ) + report = { + "schema_version": 3, "label": args.label, "base_url": args.base_url, + "model": args.model, "max_tokens": args.max_tokens, + "temperature": args.temperature, "seed": args.seed, + "ignore_eos": args.ignore_eos, "prompt_offset": args.prompt_offset, + "prompt_file_sha256": hashlib.sha256(args.prompt_file.read_bytes()).hexdigest(), + "server_metadata": metadata, "levels": results, + **client_provenance(), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(markdown(report), end="") + bad = any( + base.level_failed(level, args.ignore_eos) + or not level["request_ids_complete"] + or (args.require_effective_prompt_telemetry + and not level["effective_prompt_token_count_complete"]) + for level in results + ) + return 1 if bad else 0 + + +def main() -> int: + try: + return run(build_parser().parse_args()) + except Exception as exc: + print(f"[bench] error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/generate_feature_prompts.py b/harness/benchmarks/concurrency/generate_feature_prompts.py new file mode 100644 index 000000000..d642b138e --- /dev/null +++ b/harness/benchmarks/concurrency/generate_feature_prompts.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 +"""Generate deterministic long-context cohorts that force feature activation.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from generate_ragged_prompts import ( + PROFILES as BASE_PROFILES, + build_profile_records, + write_records, +) + + +PROFILES = { + **BASE_PROFILES, + # These word counts were chosen after observing 38,130--44,856 tokens + # with the development Qwen GGUF tokenizer. Word count is never treated as + # activation proof: runtime wire/log telemetry is checked against the + # recorded PFlash threshold. The stable word bank keeps hashes distinct. + "compression": (34000, 36000, 38000, 40000), + # The development Qwen GGUF tokenizer produced 13,463--20,190 tokens here. + # Runtime effective-token and paging telemetry, not this estimate, proves + # KVFlash pressure. + "kv-pressure": (12000, 14000, 16000, 18000), +} + + +def build_records(profile: str) -> list[dict[str, object]]: + activation_target = ( + "pflash-auto" if profile == "compression" + else "kvflash-pressure" if profile == "kv-pressure" + else "none" + ) + return build_profile_records( + profile, PROFILES, + lambda _profile: {"activation_target": activation_target}, + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--profile", choices=sorted(PROFILES), required=True) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + records = build_records(args.profile) + try: + write_records(args.out, records) + except FileExistsError as exc: + parser.error(str(exc)) + print(f"wrote {len(records)} prompts to {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/generate_ragged_prompts.py b/harness/benchmarks/concurrency/generate_ragged_prompts.py new file mode 100755 index 000000000..ca0b180a7 --- /dev/null +++ b/harness/benchmarks/concurrency/generate_ragged_prompts.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +"""Generate a small deterministic ragged-prompt manifest for concurrency runs.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Callable, Mapping + + +PROFILES = { + "short": (250, 350, 450, 550), + "medium": (650, 850, 1150, 1350), + "long": (2000, 2600, 3400, 4000), +} + +WORD_BANK = ( + "systems engineers compare latency throughput scheduling memory kernels queues " + "batches requests tokens caches pages attention arithmetic bandwidth occupancy " + "profiling measurement fairness reproducibility workloads concurrency admission " + "prefill decoding evidence tradeoffs implementation validation production service" +).split() + + +def prompt_text(profile: str, cohort: str, index: int, target_words: int) -> str: + prefix = ( + f"Ragged benchmark {profile} cohort {cohort} request {index}. " + "Write a structured engineering analysis of the following observations, " + "including assumptions, likely bottlenecks, and a concise conclusion." + ).split() + words = list(prefix) + cursor = (index * 7 + target_words) % len(WORD_BANK) + while len(words) < target_words: + words.append(WORD_BANK[cursor % len(WORD_BANK)]) + cursor += 1 + return " ".join(words[:target_words]) + + +ExtraFields = Callable[[str], Mapping[str, object]] + + +def build_profile_records( + profile: str, + profiles: Mapping[str, tuple[int, ...]], + extra_fields: ExtraFields | None = None, +) -> list[dict[str, object]]: + """Build the standard disjoint C1/C4/C8/C16 cohort layout.""" + strata = profiles[profile] + if not strata: + raise ValueError(f"profile {profile!r} has no length strata") + layout = [ + ("c1", [sum(strata) // len(strata)]), + ("c4", list(strata)), + ("c8", list(strata) * 2), + ("c16", list(strata) * 4), + ] + records: list[dict[str, object]] = [] + for cohort, targets in layout: + for target in targets: + index = len(records) + record: dict[str, object] = { + "id": f"{profile}-{index:02d}", + "cohort": cohort, + "stratum": strata.index(target) if target in strata else "mean", + "target_words": target, + "prompt": prompt_text(profile, cohort, index, target), + } + if extra_fields is not None: + additions = dict(extra_fields(profile)) + overlap = record.keys() & additions.keys() + if overlap: + raise ValueError( + "extra profile fields may not replace standard fields: " + f"{sorted(overlap)}" + ) + record.update(additions) + records.append(record) + return records + + +def build_records(profile: str) -> list[dict[str, object]]: + return build_profile_records(profile, PROFILES) + + +def write_records(path: Path, records: list[dict[str, object]]) -> None: + if path.exists(): + raise FileExistsError(f"refusing to overwrite {path}") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join(json.dumps(row, sort_keys=True) + "\n" for row in records), + encoding="utf-8", + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--profile", choices=sorted(PROFILES), required=True) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + records = build_records(args.profile) + try: + write_records(args.out, records) + except FileExistsError as exc: + parser.error(str(exc)) + print(f"wrote {len(records)} prompts to {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/record_feature_runtime.py b/harness/benchmarks/concurrency/record_feature_runtime.py new file mode 100644 index 000000000..5665a4961 --- /dev/null +++ b/harness/benchmarks/concurrency/record_feature_runtime.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Add startup-observed pool dimensions to one feature-case metadata file.""" + +from __future__ import annotations + +import argparse +import json +import re +import tempfile +from pathlib import Path +from typing import Any + + +KVFLASH_POOL_RE = re.compile( + r"\[parallel-kvflash\] physical resident pool (?P\d+) tokens; " + r"logical per-slot cap (?P\d+) across (?P\d+) slots" +) +PAGED_POOL_RE = re.compile( + r"\[paged-attention\] (?P\d+) physical blocks x " + r"(?P\d+) tokens \((?P\d+) pool tokens, " + r"per-sequence max_ctx (?P\d+)\)" +) + + +def _one_consistent(matches: list[dict[str, int]], label: str) -> dict[str, int] | None: + if not matches: + return None + first = matches[0] + if any(row != first for row in matches[1:]): + raise ValueError(f"conflicting {label} startup markers: {matches}") + return first + + +def observe_startup(log_text: str) -> dict[str, Any]: + kvflash = _one_consistent( + [ + {key: int(value) for key, value in match.groupdict().items()} + for match in KVFLASH_POOL_RE.finditer(log_text) + ], + "KVFlash pool", + ) + paged = _one_consistent( + [ + {key: int(value) for key, value in match.groupdict().items()} + for match in PAGED_POOL_RE.finditer(log_text) + ], + "paged pool", + ) + if paged and paged["blocks"] * paged["block_size"] != paged["tokens"]: + raise ValueError("paged-attention startup marker has inconsistent dimensions") + if kvflash and paged: + if kvflash["tokens"] != paged["tokens"]: + raise ValueError("KVFlash and paged-attention startup pool sizes disagree") + if kvflash["max_ctx"] != paged["max_ctx"]: + raise ValueError("KVFlash and paged-attention logical max_ctx values disagree") + + return { + "kvflash_active": kvflash is not None, + "physical_kv_pool_tokens": ( + kvflash["tokens"] if kvflash else paged["tokens"] if paged else None + ), + "physical_kv_pool_blocks": paged["blocks"] if paged else None, + "kv_block_size_tokens": paged["block_size"] if paged else None, + "logical_per_slot_max_ctx": ( + kvflash["max_ctx"] if kvflash else paged["max_ctx"] if paged else None + ), + "configured_slots": kvflash["slots"] if kvflash else None, + "proof_sources": { + "kvflash_pool_startup_marker": kvflash is not None, + "paged_pool_startup_marker": paged is not None, + }, + } + + +def update_metadata(metadata: dict[str, Any], log_text: str) -> dict[str, Any]: + observed = observe_startup(log_text) + feature_config = metadata.get("feature_config") or {} + kvflash_mode = feature_config.get("kvflash") + kvflash_requested = isinstance(kvflash_mode, str) and kvflash_mode not in ( + "", "off", "0", + ) + if kvflash_requested and not observed["kvflash_active"]: + raise ValueError( + "KVFlash metadata is enabled but its physical-pool startup marker is missing" + ) + if observed["physical_kv_pool_tokens"] is None: + raise ValueError("paged physical-pool startup marker is missing") + + result = dict(metadata) + result["schema_version"] = max(3, int(result.get("schema_version", 0))) + result["runtime_observed"] = observed + return result + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--metadata", type=Path, required=True) + parser.add_argument("--server-log", type=Path, required=True) + args = parser.parse_args() + + metadata = json.loads(args.metadata.read_text(encoding="utf-8")) + updated = update_metadata( + metadata, + args.server_log.read_text(encoding="utf-8", errors="replace"), + ) + # Replace atomically so a killed run never leaves half-written metadata. + with tempfile.NamedTemporaryFile( + "w", encoding="utf-8", dir=args.metadata.parent, + prefix=f".{args.metadata.name}.", delete=False, + ) as handle: + json.dump(updated, handle, indent=2, sort_keys=True) + handle.write("\n") + temporary = Path(handle.name) + temporary.replace(args.metadata) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh new file mode 100755 index 000000000..b6f76fe95 --- /dev/null +++ b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh @@ -0,0 +1,194 @@ +#!/usr/bin/env bash +# Paired, fresh-process Qwen3.6 concurrency benchmark for Lucebox and llama.cpp. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +REPO="${REPO:-$(cd -- "$SCRIPT_DIR/../../.." && pwd -P)}" +CLIENT="${CLIENT:-$SCRIPT_DIR/concurrent_benchmark.py}" +GENERATOR="${GENERATOR:-$SCRIPT_DIR/generate_ragged_prompts.py}" +SUMMARIZER="${SUMMARIZER:-$SCRIPT_DIR/summarize_concurrency.py}" + +MODEL="${MODEL:-}" +LUCE_SERVER_BIN="${LUCE_SERVER_BIN:-$REPO/server/build-hip/dflash_server}" +LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-$(command -v llama-server 2>/dev/null || true)}" +OUT="${OUT:-$REPO/.harness-runs/qwen36-concurrency-$(date -u +%Y%m%dT%H%M%SZ)}" +REPEATS="${REPEATS:-1}" +WORKLOADS="${WORKLOADS:-short,medium,long}" +VARIANTS="${VARIANTS:-luce-k8,luce-k1,llama}" +CLIENTS="${CLIENTS:-1,4,8,16}" +PORT="${PORT:-18114}" +COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-3}" +HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-600}" +MAX_TOKENS="${MAX_TOKENS:-64}" +WARMUP_TOKENS="${WARMUP_TOKENS:-8}" +SLOTS=16 + +usage() { + cat <<'EOF' +Usage: MODEL=/path/model.gguf [REPEATS=5] run_qwen36_concurrency.sh + +Runs fresh-server, same-concurrency warmup + measurement cases for luce-k8, +luce-k1, and llama at C=1/4/8/16. Defaults to one repeat for screening; use at +least five paired repeats for publication. For a decode-heavy comparison, set +WORKLOADS=short MAX_TOKENS=256 VARIANTS=luce-k8,llama. OUT must not already +exist. +EOF +} + +if [[ "${1:-}" == "--help" ]]; then usage; exit 0; fi +if [[ $# -ne 0 ]]; then usage >&2; exit 2; fi +for cmd in python3 curl sha256sum; do command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; }; done +[[ -r "$MODEL" ]] || { echo "set MODEL to a readable GGUF" >&2; exit 2; } +[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +[[ -x "$LLAMA_SERVER_BIN" ]] || { echo "missing llama.cpp server: $LLAMA_SERVER_BIN" >&2; exit 2; } +[[ "$REPEATS" =~ ^[1-9][0-9]*$ ]] || { echo "REPEATS must be positive" >&2; exit 2; } +[[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } +ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ + | grep -v '^LUCE_SERVER_BIN=' || true)" +if [[ -n "$ambient_tuning" ]]; then + echo "refusing ambient GPU/backend tuning variables:" >&2 + echo "$ambient_tuning" >&2 + exit 2 +fi +MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" + +IFS=, read -r -a workload_list <<< "$WORKLOADS" +IFS=, read -r -a variant_list <<< "$VARIANTS" +IFS=, read -r -a client_list <<< "$CLIENTS" +declare -A prompt_offsets=([1]=0 [4]=1 [8]=5 [16]=13) +for c in "${client_list[@]}"; do + [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } +done +for v in "${variant_list[@]}"; do + [[ "$v" == luce-k8 || "$v" == luce-k1 || "$v" == llama ]] || { echo "unknown variant $v" >&2; exit 2; } +done + +mkdir -p "$OUT/prompts" +for workload in "${workload_list[@]}"; do + python3 "$GENERATOR" --profile "$workload" --out "$OUT/prompts/$workload.jsonl" +done + +server_pid="" +stop_server() { + if [[ -n "$server_pid" ]] && kill -0 "$server_pid" 2>/dev/null; then + kill "$server_pid" 2>/dev/null || true + for _ in $(seq 1 30); do + kill -0 "$server_pid" 2>/dev/null || break + sleep 1 + done + kill -9 "$server_pid" 2>/dev/null || true + wait "$server_pid" 2>/dev/null || true + fi + server_pid="" +} +trap stop_server EXIT INT TERM + +wait_health() { + local deadline=$((SECONDS + HEALTH_TIMEOUT_SECONDS)) + while (( SECONDS < deadline )); do + kill -0 "$server_pid" 2>/dev/null || return 1 + curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && return 0 + sleep 1 + done + return 1 +} + +write_metadata() { + local path="$1" variant="$2" workload="$3" clients="$4" repeat="$5" binary="$6" max_prefills="$7" command_file="$8" + python3 -c 'import hashlib,json,pathlib,subprocess,sys +p,variant,workload,clients,repeat,binary,max_prefills,cmd_file,model_sha,prompts,repo=sys.argv[1:] +digest=lambda x: hashlib.sha256(pathlib.Path(x).read_bytes()).hexdigest() +libs={} +for line in subprocess.run(["ldd",binary],text=True,capture_output=True).stdout.splitlines(): + fields=line.replace("=>"," ").split() + paths=[x for x in fields if x.startswith("/") and pathlib.Path(x).is_file()] + for lib in paths: libs[str(pathlib.Path(lib).resolve())]=digest(lib) +lucebox_git_head=subprocess.run(["git","-C",repo,"rev-parse","HEAD"],text=True,capture_output=True).stdout.strip() or None +server_version=None +if variant == "llama": + version=subprocess.run([binary,"--version"],text=True,capture_output=True,timeout=30) + server_version="\n".join(x.strip() for x in (version.stdout,version.stderr) if x.strip()) or None + if version.returncode != 0 or server_version is None: + raise RuntimeError(f"cannot identify llama.cpp source version from {binary} --version") +obj={"variant":variant,"workload":workload,"clients":int(clients),"repeat":int(repeat), + "max_concurrent_prefills":int(max_prefills),"server_binary":str(pathlib.Path(binary).resolve()), + "server_binary_sha256":digest(binary),"model_sha256":model_sha, + "prompt_file_sha256":digest(prompts),"server_command":pathlib.Path(cmd_file).read_text().strip(), + "resolved_shared_library_sha256":libs, + "lucebox_git_head":lucebox_git_head if variant != "llama" else None, + "server_version":server_version} +pathlib.Path(p).write_text(json.dumps(obj,indent=2,sort_keys=True)+"\n")' \ + "$path" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$command_file" "$MODEL_SHA256" "$OUT/prompts/$workload.jsonl" "$REPO" +} + +run_case() { + local repeat="$1" workload="$2" clients="$3" variant="$4" + local max_ctx timeout capacity max_prefills binary model_id + if [[ "$workload" == long ]]; then + max_ctx=8192; timeout=1800 + else + max_ctx=4096; timeout=1200 + fi + capacity=$((SLOTS * max_ctx)) + local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" + mkdir -p "$case_dir" + local -a command + if [[ "$variant" == llama ]]; then + binary="$LLAMA_SERVER_BIN"; model_id=qwen36-llama; max_prefills=0 + command=("$binary" -m "$MODEL" -ngl all --parallel "$SLOTS" -c "$capacity" + -b 2048 -ub 512 --cont-batching --no-context-shift --no-mmap -fa on + -ctk q4_0 -ctv q4_0 --no-cache-prompt --host 127.0.0.1 --port "$PORT" --alias "$model_id") + else + binary="$LUCE_SERVER_BIN"; model_id=qwen36-luce + [[ "$variant" == luce-k8 ]] && max_prefills=8 || max_prefills=1 + command=("$binary" "$MODEL" --target-device hip:0 --paged-attention + --max-concurrency "$SLOTS" --kv-pool-tokens "$capacity" --max-ctx "$max_ctx" + --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 + --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 + --host 127.0.0.1 --port "$PORT" --model-name "$model_id") + fi + printf '%q ' "${command[@]}" > "$case_dir/server-command.txt"; printf '\n' >> "$case_dir/server-command.txt" + write_metadata "$case_dir/server-metadata.json" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$case_dir/server-command.txt" + + echo "[run] $workload C=$clients repeat=$repeat variant=$variant" + if [[ "$variant" == llama ]]; then + "${command[@]}" > "$case_dir/server.log" 2>&1 & + else + env DFLASH_MIN_TOKENS="$WARMUP_TOKENS" DFLASH_MAX_CONCURRENT_PREFILLS="$max_prefills" \ + "${command[@]}" > "$case_dir/server.log" 2>&1 & + fi + server_pid=$! + if ! wait_health; then tail -n 80 "$case_dir/server.log" >&2 || true; return 1; fi + + local offset="${prompt_offsets[$clients]}" prompts="$OUT/prompts/$workload.jsonl" + python3 "$CLIENT" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" \ + --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" \ + --require-distinct-prompts --max-tokens "$WARMUP_TOKENS" --temperature 0 \ + --ignore-eos --timeout "$timeout" --cooldown 0 --out "$case_dir/warmup.json" \ + --label "$variant $workload C=$clients warmup" > "$case_dir/warmup.txt" + python3 "$CLIENT" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" \ + --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" \ + --require-distinct-prompts --max-tokens "$MAX_TOKENS" --temperature 0 \ + --ignore-eos --timeout "$timeout" --cooldown 0 \ + --server-metadata-json "$case_dir/server-metadata.json" --out "$case_dir/bench.json" \ + --label "$variant $workload C=$clients repeat=$repeat" | tee "$case_dir/bench.txt" + stop_server + sleep "$COOLDOWN_SECONDS" +} + +for ((repeat=1; repeat<=REPEATS; repeat++)); do + for workload in "${workload_list[@]}"; do + for c_index in "${!client_list[@]}"; do + clients="${client_list[$c_index]}" + # Rotate start variant by case so one engine is not always hot or cold. + shift_by=$(((repeat + c_index) % ${#variant_list[@]})) + for ((i=0; i<${#variant_list[@]}; i++)); do + variant="${variant_list[$(((i + shift_by) % ${#variant_list[@]}))]}" + run_case "$repeat" "$workload" "$clients" "$variant" + done + done + done +done + +python3 "$SUMMARIZER" "$OUT" --out "$OUT/summary.md" +echo "[run] complete: $OUT" diff --git a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh new file mode 100755 index 000000000..e4b23bca1 --- /dev/null +++ b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh @@ -0,0 +1,362 @@ +#!/usr/bin/env bash +# Fresh-process Qwen3.6 concurrent feature ablations with fail-closed proof. +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +REPO="${REPO:-$(cd -- "$SCRIPT_DIR/../../.." && pwd -P)}" +CLIENT="${CLIENT:-$SCRIPT_DIR/feature_concurrent_benchmark.py}" +GENERATOR="${GENERATOR:-$SCRIPT_DIR/generate_feature_prompts.py}" +SUMMARIZER="${SUMMARIZER:-$SCRIPT_DIR/summarize_feature_matrix.py}" +PROOF_TOOL="${PROOF_TOOL:-$SCRIPT_DIR/verify_feature_metrics.py}" +METADATA_TOOL="${METADATA_TOOL:-$SCRIPT_DIR/write_feature_metadata.py}" +RUNTIME_METADATA_TOOL="${RUNTIME_METADATA_TOOL:-$SCRIPT_DIR/record_feature_runtime.py}" + +MODEL="${MODEL:-}" +DRAFT_MODEL="${DRAFT_MODEL:-}" +PREFILL_DRAFTER="${PREFILL_DRAFTER:-}" +LUCE_SERVER_BIN="${LUCE_SERVER_BIN:-$REPO/server/build-hip/dflash_server}" +LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-$(command -v llama-server 2>/dev/null || true)}" +OUT="${OUT:-$REPO/.harness-runs/qwen36-feature-matrix-$(date -u +%Y%m%dT%H%M%SZ)}" +REPEATS="${REPEATS:-1}" +WORKLOADS="${WORKLOADS:-short,compression}" +VARIANTS="${VARIANTS:-ar,ddtree,pflash,kvflash,full}" +CLIENTS="${CLIENTS:-4}" +PORT="${PORT:-18114}" +COOLDOWN_SECONDS="${COOLDOWN_SECONDS:-3}" +HEALTH_TIMEOUT_SECONDS="${HEALTH_TIMEOUT_SECONDS:-900}" +MAX_TOKENS="${MAX_TOKENS:-64}" +WARMUP_TOKENS="${WARMUP_TOKENS:-8}" +SLOTS="${SLOTS:-16}" +MAX_CONCURRENT_PREFILLS="${MAX_CONCURRENT_PREFILLS:-8}" + +# The requested Strix Halo configuration. Every value is serialized into case +# metadata; no performance-affecting DFLASH variable is inherited implicitly. +TARGET_DEVICE="${TARGET_DEVICE:-hip:0}" +DRAFT_DEVICE="${DRAFT_DEVICE:-hip:0}" +DDTREE_BUDGET="${DDTREE_BUDGET:-22}" +DRAFT_RESIDENCY="${DRAFT_RESIDENCY:-persistent}" +PREFILL_COMPRESSION="${PREFILL_COMPRESSION:-auto}" +PREFILL_THRESHOLD="${PREFILL_THRESHOLD:-32000}" +PREFILL_KEEP_RATIO="${PREFILL_KEEP_RATIO:-0.05}" +KVFLASH_MODE="${KVFLASH_MODE:-auto}" +KVFLASH_MAX_POOL_TOKENS="${KVFLASH_MAX_POOL_TOKENS:-8192}" + +usage() { + cat <<'EOF' +Usage: + MODEL=/path/Qwen3.6-27B-Q4_K_M.gguf \ + DRAFT_MODEL=/path/dflash-draft-3.6-q4_k_m.gguf \ + PREFILL_DRAFTER=/path/Qwen3-0.6B-BF16.gguf \ + harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +The default is a bounded C4 smoke matrix with independently selectable +ar, ddtree, pflash, kvflash, and full rows. The full row is the requested +Strix Halo configuration: target/draft hip:0, DDTree +budget 22, persistent PFlash auto, and KVFlash auto. The long-context profiles +are intended to cross the recorded 32K PFlash and 8K KV-residency thresholds; +word count is not treated as proof. Per-request wire/log token counts and +activation telemetry fail the case if an "auto" feature did not execute. + +llama is optional: include it explicitly with VARIANTS=ar,ddtree,llama and set +LLAMA_SERVER_BIN. For publication, set CLIENTS=1,4,8,16 and REPEATS=5. +OUT must not already exist. +EOF +} + +if [[ "${1:-}" == "--help" ]]; then usage; exit 0; fi +if [[ $# -ne 0 ]]; then usage >&2; exit 2; fi +for cmd in python3 curl sha256sum awk; do + command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; } +done +[[ -r "$MODEL" ]] || { echo "set MODEL to a readable target GGUF" >&2; exit 2; } +[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +[[ "$REPEATS" =~ ^[1-9][0-9]*$ ]] || { echo "REPEATS must be positive" >&2; exit 2; } +[[ "$SLOTS" =~ ^[1-9][0-9]*$ ]] || { echo "SLOTS must be positive" >&2; exit 2; } +[[ "$MAX_CONCURRENT_PREFILLS" =~ ^[1-9][0-9]*$ ]] || { echo "MAX_CONCURRENT_PREFILLS must be positive" >&2; exit 2; } +[[ "$DDTREE_BUDGET" =~ ^[1-9][0-9]*$ ]] || { echo "DDTREE_BUDGET must be positive" >&2; exit 2; } +[[ "$PREFILL_THRESHOLD" =~ ^[1-9][0-9]*$ ]] || { echo "PREFILL_THRESHOLD must be positive" >&2; exit 2; } +[[ "$KVFLASH_MAX_POOL_TOKENS" =~ ^[1-9][0-9]*$ ]] || { echo "KVFLASH_MAX_POOL_TOKENS must be positive" >&2; exit 2; } +[[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } + +ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ + | grep -v '^LUCE_SERVER_BIN=' || true)" +if [[ -n "$ambient_tuning" ]]; then + echo "refusing ambient GPU/backend tuning variables:" >&2 + echo "$ambient_tuning" >&2 + exit 2 +fi + +IFS=, read -r -a workload_list <<< "$WORKLOADS" +IFS=, read -r -a variant_list <<< "$VARIANTS" +IFS=, read -r -a client_list <<< "$CLIENTS" +reject_duplicates() { + local list_name="$1" value + shift + local -A seen=() + for value in "$@"; do + if [[ -n "${seen[$value]+yes}" ]]; then + echo "$list_name contains duplicate entry: $value" >&2 + return 1 + fi + seen["$value"]=1 + done +} +reject_duplicates CLIENTS "${client_list[@]}" || exit 2 +reject_duplicates VARIANTS "${variant_list[@]}" || exit 2 +declare -A prompt_offsets=([1]=0 [4]=1 [8]=5 [16]=13) +for c in "${client_list[@]}"; do + [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } + (( c <= SLOTS )) || { echo "CLIENTS=$c exceeds SLOTS=$SLOTS" >&2; exit 2; } +done +for v in "${variant_list[@]}"; do + case "$v" in + ar|ddtree|pflash|kvflash|full|llama) ;; + *) echo "unknown variant $v" >&2; exit 2 ;; + esac +done + +contains_variant() { + local needle="$1" value + for value in "${variant_list[@]}"; do [[ "$value" == "$needle" ]] && return 0; done + return 1 +} +if contains_variant ddtree || contains_variant full; then + [[ -r "$DRAFT_MODEL" ]] || { echo "DRAFT_MODEL is required for DDTree/full" >&2; exit 2; } +fi +if contains_variant pflash || contains_variant kvflash || contains_variant full; then + [[ -r "$PREFILL_DRAFTER" ]] || { + echo "PREFILL_DRAFTER is required for PFlash and drafter-scored KVFlash" >&2 + exit 2 + } +fi +if contains_variant llama; then + [[ -x "$LLAMA_SERVER_BIN" ]] || { echo "llama requested but LLAMA_SERVER_BIN is missing" >&2; exit 2; } +fi + +MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" +DRAFT_MODEL_SHA256="" +PREFILL_DRAFTER_SHA256="" +[[ -n "$DRAFT_MODEL" ]] && DRAFT_MODEL_SHA256="$(sha256sum "$DRAFT_MODEL" | awk '{print $1}')" +[[ -n "$PREFILL_DRAFTER" ]] && PREFILL_DRAFTER_SHA256="$(sha256sum "$PREFILL_DRAFTER" | awk '{print $1}')" + +mkdir -p "$OUT/prompts" +for workload in "${workload_list[@]}"; do + python3 "$GENERATOR" --profile "$workload" --out "$OUT/prompts/$workload.jsonl" +done + +server_pid="" +stop_server() { + if [[ -n "$server_pid" ]] && kill -0 "$server_pid" 2>/dev/null; then + kill "$server_pid" 2>/dev/null || true + for _ in $(seq 1 30); do + kill -0 "$server_pid" 2>/dev/null || break + sleep 1 + done + kill -9 "$server_pid" 2>/dev/null || true + wait "$server_pid" 2>/dev/null || true + fi + server_pid="" +} +trap stop_server EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +wait_health() { + local deadline=$((SECONDS + HEALTH_TIMEOUT_SECONDS)) + while (( SECONDS < deadline )); do + kill -0 "$server_pid" 2>/dev/null || return 1 + curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && return 0 + sleep 1 + done + return 1 +} + +port_is_available() { + python3 - "$PORT" <<'PY' +import socket +import sys + +port = int(sys.argv[1]) +with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind(("127.0.0.1", port)) + except OSError as exc: + print(f"PORT {port} is unavailable: {exc}", file=sys.stderr) + sys.exit(1) +PY +} + +# Auto modes cannot be proven on sub-threshold inputs. These skips are listed +# explicitly and are not emitted as successful benchmark rows. +case_applicable() { + local variant="$1" workload="$2" + case "$variant" in + pflash|full) [[ "$workload" == compression ]] ;; + kvflash) [[ "$workload" == compression || "$workload" == kv-pressure ]] ;; + *) return 0 ;; + esac +} + +workload_limits() { + case "$1" in + short) echo "4096 1200" ;; + medium) echo "8192 1800" ;; + long) echo "16384 2400" ;; + kv-pressure) echo "32768 3600" ;; + compression) echo "65536 5400" ;; + *) echo "unknown workload $1" >&2; return 1 ;; + esac +} + +run_case() { + local repeat="$1" workload="$2" clients="$3" variant="$4" + local max_ctx timeout + read -r max_ctx timeout <<< "$(workload_limits "$workload")" + local capacity=$((SLOTS * max_ctx)) + local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" + mkdir -p "$case_dir" + + local binary model_id max_prefills + local -a command launch_env metadata expected + if [[ "$variant" == llama ]]; then + binary="$LLAMA_SERVER_BIN"; model_id=qwen36-llama; max_prefills=0 + command=("$binary" -m "$MODEL" -ngl all --parallel "$SLOTS" -c "$capacity" + -b 2048 -ub 512 --cont-batching --no-context-shift --no-mmap -fa on + -ctk q4_0 -ctv q4_0 --no-cache-prompt --host 127.0.0.1 --port "$PORT" --alias "$model_id") + launch_env=() + else + binary="$LUCE_SERVER_BIN"; model_id=qwen36-luce; max_prefills="$MAX_CONCURRENT_PREFILLS" + command=("$binary" "$MODEL" --target-device "$TARGET_DEVICE" --paged-attention + --max-concurrency "$SLOTS" --kv-pool-tokens "$capacity" --max-ctx "$max_ctx" + --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 + --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 + --host 127.0.0.1 --port "$PORT" --model-name "$model_id") + launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills") + if [[ "$variant" == ddtree || "$variant" == full ]]; then + command+=(--draft "$DRAFT_MODEL" --draft-device "$DRAFT_DEVICE" + --ddtree --ddtree-budget "$DDTREE_BUDGET") + expected+=(--expect ddtree) + fi + if [[ "$variant" == pflash || "$variant" == full ]]; then + if [[ "$variant" == pflash ]]; then command+=(--draft-device "$DRAFT_DEVICE"); fi + command+=(--prefill-compression "$PREFILL_COMPRESSION" + --prefill-threshold "$PREFILL_THRESHOLD" + --prefill-keep-ratio "$PREFILL_KEEP_RATIO" + --prefill-drafter "$PREFILL_DRAFTER" + --draft-residency "$DRAFT_RESIDENCY") + expected+=(--expect pflash) + fi + if [[ "$variant" == kvflash || "$variant" == full ]]; then + # --prefill-drafter also selects the default drafter-scored KVFlash + # residency policy; it does not enable PFlash when compression is off. + if [[ "$variant" == kvflash ]]; then + command+=(--prefill-drafter "$PREFILL_DRAFTER") + fi + command+=(--kvflash "$KVFLASH_MODE") + launch_env+=("DFLASH_KVFLASH_MAX_POOL=$KVFLASH_MAX_POOL_TOKENS") + expected+=(--expect kvflash) + fi + fi + + if ((${#launch_env[@]})); then + printf 'env ' > "$case_dir/server-command.txt" + printf '%q ' "${launch_env[@]}" "${command[@]}" >> "$case_dir/server-command.txt" + else + printf '%q ' "${command[@]}" > "$case_dir/server-command.txt" + fi + printf '\n' >> "$case_dir/server-command.txt" + + metadata=(python3 "$METADATA_TOOL" --out "$case_dir/server-metadata.json" + --variant "$variant" --workload "$workload" --clients "$clients" --repeat "$repeat" + --binary "$binary" --model "$MODEL" --model-sha256 "$MODEL_SHA256" + --prompt-file "$OUT/prompts/$workload.jsonl" + --command-file "$case_dir/server-command.txt" --repo "$REPO" + --max-concurrent-prefills "$max_prefills") + if [[ "$variant" != llama ]]; then + metadata+=(--target-device "$TARGET_DEVICE") + local item + for item in "${launch_env[@]}"; do metadata+=(--launch-env "$item"); done + fi + if [[ "$variant" == ddtree || "$variant" == full ]]; then + metadata+=(--draft-device "$DRAFT_DEVICE" --draft-model "$DRAFT_MODEL" + --draft-model-sha256 "$DRAFT_MODEL_SHA256" --ddtree --ddtree-budget "$DDTREE_BUDGET") + fi + if [[ "$variant" == pflash || "$variant" == full ]]; then + metadata+=(--draft-device "$DRAFT_DEVICE" --prefill-compression "$PREFILL_COMPRESSION" + --prefill-threshold "$PREFILL_THRESHOLD" --prefill-keep-ratio "$PREFILL_KEEP_RATIO" + --prefill-drafter "$PREFILL_DRAFTER" + --prefill-drafter-sha256 "$PREFILL_DRAFTER_SHA256" + --draft-residency "$DRAFT_RESIDENCY") + fi + if [[ "$variant" == kvflash || "$variant" == full ]]; then + metadata+=(--kvflash "$KVFLASH_MODE" + --kvflash-max-pool-tokens "$KVFLASH_MAX_POOL_TOKENS" + --kvflash-scorer-drafter "$PREFILL_DRAFTER" + --kvflash-scorer-drafter-sha256 "$PREFILL_DRAFTER_SHA256") + if [[ "$variant" == kvflash ]]; then + # Record the literal server flag too, while keeping compression off. + metadata+=(--prefill-drafter "$PREFILL_DRAFTER" + --prefill-drafter-sha256 "$PREFILL_DRAFTER_SHA256") + fi + fi + "${metadata[@]}" + + echo "[run] $workload C=$clients repeat=$repeat variant=$variant" + port_is_available || return 1 + if ((${#launch_env[@]})); then + env "${launch_env[@]}" "${command[@]}" > "$case_dir/server.log" 2>&1 & + else + "${command[@]}" > "$case_dir/server.log" 2>&1 & + fi + server_pid=$! + if ! wait_health; then tail -n 120 "$case_dir/server.log" >&2 || true; return 1; fi + if [[ "$variant" != llama ]]; then + python3 "$RUNTIME_METADATA_TOOL" --metadata "$case_dir/server-metadata.json" \ + --server-log "$case_dir/server.log" + fi + + local offset="${prompt_offsets[$clients]}" prompts="$OUT/prompts/$workload.jsonl" + local -a common_client=(--base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" + --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" + --require-distinct-prompts --temperature 0 --ignore-eos --timeout "$timeout" --cooldown 0) + local -a telemetry_arg=() + [[ "$variant" != llama ]] && telemetry_arg+=(--require-effective-prompt-telemetry) + python3 "$CLIENT" "${common_client[@]}" --max-tokens "$WARMUP_TOKENS" + "${telemetry_arg[@]}" --out "$case_dir/warmup.json" + --label "$variant $workload C=$clients warmup" > "$case_dir/warmup.txt" + python3 "$CLIENT" "${common_client[@]}" --max-tokens "$MAX_TOKENS" + "${telemetry_arg[@]}" --server-metadata-json "$case_dir/server-metadata.json" + --out "$case_dir/bench.json" --label "$variant $workload C=$clients repeat=$repeat" \ + | tee "$case_dir/bench.txt" + stop_server + + if [[ "$variant" != llama ]]; then + python3 "$PROOF_TOOL" --bench "$case_dir/bench.json" --server-log "$case_dir/server.log" + "${expected[@]}" --out "$case_dir/feature-proof.json" + fi + sleep "$COOLDOWN_SECONDS" +} + +active_cases=0 +for ((repeat=1; repeat<=REPEATS; repeat++)); do + for workload in "${workload_list[@]}"; do + for c_index in "${!client_list[@]}"; do + clients="${client_list[$c_index]}" + shift_by=$(((repeat + c_index) % ${#variant_list[@]})) + for ((i=0; i<${#variant_list[@]}; i++)); do + variant="${variant_list[$(((i + shift_by) % ${#variant_list[@]}))]}" + if ! case_applicable "$variant" "$workload"; then + echo "[skip] $variant requires an activation workload; workload=$workload" + continue + fi + active_cases=$((active_cases + 1)) + run_case "$repeat" "$workload" "$clients" "$variant" + done + done + done +done +(( active_cases > 0 )) || { echo "no applicable benchmark cases" >&2; exit 2; } + +python3 "$SUMMARIZER" "$OUT" --out "$OUT/summary.md" +echo "[run] complete: $OUT" diff --git a/harness/benchmarks/concurrency/summarize_concurrency.py b/harness/benchmarks/concurrency/summarize_concurrency.py new file mode 100755 index 000000000..6f5130f41 --- /dev/null +++ b/harness/benchmarks/concurrency/summarize_concurrency.py @@ -0,0 +1,219 @@ +#!/usr/bin/env python3 +"""Summarize paired Lucebox/llama.cpp concurrency benchmark reports.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections import defaultdict +from pathlib import Path + + +def load_reports(root: Path) -> list[dict]: + reports = [] + for path in sorted(root.rglob("bench.json")): + report = json.loads(path.read_text(encoding="utf-8")) + meta = report.get("server_metadata") or {} + if len(report.get("levels", [])) != 1: + raise ValueError(f"{path}: expected exactly one client level") + level = report["levels"][0] + if ( + level.get("failures") + or not level.get("token_count_complete") + or not level.get("prompt_token_count_complete") + ): + raise ValueError(f"{path}: failed or incomplete token accounting") + if report.get("ignore_eos") and level.get("fixed_token_workload_valid") is not True: + raise ValueError(f"{path}: fixed-token validation failed") + reports.append({"path": path, "report": report, "level": level, "meta": meta}) + if not reports: + raise ValueError(f"{root}: no bench.json files found") + return reports + + +def median(values: list[float]) -> float: + return statistics.median(values) + + +def complete_median(values: list[float | None]) -> float | None: + """Return a median only when every repeat measured the metric.""" + if not values or any(value is None for value in values): + return None + return median([value for value in values if value is not None]) + + +def run_signature(item: dict) -> tuple[object, ...]: + report, meta = item["report"], item["meta"] + max_tokens = report.get("max_tokens") + ignore_eos = report.get("ignore_eos") + temperature = report.get("temperature") + seed = report.get("seed") + model_sha256 = meta.get("model_sha256") + if ( + type(max_tokens) is not int or max_tokens <= 0 + or not isinstance(ignore_eos, bool) + or type(temperature) not in (int, float) + or type(seed) is not int + or not isinstance(model_sha256, str) or not model_sha256 + ): + raise ValueError("incomplete run metadata") + return max_tokens, ignore_eos, temperature, seed, model_sha256 + + +def report_key(item: dict) -> tuple[str, int, str]: + meta, level = item["meta"], item["level"] + workload = meta.get("workload") + variant = meta.get("variant") + clients = level.get("clients") + if not isinstance(workload, str) or not workload: + raise ValueError("incomplete report metadata: missing workload") + if not isinstance(variant, str) or not variant: + raise ValueError("incomplete report metadata: missing variant") + if type(clients) is not int or clients <= 0: + raise ValueError("incomplete report metadata: invalid clients") + return workload, clients, variant + + +def summarize(reports: list[dict]) -> str: + grouped: dict[tuple[str, int, str], list[dict]] = defaultdict(list) + for item in reports: + grouped[report_key(item)].append(item) + for key, items in grouped.items(): + repeats = [item["meta"].get("repeat") for item in items] + if any(type(repeat) is not int or repeat <= 0 for repeat in repeats): + raise ValueError(f"{key}: invalid or missing repeat") + if len(repeats) != len(set(repeats)): + raise ValueError(f"{key}: duplicate repeat") + if len({run_signature(item) for item in items}) != 1: + raise ValueError(f"{key}: incompatible run metadata") + + lines = [ + "# Concurrency benchmark summary", "", + "Aggregate output goodput includes queueing, prefill, and decode. " + "Output-window goodput starts at the first observed output and is decode-facing, " + "but it can include staggered prefill. Prompt tok/s to first token includes " + "admission and TTFT.", "", + "| Workload | C | Variant | Repeats | Output goodput tok/s | " + "Output-window tok/s | Request decode tok/s | Prompt tok/s to first | " + "TTFT max s | Stable output | vs llama | Decode vs llama | K8 vs K1 |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | " + ":---: | ---: | ---: | ---: |", + ] + for workload, clients, variant in sorted(grouped): + items = grouped[(workload, clients, variant)] + prompt_digests = [ + item["level"].get("selected_prompt_set_sha256") for item in items + ] + if not all(isinstance(value, str) and value for value in prompt_digests): + raise ValueError( + f"{workload} C={clients} {variant}: missing selected prompt set hash" + ) + hashes = set(prompt_digests) + if len(hashes) != 1: + raise ValueError(f"{workload} C={clients} {variant}: prompt sets differ") + goodput_values = [item["level"].get("aggregate_tok_s") for item in items] + if any(type(value) not in (int, float) for value in goodput_values): + raise ValueError( + f"{workload} C={clients} {variant}: missing aggregate token rate" + ) + goodput = median(goodput_values) + output_window = complete_median([ + item["level"].get("output_window_tok_s") for item in items + ]) + request_decode = complete_median([ + item["level"].get("request_decode_tok_s_median") for item in items + ]) + prompt_rate = complete_median([ + item["level"].get("prompt_tokens_per_s_to_first_token") for item in items + ]) + ttft = complete_median([ + item["level"].get("ttft_max_s") for item in items + ]) + output_digests = [ + item["level"].get("selected_output_set_sha256") for item in items + ] + output_digests_complete = all( + isinstance(value, str) and bool(value) for value in output_digests + ) + output_hashes = { + value for value in output_digests if isinstance(value, str) + } + stable = ( + "n/a" if len(items) < 2 or not output_digests_complete + else "yes" if len(output_hashes) == 1 + else "NO" + ) + + def delta(other: str, metric: str) -> str: + peers = grouped.get((workload, clients, other), []) + if not peers: + return "n/a" + peer_hashes = {p["level"]["selected_prompt_set_sha256"] for p in peers} + if peer_hashes != hashes: + raise ValueError(f"{workload} C={clients}: {variant}/{other} prompts differ") + if {run_signature(item) for item in items} != { + run_signature(peer) for peer in peers + }: + raise ValueError( + f"{workload} C={clients}: {variant}/{other} run metadata differs" + ) + by_repeat = {int(item["meta"]["repeat"]): item for item in items} + peers_by_repeat = {int(item["meta"]["repeat"]): item for item in peers} + if by_repeat.keys() != peers_by_repeat.keys(): + raise ValueError( + f"{workload} C={clients}: {variant}/{other} repeat sets differ" + ) + ratios = [] + for repeat in sorted(by_repeat): + value = by_repeat[repeat]["level"].get(metric) + base = peers_by_repeat[repeat]["level"].get(metric) + if value is None or base is None: + return "n/a" + if base <= 0: + raise ValueError( + f"{workload} C={clients} repeat={repeat}: " + f"non-positive {other} {metric}" + ) + ratios.append(value / base - 1.0) + return f"{median(ratios) * 100:+.1f}%" + + vs_llama = ( + delta("llama", "aggregate_tok_s") + if variant == "luce-k8" else "—" + ) + decode_vs_llama = ( + delta("llama", "output_window_tok_s") + if variant == "luce-k8" else "—" + ) + vs_k1 = ( + delta("luce-k1", "aggregate_tok_s") + if variant == "luce-k8" else "—" + ) + output_window_text = f"{output_window:.2f}" if output_window is not None else "n/a" + request_decode_text = f"{request_decode:.2f}" if request_decode is not None else "n/a" + prompt_rate_text = f"{prompt_rate:.2f}" if prompt_rate is not None else "n/a" + ttft_text = f"{ttft:.3f}" if ttft is not None else "n/a" + lines.append( + f"| {workload} | {clients} | {variant} | {len(items)} | {goodput:.2f} | " + f"{output_window_text} | {request_decode_text} | {prompt_rate_text} | " + f"{ttft_text} | {stable} | {vs_llama} | {decode_vs_llama} | {vs_k1} |" + ) + lines.append("") + return "\n".join(lines) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("root", type=Path) + parser.add_argument("--out", type=Path) + args = parser.parse_args() + text = summarize(load_reports(args.root)) + if args.out: + args.out.write_text(text + "\n", encoding="utf-8") + print(text) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/summarize_feature_matrix.py b/harness/benchmarks/concurrency/summarize_feature_matrix.py new file mode 100755 index 000000000..b5b29c01e --- /dev/null +++ b/harness/benchmarks/concurrency/summarize_feature_matrix.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +"""Summarize Qwen3.6 concurrent feature ablations and activation proof.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from collections import defaultdict +from pathlib import Path + + +def load_reports(root: Path) -> list[dict]: + reports = [] + for path in sorted(root.rglob("bench.json")): + report = json.loads(path.read_text(encoding="utf-8")) + meta = report.get("server_metadata") or {} + if len(report.get("levels", [])) != 1: + raise ValueError(f"{path}: expected exactly one client level") + level = report["levels"][0] + variant = str(meta.get("variant", "")) + if ( + level.get("failures") + or not level.get("token_count_complete") + or not level.get("prompt_token_count_complete") + or (variant != "llama" + and not level.get("effective_prompt_token_count_complete")) + ): + raise ValueError(f"{path}: failed or incomplete token accounting") + if report.get("ignore_eos") and level.get("fixed_token_workload_valid") is not True: + raise ValueError(f"{path}: fixed-token validation failed") + proof_path = path.with_name("feature-proof.json") + proof = None + if variant != "llama": + if not proof_path.is_file(): + raise ValueError(f"{path}: missing feature-proof.json") + proof = json.loads(proof_path.read_text(encoding="utf-8")) + if proof.get("valid") is not True: + raise ValueError(f"{proof_path}: activation proof failed") + expected_by_variant = { + "ar": [], "ddtree": ["ddtree"], "pflash": ["pflash"], + "kvflash": ["kvflash"], "full": ["ddtree", "kvflash", "pflash"], + } + if variant not in expected_by_variant: + raise ValueError(f"{path}: unknown Lucebox variant {variant!r}") + if proof.get("expected_features") != expected_by_variant[variant]: + raise ValueError( + f"{proof_path}: expected_features does not match variant {variant}" + ) + reports.append({ + "path": path, "report": report, "level": level, + "meta": meta, "proof": proof, + }) + if not reports: + raise ValueError(f"{root}: no bench.json files found") + return reports + + +def median(values: list[float]) -> float: + return statistics.median(values) + + +def fmt(value: float | None, digits: int = 2) -> str: + return f"{value:.{digits}f}" if value is not None else "n/a" + + +def complete_median(values: list[float | None]) -> float | None: + """Return a median only when every repeat measured the metric.""" + if not values or any(value is None for value in values): + return None + return median([value for value in values if value is not None]) + + +def run_signature(item: dict) -> tuple[object, ...]: + report, meta = item["report"], item["meta"] + max_tokens = report.get("max_tokens") + ignore_eos = report.get("ignore_eos") + model_sha256 = meta.get("model_sha256") + if ( + type(max_tokens) is not int or max_tokens <= 0 + or not isinstance(ignore_eos, bool) + or not isinstance(model_sha256, str) or not model_sha256 + ): + raise ValueError("incomplete run metadata") + return max_tokens, ignore_eos, model_sha256 + + +def output_stability(items: list[dict]) -> str: + output_digests = [ + item["level"].get("selected_output_set_sha256") for item in items + ] + complete = all(isinstance(value, str) and bool(value) for value in output_digests) + hashes = {value for value in output_digests if isinstance(value, str)} + return ( + "n/a" if len(items) < 2 or not complete + else "yes" if len(hashes) == 1 + else "NO" + ) + + +def summarize(reports: list[dict]) -> str: + grouped: dict[tuple[str, int, str], list[dict]] = defaultdict(list) + for item in reports: + meta, level = item["meta"], item["level"] + key = (str(meta["workload"]), int(level["clients"]), str(meta["variant"])) + grouped[key].append(item) + for key, items in grouped.items(): + repeats = [int(item["meta"]["repeat"]) for item in items] + if len(repeats) != len(set(repeats)): + raise ValueError(f"{key}: duplicate repeat") + if len({run_signature(item) for item in items}) != 1: + raise ValueError(f"{key}: incompatible run metadata") + + lines = [ + "# Qwen3.6 concurrent feature matrix", "", + "Every Lucebox row is included only after request-correlated server telemetry " + "proves its requested features executed. Throughput is the median across fresh-process repeats.", + "", + "| Workload | C | Variant | N | Output goodput | Output-window | vs AR | " + "Effective/wire | DDTree accepted/step | Target forwards | KV in/out | " + "PFlash requests | TTFT max s | Stable output |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | " + ":--- | ---: | ---: | :---: |", + ] + for workload, clients, variant in sorted(grouped): + items = grouped[(workload, clients, variant)] + prompt_hashes = {item["level"]["selected_prompt_set_sha256"] for item in items} + if len(prompt_hashes) != 1: + raise ValueError(f"{workload} C={clients} {variant}: prompt sets differ") + goodput = median([item["level"]["aggregate_tok_s"] for item in items]) + window = complete_median([ + item["level"].get("output_window_tok_s") for item in items + ]) + ratio = complete_median([ + item["level"].get("effective_to_wire_prompt_ratio") for item in items + ]) + ttft = complete_median([ + item["level"].get("ttft_max_s") for item in items + ]) + stable = output_stability(items) + + peers = grouped.get((workload, clients, "ar"), []) + vs_ar = "—" if variant == "ar" else "n/a" + if variant != "ar" and peers: + peer_hashes = {p["level"]["selected_prompt_set_sha256"] for p in peers} + if peer_hashes != prompt_hashes: + raise ValueError(f"{workload} C={clients}: {variant}/ar prompts differ") + if {run_signature(item) for item in items} != { + run_signature(peer) for peer in peers + }: + raise ValueError( + f"{workload} C={clients}: {variant}/ar run metadata differs" + ) + by_repeat = {int(item["meta"]["repeat"]): item for item in items} + peers_by_repeat = {int(item["meta"]["repeat"]): item for item in peers} + if by_repeat.keys() != peers_by_repeat.keys(): + raise ValueError( + f"{workload} C={clients}: {variant}/ar repeat sets differ" + ) + if stable != "NO" and output_stability(peers) != "NO": + ratios = [] + for repeat in sorted(by_repeat): + value = by_repeat[repeat]["level"].get("aggregate_tok_s") + base = peers_by_repeat[repeat]["level"].get("aggregate_tok_s") + if value is None or base is None or base <= 0: + raise ValueError( + f"{workload} C={clients} repeat={repeat}: invalid AR goodput" + ) + ratios.append(value / base - 1.0) + vs_ar = f"{median(ratios) * 100:+.1f}%" + + proofs = [item["proof"] for item in items if item["proof"] is not None] + aggregates = [p["aggregate"] for p in proofs] + steps = sum(a["ddtree_steps"] for a in aggregates) + accepted = sum(a["ddtree_accepted_tokens"] for a in aggregates) + accepted_per_step = accepted / steps if steps else None + target_forwards = median([a["target_forwards"] for a in aggregates]) if aggregates else None + page_ins = median([a["kvflash_page_ins"] for a in aggregates]) if aggregates else None + page_outs = median([a["kvflash_page_outs"] for a in aggregates]) if aggregates else None + pflash_requests = median([a["pflash_applied_requests"] for a in aggregates]) if aggregates else None + kv_text = ( + f"{page_ins:.0f}/{page_outs:.0f}" + if page_ins is not None and page_outs is not None else "n/a" + ) + lines.append( + f"| {workload} | {clients} | {variant} | {len(items)} | {goodput:.2f} | " + f"{fmt(window)} | {vs_ar} | {fmt(ratio, 3)} | {fmt(accepted_per_step)} | " + f"{fmt(target_forwards, 0)} | {kv_text} | {fmt(pflash_requests, 0)} | " + f"{fmt(ttft, 3)} | {stable} |" + ) + lines.append("") + return "\n".join(lines) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("root", type=Path) + parser.add_argument("--out", type=Path) + args = parser.parse_args() + text = summarize(load_reports(args.root)) + if args.out: + args.out.write_text(text + "\n", encoding="utf-8") + print(text) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/test_concurrency_tools.py b/harness/benchmarks/concurrency/test_concurrency_tools.py new file mode 100644 index 000000000..1aa8dc2d8 --- /dev/null +++ b/harness/benchmarks/concurrency/test_concurrency_tools.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 +"""Unit tests for the deterministic prompt generator and compact summarizer.""" + +from __future__ import annotations + +import importlib.util +import json +import tempfile +import unittest +from pathlib import Path + + +HERE = Path(__file__).parent + + +def load(name: str): + path = HERE / f"{name}.py" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +generator = load("generate_ragged_prompts") +summarizer = load("summarize_concurrency") + + +class PromptGeneratorTests(unittest.TestCase): + def test_cohorts_are_disjoint_ragged_and_mean_matched(self) -> None: + records = generator.build_records("short") + self.assertEqual(len(records), 29) + self.assertEqual( + [row["cohort"] for row in records], + ["c1"] + ["c4"] * 4 + ["c8"] * 8 + ["c16"] * 16, + ) + self.assertEqual(len({row["prompt"] for row in records}), 29) + by_cohort = { + cohort: [row for row in records if row["cohort"] == cohort] + for cohort in ("c1", "c4", "c8", "c16") + } + means = { + cohort: sum(row["target_words"] for row in rows) / len(rows) + for cohort, rows in by_cohort.items() + } + self.assertEqual(len(set(means.values())), 1) + for cohort in ("c4", "c8", "c16"): + self.assertEqual(len({row["target_words"] for row in by_cohort[cohort]}), 4) + for row in records: + self.assertEqual(len(row["prompt"].split()), row["target_words"]) + + +class SummarizerTests(unittest.TestCase): + @staticmethod + def item( + variant: str, + goodput: float, + output_window: float | None = None, + *, + repeat: int = 1, + output_hash: str | None = "same-outputs", + ) -> dict: + return { + "report": { + "max_tokens": 256, "ignore_eos": True, + "temperature": 0.0, "seed": 1, + }, + "meta": { + "workload": "short", "variant": variant, "repeat": repeat, + "model_sha256": "a" * 64, + }, + "level": { + "clients": 8, + "aggregate_tok_s": goodput, + "output_window_tok_s": output_window if output_window is not None else goodput, + "request_decode_tok_s_median": goodput / 8, + "prompt_tokens_per_s_to_first_token": 100.0, + "ttft_max_s": 2.0, + "selected_prompt_set_sha256": "same-prompts", + "selected_output_set_sha256": output_hash, + }, + } + + def test_summary_reports_product_and_packing_deltas(self) -> None: + text = summarizer.summarize([ + self.item("luce-k8", 20.0), + self.item("luce-k1", 10.0), + self.item("llama", 8.0), + ]) + self.assertIn("+150.0%", text) + self.assertIn("+100.0%", text) + self.assertIn("Decode vs llama", text) + + def test_summary_uses_same_repeat_ratios(self) -> None: + reports = [] + for repeat, luce, llama in ( + (1, 10.0, 1.0), + (2, 20.0, 90.0), + (3, 100.0, 50.0), + ): + reports.extend([ + self.item("luce-k8", luce, repeat=repeat), + self.item("llama", llama, repeat=repeat), + ]) + text = summarizer.summarize(reports) + luce_row = next(line for line in text.splitlines() if "| luce-k8 |" in line) + self.assertIn("+100.0%", luce_row) + self.assertNotIn("-60.0%", luce_row) + + def test_summary_rejects_mismatched_repeat_sets(self) -> None: + reports = [ + self.item("luce-k8", 20.0, repeat=1), + self.item("luce-k8", 22.0, repeat=2), + self.item("llama", 10.0, repeat=1), + ] + with self.assertRaisesRegex(ValueError, "repeat sets differ"): + summarizer.summarize(reports) + + def test_single_repeat_does_not_claim_stability(self) -> None: + text = summarizer.summarize([self.item("llama", 8.0)]) + row = next(line for line in text.splitlines() if "| llama |" in line) + self.assertEqual(row.split("|")[10].strip(), "n/a") + + def test_multiple_repeats_report_output_stability(self) -> None: + stable = summarizer.summarize([ + self.item("llama", 8.0, repeat=1), + self.item("llama", 9.0, repeat=2), + ]) + stable_row = next(line for line in stable.splitlines() if "| llama |" in line) + self.assertEqual(stable_row.split("|")[10].strip(), "yes") + + unstable = summarizer.summarize([ + self.item("llama", 8.0, repeat=1, output_hash="first"), + self.item("llama", 9.0, repeat=2, output_hash="second"), + ]) + unstable_row = next(line for line in unstable.splitlines() if "| llama |" in line) + self.assertEqual(unstable_row.split("|")[10].strip(), "NO") + + def test_missing_output_digest_does_not_claim_stability(self) -> None: + reports = [ + self.item("llama", 8.0, repeat=1, output_hash=None), + self.item("llama", 9.0, repeat=2, output_hash=None), + ] + text = summarizer.summarize(reports) + row = next(line for line in text.splitlines() if "| llama |" in line) + self.assertEqual(row.split("|")[10].strip(), "n/a") + + def test_incomplete_repeat_metrics_are_reported_as_na(self) -> None: + first = self.item("llama", 8.0, repeat=1) + second = self.item("llama", 9.0, repeat=2) + for metric in ( + "output_window_tok_s", "request_decode_tok_s_median", + "prompt_tokens_per_s_to_first_token", "ttft_max_s", + ): + second["level"][metric] = None + text = summarizer.summarize([first, second]) + row = next(line for line in text.splitlines() if "| llama |" in line) + self.assertEqual( + [row.split("|")[column].strip() for column in range(6, 10)], + ["n/a"] * 4, + ) + + def test_comparison_rejects_incompatible_run_metadata(self) -> None: + fields = ( + ("report", "max_tokens", 64), + ("report", "ignore_eos", False), + ("report", "temperature", 0.5), + ("report", "seed", 2), + ("meta", "model_sha256", "b" * 64), + ) + for container, key, value in fields: + with self.subTest(key=key): + luce = self.item("luce-k8", 20.0) + llama = self.item("llama", 10.0) + luce[container][key] = value + with self.assertRaisesRegex(ValueError, "run metadata differs"): + summarizer.summarize([luce, llama]) + + def test_summary_rejects_incomplete_run_metadata(self) -> None: + item = self.item("llama", 10.0) + del item["report"]["max_tokens"] + with self.assertRaisesRegex(ValueError, "incomplete run metadata"): + summarizer.summarize([item]) + + def test_summary_reports_descriptive_errors_for_missing_fields(self) -> None: + fields = ( + ("meta", "workload", "workload"), + ("level", "clients", "clients"), + ("level", "aggregate_tok_s", "aggregate token rate"), + ("level", "selected_prompt_set_sha256", "prompt set hash"), + ) + for container, key, message in fields: + with self.subTest(key=key): + item = self.item("llama", 10.0) + del item[container][key] + with self.assertRaisesRegex(ValueError, message): + summarizer.summarize([item]) + + def test_load_reports_rejects_missing_prompt_usage(self) -> None: + report = { + "ignore_eos": True, + "server_metadata": {"workload": "short", "variant": "llama", "repeat": 1}, + "levels": [{ + "failures": 0, + "token_count_complete": True, + "prompt_token_count_complete": False, + "fixed_token_workload_valid": True, + }], + } + with tempfile.TemporaryDirectory() as root: + path = Path(root) / "bench.json" + path.write_text(json.dumps(report), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "incomplete token accounting"): + summarizer.load_reports(Path(root)) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_concurrent_benchmark.py new file mode 100644 index 000000000..c0455edc3 --- /dev/null +++ b/harness/benchmarks/concurrency/test_concurrent_benchmark.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Focused tests for concurrent_benchmark.py.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import unittest +from pathlib import Path +from unittest import mock + + +SCRIPT = Path(__file__).with_name("concurrent_benchmark.py") +SPEC = importlib.util.spec_from_file_location("concurrent_benchmark", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +benchmark = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(benchmark) + + +class BenchmarkTests(unittest.TestCase): + def test_sse_parser_handles_events_and_done(self) -> None: + lines = [ + b'data: {"choices":[{"delta":{"content":"hi"}}]}\n', b"\n", + b"data: [DONE]\n", b"\n", + ] + self.assertEqual( + list(benchmark.iter_sse_data(lines)), + ['{"choices":[{"delta":{"content":"hi"}}]}', "[DONE]"], + ) + + def test_prompt_selection_never_wraps(self) -> None: + self.assertEqual(benchmark.request_prompts(["a", "b", "c"], 2, 1), ["b", "c"]) + with self.assertRaisesRegex(ValueError, "refusing to reuse"): + benchmark.request_prompts(["a", "b"], 2, 1) + + def test_level_uses_exact_usage_and_first_token_window(self) -> None: + def fake_request(_args: argparse.Namespace, prompt: str) -> dict: + start, first, end, prompt_tokens = { + "first": (10.0, 12.0, 14.0, 10), + "second": (10.25, 11.25, 15.0, 30), + }[prompt] + return { + "t_start": start, "t_first": first, "t_end": end, + "duration_s": end - start, "ttft_s": first - start, + "decode_duration_s": end - first, + "completion_tokens": 8, "prompt_tokens": prompt_tokens, + "finish_reason": "length", "error": None, + "content_sha256": benchmark.sha256_text(prompt + " output"), + "reasoning_content_sha256": benchmark.sha256_text(""), + "content_chars": 6, "reasoning_content_chars": 0, + "request_output_tok_s": 8 / (end - start), + "request_decode_tok_s": 7 / (end - first), + } + + args = argparse.Namespace(max_tokens=8, ignore_eos=True, timeout=2.0) + with mock.patch.object(benchmark, "stream_request", side_effect=fake_request): + level = benchmark.run_level(2, args, ["first", "second"], 0) + self.assertEqual(level["completion_tokens_total"], 16) + self.assertEqual(level["prompt_tokens_total"], 40) + self.assertTrue(level["fixed_token_workload_valid"]) + # Assert independently known windows before their derived rates. + self.assertEqual(level["wall_s"], 5.0) + self.assertEqual(level["output_window_s"], 3.75) + self.assertEqual(level["prompt_to_first_token_s"], 2.0) + self.assertAlmostEqual(level["aggregate_tok_s"], 3.2) + self.assertAlmostEqual( + level["output_window_tok_s"], 16 / 3.75, + ) + self.assertAlmostEqual( + level["request_decode_tok_s_median"], (3.5 + 7 / 3.75) / 2, + ) + self.assertAlmostEqual(level["prompt_tokens_per_s_to_first_token"], 20.0) + + def test_stream_request_keeps_usage_separate_from_sse_chunks(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"choices":[{"delta":{"content":"one chunk"}}]}\n', b"\n", + b'data: {"choices":[{"delta":{},"finish_reason":"length"}]}\n', b"\n", + b'data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":64}}\n', b"\n", + b"data: [DONE]\n", b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=64, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + record = benchmark.stream_request(args, "prompt") + self.assertEqual(record["completion_tokens"], 64) + self.assertEqual(record["prompt_tokens"], 12) + self.assertTrue(record["done_received"]) + self.assertIsNone(record["error"]) + self.assertIsNotNone(record["request_decode_tok_s"]) + self.assertEqual(record["content_sha256"], benchmark.sha256_text("one chunk")) + + def test_stream_request_rejects_clean_eof_without_done(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"choices":[{"delta":{"content":"partial"}}]}\n', b"\n", + b'data: {"choices":[{"delta":{},"finish_reason":"length"}]}\n', b"\n", + b'data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":64}}\n', b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=64, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + record = benchmark.stream_request(args, "prompt") + self.assertFalse(record["done_received"]) + self.assertIn("before [DONE]", record["error"]) + + def test_missing_prompt_usage_fails_level(self) -> None: + level = { + "failures": 0, + "token_count_complete": True, + "prompt_token_count_complete": False, + "fixed_token_workload_valid": True, + } + self.assertTrue(benchmark.level_failed(level, ignore_eos=True)) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py new file mode 100644 index 000000000..d582f6e32 --- /dev/null +++ b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Focused tests for request-correlated feature benchmark telemetry.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import sys +import unittest +from pathlib import Path +from unittest import mock + + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +SCRIPT = HERE / "feature_concurrent_benchmark.py" +SPEC = importlib.util.spec_from_file_location("feature_concurrent_benchmark", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +benchmark = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(benchmark) + + +class FeatureBenchmarkTests(unittest.TestCase): + def test_stream_request_captures_id_and_effective_prompt(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"id":"chatcmpl-42","choices":[{"delta":{"content":"x"}}]}\n', b"\n", + b'data: {"id":"chatcmpl-42","choices":[{"delta":{},' + b'"finish_reason":"length"}]}\n', b"\n", + b'data: {"id":"chatcmpl-42","choices":[],"usage":' + b'{"prompt_tokens":40000,"completion_tokens":8,"timings":' + b'{"effective_prompt_tokens":2000,"prefilled_tokens":2000,' + b'"cached_prefix_tokens":0,"cache_hit":false,"prefill_ms":12.5,' + b'"decode_ms":20.0,"decode_tokens_per_sec":400.0}}}\n', b"\n", + b"data: [DONE]\n", b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=8, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + row = benchmark.stream_request(args, "prompt") + self.assertEqual(row["request_id"], "chatcmpl-42") + self.assertEqual(row["prompt_tokens"], 40000) + self.assertEqual(row["effective_prompt_tokens"], 2000) + self.assertEqual(row["server_prefill_ms"], 12.5) + self.assertTrue(row["done_received"]) + self.assertIsNone(row["error"]) + + def test_clean_eof_without_done_is_rejected(self) -> None: + class Response: + def __enter__(self): return self + def __exit__(self, *_args): return None + def __iter__(self): + return iter([ + b'data: {"id":"chatcmpl-cut","choices":[{"delta":{"content":"x"},' + b'"finish_reason":"length"}]}\n', b"\n", + ]) + + args = argparse.Namespace( + model="m", max_tokens=8, temperature=0.0, seed=1, ignore_eos=True, + api_key="", base_url="http://localhost/v1", timeout=2.0, + ) + with mock.patch.object(benchmark.urllib.request, "urlopen", return_value=Response()): + row = benchmark.stream_request(args, "prompt") + self.assertIn("before [DONE]", row["error"]) + + def test_enrich_level_reports_compression_ratio(self) -> None: + level = { + "prompt_tokens_total": 40000, + "requests_detail": [{ + "request_id": "r1", "error": None, + "effective_prompt_tokens": 2000, + "server_prefill_ms": 10.0, "server_decode_ms": 20.0, + "server_decode_tokens_per_sec": 100.0, + }], + } + benchmark.enrich_level(level) + self.assertTrue(level["request_ids_complete"]) + self.assertTrue(level["effective_prompt_token_count_complete"]) + self.assertEqual(level["effective_prompt_tokens_total"], 2000) + self.assertEqual(level["effective_to_wire_prompt_ratio"], 0.05) + + def test_duplicate_request_ids_are_not_complete(self) -> None: + level = { + "prompt_tokens_total": 20, + "requests_detail": [ + {"request_id": "same", "error": None, "effective_prompt_tokens": 10}, + {"request_id": "same", "error": None, "effective_prompt_tokens": 10}, + ], + } + benchmark.enrich_level(level) + self.assertFalse(level["request_ids_complete"]) + + def test_boolean_wire_counts_are_not_accepted_as_integers(self) -> None: + level = { + "prompt_tokens_total": 1, + "requests_detail": [{ + "request_id": "r1", "error": None, + "effective_prompt_tokens": True, + }], + } + benchmark.enrich_level(level) + self.assertFalse(level["effective_prompt_token_count_complete"]) + + def test_client_provenance_records_exact_argv_and_source_digest(self) -> None: + argv = ["python3", "feature_concurrent_benchmark.py", "--clients", "4"] + result = benchmark.client_provenance(argv) + self.assertEqual(result["client_argv"], argv) + self.assertEqual(result["client_script"], str(SCRIPT.resolve())) + self.assertEqual( + result["client_script_sha256"], + hashlib.sha256(SCRIPT.read_bytes()).hexdigest(), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_metadata.py b/harness/benchmarks/concurrency/test_feature_metadata.py new file mode 100644 index 000000000..7349fa21a --- /dev/null +++ b/harness/benchmarks/concurrency/test_feature_metadata.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""Tests for literal feature flags and reproducibility metadata.""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + + +HERE = Path(__file__).parent +SCRIPT = HERE / "write_feature_metadata.py" +SPEC = importlib.util.spec_from_file_location("write_feature_metadata", SCRIPT) +assert SPEC is not None and SPEC.loader is not None +metadata = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(metadata) +RUNTIME_SCRIPT = HERE / "record_feature_runtime.py" +RUNTIME_SPEC = importlib.util.spec_from_file_location( + "record_feature_runtime", RUNTIME_SCRIPT, +) +assert RUNTIME_SPEC is not None and RUNTIME_SPEC.loader is not None +runtime_metadata = importlib.util.module_from_spec(RUNTIME_SPEC) +RUNTIME_SPEC.loader.exec_module(runtime_metadata) + + +class FeatureMetadataTests(unittest.TestCase): + def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + target = root / "target.gguf" + draft = root / "draft.gguf" + prefill = root / "prefill.gguf" + prompts = root / "prompts.jsonl" + command = root / "command.txt" + out = root / "metadata.json" + target.write_bytes(b"target") + draft.write_bytes(b"draft") + prefill.write_bytes(b"prefill") + prompts.write_text('{"prompt":"p"}\n', encoding="utf-8") + command.write_text("server --target-device hip:0\n", encoding="utf-8") + argv = [ + str(SCRIPT), "--out", str(out), "--variant", "full", + "--workload", "compression", "--clients", "4", "--repeat", "1", + "--binary", "/bin/true", "--model", str(target), + "--model-sha256", hashlib.sha256(target.read_bytes()).hexdigest(), + "--prompt-file", str(prompts), "--command-file", str(command), + "--repo", str(HERE.parents[2]), "--max-concurrent-prefills", "8", + "--target-device", "hip:0", "--draft-device", "hip:0", + "--draft-model", str(draft), "--draft-model-sha256", "draft-sha", + "--ddtree", "--ddtree-budget", "22", + "--prefill-compression", "auto", "--prefill-threshold", "32000", + "--prefill-keep-ratio", "0.05", "--prefill-drafter", str(prefill), + "--prefill-drafter-sha256", "prefill-sha", + "--draft-residency", "persistent", "--kvflash", "auto", + "--kvflash-max-pool-tokens", "8192", + "--kvflash-scorer-drafter", str(prefill), + "--kvflash-scorer-drafter-sha256", "prefill-sha", + ] + with mock.patch.object(sys, "argv", argv): + self.assertEqual(metadata.main(), 0) + result = json.loads(out.read_text(encoding="utf-8")) + self.assertEqual(result["model_sha256"], hashlib.sha256(b"target").hexdigest()) + self.assertTrue(result["server_binary_sha256"]) + self.assertTrue(result["git_head"]) + self.assertEqual(result["feature_config"]["draft_model_sha256"], "draft-sha") + self.assertEqual(result["feature_config"]["prefill_drafter_sha256"], "prefill-sha") + self.assertEqual( + result["feature_config"]["kvflash_scorer_drafter"], + str(prefill.resolve()), + ) + self.assertEqual( + result["feature_config"]["kvflash_scorer_drafter_sha256"], + "prefill-sha", + ) + self.assertEqual(result["literal_screenshot_flags"], [ + "--target-device", "hip:0", + "--draft-device", "hip:0", + "--ddtree", + "--ddtree-budget", "22", + "--draft-residency", "persistent", + "--prefill-compression", "auto", + "--prefill-drafter", str(prefill.resolve()), + "--kvflash", "auto", + ]) + self.assertIsNone(result["runtime_observed"]) + + def test_ldd_failure_is_fatal(self) -> None: + failed = mock.Mock(returncode=1, stdout="", stderr="not a dynamic executable") + with mock.patch.object(metadata.subprocess, "run", return_value=failed): + with self.assertRaisesRegex(RuntimeError, "ldd failed"): + metadata.resolved_libraries(Path("/tmp/server")) + + def test_unresolved_shared_library_is_fatal(self) -> None: + unresolved = mock.Mock( + returncode=0, stdout="libmissing.so => not found\n", stderr="", + ) + with mock.patch.object(metadata.subprocess, "run", return_value=unresolved): + with self.assertRaisesRegex(RuntimeError, "unresolved libraries"): + metadata.resolved_libraries(Path("/tmp/server")) + + def test_git_revision_failure_and_empty_output_are_fatal(self) -> None: + for result, message in ( + (mock.Mock(returncode=128, stdout="", stderr="not a repository"), + "git rev-parse failed"), + (mock.Mock(returncode=0, stdout="\n", stderr=""), + "empty revision"), + ): + with self.subTest(message=message): + with mock.patch.object(metadata.subprocess, "run", return_value=result): + with self.assertRaisesRegex(RuntimeError, message): + metadata.repository_head(Path("/tmp/repo")) + + def test_runtime_records_actual_kvflash_pool_from_startup(self) -> None: + original = { + "schema_version": 3, + "feature_config": { + "kvflash": "auto", + "kvflash_max_pool_tokens": 16384, + }, + } + log = "\n".join(( + "[parallel-kvflash] physical resident pool 8192 tokens; " + "logical per-slot cap 65536 across 16 slots " + "(--kv-pool-tokens does not expand resident VRAM)", + "[paged-attention] 512 physical blocks x 16 tokens " + "(8192 pool tokens, per-sequence max_ctx 65536)", + )) + result = runtime_metadata.update_metadata(original, log) + observed = result["runtime_observed"] + self.assertTrue(observed["kvflash_active"]) + self.assertEqual(observed["physical_kv_pool_tokens"], 8192) + self.assertEqual(observed["physical_kv_pool_blocks"], 512) + self.assertEqual(observed["kv_block_size_tokens"], 16) + self.assertEqual(observed["logical_per_slot_max_ctx"], 65536) + self.assertEqual(observed["configured_slots"], 16) + + def test_runtime_rejects_enabled_kvflash_without_marker(self) -> None: + original = {"feature_config": {"kvflash": "auto"}} + with self.assertRaisesRegex(ValueError, "startup marker is missing"): + runtime_metadata.update_metadata( + original, + "[paged-attention] 512 physical blocks x 16 tokens " + "(8192 pool tokens, per-sequence max_ctx 65536)", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py new file mode 100644 index 000000000..85c2294e6 --- /dev/null +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -0,0 +1,556 @@ +#!/usr/bin/env python3 +"""Tests for pressure prompts, activation proof, and feature summaries.""" + +from __future__ import annotations + +import importlib.util +import json +import os +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) + + +def load(name: str): + path = HERE / f"{name}.py" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +generator = load("generate_feature_prompts") +proof = load("verify_feature_metrics") +summary = load("summarize_feature_matrix") + + +def report( + variant: str = "full", workload: str = "compression", + effective: tuple[int, int] = (2000, 2100), +) -> dict: + return { + "server_metadata": { + "variant": variant, + "workload": workload, + "feature_config": { + "prefill_compression": "auto", + "prefill_threshold": 32000, + "kvflash": "auto", + "kvflash_max_pool_tokens": 8192, + "kvflash_scorer_drafter": "/models/Qwen3-0.6B-BF16.gguf", + "kvflash_scorer_drafter_sha256": "a" * 64, + }, + "runtime_observed": { + "kvflash_active": True, + "physical_kv_pool_tokens": 8192, + }, + }, + "levels": [{ + "requests_detail": [ + {"request_id": "r1", "error": None, + "prompt_tokens": 40000, + "effective_prompt_tokens": effective[0]}, + {"request_id": "r2", "error": None, + "prompt_tokens": 40000, + "effective_prompt_tokens": effective[1]}, + ], + }], + } + + +def metric(request_id: str, effective: int, page_outs: int = 1) -> dict: + return { + "request_id": request_id, + "effective_prompt_tokens": effective, + "ddtree_steps": 3, + "ddtree_suspensions": 0, + # Zero acceptance is legitimate and must not invalidate execution proof. + "ddtree_accepted_tokens": 0, + "target_forwards": 3, + "kvflash_page_ins": 0, + "kvflash_page_outs": page_outs, + "kvflash_resident_blocks": 8, + "kvflash_reselects": 1, + "pflash_applied": True, + "pflash_input_tokens": 40000, + "pflash_output_tokens": effective, + } + + +class FeaturePromptTests(unittest.TestCase): + def test_activation_profiles_are_disjoint_and_above_thresholds(self) -> None: + compression = generator.build_records("compression") + pressure = generator.build_records("kv-pressure") + self.assertEqual(len(compression), 29) + self.assertEqual(len(pressure), 29) + self.assertEqual(len({row["prompt"] for row in compression}), 29) + self.assertEqual(len({row["prompt"] for row in pressure}), 29) + self.assertTrue( + {row["prompt"] for row in compression}.isdisjoint( + row["prompt"] for row in pressure + ) + ) + self.assertGreaterEqual(min(row["target_words"] for row in compression), 34000) + self.assertGreaterEqual(min(row["target_words"] for row in pressure), 12000) + self.assertTrue(all(row["activation_target"] == "pflash-auto" for row in compression)) + + +class FeatureRunnerShellTests(unittest.TestCase): + def run_invalid_matrix( + self, tmp: str, **overrides: str, + ) -> subprocess.CompletedProcess[str]: + model = Path(tmp) / "model.gguf" + model.touch() + env = { + key: value for key, value in os.environ.items() + if not key.startswith(("GGML_", "DFLASH_", "LUCE_", "HIP_", "ROCR_", "HSA_")) + and key not in ("LD_PRELOAD", "LD_LIBRARY_PATH") + } + env.update({ + "MODEL": str(model), + "LUCE_SERVER_BIN": "/bin/true", + "OUT": str(Path(tmp) / "out"), + "VARIANTS": "ar", + **overrides, + }) + return subprocess.run( + ["bash", str(HERE / "run_qwen36_feature_matrix.sh")], + env=env, capture_output=True, text=True, check=False, + ) + + def test_client_and_proof_invocations_are_array_backed(self) -> None: + runner = (HERE / "run_qwen36_feature_matrix.sh").read_text( + encoding="utf-8", + ) + client_lines = [ + line.strip() for line in runner.splitlines() + if 'python3 "$CLIENT"' in line + ] + self.assertEqual(client_lines, [ + 'python3 "$CLIENT" "${common_client[@]}"', + 'python3 "$CLIENT" "${common_client[@]}"', + ]) + self.assertIn('local -a warmup_cmd=(', runner) + self.assertIn('local -a benchmark_cmd=(', runner) + self.assertIn( + '"${warmup_cmd[@]}" > "$case_dir/warmup.txt"', + runner, + ) + self.assertIn( + '"${benchmark_cmd[@]}" | tee "$case_dir/bench.txt"', + runner, + ) + + proof_lines = [ + line.strip() for line in runner.splitlines() + if 'python3 "$PROOF_TOOL"' in line + ] + self.assertEqual(proof_lines, ['python3 "$PROOF_TOOL"']) + self.assertIn('local -a proof_cmd=(', runner) + self.assertIn('"${proof_cmd[@]}"', runner) + + def test_signals_exit_and_launch_environment_is_recorded(self) -> None: + runner = (HERE / "run_qwen36_feature_matrix.sh").read_text(encoding="utf-8") + self.assertIn("trap stop_server EXIT", runner) + self.assertIn("trap 'exit 130' INT", runner) + self.assertIn("trap 'exit 143' TERM", runner) + self.assertNotIn("trap stop_server EXIT INT TERM", runner) + self.assertIn("'env ' > \"$case_dir/server-command.txt\"", runner) + self.assertIn('"${launch_env[@]}" "${command[@]}"', runner) + + def test_duplicate_clients_are_rejected_before_artifacts_are_created(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = self.run_invalid_matrix(tmp, CLIENTS="4,4") + self.assertEqual(result.returncode, 2, result.stderr) + self.assertIn("CLIENTS contains duplicate entry: 4", result.stderr) + self.assertFalse((Path(tmp) / "out").exists()) + + def test_duplicate_variants_are_rejected_before_artifacts_are_created(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = self.run_invalid_matrix(tmp, VARIANTS="ar,ar") + self.assertEqual(result.returncode, 2, result.stderr) + self.assertIn("VARIANTS contains duplicate entry: ar", result.stderr) + self.assertFalse((Path(tmp) / "out").exists()) + + +class FeatureProofTests(unittest.TestCase): + def test_full_below_pool_passes_without_page_traffic(self) -> None: + rows = [ + metric("warmup", 50), + metric("r1", 2000, page_outs=0), + metric("r2", 2100, page_outs=0), + ] + result = proof.verify( + report(), rows, {"ddtree", "pflash", "kvflash"}, + ) + self.assertTrue(result["valid"], result["errors"]) + self.assertFalse(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "compressed-prompt-fits-pool", + ) + self.assertEqual(result["aggregate"]["ddtree_accepted_tokens"], 0) + self.assertEqual(result["aggregate"]["ddtree_suspensions"], 0) + self.assertEqual(result["matched_metric_count"], 2) + + def test_ddtree_suspensions_required_and_aggregated(self) -> None: + missing = metric("r1", 2000) + del missing["ddtree_suspensions"] + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "server.log" + path.write_text( + proof.PREFIX + json.dumps(missing) + "\n", + encoding="utf-8", + ) + with self.assertRaisesRegex( + ValueError, "missing telemetry keys.*ddtree_suspensions", + ): + proof.parse_markers(path) + + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[1]["ddtree_suspensions"] = 1 + result = proof.verify( + report(), rows, {"ddtree", "pflash", "kvflash"}, + ) + self.assertTrue(result["valid"], result["errors"]) + self.assertEqual(result["aggregate"]["ddtree_suspensions"], 1) + self.assertEqual( + [row["ddtree_suspensions"] for row in result["requests"]], + [0, 1], + ) + + def test_ddtree_suspensions_must_be_binary_per_request(self) -> None: + raw_two = metric("r1", 2000) + raw_two["ddtree_suspensions"] = 2 + with self.assertRaisesRegex( + ValueError, "ddtree_suspensions must be 0 or 1 per request", + ): + proof.aggregate_rows([raw_two]) + + def test_duplicate_terminal_telemetry_is_rejected(self) -> None: + row = metric("r1", 2000) + with self.assertRaisesRegex(ValueError, "duplicate telemetry request ID r1"): + proof.aggregate_rows([row, dict(row)]) + + def test_boolean_telemetry_counts_are_rejected(self) -> None: + integer_fields = (*proof.COUNTERS, "effective_prompt_tokens", + "kvflash_resident_blocks", "pflash_input_tokens", + "pflash_output_tokens") + for key in integer_fields: + with self.subTest(key=key): + row = metric("r1", 2000) + row[key] = True + with self.assertRaisesRegex(ValueError, key): + proof.aggregate_rows([row]) + + def test_boolean_wire_and_metadata_counts_do_not_prove_features(self) -> None: + rows = [metric("r1", 2000), metric("r2", 2100)] + wire = report() + wire["levels"][0]["requests_detail"][0]["effective_prompt_tokens"] = True + result = proof.verify(wire, rows, set()) + self.assertFalse(result["valid"]) + self.assertIn("missing usage.timings.effective_prompt_tokens", result["errors"][0]) + + threshold = report() + threshold["server_metadata"]["feature_config"]["prefill_threshold"] = True + result = proof.verify(threshold, rows, {"pflash"}) + self.assertFalse(result["valid"]) + self.assertIn("positive recorded token threshold", "\n".join(result["errors"])) + + pool = report() + pool["server_metadata"]["runtime_observed"]["physical_kv_pool_tokens"] = True + result = proof.verify(pool, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + self.assertIn("positive startup-observed physical pool", "\n".join(result["errors"])) + + requested = report() + requested["server_metadata"]["feature_config"]["kvflash_max_pool_tokens"] = True + result = proof.verify(requested, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + self.assertIn("recorded pool-token cap", "\n".join(result["errors"])) + + def test_full_above_pool_requires_page_traffic(self) -> None: + rows = [ + metric("r1", 9000, page_outs=0), + metric("r2", 9100, page_outs=0), + ] + result = proof.verify( + report(effective=(9000, 9100)), rows, + {"ddtree", "pflash", "kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertIn( + "no page-in/page-out", "\n".join(result["errors"]) + ) + + def test_full_aggregate_effective_demand_requires_page_traffic(self) -> None: + rows = [ + metric("r1", 5000, page_outs=0), + metric("r2", 5000, page_outs=0), + ] + result = proof.verify( + report(effective=(5000, 5000)), rows, + {"ddtree", "pflash", "kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "effective-prompt-exceeds-pool", + ) + self.assertIn("no page-in/page-out", "\n".join(result["errors"])) + + def test_full_effective_demand_is_evaluated_per_client_level(self) -> None: + input_report = report(effective=(5000, 5000)) + requests = input_report["levels"][0]["requests_detail"] + input_report["levels"] = [ + {"requests_detail": [requests[0]]}, + {"requests_detail": [requests[1]]}, + ] + rows = [ + metric("r1", 5000, page_outs=0), + metric("r2", 5000, page_outs=0), + ] + result = proof.verify( + input_report, rows, {"ddtree", "pflash", "kvflash"}, + ) + self.assertTrue(result["valid"], result["errors"]) + self.assertFalse(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "compressed-prompt-fits-pool", + ) + + def test_unknown_kvflash_variant_fails_closed(self) -> None: + rows = [ + metric("r1", 2000, page_outs=0), + metric("r2", 2100, page_outs=0), + ] + result = proof.verify( + report(variant="typo"), rows, {"kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertEqual(result["kvflash_page_traffic_reason"], "unknown-variant") + self.assertIn("no page-in/page-out", "\n".join(result["errors"])) + + def test_pflash_auto_requires_measured_input_above_threshold(self) -> None: + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[0]["pflash_input_tokens"] = 31999 + result = proof.verify(report(), rows, {"pflash"}) + self.assertFalse(result["valid"]) + self.assertIn( + "did not reach its recorded token threshold", + "\n".join(result["errors"]), + ) + + def test_kvflash_only_cannot_succeed_without_page_traffic(self) -> None: + rows = [ + metric("r1", 2000, page_outs=0), + metric("r2", 2100, page_outs=0), + ] + result = proof.verify( + report(variant="kvflash"), rows, {"kvflash"}, + ) + self.assertFalse(result["valid"]) + self.assertTrue(result["kvflash_page_traffic_required"]) + self.assertEqual( + result["kvflash_page_traffic_reason"], + "kvflash-only-ablation", + ) + self.assertIn( + "no page-in/page-out", "\n".join(result["errors"]) + ) + + def test_kvflash_requires_explicit_hashed_scorer(self) -> None: + input_report = report(variant="kvflash") + config = input_report["server_metadata"]["feature_config"] + config["kvflash_scorer_drafter"] = None + config["kvflash_scorer_drafter_sha256"] = None + rows = [metric("r1", 9000), metric("r2", 9100)] + result = proof.verify(input_report, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + text = "\n".join(result["errors"]) + self.assertIn("no explicit scorer drafter", text) + self.assertIn("not a valid SHA-256 digest", text) + + def test_kvflash_rejects_malformed_scorer_hash(self) -> None: + rows = [metric("r1", 9000), metric("r2", 9100)] + for digest in ("not-a-hash", "g" * 64, "a" * 63, "a" * 65): + with self.subTest(digest=digest): + input_report = report(variant="kvflash") + input_report["server_metadata"]["feature_config"][ + "kvflash_scorer_drafter_sha256" + ] = digest + result = proof.verify(input_report, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + self.assertIn( + "not a valid SHA-256 digest", "\n".join(result["errors"]) + ) + + def test_kvflash_requires_startup_observed_pool(self) -> None: + input_report = report(variant="kvflash") + input_report["server_metadata"]["runtime_observed"] = None + rows = [metric("r1", 9000), metric("r2", 9100)] + result = proof.verify(input_report, rows, {"kvflash"}) + self.assertFalse(result["valid"]) + text = "\n".join(result["errors"]) + self.assertIn("startup marker was not recorded", text) + self.assertIn("no positive startup-observed physical pool", text) + + def test_requested_features_cannot_succeed_silently(self) -> None: + rows = [metric("r1", 9000, page_outs=0), metric("r2", 9100, page_outs=0)] + rows[0]["ddtree_steps"] = 0 + rows[1]["pflash_applied"] = False + result = proof.verify( + report(effective=(9000, 9100)), rows, + {"ddtree", "pflash", "kvflash"}, + ) + self.assertFalse(result["valid"]) + text = "\n".join(result["errors"]) + self.assertIn("ddtree_steps is zero", text) + self.assertIn("pflash_applied is false", text) + self.assertIn("no page-in/page-out", text) + + def test_pflash_input_mismatch_with_wire_tokens_fails(self) -> None: + rows = [metric("r1", 2000), metric("r2", 2100)] + rows[0]["pflash_input_tokens"] = 39999 + result = proof.verify(report(), rows, {"pflash"}) + self.assertFalse(result["valid"]) + self.assertIn( + "pflash_input_tokens=39999 does not match wire value 40000", + "\n".join(result["errors"]), + ) + + def test_effective_prompt_mismatch_fails(self) -> None: + rows = [metric("r1", 1999), metric("r2", 2100)] + result = proof.verify(report(), rows, set()) + self.assertFalse(result["valid"]) + self.assertIn("does not match wire value", result["errors"][0]) + + def test_measured_request_requires_explicit_error_status(self) -> None: + input_report = report() + del input_report["levels"][0]["requests_detail"][0]["error"] + with self.assertRaisesRegex(ValueError, "explicit error status"): + proof.measured_requests(input_report) + + +class FeatureSummaryTests(unittest.TestCase): + @staticmethod + def item( + variant: str, goodput: float, *, repeat: int = 1, + output_hash: str | None = "same-output", + ) -> dict: + return { + "report": {"max_tokens": 256, "ignore_eos": True}, + "meta": { + "workload": "compression", "variant": variant, "repeat": repeat, + "model_sha256": "a" * 64, + }, + "level": { + "clients": 8, + "aggregate_tok_s": goodput, + "output_window_tok_s": goodput, + "effective_to_wire_prompt_ratio": 0.05 if variant == "full" else 1.0, + "ttft_max_s": 2.0, + "selected_prompt_set_sha256": "same-prompts", + "selected_output_set_sha256": output_hash, + }, + "proof": { + "aggregate": { + "ddtree_steps": 4 if variant == "full" else 0, + "ddtree_suspensions": 1 if variant == "full" else 0, + "ddtree_accepted_tokens": 8 if variant == "full" else 0, + "target_forwards": 4, + "kvflash_page_ins": 2 if variant == "full" else 0, + "kvflash_page_outs": 3 if variant == "full" else 0, + "pflash_applied_requests": 8 if variant == "full" else 0, + }, + }, + } + + def test_summary_compares_feature_row_to_ar(self) -> None: + text = summary.summarize([self.item("ar", 10.0), self.item("full", 12.0)]) + self.assertIn("+20.0%", text) + self.assertIn("2.00", text) + self.assertIn("DDTree steps/susp.", text) + self.assertIn("| 4/1 | 4 | 2/3 |", text) + + def test_feature_row_without_ar_control_reports_na(self) -> None: + text = summary.summarize([self.item("full", 12.0)]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[7].strip(), "n/a") + + def test_missing_output_digest_does_not_claim_stability(self) -> None: + first = self.item("full", 12.0, output_hash=None) + second = self.item("full", 13.0, repeat=2, output_hash=None) + text = summary.summarize([first, second]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[15].strip(), "n/a") + + def test_incomplete_repeat_metrics_are_reported_as_na(self) -> None: + first = self.item("full", 12.0) + second = self.item("full", 13.0, repeat=2) + second["level"]["output_window_tok_s"] = None + second["level"]["effective_to_wire_prompt_ratio"] = None + text = summary.summarize([first, second]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[6].strip(), "n/a") + self.assertEqual(row.split("|")[8].strip(), "n/a") + + def test_incomplete_repeat_ttft_is_reported_as_na(self) -> None: + first = self.item("full", 12.0) + second = self.item("full", 13.0, repeat=2) + second["level"]["ttft_max_s"] = None + text = summary.summarize([first, second]) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[14].strip(), "n/a") + + def test_summary_rejects_incompatible_repeat_metadata(self) -> None: + first = self.item("full", 12.0) + second = self.item("full", 13.0, repeat=2) + second["report"]["max_tokens"] = 128 + with self.assertRaisesRegex(ValueError, "incompatible run metadata"): + summary.summarize([first, second]) + + def test_summary_rejects_incompatible_ar_control_metadata(self) -> None: + ar = self.item("ar", 10.0) + feature = self.item("full", 12.0) + feature["meta"]["model_sha256"] = "b" * 64 + with self.assertRaisesRegex(ValueError, "run metadata differs"): + summary.summarize([ar, feature]) + + def test_unstable_feature_row_suppresses_ar_delta(self) -> None: + reports = [ + self.item("ar", 10.0, repeat=1), + self.item("ar", 10.0, repeat=2), + self.item("full", 12.0, repeat=1, output_hash="first"), + self.item("full", 13.0, repeat=2, output_hash="second"), + ] + text = summary.summarize(reports) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[7].strip(), "n/a") + self.assertEqual(row.split("|")[15].strip(), "NO") + + def test_unstable_ar_control_suppresses_feature_delta(self) -> None: + reports = [ + self.item("ar", 10.0, repeat=1, output_hash="first"), + self.item("ar", 10.0, repeat=2, output_hash="second"), + self.item("full", 12.0, repeat=1), + self.item("full", 13.0, repeat=2), + ] + text = summary.summarize(reports) + row = next(line for line in text.splitlines() if "| full |" in line) + self.assertEqual(row.split("|")[7].strip(), "n/a") + + +if __name__ == "__main__": + unittest.main() diff --git a/harness/benchmarks/concurrency/verify_feature_metrics.py b/harness/benchmarks/concurrency/verify_feature_metrics.py new file mode 100644 index 000000000..b8faf043b --- /dev/null +++ b/harness/benchmarks/concurrency/verify_feature_metrics.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +"""Fail closed unless server telemetry proves requested features executed.""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from pathlib import Path +from typing import Any + + +PREFIX = "[concurrency-metrics] " +COUNTERS = ( + "ddtree_steps", "ddtree_suspensions", "ddtree_accepted_tokens", + "target_forwards", "kvflash_page_ins", "kvflash_page_outs", + "kvflash_reselects", +) +REQUIRED_KEYS = ( + "request_id", "effective_prompt_tokens", *COUNTERS, + "kvflash_resident_blocks", "pflash_applied", "pflash_input_tokens", + "pflash_output_tokens", +) + + +def parse_markers(path: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for line_no, line in enumerate(path.read_text(encoding="utf-8", errors="replace").splitlines(), 1): + marker = line.find(PREFIX) + if marker < 0: + continue + raw = line[marker + len(PREFIX):] + try: + row = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError(f"{path}:{line_no}: invalid concurrency metric JSON: {exc}") from exc + if not isinstance(row, dict): + raise ValueError(f"{path}:{line_no}: concurrency metric must be an object") + missing = [key for key in REQUIRED_KEYS if key not in row] + if missing: + raise ValueError(f"{path}:{line_no}: missing telemetry keys {missing}") + rows.append(row) + return rows + + +def measured_requests(report: dict[str, Any]) -> dict[str, dict[str, Any]]: + requests: dict[str, dict[str, Any]] = {} + for level in report.get("levels") or []: + for row in level.get("requests_detail") or []: + if "error" not in row: + raise ValueError("bench report request lacks an explicit error status") + if row["error"] is not None: + continue + request_id = row.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise ValueError("bench report lacks a request ID for a successful request") + if request_id in requests: + raise ValueError(f"duplicate measured request ID {request_id}") + requests[request_id] = row + if not requests: + raise ValueError("bench report has no successful measured requests") + return requests + + +def aggregate_rows(rows: list[dict[str, Any]]) -> dict[str, dict[str, Any]]: + aggregate: dict[str, dict[str, Any]] = {} + for row in rows: + request_id = row.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise ValueError("telemetry request_id must be a non-empty string") + if request_id in aggregate: + raise ValueError(f"duplicate telemetry request ID {request_id}") + dst = { + "request_id": request_id, + "effective_prompt_tokens": row["effective_prompt_tokens"], + "kvflash_resident_blocks": row["kvflash_resident_blocks"], + "pflash_applied": False, + "pflash_input_tokens": row["pflash_input_tokens"], + "pflash_output_tokens": row["pflash_output_tokens"], + **{key: 0 for key in COUNTERS}, + } + aggregate[request_id] = dst + for key in COUNTERS: + value = row[key] + if key == "ddtree_suspensions": + if type(value) is not int or value not in (0, 1): + raise ValueError( + f"{request_id}: ddtree_suspensions must be 0 or 1 " + "per request" + ) + elif type(value) is not int or value < 0: + raise ValueError(f"{request_id}: {key} must be a non-negative integer") + dst[key] += value + for key in ( + "effective_prompt_tokens", "kvflash_resident_blocks", + "pflash_input_tokens", "pflash_output_tokens", + ): + value = row[key] + if type(value) is not int or value < 0: + raise ValueError(f"{request_id}: {key} must be a non-negative integer") + dst[key] = value + if not isinstance(row["pflash_applied"], bool): + raise ValueError(f"{request_id}: pflash_applied must be boolean") + dst["pflash_applied"] = dst["pflash_applied"] or row["pflash_applied"] + for request_id, row in aggregate.items(): + if row["ddtree_suspensions"] not in (0, 1): + raise ValueError( + f"{request_id}: ddtree_suspensions must be 0 or 1 per request" + ) + return aggregate + + +def verify( + report: dict[str, Any], markers: list[dict[str, Any]], expected: set[str], +) -> dict[str, Any]: + measured = measured_requests(report) + all_rows = aggregate_rows(markers) + rows = {request_id: all_rows[request_id] for request_id in measured if request_id in all_rows} + errors: list[str] = [] + missing = sorted(set(measured) - set(rows)) + if missing: + errors.append(f"missing concurrency telemetry for {len(missing)} measured request(s): {missing}") + + for request_id, measured_row in measured.items(): + metric = rows.get(request_id) + if metric is None: + continue + wire_effective = measured_row.get("effective_prompt_tokens") + if type(wire_effective) is not int or wire_effective < 0: + errors.append(f"{request_id}: missing usage.timings.effective_prompt_tokens") + elif metric["effective_prompt_tokens"] != wire_effective: + errors.append( + f"{request_id}: log effective_prompt_tokens={metric['effective_prompt_tokens']} " + f"does not match wire value {wire_effective}" + ) + if "pflash" in expected: + wire_input = measured_row.get("prompt_tokens") + if type(wire_input) is not int or wire_input < 0: + errors.append(f"{request_id}: missing usage.prompt_tokens") + elif metric["pflash_input_tokens"] != wire_input: + errors.append( + f"{request_id}: log pflash_input_tokens=" + f"{metric['pflash_input_tokens']} does not match wire value " + f"{wire_input}" + ) + if "ddtree" in expected: + if metric["ddtree_steps"] <= 0: + errors.append(f"{request_id}: DDTree requested but ddtree_steps is zero") + if metric["target_forwards"] <= 0: + errors.append(f"{request_id}: DDTree requested but target_forwards is zero") + if "pflash" in expected: + if metric["pflash_applied"] is not True: + errors.append(f"{request_id}: PFlash requested but pflash_applied is false") + if not (0 < metric["pflash_output_tokens"] < metric["pflash_input_tokens"]): + errors.append( + f"{request_id}: PFlash did not reduce prompt tokens " + f"({metric['pflash_input_tokens']} -> {metric['pflash_output_tokens']})" + ) + + totals = {key: sum(row[key] for row in rows.values()) for key in COUNTERS} + resident = [row["kvflash_resident_blocks"] for row in rows.values()] + metadata = report.get("server_metadata") or {} + feature_config = metadata.get("feature_config") or {} + variant = str(metadata.get("variant") or "") + workload = str(metadata.get("workload") or "") + pflash_mode = feature_config.get("prefill_compression") + pflash_threshold = feature_config.get("prefill_threshold") + if "pflash" in expected: + if not isinstance(pflash_mode, str) or pflash_mode in ("", "off", "0"): + errors.append("PFlash requested but metadata does not prove it was enabled") + if pflash_mode == "auto": + if type(pflash_threshold) is not int or pflash_threshold <= 0: + errors.append( + "PFlash auto row lacks a positive recorded token threshold" + ) + else: + below_threshold = sorted( + request_id for request_id, row in rows.items() + if row["pflash_input_tokens"] < pflash_threshold + ) + if below_threshold: + errors.append( + "PFlash auto input did not reach its recorded token threshold " + f"for request(s): {below_threshold}" + ) + kvflash_mode = feature_config.get("kvflash") + requested_pool_tokens = feature_config.get("kvflash_max_pool_tokens") + runtime_observed = metadata.get("runtime_observed") or {} + pool_tokens = runtime_observed.get("physical_kv_pool_tokens") + kvflash_page_traffic_required = False + kvflash_page_traffic_reason = "not-requested" + if "kvflash" in expected: + if not isinstance(kvflash_mode, str) or kvflash_mode in ("", "off", "0"): + errors.append("KVFlash requested but metadata does not prove it was enabled") + if (type(requested_pool_tokens) is not int + or requested_pool_tokens <= 0): + errors.append( + "KVFlash requested but its recorded pool-token cap is not a positive integer") + scorer_drafter = feature_config.get("kvflash_scorer_drafter") + scorer_sha256 = feature_config.get("kvflash_scorer_drafter_sha256") + if not isinstance(scorer_drafter, str) or not scorer_drafter: + errors.append( + "KVFlash requested but no explicit scorer drafter was recorded" + ) + if ( + not isinstance(scorer_sha256, str) + or re.fullmatch(r"[0-9a-fA-F]{64}", scorer_sha256) is None + ): + errors.append( + "KVFlash requested but the scorer drafter hash is not a valid SHA-256 digest" + ) + if runtime_observed.get("kvflash_active") is not True: + errors.append( + "KVFlash requested but its physical-pool startup marker was not recorded" + ) + if type(pool_tokens) is not int or pool_tokens <= 0: + errors.append( + "KVFlash requested but no positive startup-observed physical pool was recorded" + ) + if not resident or max(resident) <= 0: + errors.append("KVFlash requested but resident block count never became positive") + + effective_demand_by_level = [] + for level in report.get("levels") or []: + demand = 0 + for request in level.get("requests_detail") or []: + request_id = request.get("request_id") + if request.get("error") is None and request_id in rows: + demand += rows[request_id]["effective_prompt_tokens"] + effective_demand_by_level.append(demand) + if variant not in ("kvflash", "full"): + # A report whose declared variant is inconsistent with a KVFlash + # proof must not pass merely because its current prompts fit. + kvflash_page_traffic_required = True + kvflash_page_traffic_reason = "unknown-variant" + elif variant == "kvflash": + kvflash_page_traffic_required = True + kvflash_page_traffic_reason = "kvflash-only-ablation" + elif workload == "kv-pressure": + kvflash_page_traffic_required = True + kvflash_page_traffic_reason = "kv-pressure-workload" + elif variant == "full": + if type(pool_tokens) is not int or pool_tokens <= 0: + errors.append( + "full KVFlash row lacks a positive recorded pool-token limit" + ) + kvflash_page_traffic_reason = "missing-pool-limit" + else: + kvflash_page_traffic_required = any( + demand > pool_tokens + for demand in effective_demand_by_level + ) + kvflash_page_traffic_reason = ( + "effective-prompt-exceeds-pool" + if kvflash_page_traffic_required + else "compressed-prompt-fits-pool" + ) + + if ( + kvflash_page_traffic_required + and totals["kvflash_page_ins"] + totals["kvflash_page_outs"] <= 0 + ): + errors.append( + "KVFlash paging was required but no page-in/page-out was observed" + ) + + return { + "schema_version": 3, + "expected_features": sorted(expected), + "valid": not errors, + "errors": errors, + "measured_request_count": len(measured), + "matched_metric_count": len(rows), + "ignored_marker_count": len(markers) - len(rows), + "kvflash_page_traffic_required": kvflash_page_traffic_required, + "kvflash_page_traffic_reason": kvflash_page_traffic_reason, + "kvflash_pool_tokens": pool_tokens if type(pool_tokens) is int else None, + "kvflash_requested_max_pool_tokens": ( + requested_pool_tokens if type(requested_pool_tokens) is int else None + ), + "pflash_threshold_tokens": ( + pflash_threshold if type(pflash_threshold) is int else None + ), + "aggregate": { + **totals, + "kvflash_resident_blocks_max": max(resident) if resident else None, + "pflash_applied_requests": sum(row["pflash_applied"] is True for row in rows.values()), + "pflash_input_tokens": sum(row["pflash_input_tokens"] for row in rows.values()), + "pflash_output_tokens": sum(row["pflash_output_tokens"] for row in rows.values()), + }, + "requests": [rows[key] for key in sorted(rows)], + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--bench", type=Path, required=True) + parser.add_argument("--server-log", type=Path, required=True) + parser.add_argument("--expect", action="append", choices=("ddtree", "pflash", "kvflash"), default=[]) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + try: + report = json.loads(args.bench.read_text(encoding="utf-8")) + result = verify(report, parse_markers(args.server_log), set(args.expect)) + except Exception as exc: + print(f"[proof] error: {exc}", file=sys.stderr) + return 2 + args.out.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") + if not result["valid"]: + for error in result["errors"]: + print(f"[proof] {error}", file=sys.stderr) + return 1 + print( + f"[proof] valid features={','.join(result['expected_features']) or 'ar'} " + f"requests={result['matched_metric_count']}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harness/benchmarks/concurrency/write_feature_metadata.py b/harness/benchmarks/concurrency/write_feature_metadata.py new file mode 100644 index 000000000..9500c41b9 --- /dev/null +++ b/harness/benchmarks/concurrency/write_feature_metadata.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Write reproducible server and feature configuration metadata for one case.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import pathlib +import subprocess + + +def digest(path: pathlib.Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def resolved_libraries(binary: pathlib.Path) -> dict[str, str]: + result = subprocess.run( + ["ldd", str(binary)], text=True, capture_output=True, + ) + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() or "no diagnostic" + raise RuntimeError(f"ldd failed for {binary}: {detail}") + unresolved = [ + line.strip() for line in result.stdout.splitlines() + if "=> not found" in line + ] + if unresolved: + raise RuntimeError(f"ldd found unresolved libraries for {binary}: {unresolved}") + libraries: dict[str, str] = {} + for line in result.stdout.splitlines(): + fields = line.replace("=>", " ").split() + paths = [pathlib.Path(value) for value in fields if value.startswith("/")] + for path in paths: + if path.is_file(): + libraries[str(path.resolve())] = digest(path) + return libraries + + +def repository_head(repo: pathlib.Path) -> str: + result = subprocess.run( + ["git", "-C", str(repo), "rev-parse", "HEAD"], + text=True, capture_output=True, + ) + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() or "no diagnostic" + raise RuntimeError(f"git rev-parse failed for {repo}: {detail}") + head = result.stdout.strip() + if not head: + raise RuntimeError(f"git rev-parse returned an empty revision for {repo}") + return head + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=pathlib.Path, required=True) + parser.add_argument("--variant", required=True) + parser.add_argument("--workload", required=True) + parser.add_argument("--clients", type=int, required=True) + parser.add_argument("--repeat", type=int, required=True) + parser.add_argument("--binary", type=pathlib.Path, required=True) + parser.add_argument("--model", type=pathlib.Path, required=True) + parser.add_argument("--model-sha256", required=True) + parser.add_argument("--prompt-file", type=pathlib.Path, required=True) + parser.add_argument("--command-file", type=pathlib.Path, required=True) + parser.add_argument("--repo", type=pathlib.Path, required=True) + parser.add_argument("--max-concurrent-prefills", type=int, required=True) + parser.add_argument("--target-device", default=None) + parser.add_argument("--draft-device", default=None) + parser.add_argument("--draft-model", type=pathlib.Path) + parser.add_argument("--draft-model-sha256") + parser.add_argument("--ddtree", action="store_true") + parser.add_argument("--ddtree-budget", type=int) + parser.add_argument("--prefill-compression", default="off") + parser.add_argument("--prefill-threshold", type=int) + parser.add_argument("--prefill-keep-ratio", type=float) + parser.add_argument("--prefill-drafter", type=pathlib.Path) + parser.add_argument("--prefill-drafter-sha256") + parser.add_argument("--draft-residency", default=None) + parser.add_argument("--kvflash", default="off") + parser.add_argument("--kvflash-max-pool-tokens", type=int) + parser.add_argument("--kvflash-scorer-drafter", type=pathlib.Path) + parser.add_argument("--kvflash-scorer-drafter-sha256") + parser.add_argument("--launch-env", action="append", default=[]) + args = parser.parse_args() + + launch_env: dict[str, str] = {} + for item in args.launch_env: + key, sep, value = item.partition("=") + if not sep or not key: + parser.error(f"bad --launch-env {item!r}; expected KEY=VALUE") + launch_env[key] = value + + libraries = resolved_libraries(args.binary) + git_head = repository_head(args.repo) + literal_flags: list[str] = [] + if args.target_device: + literal_flags += ["--target-device", args.target_device] + if args.draft_device: + literal_flags += ["--draft-device", args.draft_device] + if args.ddtree: + literal_flags += ["--ddtree"] + if args.ddtree_budget is not None: + literal_flags += ["--ddtree-budget", str(args.ddtree_budget)] + if args.draft_residency: + literal_flags += ["--draft-residency", args.draft_residency] + if args.prefill_compression != "off": + literal_flags += ["--prefill-compression", args.prefill_compression] + if args.prefill_drafter: + literal_flags += ["--prefill-drafter", str(args.prefill_drafter.resolve())] + if args.kvflash != "off": + literal_flags += ["--kvflash", args.kvflash] + + obj = { + "schema_version": 3, + "variant": args.variant, + "workload": args.workload, + "clients": args.clients, + "repeat": args.repeat, + "max_concurrent_prefills": args.max_concurrent_prefills, + "server_binary": str(args.binary.resolve()), + "server_binary_sha256": digest(args.binary), + "model": str(args.model.resolve()), + "model_sha256": args.model_sha256, + "prompt_file_sha256": digest(args.prompt_file), + "server_command": args.command_file.read_text(encoding="utf-8").strip(), + "launch_environment": launch_env, + "resolved_shared_library_sha256": libraries, + "git_head": git_head, + "literal_screenshot_flags": literal_flags, + # Populated from fail-closed startup markers after the server is healthy. + "runtime_observed": None, + "feature_config": { + "target_device": args.target_device, + "draft_device": args.draft_device, + "draft_model": str(args.draft_model.resolve()) if args.draft_model else None, + "draft_model_sha256": args.draft_model_sha256, + "ddtree": args.ddtree, + "ddtree_budget": args.ddtree_budget, + "prefill_compression": args.prefill_compression, + "prefill_threshold": args.prefill_threshold, + "prefill_keep_ratio": args.prefill_keep_ratio, + "prefill_drafter": ( + str(args.prefill_drafter.resolve()) if args.prefill_drafter else None + ), + "prefill_drafter_sha256": args.prefill_drafter_sha256, + "draft_residency": args.draft_residency, + "kvflash": args.kvflash, + "kvflash_max_pool_tokens": args.kvflash_max_pool_tokens, + "kvflash_scorer_drafter": ( + str(args.kvflash_scorer_drafter.resolve()) + if args.kvflash_scorer_drafter else None + ), + "kvflash_scorer_drafter_sha256": ( + args.kvflash_scorer_drafter_sha256 + ), + }, + } + args.out.write_text(json.dumps(obj, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 631a9e176cba00fe3df8993e3219b039899cdbb8 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 12 Aug 2026 23:09:18 +0000 Subject: [PATCH 03/18] fix(qwen36): harden feature smoke paths --- .../concurrency/run_qwen36_feature_matrix.sh | 33 +++-- .../qwen35/concurrency/qwen35_seq_engine.cpp | 10 +- server/src/qwen35/graph_builders.cpp | 75 +++++++++-- server/src/qwen35/graph_builders.h | 36 ++++++ server/test/test_recurrent_snapshot.cpp | 116 ++++++++++++++++++ 5 files changed, 248 insertions(+), 22 deletions(-) diff --git a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh index e4b23bca1..d0ea1786a 100755 --- a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +++ b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh @@ -322,18 +322,33 @@ run_case() { --require-distinct-prompts --temperature 0 --ignore-eos --timeout "$timeout" --cooldown 0) local -a telemetry_arg=() [[ "$variant" != llama ]] && telemetry_arg+=(--require-effective-prompt-telemetry) - python3 "$CLIENT" "${common_client[@]}" --max-tokens "$WARMUP_TOKENS" - "${telemetry_arg[@]}" --out "$case_dir/warmup.json" - --label "$variant $workload C=$clients warmup" > "$case_dir/warmup.txt" - python3 "$CLIENT" "${common_client[@]}" --max-tokens "$MAX_TOKENS" - "${telemetry_arg[@]}" --server-metadata-json "$case_dir/server-metadata.json" - --out "$case_dir/bench.json" --label "$variant $workload C=$clients repeat=$repeat" \ - | tee "$case_dir/bench.txt" + local -a warmup_cmd=( + python3 "$CLIENT" "${common_client[@]}" + --max-tokens "$WARMUP_TOKENS" "${telemetry_arg[@]}" + --out "$case_dir/warmup.json" + --label "$variant $workload C=$clients warmup" + ) + "${warmup_cmd[@]}" > "$case_dir/warmup.txt" + + local -a benchmark_cmd=( + python3 "$CLIENT" "${common_client[@]}" + --max-tokens "$MAX_TOKENS" "${telemetry_arg[@]}" + --server-metadata-json "$case_dir/server-metadata.json" + --out "$case_dir/bench.json" + --label "$variant $workload C=$clients repeat=$repeat" + ) + "${benchmark_cmd[@]}" | tee "$case_dir/bench.txt" stop_server if [[ "$variant" != llama ]]; then - python3 "$PROOF_TOOL" --bench "$case_dir/bench.json" --server-log "$case_dir/server.log" - "${expected[@]}" --out "$case_dir/feature-proof.json" + local -a proof_cmd=( + python3 "$PROOF_TOOL" + --bench "$case_dir/bench.json" + --server-log "$case_dir/server.log" + "${expected[@]}" + --out "$case_dir/feature-proof.json" + ) + "${proof_cmd[@]}" fi sleep "$COOLDOWN_SECONDS" } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 545e7d9bf..10e50450a 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -339,8 +339,14 @@ std::optional Qwen35SeqEngine::step_ddtree( sizeof(int32_t) * parents.size()); ggml_backend_tensor_set(tree_sg.tree_sizes, sizes.data(), 0, sizeof(int32_t) * sizes.size()); - ggml_backend_tensor_set(tree_sg.active_slot_ids, tree_slots.data(), 0, - sizeof(int32_t) * tree_slots.size()); + // Mapped-tree DeltaNet uses active_slot_ids only as a topology marker; + // gallocr may therefore optimize away its backing buffer. The actual + // state/attention mappings below are live graph inputs and remain + // mandatory. Upload the marker only if a future topology consumes it. + if (detail::target_paged_tree_active_slots_need_upload(tree_sg)) { + ggml_backend_tensor_set(tree_sg.active_slot_ids, tree_slots.data(), 0, + sizeof(int32_t) * tree_slots.size()); + } ggml_backend_tensor_set(tree_sg.state_slot_ids, tree_state_slots.data(), 0, sizeof(int32_t) * tree_state_slots.size()); ggml_backend_tensor_set(tree_sg.paged_query_seq_ids, query_slots.data(), 0, diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 0ca65ecdb..f96349af3 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -5,10 +5,56 @@ #include #include #include +#include #include namespace dflash::common { +bool detail::target_graph_capacity_for_parallel_segments( + int n_parallel_segments, + size_t & capacity) { + static constexpr int k_max_parallel_segments = 64; + static constexpr size_t k_base_capacity = 16384; + static constexpr int k_segments_per_capacity = 8; + static constexpr size_t k_max_capacity = + k_base_capacity * + (k_max_parallel_segments / k_segments_per_capacity); + + if (n_parallel_segments < 0 || + n_parallel_segments > k_max_parallel_segments) { + return false; + } + const int64_t scale = std::max( + 1, ((int64_t)n_parallel_segments + + k_segments_per_capacity - 1) / + k_segments_per_capacity); + if ((uint64_t)scale > + std::numeric_limits::max() / k_base_capacity) { + return false; + } + const size_t computed = k_base_capacity * (size_t)scale; + if (computed > k_max_capacity) return false; + capacity = computed; + return true; +} + +bool detail::target_paged_tree_graph_capacity( + int tree_width, + int n_tree_seqs, + size_t & capacity) { + static constexpr int tree_buckets[] = { + 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, + }; + if (tree_width < 1 || tree_width > 256 || + std::find(std::begin(tree_buckets), std::end(tree_buckets), + n_tree_seqs) == std::end(tree_buckets) || + (int64_t)tree_width * n_tree_seqs > INT32_MAX) { + return false; + } + return target_graph_capacity_for_parallel_segments( + n_tree_seqs, capacity); +} + bool detail::validate_target_paged_tree_layout( const TargetCache & cache, int tree_width, @@ -16,12 +62,9 @@ bool detail::validate_target_paged_tree_layout( int paged_max_kv_len, int tree_scratch_base, int tree_scratch_stride) { - static constexpr int tree_buckets[] = { - 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, - }; - if (tree_width < 1 || - std::find(std::begin(tree_buckets), std::end(tree_buckets), - n_tree_seqs) == std::end(tree_buckets) || + size_t graph_capacity = 0; + if (!target_paged_tree_graph_capacity( + tree_width, n_tree_seqs, graph_capacity) || cache.n_seq_slots <= 1 || !cache.paged_block_table || !cache.paged_kv_seq_lens || paged_max_kv_len < 1 || tree_scratch_base <= 0 || @@ -43,8 +86,7 @@ bool detail::validate_target_paged_tree_layout( (int64_t)tree_scratch_base + (int64_t)(cache.n_seq_slots - 1) * tree_scratch_stride + tree_width; - return scratch_end <= physical_kv_rows && - (int64_t)tree_width * n_tree_seqs <= INT32_MAX; + return scratch_end <= physical_kv_rows; } // ── build_layer_step ──────────────────────────────────────────── @@ -370,6 +412,11 @@ bool build_target_step( } if (segment_total != n_prefill_tokens) return false; if (n_logits_rows > 0 && n_prefill_tokens == 0) return false; + size_t graph_capacity = 0; + if (!detail::target_graph_capacity_for_parallel_segments( + n_prefill_segments, graph_capacity)) { + return false; + } // Persistent thread_local arena: rebuilt step graphs land at identical // addresses, keeping the ggml-cuda CUDA-graph cache key (nodes[0]) and @@ -535,7 +582,7 @@ bool build_target_step( ggml_set_input(sg.target_feat_rows); } - sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); + sg.gf = ggml_new_graph_custom(sg.ctx, graph_capacity, false); // Step-invariant KV write: only when topology can't vary per step. // DFLASH_QWEN35_NO_KVPAD=1 restores the legacy cpy append + exact-length @@ -696,6 +743,11 @@ bool build_target_step_paged_tree( tree_scratch_base, tree_scratch_stride)) { return false; } + size_t graph_capacity = 0; + if (!detail::target_paged_tree_graph_capacity( + tree_width, n_tree_seqs, graph_capacity)) { + return false; + } const int n_tokens = tree_width * n_tree_seqs; ggml_init_params ip{}; @@ -748,7 +800,7 @@ bool build_target_step_paged_tree( ggml_set_input(input.tensor); } - sg.gf = ggml_new_graph_custom(sg.ctx, 16384, false); + sg.gf = ggml_new_graph_custom(sg.ctx, graph_capacity, false); QwenGraphInputs gi{}; gi.inp_embed = sg.inp_embed; gi.positions = sg.positions; @@ -783,7 +835,8 @@ bool build_target_step_paged_tree( sg.alloc = ggml_gallocr_new( ggml_backend_get_default_buffer_type(backend)); } - return ggml_gallocr_alloc_graph(sg.alloc, sg.gf); + return ggml_gallocr_alloc_graph(sg.alloc, sg.gf) && + detail::target_paged_tree_uploads_ready(sg); } diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index fac406863..cbe58b786 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -25,6 +25,20 @@ namespace dflash::common { namespace detail { +// Qwen's recurrent graph duplicates one small subgraph per ragged sequence. +// Return a graph capacity that covers every supported concurrent bucket while +// keeping the legacy allocation for the common <= 8-sequence case. +bool target_graph_capacity_for_parallel_segments( + int n_parallel_segments, + size_t & capacity); + +// Checked packed-tree shape/capacity contract. The public DDTree budget allows +// at most 255 children plus the root, and concurrent serving at most 64 slots. +bool target_paged_tree_graph_capacity( + int tree_width, + int n_tree_seqs, + size_t & capacity); + // Model-free validation shared by the packed-tree builder and its shape // tests. paged_max_kv_len is a logical launch bound and may exceed the // bounded physical K/V pool; only the per-slot scratch slabs must fit in the @@ -37,6 +51,28 @@ bool validate_target_paged_tree_layout( int tree_scratch_base, int tree_scratch_stride); +// `active_slot_ids` is a topology marker in mapped-tree graphs. It may be +// optimized out by gallocr because the actual recurrent and attention row +// mappings are carried by state_slot_ids and paged_query_seq_ids. Every other +// tensor listed here is read by a graph node and must have backend storage +// before the engine uploads metadata. +inline bool target_paged_tree_uploads_ready(const StepGraph & sg) { + const auto allocated = [](const ggml_tensor * tensor) { + return tensor && tensor->buffer; + }; + return sg.active_slot_ids && + allocated(sg.inp_embed) && allocated(sg.positions) && + allocated(sg.parent_ids) && allocated(sg.tree_sizes) && + allocated(sg.state_slot_ids) && + allocated(sg.paged_query_seq_ids) && + allocated(sg.kv_write_rows); +} + +inline bool target_paged_tree_active_slots_need_upload( + const StepGraph & sg) { + return sg.active_slot_ids && sg.active_slot_ids->buffer; +} + } // namespace detail // Layer-segmented prefill: process one target layer for chunk_start..chunk_start+n_tokens. diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index 05e8a59cb..b0bb0475b 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -12,6 +12,7 @@ using namespace CppUnitTestFramework; using dflash::common::TargetCache; +using dflash::common::StepGraph; using dflash::common::restore_ssm_state; using dflash::common::snapshot_ssm_state; @@ -33,6 +34,90 @@ static std::vector get_tensor(const ggml_tensor * tensor) { return values; } +TEST_CASE(RecurrentSnapshotFixture, hardens_feature_smoke_paths) { + size_t graph_capacity = 0; + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 0, graph_capacity) && graph_capacity == 16384); + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 8, graph_capacity) && graph_capacity == 16384); + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 16, graph_capacity) && graph_capacity == 32768); + CHECK(dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 64, graph_capacity) && graph_capacity == 131072); + CHECK(!dflash::common::detail:: + target_graph_capacity_for_parallel_segments( + 65, graph_capacity)); + CHECK(dflash::common::detail::target_paged_tree_graph_capacity( + 23, 16, graph_capacity) && graph_capacity == 32768); + CHECK(!dflash::common::detail::target_paged_tree_graph_capacity( + 257, 16, graph_capacity)); + + // Mapped-tree active_slot_ids is a topology marker and can legitimately + // be left without gallocr storage. Required state/query/write metadata + // remains allocated and uploadable. + { + ggml_backend_t tree_backend = ggml_backend_cpu_init(); + CHECK(tree_backend != nullptr); + ggml_init_params marker_params{}; + marker_params.mem_size = 4 * ggml_tensor_overhead(); + marker_params.no_alloc = true; + ggml_context * marker_ctx = ggml_init(marker_params); + ggml_init_params live_params{}; + live_params.mem_size = 16 * ggml_tensor_overhead(); + live_params.no_alloc = true; + ggml_context * live_ctx = ggml_init(live_params); + CHECK(marker_ctx != nullptr); + CHECK(live_ctx != nullptr); + if (tree_backend && marker_ctx && live_ctx) { + StepGraph tree; + tree.active_slot_ids = + ggml_new_tensor_1d(marker_ctx, GGML_TYPE_I32, 2); + ggml_tensor * unallocated_state_ids = + ggml_new_tensor_1d(marker_ctx, GGML_TYPE_I32, 2); + tree.inp_embed = ggml_new_tensor_2d( + live_ctx, GGML_TYPE_F32, 4, 4); + tree.positions = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 16); + tree.parent_ids = + ggml_new_tensor_2d(live_ctx, GGML_TYPE_I32, 2, 2); + tree.tree_sizes = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 2); + tree.state_slot_ids = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 2); + tree.paged_query_seq_ids = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 4); + tree.kv_write_rows = + ggml_new_tensor_2d(live_ctx, GGML_TYPE_I64, 4, 1); + ggml_backend_buffer_t live_buffer = + ggml_backend_alloc_ctx_tensors(live_ctx, tree_backend); + CHECK(live_buffer != nullptr); + if (live_buffer) { + CHECK(tree.active_slot_ids->buffer == nullptr); + CHECK(dflash::common::detail:: + target_paged_tree_uploads_ready(tree)); + CHECK(!dflash::common::detail:: + target_paged_tree_active_slots_need_upload(tree)); + + const int32_t state_ids[] = {0, 1}; + ggml_backend_tensor_set(tree.state_slot_ids, state_ids, 0, + sizeof(state_ids)); + tree.state_slot_ids = unallocated_state_ids; + CHECK(!dflash::common::detail:: + target_paged_tree_uploads_ready(tree)); + ggml_backend_buffer_free(live_buffer); + } + } + if (live_ctx) ggml_free(live_ctx); + if (marker_ctx) ggml_free(marker_ctx); + if (tree_backend) ggml_backend_free(tree_backend); + } + +} + TEST_CASE(RecurrentSnapshotFixture, validates_paged_tree_layout) { // The packed-tree launch length is logical. KVFlash may keep a much // smaller physical resident pool, provided every tree scratch slab still @@ -70,6 +155,37 @@ TEST_CASE(RecurrentSnapshotFixture, snapshot_and_restore_recurrent_state) { CHECK(backend != nullptr); if (!backend) SKIP("CPU backend is unavailable"); + // C16 x width-23 selects a 32K graph. Prove that graph traversal and + // gallocr can cross the old 16K hard ceiling without asserting. + { + CHECK(dflash::common::detail::target_paged_tree_graph_capacity( + 23, 16, graph_capacity)); + ggml_init_params graph_params{}; + graph_params.mem_size = 32 * 1024 * 1024; + graph_params.no_alloc = true; + ggml_context * graph_ctx = ggml_init(graph_params); + CHECK(graph_ctx != nullptr); + if (graph_ctx) { + ggml_tensor * input = + ggml_new_tensor_1d(graph_ctx, GGML_TYPE_F32, 1); + ggml_set_input(input); + ggml_cgraph * graph = ggml_new_graph_custom( + graph_ctx, graph_capacity, false); + for (int i = 0; i < 16385; ++i) { + ggml_build_forward_expand( + graph, ggml_dup(graph_ctx, input)); + } + CHECK(ggml_graph_n_nodes(graph) == 16385); + ggml_gallocr_t graph_alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + CHECK(graph_alloc != nullptr); + CHECK(graph_alloc && + ggml_gallocr_alloc_graph(graph_alloc, graph)); + if (graph_alloc) ggml_gallocr_free(graph_alloc); + ggml_free(graph_ctx); + } + } + ggml_init_params params{}; params.mem_size = 8 * ggml_tensor_overhead(); params.no_alloc = true; From 5f7c5df392bffe9316e0e0a56244c7466e0b616d Mon Sep 17 00:00:00 2001 From: Graffioh Date: Thu, 13 Aug 2026 00:04:36 +0000 Subject: [PATCH 04/18] feat(qwen36): tune concurrent serving for Strix Halo --- .../benchmarks/concurrency/FEATURE_MATRIX.md | 30 +++- .../concurrency/STRIX_HALO_RESULTS.md | 134 ++++++++++++++++++ .../concurrency/concurrent_benchmark.py | 11 +- .../feature_concurrent_benchmark.py | 8 +- .../concurrency/record_feature_runtime.py | 2 +- .../concurrency/run_qwen36_concurrency.sh | 78 ++++++++-- .../concurrency/run_qwen36_feature_matrix.sh | 32 ++++- .../concurrency/summarize_concurrency.py | 30 ++-- .../concurrency/summarize_feature_matrix.py | 39 +++-- .../concurrency/test_concurrency_tools.py | 48 +++++++ .../concurrency/test_concurrent_benchmark.py | 16 +++ .../test_feature_concurrent_benchmark.py | 26 ++++ .../concurrency/test_feature_metadata.py | 45 +++++- .../concurrency/test_feature_tools.py | 36 ++++- .../concurrency/write_feature_metadata.py | 53 ++++++- .../common/concurrency/paged_kv_residency.cpp | 81 ++++++++--- .../common/concurrency/paged_kv_residency.h | 1 + .../concurrency/qwen_paged_kv_transfer.cpp | 6 +- server/src/common/concurrency/seq_engine.h | 5 + .../qwen35/concurrency/qwen35_seq_engine.cpp | 42 +++++- .../concurrency/qwen35_slot_manager.cpp | 42 +++++- .../qwen35/concurrency/qwen35_slot_manager.h | 26 +++- server/src/qwen35/qwen35_backend.cpp | 7 + server/src/server/scheduler.cpp | 3 + server/test/test_ddtree_path.cpp | 24 ++-- server/test/test_paged_kv_residency.cpp | 105 +++++++++++++- server/test/test_seq_batch_plan.cpp | 13 ++ server/test/test_seq_slot_manager.cpp | 114 +++++++++++++++ 28 files changed, 955 insertions(+), 102 deletions(-) create mode 100644 harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index 7a1885f18..62fce6a46 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -1,5 +1,8 @@ # Qwen3.6 concurrent feature matrix +The bounded Strix Halo measurements collected for the draft implementation are +recorded in [`STRIX_HALO_RESULTS.md`](STRIX_HALO_RESULTS.md). + `run_qwen36_feature_matrix.sh` extends the PR #596 protocol with feature ablations for the complete Strix Halo configuration: @@ -19,12 +22,25 @@ Run the default bounded C4 screening repeat (seven applicable fresh-server cases ```bash MODEL=/opt/models/Qwen3.6-27B-Q4_K_M.gguf \ -DRAFT_MODEL=/opt/models/draft/dflash-draft-3.6-q4_k_m.gguf \ +DRAFT_MODEL=/opt/models/draft/dflash-draft-3.6-q8_0.gguf \ PREFILL_DRAFTER=/opt/models/Qwen3-0.6B-BF16.gguf \ REPEATS=1 \ harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh ``` +For the canonical AMD Strix Halo recipe, use the published Q8_0 3.6 drafter and make the tuning explicit: + +```bash +MODEL=/opt/models/Qwen3.6-27B-Q4_K_M.gguf \ +DRAFT_MODEL=/opt/models/draft/dflash-draft-3.6-q8_0.gguf \ +PREFILL_DRAFTER=/opt/models/Qwen3-0.6B-BF16.gguf \ +DRAFT_SWA=2048 PREFILL_UBATCH=512 DDTREE_ADAPTIVE=0 \ +VARIANTS=ddtree,pflash,kvflash,full CLIENTS=1,4,8,16 REPEATS=5 \ +harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +`DDTREE_ADAPTIVE=0` matches the blog's continuous DDTree probe policy; leave it at the default `1` when measuring the concurrent engine's adaptive fallback policy. The concurrent path now records a startup `[parallel-ddtree]` marker and per-request `ddtree_steps`; these are the proof that DDTree actually ran. + On a 128 GiB Strix Halo host, budget roughly 45–90 minutes for this smoke run; the long-context AR controls dominate and actual time depends on the build. Every row remains independently selectable through `VARIANTS`. For example: @@ -77,8 +93,9 @@ The server must write one JSON object per completed request with this prefix: ``` Required fields are `effective_prompt_tokens`, `ddtree_steps`, -`ddtree_accepted_tokens`, `target_forwards`, `kvflash_page_ins`, -`kvflash_page_outs`, `kvflash_resident_blocks`, `kvflash_reselects`, +`ddtree_suspensions`, `ddtree_accepted_tokens`, `target_forwards`, +`kvflash_page_ins`, `kvflash_page_outs`, `kvflash_resident_blocks`, +`kvflash_reselects`, `pflash_applied`, `pflash_input_tokens`, and `pflash_output_tokens`. The proof tool correlates log objects with measured SSE request IDs and also @@ -87,6 +104,10 @@ checks the log's effective token count against unless: - DDTree has positive step and target-forward counts. Acceptance may be zero. + The required per-request suspension counter must be either zero or one. It + records adaptive fallback activation, but does not prove whether AR work ran + before or after the suspension; temporal claims require direct ordered-log + evidence. - PFlash reports `pflash_applied=true`, a smaller output prompt, and (in auto mode) an input token count at or above the recorded activation threshold. - KVFlash always records an explicit hashed scorer drafter, reports its @@ -108,6 +129,3 @@ binary/shared-library/target/draft/PFlash-and-KV-scorer hashes, the ordered `literal_screenshot_flags` array, all feature values, raw request report, server log, and `feature-proof.json`. The summary refuses to include a Lucebox row whose proof is missing or invalid. - - - diff --git a/harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md b/harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md new file mode 100644 index 000000000..c42eea4f9 --- /dev/null +++ b/harness/benchmarks/concurrency/STRIX_HALO_RESULTS.md @@ -0,0 +1,134 @@ +# Qwen3.6 concurrent feature results — Strix Halo + +- Date: 2026-08-13 +- Implementation: `568fbac03b498d53d6efc0b2ab5893044543a321` +- Stack base: PR #595 head `a90ffe45c1d4ad58f5f73c4107571d3cf6c51bfd` + +These are bounded engineering measurements for the draft PR, not the +five-repeat publication matrix described in `FEATURE_MATRIX.md`. The paired +AR/DDTree screen has three fresh-process repeats. The long-context activation +rows have one fresh-process repeat per concurrency level because they are much +more expensive; treat their throughput as screening data. + +## System and artifacts + +- AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`), 128 GiB unified memory. +- ROCm runtime 7.2.4. +- Release HIP build for `gfx1151` with + `DFLASH27B_HIP_SM80_EQUIV=ON`. +- Server SHA-256: + `c77b4d2c7d1505fcc751600a6603cd65e51514b685bc66c4f7d33cd64a87c8a6`. +- Target SHA-256: + `5ed60d0af4650a854b1755bd392f9aef4872643dc25a254bc68043fa638392a0`. +- Decode draft SHA-256: + `e2500e90165a0f8e7b52c9882c29ed1fa391c60b300ff11b817bf10e31fa092e`. +- PFlash/KV scorer drafter SHA-256: + `f9c9f1d3c1e21755b82d4e165f88dbbbd4355646d632fb5d6cef7c66ed4ee04e`. + +Every case started a fresh server, discarded an 8-token same-concurrency +warmup, then requested exactly 64 output tokens per request with temperature +zero, seed one, and EOS ignored. Prompts were deterministic, disjoint across +concurrency levels, and identical between paired variants. The runner rotated +variant order across repeats. + +`Output-window` counts all completion tokens from the earliest first output to +the last completion. `Goodput` counts completion tokens over the whole level, +including TTFT. Every reported row passed exact token accounting and the +request-ID-correlated feature proof. + +The retained screening artifacts contain maximum TTFT but not median TTFT. +Their max-only columns below are an explicit screening exception, not a +protocol-complete publication result; a publication rerun must report both. + +## Paired AR and adaptive DDTree + +The DDTree configuration adds the local decode draft, budget 22, and target and +draft placement on `hip:0`. Values are medians over three fresh-process +repeats. + +| C | Variant | N | Goodput tok/s | Output-window tok/s | vs AR goodput | Accepted/step | Steps/suspensions | Max TTFT s | Output hashes stable | +| ---: | :--- | ---: | ---: | ---: | ---: | ---: | :--- | ---: | :---: | +| 1 | AR | 3 | 9.41 | 12.57 | — | — | 0/0 | 1.707 | yes | +| 1 | DDTree | 3 | 9.00 | 11.85 | -4.4% | 1.00 | 1/1 | 1.715 | yes | +| 4 | AR | 3 | 20.46 | 36.21 | — | — | 0/0 | 5.508 | no | +| 4 | DDTree | 3 | 19.43 | 33.61 | n/a | 3.08 | 4/4 | 5.533 | no | +| 8 | AR | 3 | 27.44 | 65.93 | — | — | 0/0 | 11.008 | no | +| 8 | DDTree | 3 | 25.63 | 56.32 | n/a | 2.79 | 8/8 | 11.076 | no | +| 16 | AR | 3 | 31.82 | 58.79 | — | — | 0/0 | 22.487 | no | +| 16 | DDTree | 3 | 29.36 | 54.63 | n/a | 2.00 | 16/16 | 22.514 | no | + +The supplied draft had weak acceptance on this cohort. The adaptive policy +sampled one real packed-tree step, then suspended the whole cohort because its +aggregate emitted yield was below six tokens per request. At C4 and above, +the raw timings are retained only to diagnose this fallback behavior; unstable +outputs do not support a performance comparison with AR. + +At C4 and above, greedy text hashes varied across fresh repeats in both the AR +control and DDTree. C1 was byte-stable. These measurements therefore establish +exact token accounting and feature execution, but do not claim bitwise text +reproducibility for concurrent batches. + +## Full screenshot configuration + +These rows enable the complete requested product configuration: + +```text +--target-device hip:0 +--draft-device hip:0 +--ddtree +--ddtree-budget 22 +--draft-residency persistent +--prefill-compression auto +--prefill-drafter /opt/models/Qwen3-0.6B-BF16.gguf +--kvflash auto +``` + +The controlled runner sets the auto PFlash threshold to 32K tokens, the keep +ratio to 0.05, and the KVFlash resident cap to 8,192 tokens. Startup telemetry +confirmed 512 physical blocks of 16 tokens, 16 configured slots, and a 65,536 +logical-token bound per slot. + +| C | N | Goodput tok/s | Output-window tok/s | Request decode tok/s | Raw prompt range | Effective prompt range | Max TTFT s | DDTree steps/susp. | KV page in/out | PFlash requests | +| ---: | ---: | ---: | ---: | ---: | :--- | :--- | ---: | :--- | :--- | ---: | +| 1 | 1 | 2.61 | 12.39 | 12.19 | 41,504 | 2,021 | 19.324 | 1/1 | 0/0 | 1 | +| 4 | 1 | 3.00 | 31.61 | 7.95 | 38,142–44,866 | 1,870–2,235 | 77.495 | 4/4 | 1/18 | 4 | +| 8 | 1 | 3.13 | 52.60 | 6.56 | 38,141–44,870 | 1,869–2,237 | 153.782 | 8/8 | 245/792 | 8 | +| 16 | 1 | 3.19 | 18.73 | 2.34 | 38,140–44,872 | 1,867–2,238 | 307.360 | 16/16 | 293/1,880 | 16 | + +All four rows proved DDTree, PFlash, and KVFlash active. PFlash retained about +4.9% of raw prompt tokens. Output-window throughput scaled through C8, then +dropped at C16 while roughly 32K effective prompt tokens shared the 8K resident +pool; the concurrent page traffic rose accordingly. + +## Feature ablations + +| Workload | C | Variant | N | Goodput tok/s | Output-window tok/s | Request decode tok/s | Effective/raw | Max TTFT s | Activation evidence | +| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | :--- | +| compression | 4 | PFlash | 1 | 3.15 | 35.54 | 8.89 | 0.049 | 74.261 | 4/4 prompts compressed, 166,016 -> 8,179 tokens | +| kv-pressure | 4 | KVFlash | 1 | 1.24 | 8.56 | 5.32 | 1.000 | 197.859 | 129 resident blocks max, 0 page-ins / 3,714 page-outs | + +The PFlash-only row uses the same C4 prompts as the full row; adding DDTree and +KVFlash reduced output-window throughput from 35.54 to 31.61 tok/s in this +single screening repeat. The KVFlash-only row deliberately disables PFlash and +uses 13,474–20,203-token histories against the 8K pool. It is an activation and +pressure test, not a recommended latency configuration. + +## Reproduction + +The exact per-case command, controlled environment, startup-observed pool, +binary/shared-library/model hashes, raw request report, server log, and +`feature-proof.json` are retained by the runner. The principal invocations were: + +```bash +WORKLOADS=short CLIENTS=1,4,8,16 VARIANTS=ar,ddtree MAX_TOKENS=64 REPEATS=3 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +WORKLOADS=compression CLIENTS=1,4,8,16 VARIANTS=full MAX_TOKENS=64 REPEATS=1 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +WORKLOADS=compression CLIENTS=4 VARIANTS=pflash MAX_TOKENS=64 REPEATS=1 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh + +WORKLOADS=kv-pressure CLIENTS=4 VARIANTS=kvflash MAX_TOKENS=64 REPEATS=1 SLOTS=16 harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +``` + +Set `MODEL`, `DRAFT_MODEL`, `PREFILL_DRAFTER`, and `LUCE_SERVER_BIN` as shown +in `FEATURE_MATRIX.md`. For publication-quality claims, rerun the documented +five-repeat 256-token matrix. diff --git a/harness/benchmarks/concurrency/concurrent_benchmark.py b/harness/benchmarks/concurrency/concurrent_benchmark.py index 773260393..4c3e85774 100755 --- a/harness/benchmarks/concurrency/concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/concurrent_benchmark.py @@ -164,11 +164,16 @@ def worker(index: int) -> None: threads = [threading.Thread(target=worker, args=(i,), daemon=True) for i in range(clients)] for thread in threads: thread.start() + deadline = time.monotonic() + args.timeout + 30 for thread in threads: - thread.join(args.timeout + 30) - completed = [record for record in records if record is not None] + thread.join(max(0.0, deadline - time.monotonic())) hung = sum(thread.is_alive() for thread in threads) - failures = hung + sum(record["error"] is not None for record in completed) + if hung: + raise TimeoutError( + f"{hung} request worker(s) exceeded the level deadline" + ) + completed = [record for record in records if record is not None] + failures = sum(record["error"] is not None for record in completed) ok = [record for record in completed if record["error"] is None] starts = [record["t_start"] for record in completed] ends = [record["t_end"] for record in completed] diff --git a/harness/benchmarks/concurrency/feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/feature_concurrent_benchmark.py index bc509c590..c9df122e1 100755 --- a/harness/benchmarks/concurrency/feature_concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/feature_concurrent_benchmark.py @@ -171,9 +171,9 @@ def markdown(report: dict[str, Any]) -> str: lines = [ f"# Concurrent feature benchmark — {report['label']}", "", "| C | Ok | Output goodput tok/s | Output-window tok/s | " - "Request decode tok/s | Wire prompt range | Effective prompt range | " - "Effective/wire | TTFT max s |", - "| ---: | ---: | ---: | ---: | ---: | :--- | :--- | ---: | ---: |", + "Request decode tok/s | Prompt tok/s to first | Wire prompt range | " + "Effective prompt range | Effective/wire | TTFT median/max s |", + "| ---: | ---: | ---: | ---: | ---: | ---: | :--- | :--- | ---: | :--- |", ] for level in report["levels"]: lines.append( @@ -181,10 +181,12 @@ def markdown(report: dict[str, Any]) -> str: f"{base.fmt(level['aggregate_tok_s'])} | " f"{base.fmt(level['output_window_tok_s'])} | " f"{base.fmt(level['request_decode_tok_s_median'])} | " + f"{base.fmt(level['prompt_tokens_per_s_to_first_token'])} | " f"{base.fmt(level['prompt_tokens_min'], '.0f')}–{base.fmt(level['prompt_tokens_max'], '.0f')} | " f"{base.fmt(level['effective_prompt_tokens_min'], '.0f')}–" f"{base.fmt(level['effective_prompt_tokens_max'], '.0f')} | " f"{base.fmt(level['effective_to_wire_prompt_ratio'], '.3f')} | " + f"{base.fmt(level['ttft_median_s'], '.3f')}/" f"{base.fmt(level['ttft_max_s'], '.3f')} |" ) return "\n".join(lines) + "\n" diff --git a/harness/benchmarks/concurrency/record_feature_runtime.py b/harness/benchmarks/concurrency/record_feature_runtime.py index 5665a4961..afb4fea6e 100644 --- a/harness/benchmarks/concurrency/record_feature_runtime.py +++ b/harness/benchmarks/concurrency/record_feature_runtime.py @@ -83,7 +83,7 @@ def update_metadata(metadata: dict[str, Any], log_text: str) -> dict[str, Any]: raise ValueError( "KVFlash metadata is enabled but its physical-pool startup marker is missing" ) - if observed["physical_kv_pool_tokens"] is None: + if not observed["proof_sources"]["paged_pool_startup_marker"]: raise ValueError("paged physical-pool startup marker is missing") result = dict(metadata) diff --git a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh index b6f76fe95..62c1e513e 100755 --- a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh +++ b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh @@ -55,6 +55,20 @@ MODEL_SHA256="$(sha256sum "$MODEL" | awk '{print $1}')" IFS=, read -r -a workload_list <<< "$WORKLOADS" IFS=, read -r -a variant_list <<< "$VARIANTS" IFS=, read -r -a client_list <<< "$CLIENTS" +reject_duplicates() { + local list_name="$1" value + shift + local -A seen=() + for value in "$@"; do + if [[ -n "${seen[$value]+yes}" ]]; then + echo "$list_name contains duplicate entry: $value" >&2 + return 1 + fi + seen["$value"]=1 + done +} +reject_duplicates CLIENTS "${client_list[@]}" || exit 2 +reject_duplicates VARIANTS "${variant_list[@]}" || exit 2 declare -A prompt_offsets=([1]=0 [4]=1 [8]=5 [16]=13) for c in "${client_list[@]}"; do [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } @@ -83,16 +97,55 @@ stop_server() { } trap stop_server EXIT INT TERM +served_model_matches() { + python3 - "$PORT" "$1" <<'PY' +import json +import sys +import urllib.request + +port, expected = sys.argv[1:] +with urllib.request.urlopen( + f"http://127.0.0.1:{port}/v1/models", timeout=2, +) as response: + payload = json.load(response) +matches = any( + isinstance(row, dict) and row.get("id") == expected + for row in payload.get("data", []) +) +raise SystemExit(0 if matches else 1) +PY +} + wait_health() { + local expected_model="$1" local deadline=$((SECONDS + HEALTH_TIMEOUT_SECONDS)) while (( SECONDS < deadline )); do kill -0 "$server_pid" 2>/dev/null || return 1 - curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && return 0 + if curl -fsS --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1 && + served_model_matches "$expected_model" >/dev/null 2>&1; then + return 0 + fi sleep 1 done return 1 } +port_is_available() { + python3 - "$PORT" <<'PY' +import socket +import sys + +port = int(sys.argv[1]) +with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + sock.bind(("127.0.0.1", port)) + except OSError as exc: + print(f"PORT {port} is unavailable: {exc}", file=sys.stderr) + raise SystemExit(1) +PY +} + write_metadata() { local path="$1" variant="$2" workload="$3" clients="$4" repeat="$5" binary="$6" max_prefills="$7" command_file="$8" python3 -c 'import hashlib,json,pathlib,subprocess,sys @@ -132,12 +185,13 @@ run_case() { capacity=$((SLOTS * max_ctx)) local case_dir="$OUT/$workload/c$clients/r$repeat/$variant" mkdir -p "$case_dir" - local -a command + local -a command launch_env if [[ "$variant" == llama ]]; then binary="$LLAMA_SERVER_BIN"; model_id=qwen36-llama; max_prefills=0 command=("$binary" -m "$MODEL" -ngl all --parallel "$SLOTS" -c "$capacity" -b 2048 -ub 512 --cont-batching --no-context-shift --no-mmap -fa on -ctk q4_0 -ctv q4_0 --no-cache-prompt --host 127.0.0.1 --port "$PORT" --alias "$model_id") + launch_env=() else binary="$LUCE_SERVER_BIN"; model_id=qwen36-luce [[ "$variant" == luce-k8 ]] && max_prefills=8 || max_prefills=1 @@ -146,19 +200,27 @@ run_case() { --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 --host 127.0.0.1 --port "$PORT" --model-name "$model_id") + launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" + "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills") + fi + if ((${#launch_env[@]})); then + printf 'env ' > "$case_dir/server-command.txt" + printf '%q ' "${launch_env[@]}" "${command[@]}" >> "$case_dir/server-command.txt" + else + printf '%q ' "${command[@]}" > "$case_dir/server-command.txt" fi - printf '%q ' "${command[@]}" > "$case_dir/server-command.txt"; printf '\n' >> "$case_dir/server-command.txt" + printf '\n' >> "$case_dir/server-command.txt" write_metadata "$case_dir/server-metadata.json" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$case_dir/server-command.txt" echo "[run] $workload C=$clients repeat=$repeat variant=$variant" - if [[ "$variant" == llama ]]; then - "${command[@]}" > "$case_dir/server.log" 2>&1 & + port_is_available || return 1 + if ((${#launch_env[@]})); then + env "${launch_env[@]}" "${command[@]}" > "$case_dir/server.log" 2>&1 & else - env DFLASH_MIN_TOKENS="$WARMUP_TOKENS" DFLASH_MAX_CONCURRENT_PREFILLS="$max_prefills" \ - "${command[@]}" > "$case_dir/server.log" 2>&1 & + "${command[@]}" > "$case_dir/server.log" 2>&1 & fi server_pid=$! - if ! wait_health; then tail -n 80 "$case_dir/server.log" >&2 || true; return 1; fi + if ! wait_health "$model_id"; then tail -n 80 "$case_dir/server.log" >&2 || true; return 1; fi local offset="${prompt_offsets[$clients]}" prompts="$OUT/prompts/$workload.jsonl" python3 "$CLIENT" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" \ diff --git a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh index d0ea1786a..c7a113916 100755 --- a/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh +++ b/harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh @@ -28,6 +28,9 @@ MAX_TOKENS="${MAX_TOKENS:-64}" WARMUP_TOKENS="${WARMUP_TOKENS:-8}" SLOTS="${SLOTS:-16}" MAX_CONCURRENT_PREFILLS="${MAX_CONCURRENT_PREFILLS:-8}" +DRAFT_SWA="${DRAFT_SWA:-2048}" +PREFILL_UBATCH="${PREFILL_UBATCH:-512}" +DDTREE_ADAPTIVE="${DDTREE_ADAPTIVE:-1}" # The requested Strix Halo configuration. Every value is serialized into case # metadata; no performance-affecting DFLASH variable is inherited implicitly. @@ -45,13 +48,14 @@ usage() { cat <<'EOF' Usage: MODEL=/path/Qwen3.6-27B-Q4_K_M.gguf \ - DRAFT_MODEL=/path/dflash-draft-3.6-q4_k_m.gguf \ + DRAFT_MODEL=/path/dflash-draft-3.6-q8_0.gguf \ PREFILL_DRAFTER=/path/Qwen3-0.6B-BF16.gguf \ + DRAFT_SWA=2048 PREFILL_UBATCH=512 \ harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh The default is a bounded C4 smoke matrix with independently selectable ar, ddtree, pflash, kvflash, and full rows. The full row is the requested -Strix Halo configuration: target/draft hip:0, DDTree +Strix Halo configuration: target/draft hip:0, Q8 DFlash SWA=2048, DDTree budget 22, persistent PFlash auto, and KVFlash auto. The long-context profiles are intended to cross the recorded 32K PFlash and 8K KV-residency thresholds; word count is not treated as proof. Per-request wire/log token counts and @@ -59,6 +63,7 @@ activation telemetry fail the case if an "auto" feature did not execute. llama is optional: include it explicitly with VARIANTS=ar,ddtree,llama and set LLAMA_SERVER_BIN. For publication, set CLIENTS=1,4,8,16 and REPEATS=5. +For the AMD blog recipe, use the Q8_0 3.6 drafter, DRAFT_SWA=2048, PREFILL_UBATCH=512, and set DDTREE_ADAPTIVE=0 to keep DDTree active for every eligible step. OUT must not already exist. EOF } @@ -69,13 +74,15 @@ for cmd in python3 curl sha256sum awk; do command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; } done [[ -r "$MODEL" ]] || { echo "set MODEL to a readable target GGUF" >&2; exit 2; } -[[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } [[ "$REPEATS" =~ ^[1-9][0-9]*$ ]] || { echo "REPEATS must be positive" >&2; exit 2; } [[ "$SLOTS" =~ ^[1-9][0-9]*$ ]] || { echo "SLOTS must be positive" >&2; exit 2; } [[ "$MAX_CONCURRENT_PREFILLS" =~ ^[1-9][0-9]*$ ]] || { echo "MAX_CONCURRENT_PREFILLS must be positive" >&2; exit 2; } [[ "$DDTREE_BUDGET" =~ ^[1-9][0-9]*$ ]] || { echo "DDTREE_BUDGET must be positive" >&2; exit 2; } [[ "$PREFILL_THRESHOLD" =~ ^[1-9][0-9]*$ ]] || { echo "PREFILL_THRESHOLD must be positive" >&2; exit 2; } [[ "$KVFLASH_MAX_POOL_TOKENS" =~ ^[1-9][0-9]*$ ]] || { echo "KVFLASH_MAX_POOL_TOKENS must be positive" >&2; exit 2; } +[[ "$DRAFT_SWA" =~ ^[0-9]+$ ]] || { echo "DRAFT_SWA must be a non-negative integer" >&2; exit 2; } +[[ "$PREFILL_UBATCH" =~ ^[1-9][0-9]*$ ]] || { echo "PREFILL_UBATCH must be positive" >&2; exit 2; } +[[ "$DDTREE_ADAPTIVE" == 0 || "$DDTREE_ADAPTIVE" == 1 ]] || { echo "DDTREE_ADAPTIVE must be 0 or 1" >&2; exit 2; } [[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ @@ -108,12 +115,17 @@ for c in "${client_list[@]}"; do [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } (( c <= SLOTS )) || { echo "CLIENTS=$c exceeds SLOTS=$SLOTS" >&2; exit 2; } done +luce_requested=0 for v in "${variant_list[@]}"; do case "$v" in - ar|ddtree|pflash|kvflash|full|llama) ;; + llama) ;; + ar|ddtree|pflash|kvflash|full) luce_requested=1 ;; *) echo "unknown variant $v" >&2; exit 2 ;; esac done +if (( luce_requested )); then + [[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } +fi contains_variant() { local needle="$1" value @@ -232,11 +244,16 @@ run_case() { --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 --host 127.0.0.1 --port "$PORT" --model-name "$model_id") - launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills") + launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills" + "DFLASH27B_DRAFT_SWA=$DRAFT_SWA" "DFLASH27B_PREFILL_UBATCH=$PREFILL_UBATCH") + if [[ "$DDTREE_ADAPTIVE" == 0 ]]; then + launch_env+=("DFLASH_DDTREE_ADAPTIVE=0") + fi if [[ "$variant" == ddtree || "$variant" == full ]]; then command+=(--draft "$DRAFT_MODEL" --draft-device "$DRAFT_DEVICE" - --ddtree --ddtree-budget "$DDTREE_BUDGET") + --ddtree --ddtree-budget "$DDTREE_BUDGET" --fast-rollback) expected+=(--expect ddtree) + if [[ "$variant" == ddtree ]]; then command+=(--draft-residency "$DRAFT_RESIDENCY"); fi fi if [[ "$variant" == pflash || "$variant" == full ]]; then if [[ "$variant" == pflash ]]; then command+=(--draft-device "$DRAFT_DEVICE"); fi @@ -280,7 +297,8 @@ run_case() { fi if [[ "$variant" == ddtree || "$variant" == full ]]; then metadata+=(--draft-device "$DRAFT_DEVICE" --draft-model "$DRAFT_MODEL" - --draft-model-sha256 "$DRAFT_MODEL_SHA256" --ddtree --ddtree-budget "$DDTREE_BUDGET") + --draft-model-sha256 "$DRAFT_MODEL_SHA256" --ddtree --ddtree-budget "$DDTREE_BUDGET" --fast-rollback) + if [[ "$variant" == ddtree ]]; then metadata+=(--draft-residency "$DRAFT_RESIDENCY"); fi fi if [[ "$variant" == pflash || "$variant" == full ]]; then metadata+=(--draft-device "$DRAFT_DEVICE" --prefill-compression "$PREFILL_COMPRESSION" diff --git a/harness/benchmarks/concurrency/summarize_concurrency.py b/harness/benchmarks/concurrency/summarize_concurrency.py index 6f5130f41..0f4374bcb 100755 --- a/harness/benchmarks/concurrency/summarize_concurrency.py +++ b/harness/benchmarks/concurrency/summarize_concurrency.py @@ -43,6 +43,19 @@ def complete_median(values: list[float | None]) -> float | None: return median([value for value in values if value is not None]) +def output_stability(items: list[dict]) -> str: + output_digests = [ + item["level"].get("selected_output_set_sha256") for item in items + ] + complete = all(isinstance(value, str) and bool(value) for value in output_digests) + hashes = {value for value in output_digests if isinstance(value, str)} + return ( + "n/a" if len(items) < 2 or not complete + else "yes" if len(hashes) == 1 + else "NO" + ) + + def run_signature(item: dict) -> tuple[object, ...]: report, meta = item["report"], item["meta"] max_tokens = report.get("max_tokens") @@ -130,20 +143,7 @@ def summarize(reports: list[dict]) -> str: ttft = complete_median([ item["level"].get("ttft_max_s") for item in items ]) - output_digests = [ - item["level"].get("selected_output_set_sha256") for item in items - ] - output_digests_complete = all( - isinstance(value, str) and bool(value) for value in output_digests - ) - output_hashes = { - value for value in output_digests if isinstance(value, str) - } - stable = ( - "n/a" if len(items) < 2 or not output_digests_complete - else "yes" if len(output_hashes) == 1 - else "NO" - ) + stable = output_stability(items) def delta(other: str, metric: str) -> str: peers = grouped.get((workload, clients, other), []) @@ -164,6 +164,8 @@ def delta(other: str, metric: str) -> str: raise ValueError( f"{workload} C={clients}: {variant}/{other} repeat sets differ" ) + if stable == "NO" or output_stability(peers) == "NO": + return "n/a" ratios = [] for repeat in sorted(by_repeat): value = by_repeat[repeat]["level"].get(metric) diff --git a/harness/benchmarks/concurrency/summarize_feature_matrix.py b/harness/benchmarks/concurrency/summarize_feature_matrix.py index b5b29c01e..b789a5fbd 100755 --- a/harness/benchmarks/concurrency/summarize_feature_matrix.py +++ b/harness/benchmarks/concurrency/summarize_feature_matrix.py @@ -75,14 +75,18 @@ def run_signature(item: dict) -> tuple[object, ...]: report, meta = item["report"], item["meta"] max_tokens = report.get("max_tokens") ignore_eos = report.get("ignore_eos") + temperature = report.get("temperature") + seed = report.get("seed") model_sha256 = meta.get("model_sha256") if ( type(max_tokens) is not int or max_tokens <= 0 or not isinstance(ignore_eos, bool) + or type(temperature) not in (int, float) + or type(seed) is not int or not isinstance(model_sha256, str) or not model_sha256 ): raise ValueError("incomplete run metadata") - return max_tokens, ignore_eos, model_sha256 + return max_tokens, ignore_eos, temperature, seed, model_sha256 def output_stability(items: list[dict]) -> str: @@ -117,14 +121,22 @@ def summarize(reports: list[dict]) -> str: "proves its requested features executed. Throughput is the median across fresh-process repeats.", "", "| Workload | C | Variant | N | Output goodput | Output-window | vs AR | " - "Effective/wire | DDTree accepted/step | Target forwards | KV in/out | " - "PFlash requests | TTFT max s | Stable output |", - "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | " - ":--- | ---: | ---: | :---: |", + "Effective/wire | DDTree accepted/step | DDTree steps/susp. | " + "Target forwards | KV in/out | PFlash requests | TTFT max s | " + "Stable output |", + "| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | ---: | " + ":--- | ---: | :--- | ---: | ---: | :---: |", ] for workload, clients, variant in sorted(grouped): items = grouped[(workload, clients, variant)] - prompt_hashes = {item["level"]["selected_prompt_set_sha256"] for item in items} + prompt_digests = [ + item["level"].get("selected_prompt_set_sha256") for item in items + ] + if not all(isinstance(value, str) and value for value in prompt_digests): + raise ValueError( + f"{workload} C={clients} {variant}: missing selected prompt set hash" + ) + prompt_hashes = set(prompt_digests) if len(prompt_hashes) != 1: raise ValueError(f"{workload} C={clients} {variant}: prompt sets differ") goodput = median([item["level"]["aggregate_tok_s"] for item in items]) @@ -174,6 +186,17 @@ def summarize(reports: list[dict]) -> str: steps = sum(a["ddtree_steps"] for a in aggregates) accepted = sum(a["ddtree_accepted_tokens"] for a in aggregates) accepted_per_step = accepted / steps if steps else None + median_steps = median( + [a["ddtree_steps"] for a in aggregates] + ) if aggregates else None + median_suspensions = median( + [a["ddtree_suspensions"] for a in aggregates] + ) if aggregates else None + ddtree_activity = ( + f"{median_steps:.0f}/{median_suspensions:.0f}" + if median_steps is not None and median_suspensions is not None + else "n/a" + ) target_forwards = median([a["target_forwards"] for a in aggregates]) if aggregates else None page_ins = median([a["kvflash_page_ins"] for a in aggregates]) if aggregates else None page_outs = median([a["kvflash_page_outs"] for a in aggregates]) if aggregates else None @@ -185,8 +208,8 @@ def summarize(reports: list[dict]) -> str: lines.append( f"| {workload} | {clients} | {variant} | {len(items)} | {goodput:.2f} | " f"{fmt(window)} | {vs_ar} | {fmt(ratio, 3)} | {fmt(accepted_per_step)} | " - f"{fmt(target_forwards, 0)} | {kv_text} | {fmt(pflash_requests, 0)} | " - f"{fmt(ttft, 3)} | {stable} |" + f"{ddtree_activity} | {fmt(target_forwards, 0)} | {kv_text} | " + f"{fmt(pflash_requests, 0)} | {fmt(ttft, 3)} | {stable} |" ) lines.append("") return "\n".join(lines) diff --git a/harness/benchmarks/concurrency/test_concurrency_tools.py b/harness/benchmarks/concurrency/test_concurrency_tools.py index 1aa8dc2d8..60bc925a9 100644 --- a/harness/benchmarks/concurrency/test_concurrency_tools.py +++ b/harness/benchmarks/concurrency/test_concurrency_tools.py @@ -50,6 +50,23 @@ def test_cohorts_are_disjoint_ragged_and_mean_matched(self) -> None: self.assertEqual(len(row["prompt"].split()), row["target_words"]) +class RunnerShellTests(unittest.TestCase): + def test_runner_guards_case_identity_and_records_launch_environment(self) -> None: + runner = (HERE / "run_qwen36_concurrency.sh").read_text( + encoding="utf-8", + ) + self.assertIn( + 'reject_duplicates CLIENTS "${client_list[@]}"', runner, + ) + self.assertIn( + 'reject_duplicates VARIANTS "${variant_list[@]}"', runner, + ) + self.assertIn('printf \'env \' > "$case_dir/server-command.txt"', runner) + self.assertIn('"${launch_env[@]}" "${command[@]}"', runner) + self.assertIn("port_is_available || return 1", runner) + self.assertIn('wait_health "$model_id"', runner) + + class SummarizerTests(unittest.TestCase): @staticmethod def item( @@ -136,6 +153,37 @@ def test_multiple_repeats_report_output_stability(self) -> None: unstable_row = next(line for line in unstable.splitlines() if "| llama |" in line) self.assertEqual(unstable_row.split("|")[10].strip(), "NO") + def test_unstable_current_variant_suppresses_deltas(self) -> None: + reports = [ + self.item("llama", 10.0, repeat=1), + self.item("llama", 10.0, repeat=2), + self.item("luce-k8", 20.0, repeat=1, output_hash="first"), + self.item("luce-k8", 22.0, repeat=2, output_hash="second"), + ] + text = summarizer.summarize(reports) + row = next(line for line in text.splitlines() if "| luce-k8 |" in line) + self.assertEqual(row.split("|")[10].strip(), "NO") + self.assertEqual(row.split("|")[11].strip(), "n/a") + self.assertEqual(row.split("|")[12].strip(), "n/a") + + def test_unstable_peer_variant_suppresses_deltas(self) -> None: + reports = [ + self.item( + "llama", 10.0, repeat=1, output_hash="first", + ), + self.item( + "llama", 10.0, repeat=2, output_hash="second", + ), + self.item("luce-k8", 20.0, repeat=1), + self.item("luce-k8", 22.0, repeat=2), + ] + text = summarizer.summarize(reports) + row = next(line for line in text.splitlines() if "| luce-k8 |" in line) + self.assertEqual(row.split("|")[10].strip(), "yes") + self.assertEqual(row.split("|")[11].strip(), "n/a") + self.assertEqual(row.split("|")[12].strip(), "n/a") + + def test_missing_output_digest_does_not_claim_stability(self) -> None: reports = [ self.item("llama", 8.0, repeat=1, output_hash=None), diff --git a/harness/benchmarks/concurrency/test_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_concurrent_benchmark.py index c0455edc3..da1251c79 100644 --- a/harness/benchmarks/concurrency/test_concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/test_concurrent_benchmark.py @@ -126,6 +126,22 @@ def test_missing_prompt_usage_fails_level(self) -> None: } self.assertTrue(benchmark.level_failed(level, ignore_eos=True)) + def test_hung_worker_aborts_level_instead_of_overlapping_next(self) -> None: + thread = mock.Mock() + thread.is_alive.return_value = True + args = argparse.Namespace(timeout=1.0) + with ( + mock.patch.object(benchmark.threading, "Thread", return_value=thread), + mock.patch.object( + benchmark.time, "monotonic", side_effect=(10.0, 50.0), + ), + ): + with self.assertRaisesRegex(TimeoutError, "exceeded the level deadline"): + benchmark.run_level(1, args, ["prompt"], 0) + thread.start.assert_called_once_with() + thread.join.assert_called_once_with(0.0) + + if __name__ == "__main__": unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py index d582f6e32..232057dcc 100644 --- a/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py +++ b/harness/benchmarks/concurrency/test_feature_concurrent_benchmark.py @@ -118,6 +118,32 @@ def test_client_provenance_records_exact_argv_and_source_digest(self) -> None: hashlib.sha256(SCRIPT.read_bytes()).hexdigest(), ) + def test_markdown_reports_prompt_rate_and_ttft_median_max(self) -> None: + text = benchmark.markdown({ + "label": "feature", + "levels": [{ + "clients": 4, + "requests_ok": 4, + "requests": 4, + "aggregate_tok_s": 12.5, + "output_window_tok_s": 20.0, + "request_decode_tok_s_median": 5.0, + "prompt_tokens_per_s_to_first_token": 123.4, + "prompt_tokens_min": 100, + "prompt_tokens_max": 400, + "effective_prompt_tokens_min": 50, + "effective_prompt_tokens_max": 200, + "effective_to_wire_prompt_ratio": 0.5, + "ttft_median_s": 1.25, + "ttft_max_s": 2.5, + }], + }) + self.assertIn("Prompt tok/s to first", text) + self.assertIn("TTFT median/max s", text) + self.assertIn("| 4 | 4/4 | 12.50 | 20.00 | 5.00 | 123.40 |", text) + self.assertIn("| 0.500 | 1.250/2.500 |", text) + + if __name__ == "__main__": unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_metadata.py b/harness/benchmarks/concurrency/test_feature_metadata.py index 7349fa21a..229685030 100644 --- a/harness/benchmarks/concurrency/test_feature_metadata.py +++ b/harness/benchmarks/concurrency/test_feature_metadata.py @@ -41,6 +41,8 @@ def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: target.write_bytes(b"target") draft.write_bytes(b"draft") prefill.write_bytes(b"prefill") + draft_sha = hashlib.sha256(draft.read_bytes()).hexdigest() + prefill_sha = hashlib.sha256(prefill.read_bytes()).hexdigest() prompts.write_text('{"prompt":"p"}\n', encoding="utf-8") command.write_text("server --target-device hip:0\n", encoding="utf-8") argv = [ @@ -51,15 +53,15 @@ def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: "--prompt-file", str(prompts), "--command-file", str(command), "--repo", str(HERE.parents[2]), "--max-concurrent-prefills", "8", "--target-device", "hip:0", "--draft-device", "hip:0", - "--draft-model", str(draft), "--draft-model-sha256", "draft-sha", - "--ddtree", "--ddtree-budget", "22", + "--draft-model", str(draft), "--draft-model-sha256", draft_sha, + "--ddtree", "--ddtree-budget", "22", "--fast-rollback", "--prefill-compression", "auto", "--prefill-threshold", "32000", "--prefill-keep-ratio", "0.05", "--prefill-drafter", str(prefill), - "--prefill-drafter-sha256", "prefill-sha", + "--prefill-drafter-sha256", prefill_sha, "--draft-residency", "persistent", "--kvflash", "auto", "--kvflash-max-pool-tokens", "8192", "--kvflash-scorer-drafter", str(prefill), - "--kvflash-scorer-drafter-sha256", "prefill-sha", + "--kvflash-scorer-drafter-sha256", prefill_sha, ] with mock.patch.object(sys, "argv", argv): self.assertEqual(metadata.main(), 0) @@ -67,21 +69,22 @@ def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: self.assertEqual(result["model_sha256"], hashlib.sha256(b"target").hexdigest()) self.assertTrue(result["server_binary_sha256"]) self.assertTrue(result["git_head"]) - self.assertEqual(result["feature_config"]["draft_model_sha256"], "draft-sha") - self.assertEqual(result["feature_config"]["prefill_drafter_sha256"], "prefill-sha") + self.assertEqual(result["feature_config"]["draft_model_sha256"], draft_sha) + self.assertEqual(result["feature_config"]["prefill_drafter_sha256"], prefill_sha) self.assertEqual( result["feature_config"]["kvflash_scorer_drafter"], str(prefill.resolve()), ) self.assertEqual( result["feature_config"]["kvflash_scorer_drafter_sha256"], - "prefill-sha", + prefill_sha, ) self.assertEqual(result["literal_screenshot_flags"], [ "--target-device", "hip:0", "--draft-device", "hip:0", "--ddtree", "--ddtree-budget", "22", + "--fast-rollback", "--draft-residency", "persistent", "--prefill-compression", "auto", "--prefill-drafter", str(prefill.resolve()), @@ -89,6 +92,24 @@ def test_full_row_records_literal_screenshot_flags_and_hashes(self) -> None: ]) self.assertIsNone(result["runtime_observed"]) + def test_model_digest_claim_must_match_referenced_file(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + model = Path(tmp) / "model.gguf" + model.write_bytes(b"model") + expected = hashlib.sha256(b"model").hexdigest() + cache: dict[Path, str] = {} + self.assertEqual( + metadata.validated_digest(model, expected, "model", cache), + expected, + ) + with self.assertRaisesRegex(ValueError, "does not match"): + metadata.validated_digest(model, "stale", "model", cache) + with self.assertRaisesRegex(ValueError, "is required"): + metadata.validated_digest(model, None, "model", cache) + with self.assertRaisesRegex(ValueError, "without a model file"): + metadata.validated_digest(None, expected, "model", cache) + + def test_ldd_failure_is_fatal(self) -> None: failed = mock.Mock(returncode=1, stdout="", stderr="not a dynamic executable") with mock.patch.object(metadata.subprocess, "run", return_value=failed): @@ -148,6 +169,16 @@ def test_runtime_rejects_enabled_kvflash_without_marker(self) -> None: "(8192 pool tokens, per-sequence max_ctx 65536)", ) + def test_runtime_rejects_kvflash_marker_without_paged_marker(self) -> None: + original = {"feature_config": {"kvflash": "auto"}} + log = ( + "[parallel-kvflash] physical resident pool 8192 tokens; " + "logical per-slot cap 65536 across 16 slots " + "(--kv-pool-tokens does not expand resident VRAM)" + ) + with self.assertRaisesRegex(ValueError, "paged physical-pool"): + runtime_metadata.update_metadata(original, log) + if __name__ == "__main__": unittest.main() diff --git a/harness/benchmarks/concurrency/test_feature_tools.py b/harness/benchmarks/concurrency/test_feature_tools.py index 85c2294e6..5e5645f7d 100644 --- a/harness/benchmarks/concurrency/test_feature_tools.py +++ b/harness/benchmarks/concurrency/test_feature_tools.py @@ -179,6 +179,17 @@ def test_duplicate_variants_are_rejected_before_artifacts_are_created(self) -> N self.assertIn("VARIANTS contains duplicate entry: ar", result.stderr) self.assertFalse((Path(tmp) / "out").exists()) + def test_llama_only_does_not_require_lucebox_binary(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = self.run_invalid_matrix( + tmp, VARIANTS="llama", LUCE_SERVER_BIN="/does/not/exist", + LLAMA_SERVER_BIN="/bin/true", REPEATS="0", + ) + self.assertEqual(result.returncode, 2, result.stderr) + self.assertIn("REPEATS must be positive", result.stderr) + self.assertNotIn("missing Lucebox server", result.stderr) + + class FeatureProofTests(unittest.TestCase): def test_full_below_pool_passes_without_page_traffic(self) -> None: @@ -450,7 +461,10 @@ def item( output_hash: str | None = "same-output", ) -> dict: return { - "report": {"max_tokens": 256, "ignore_eos": True}, + "report": { + "max_tokens": 256, "ignore_eos": True, + "temperature": 0.0, "seed": 1, + }, "meta": { "workload": "compression", "variant": variant, "repeat": repeat, "model_sha256": "a" * 64, @@ -528,6 +542,26 @@ def test_summary_rejects_incompatible_ar_control_metadata(self) -> None: with self.assertRaisesRegex(ValueError, "run metadata differs"): summary.summarize([ar, feature]) + def test_summary_rejects_sampling_mismatch(self) -> None: + for field, value in (("temperature", 0.5), ("seed", 2)): + with self.subTest(field=field): + ar = self.item("ar", 10.0) + feature = self.item("full", 12.0) + feature["report"][field] = value + with self.assertRaisesRegex(ValueError, "run metadata differs"): + summary.summarize([ar, feature]) + + def test_summary_rejects_missing_prompt_hash(self) -> None: + for value in (None, ""): + with self.subTest(value=value): + feature = self.item("full", 12.0) + feature["level"]["selected_prompt_set_sha256"] = value + with self.assertRaisesRegex( + ValueError, "missing selected prompt set hash", + ): + summary.summarize([feature]) + + def test_unstable_feature_row_suppresses_ar_delta(self) -> None: reports = [ self.item("ar", 10.0, repeat=1), diff --git a/harness/benchmarks/concurrency/write_feature_metadata.py b/harness/benchmarks/concurrency/write_feature_metadata.py index 9500c41b9..01e40146f 100644 --- a/harness/benchmarks/concurrency/write_feature_metadata.py +++ b/harness/benchmarks/concurrency/write_feature_metadata.py @@ -14,6 +14,27 @@ def digest(path: pathlib.Path) -> str: return hashlib.sha256(path.read_bytes()).hexdigest() +def validated_digest( + path: pathlib.Path | None, + claimed: str | None, + label: str, + cache: dict[pathlib.Path, str], +) -> str | None: + if path is None: + if claimed is not None: + raise ValueError(f"{label} SHA-256 supplied without a model file") + return None + if not isinstance(claimed, str) or not claimed: + raise ValueError(f"{label} SHA-256 is required") + resolved = path.resolve() + if resolved not in cache: + cache[resolved] = digest(resolved) + actual = cache[resolved] + if claimed != actual: + raise ValueError(f"{label} SHA-256 does not match {resolved}") + return actual + + def resolved_libraries(binary: pathlib.Path) -> dict[str, str]: result = subprocess.run( ["ldd", str(binary)], text=True, capture_output=True, @@ -70,6 +91,7 @@ def main() -> int: parser.add_argument("--draft-model", type=pathlib.Path) parser.add_argument("--draft-model-sha256") parser.add_argument("--ddtree", action="store_true") + parser.add_argument("--fast-rollback", action="store_true") parser.add_argument("--ddtree-budget", type=int) parser.add_argument("--prefill-compression", default="off") parser.add_argument("--prefill-threshold", type=int) @@ -91,6 +113,24 @@ def main() -> int: parser.error(f"bad --launch-env {item!r}; expected KEY=VALUE") launch_env[key] = value + digest_cache: dict[pathlib.Path, str] = {} + model_sha256 = validated_digest( + args.model, args.model_sha256, "target model", digest_cache, + ) + draft_model_sha256 = validated_digest( + args.draft_model, args.draft_model_sha256, "draft model", digest_cache, + ) + prefill_drafter_sha256 = validated_digest( + args.prefill_drafter, args.prefill_drafter_sha256, + "prefill drafter", digest_cache, + ) + kvflash_scorer_drafter_sha256 = validated_digest( + args.kvflash_scorer_drafter, + args.kvflash_scorer_drafter_sha256, + "KVFlash scorer drafter", + digest_cache, + ) + libraries = resolved_libraries(args.binary) git_head = repository_head(args.repo) literal_flags: list[str] = [] @@ -102,6 +142,8 @@ def main() -> int: literal_flags += ["--ddtree"] if args.ddtree_budget is not None: literal_flags += ["--ddtree-budget", str(args.ddtree_budget)] + if args.fast_rollback: + literal_flags += ["--fast-rollback"] if args.draft_residency: literal_flags += ["--draft-residency", args.draft_residency] if args.prefill_compression != "off": @@ -121,7 +163,7 @@ def main() -> int: "server_binary": str(args.binary.resolve()), "server_binary_sha256": digest(args.binary), "model": str(args.model.resolve()), - "model_sha256": args.model_sha256, + "model_sha256": model_sha256, "prompt_file_sha256": digest(args.prompt_file), "server_command": args.command_file.read_text(encoding="utf-8").strip(), "launch_environment": launch_env, @@ -134,8 +176,9 @@ def main() -> int: "target_device": args.target_device, "draft_device": args.draft_device, "draft_model": str(args.draft_model.resolve()) if args.draft_model else None, - "draft_model_sha256": args.draft_model_sha256, + "draft_model_sha256": draft_model_sha256, "ddtree": args.ddtree, + "fast_rollback": args.fast_rollback, "ddtree_budget": args.ddtree_budget, "prefill_compression": args.prefill_compression, "prefill_threshold": args.prefill_threshold, @@ -143,7 +186,7 @@ def main() -> int: "prefill_drafter": ( str(args.prefill_drafter.resolve()) if args.prefill_drafter else None ), - "prefill_drafter_sha256": args.prefill_drafter_sha256, + "prefill_drafter_sha256": prefill_drafter_sha256, "draft_residency": args.draft_residency, "kvflash": args.kvflash, "kvflash_max_pool_tokens": args.kvflash_max_pool_tokens, @@ -151,9 +194,7 @@ def main() -> int: str(args.kvflash_scorer_drafter.resolve()) if args.kvflash_scorer_drafter else None ), - "kvflash_scorer_drafter_sha256": ( - args.kvflash_scorer_drafter_sha256 - ), + "kvflash_scorer_drafter_sha256": kvflash_scorer_drafter_sha256, }, } args.out.write_text(json.dumps(obj, indent=2, sort_keys=True) + "\n", encoding="utf-8") diff --git a/server/src/common/concurrency/paged_kv_residency.cpp b/server/src/common/concurrency/paged_kv_residency.cpp index 43c8443b0..c0af73b12 100644 --- a/server/src/common/concurrency/paged_kv_residency.cpp +++ b/server/src/common/concurrency/paged_kv_residency.cpp @@ -243,10 +243,13 @@ PagedKvResidencyStatus PagedKvResidencyManager::prepare_append( } if (token_count == 0) return finish_transfers(PagedKvResidencyStatus::Ok); + BlockState * append_head = nullptr; bool restore_partial_head = false; uint32_t partial_head = 0; if (snapshot.kv_seq_len % pool_.block_size() != 0) { partial_head = snapshot.kv_seq_len / pool_.block_size(); + append_head = &sequences_[handle.slot].blocks[partial_head]; + append_head->reservation_pending = true; restore_partial_head = snapshot.block_table[partial_head] == PAGED_KV_COLD_BLOCK; } @@ -265,15 +268,22 @@ PagedKvResidencyStatus PagedKvResidencyManager::prepare_append( const auto room = make_room( handle, globally_needed + (restore_partial_head ? 1u : 0u), additional_blocks + (restore_partial_head ? 1u : 0u)); - if (room != PagedKvResidencyStatus::Ok) return finish_transfers(room); + if (room != PagedKvResidencyStatus::Ok) { + if (append_head) append_head->reservation_pending = false; + return room; + } if (restore_partial_head) { const auto restored = restore_block_async( handle, partial_head); if (restored != PagedKvResidencyStatus::Ok) { - return finish_transfers(restored); + const auto finished = finish_transfers(restored); + if (append_head) append_head->reservation_pending = false; + return finished; } } - return finish_transfers(PagedKvResidencyStatus::Ok); + const auto finished = finish_transfers(PagedKvResidencyStatus::Ok); + if (append_head) append_head->reservation_pending = false; + return finished; } PagedKvResidentAppendResult PagedKvResidencyManager::append( @@ -305,18 +315,26 @@ PagedKvResidencyStatus PagedKvResidencyManager::observe_append( !state.blocks[remap.logical_block].host_valid) { return finish_transfers(PagedKvResidencyStatus::HostCopyMissing); } + BlockState & block = state.blocks[remap.logical_block]; + // A callback may reject only after queuing a prefix of a multi-tensor + // copy. Quarantine the destination before invoking it either way. + transfers_pending_ = true; + block.page_in_pending = true; bool queued = false; try { queued = transfers_.copy_in_async( handle, remap.logical_block, remap.physical_block, - state.blocks[remap.logical_block].host, config_.block_bytes); + block.host, config_.block_bytes); } catch (...) { queued = false; } - if (!queued) return finish_transfers(PagedKvResidencyStatus::TransferFailed); - transfers_pending_ = true; - stats_.page_ins++; - stats_.moved_bytes += config_.block_bytes; + const PendingTransfer transfer{ + handle, remap.logical_block, remap.physical_block}; + if (!queued) { + pending_page_in_rollbacks_.push_back(transfer); + return finish_transfers(PagedKvResidencyStatus::TransferFailed); + } + pending_page_ins_.push_back(transfer); } const auto synced = finish_transfers(PagedKvResidencyStatus::Ok); if (synced != PagedKvResidencyStatus::Ok) return synced; @@ -424,8 +442,34 @@ PagedKvResidencyStatus PagedKvResidencyManager::synchronize_before_read() { stats_.page_ins++; stats_.moved_bytes += config_.block_bytes; } + for (const PendingTransfer & transfer : pending_page_in_rollbacks_) { + if (validate_registered(transfer.handle) != + PagedKvResidencyStatus::Ok) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + BlockState & block = + sequences_[transfer.handle.slot].blocks[transfer.logical_block]; + PagedKvSequenceSnapshot snapshot; + if (pool_.sequence(transfer.handle, snapshot) != PagedKvStatus::Ok || + transfer.logical_block >= snapshot.block_table.size() || + snapshot.block_table[transfer.logical_block] != + transfer.physical_block) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + uint32_t released = PAGED_KV_COLD_BLOCK; + if (pool_.page_out_block( + transfer.handle, transfer.logical_block, released) != + PagedKvStatus::Ok || released != transfer.physical_block) { + result = PagedKvResidencyStatus::InconsistentPoolState; + continue; + } + block.page_in_pending = false; + } pending_page_outs_.clear(); pending_page_ins_.clear(); + pending_page_in_rollbacks_.clear(); if (result != PagedKvResidencyStatus::Ok) return result; return PagedKvResidencyStatus::Ok; } @@ -630,6 +674,10 @@ PagedKvResidencyStatus PagedKvResidencyManager::evict_block_async( const uint32_t physical = snapshot.block_table[logical_block]; bool queued = false; + // False may mean a multi-tensor callback queued only a prefix. Invalidate + // the old host image and require a barrier before any later operation. + transfers_pending_ = true; + block.host_valid = false; try { queued = transfers_.copy_out_async( handle, logical_block, physical, block.host, @@ -638,8 +686,6 @@ PagedKvResidencyStatus PagedKvResidencyManager::evict_block_async( queued = false; } if (!queued) return PagedKvResidencyStatus::TransferFailed; - transfers_pending_ = true; - block.host_valid = false; block.page_out_pending = true; pending_page_outs_.push_back({handle, logical_block, physical}); return PagedKvResidencyStatus::Ok; @@ -673,6 +719,8 @@ PagedKvResidencyStatus PagedKvResidencyManager::restore_block_async( if (page_status != PagedKvStatus::Ok) return from_pool_status(page_status); bool queued = false; + transfers_pending_ = true; + block.page_in_pending = true; try { queued = transfers_.copy_in_async( handle, logical_block, physical, block.host, @@ -680,17 +728,14 @@ PagedKvResidencyStatus PagedKvResidencyManager::restore_block_async( } catch (...) { queued = false; } + const PendingTransfer transfer{handle, logical_block, physical}; if (!queued) { - uint32_t released = PAGED_KV_COLD_BLOCK; - if (pool_.page_out_block(handle, logical_block, released) != - PagedKvStatus::Ok) { - return PagedKvResidencyStatus::InconsistentPoolState; - } + // The callback may have queued a prefix. Keep the physical mapping + // quarantined until finish_transfers() proves the stream is drained. + pending_page_in_rollbacks_.push_back(transfer); return PagedKvResidencyStatus::TransferFailed; } - transfers_pending_ = true; - block.page_in_pending = true; - pending_page_ins_.push_back({handle, logical_block, physical}); + pending_page_ins_.push_back(transfer); block.last_use = ++clock_; return PagedKvResidencyStatus::Ok; } diff --git a/server/src/common/concurrency/paged_kv_residency.h b/server/src/common/concurrency/paged_kv_residency.h index dc6f67cf5..ac2d982b0 100644 --- a/server/src/common/concurrency/paged_kv_residency.h +++ b/server/src/common/concurrency/paged_kv_residency.h @@ -255,6 +255,7 @@ class PagedKvResidencyManager { bool transfer_barrier_failed_ = false; std::vector pending_page_outs_; std::vector pending_page_ins_; + std::vector pending_page_in_rollbacks_; }; } // namespace dflash::common diff --git a/server/src/common/concurrency/qwen_paged_kv_transfer.cpp b/server/src/common/concurrency/qwen_paged_kv_transfer.cpp index 111e6a36f..d9a2b74cf 100644 --- a/server/src/common/concurrency/qwen_paged_kv_transfer.cpp +++ b/server/src/common/concurrency/qwen_paged_kv_transfer.cpp @@ -212,9 +212,9 @@ struct QwenPagedKvResidencyTransfer::State { } } if (status != cudaSuccess) { - // A callback that returns false is not marked pending by the - // residency manager. Drain any prefix already queued here so - // its pinned buffer can still be released safely. + // Drain any prefix already queued here. The residency manager + // also treats a false return as pending, so a failed drain is + // quarantined and retried before ownership can be released. (void)cudaGetLastError(); if (cudaStreamSynchronize(stream) != cudaSuccess) { stream_may_reference_host = true; diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index fa2909da3..6371e2552 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -202,6 +202,7 @@ class SeqEngine { // until retirement and emits one machine-readable proof record. uint64_t ddtree_steps = 0; uint64_t ddtree_accepted_tokens = 0; + uint64_t ddtree_suspensions = 0; uint64_t target_forwards = 0; uint64_t kvflash_page_ins = 0; uint64_t kvflash_page_outs = 0; @@ -329,6 +330,8 @@ inline std::string validate_step_result( return "failed decode has no diagnostic"; if (output.token >= 0 || !output.committed_tokens.empty()) return "failed decode exposes token payload"; + if (output.ddtree_suspensions != 0) + return "failed decode carries DDTree suspension telemetry"; } else { if (output.token < 0) return "successful decode has no pending token"; @@ -342,6 +345,8 @@ inline std::string validate_step_result( output.committed_tokens.end(), [](int32_t token) { return token < 0; })) return "decode output burst contains an invalid token"; + if (output.ddtree_suspensions > output.ddtree_steps) + return "DDTree suspension has no successful DDTree step"; } decode_seen[(size_t)output.slot] = 1; } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 10e50450a..52949acfc 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -154,6 +154,7 @@ bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { if (!in.allow_speculation || in.slot < 0 || in.slot >= slots_.slot_count() || !slots_.slot(in.slot).decoding() || + !slots_.ddtree_speculation_allowed(in.slot) || slots_.slot(in.slot).sampler.needs_logit_processing() || slots_.slot(in.slot).cur_pos < 1 || slots_.slot(in.slot).cur_pos >= slots_.max_context()) { @@ -234,6 +235,13 @@ std::optional Qwen35SeqEngine::step_ddtree( GGML_STATUS_SUCCESS) { return proposal_fallback(); } + // The draft and target backends own separate HIP streams even when + // both are placed on hip:0. Projection consumes the draft hidden + // state on the target stream, so establish the producer/consumer + // ordering explicitly before the cross-backend tensor copy. Without + // this barrier the first tree proposal can race stale hidden rows and + // collapse acceptance to the one-token fallback. + ggml_backend_synchronize(b_.draft_backend_); ggml_backend_tensor_copy( draft->hidden_states, b_.proj_sg_.hidden_input); if (ggml_backend_graph_compute( @@ -524,19 +532,51 @@ std::optional Qwen35SeqEngine::step_ddtree( return result; } } + + uint64_t cohort_emitted = 0; + for (const Proposal & p : proposals) { + // accepted contains the replay root plus accepted children. The + // output emits those children plus one separately computed pending + // scalar, so accepted.size() is this request's emitted yield. + cohort_emitted += (uint64_t)p.accepted.size(); + } + const bool suspend_cohort = + Qwen35SlotManager::ddtree_cohort_should_suspend( + cohort_emitted, active); + std::vector newly_suspended((size_t)slots_.slot_count(), false); + for (const Proposal & p : proposals) { + newly_suspended[(size_t)p.slot] = + slots_.record_ddtree_sample(p.slot, suspend_cohort); + } + result.decode.reserve((size_t)active); for (Proposal & p : proposals) { DecodeOutput out; out.slot = p.slot; out.token = p.bonus; out.ddtree_steps = 1; - out.ddtree_accepted_tokens = p.accepted.size() - 1; + const int accepted_children = (int)p.accepted.size() - 1; + out.ddtree_accepted_tokens = (uint64_t)accepted_children; out.target_forwards = 2; for (size_t i = 1; i < p.accepted.size(); ++i) { const int dfs = p.accepted[i]; out.committed_tokens.push_back( p.tree.token_ids[(size_t)dfs - 1]); } + if (newly_suspended[(size_t)p.slot]) { + out.ddtree_suspensions = 1; + const Qwen35Slot & seq = slots_.slot(p.slot); + std::fprintf(stderr, + "[parallel-ddtree] adaptive suspend request=%llu slot=%d " + "sample=%llu emitted=%d accepted_children=%d " + "cohort_emitted=%llu cohort_size=%d target_forwards=2 " + "floor=%d\n", + (unsigned long long)seq.request_id, p.slot, + (unsigned long long)seq.ddtree_sampled_steps, + accepted_children + 1, accepted_children, + (unsigned long long)cohort_emitted, active, + Qwen35SlotManager::kDdtreeMinEmittedTokens); + } attach_residency_telemetry(out); result.decode.push_back(std::move(out)); } diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index f740a283b..e0e64170d 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -2,6 +2,7 @@ #include #include +#include namespace dflash::common { @@ -138,9 +139,15 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( return r; } + // Retirement keeps failed copy-stream ownership quarantined. Retry those + // barriers before deciding whether a sequence slot is actually available. + for (int i = 0; i < (int)slots_.size(); ++i) { + if (slots_[(size_t)i].retiring()) retire(i); + } + int slot = -1; for (int i = 0; i < (int)slots_.size(); i++) { - if (!slots_[(size_t)i].active()) { slot = i; break; } + if (slots_[(size_t)i].phase == Qwen35SlotPhase::free) { slot = i; break; } } if (slot < 0) { r.status = AdmitStatus::busy; @@ -196,6 +203,9 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( Qwen35Slot & s = slots_[(size_t)slot]; s.phase = Qwen35SlotPhase::prefill; + s.request_id = request_id; + s.ddtree_suspended = false; + s.ddtree_sampled_steps = 0; s.handle = handle; s.cur_pos = 0; s.prompt_len = prompt_len; @@ -214,6 +224,33 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( return r; } +bool Qwen35SlotManager::ddtree_speculation_allowed(int slot) const { + return is_active(slot) && !slots_[(size_t)slot].ddtree_suspended; +} + +bool Qwen35SlotManager::ddtree_cohort_should_suspend( + uint64_t total_emitted, int active) { + const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); + if (adaptive && std::atoi(adaptive) == 0) { + return false; + } + if (active <= 0) return false; + return total_emitted < + (uint64_t)active * (uint64_t)kDdtreeMinEmittedTokens; +} + +bool Qwen35SlotManager::record_ddtree_sample( + int slot, bool suspend_cohort) { + if (!is_active(slot)) return false; + Qwen35Slot & s = slots_[(size_t)slot]; + // A suspended request must never pay for another probe. + if (s.ddtree_suspended) return false; + ++s.ddtree_sampled_steps; + if (!suspend_cohort) return false; + s.ddtree_suspended = true; + return true; +} + Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( int slot, int n_tokens) { PrefillChunk out; @@ -454,7 +491,7 @@ void Qwen35SlotManager::take_residency_telemetry( void Qwen35SlotManager::retire(int slot) { if (slot < 0 || slot >= (int)slots_.size()) return; Qwen35Slot & s = slots_[(size_t)slot]; - if (!s.active()) return; + if (s.phase == Qwen35SlotPhase::free) return; if (residency_) { const PagedKvResidencyStatus resident_status = residency_->forget_sequence(s.handle); @@ -467,6 +504,7 @@ void Qwen35SlotManager::retire(int slot) { // A failed copy-stream barrier leaves physical pages in flight. // Keep the slot and its pool handle intact so a later retirement // can retry forget_sequence without recycling those pages. + s.phase = Qwen35SlotPhase::retiring; return; } } diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index dd077496c..9a84fb3be 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -32,10 +32,12 @@ enum class Qwen35SlotPhase { free, prefill, decode, + retiring, }; struct Qwen35Slot { Qwen35SlotPhase phase = Qwen35SlotPhase::free; + uint64_t request_id = 0; PagedKvSequenceHandle handle; // Prompt tokens are the immutable prefix of sample_history. Decode tokens // append to the same allocation, avoiding a second full prompt copy. @@ -64,13 +66,28 @@ struct Qwen35Slot { : 0; } - bool active() const { return phase != Qwen35SlotPhase::free; } + // Real packed-tree samples are counted per request. A low aggregate-yield + // cohort sample suspends every participating request; ordinary AR keeps + // every target cache and feature-ring row current. + bool ddtree_suspended = false; + uint64_t ddtree_sampled_steps = 0; + + bool active() const { + return phase == Qwen35SlotPhase::prefill || + phase == Qwen35SlotPhase::decode; + } bool prefilling() const { return phase == Qwen35SlotPhase::prefill; } bool decoding() const { return phase == Qwen35SlotPhase::decode; } + bool retiring() const { return phase == Qwen35SlotPhase::retiring; } }; class Qwen35SlotManager { public: + // Packed DDTree pays for verify + accepted-path replay. Requiring six + // emitted tokens makes continuation earn at least three tokens per target + // forward before accounting for its additional draft/tree work. + static constexpr int kDdtreeMinEmittedTokens = 6; + // `max_ctx` is the per-sequence logical bound; slot count comes from the // pool's max_sequences. The pool must outlive the manager. Qwen35SlotManager(PagedKvPool & pool, int max_ctx, @@ -139,6 +156,13 @@ class Qwen35SlotManager { std::string * error = nullptr); void take_residency_telemetry(int slot, SeqEngine::DecodeOutput & out); + bool ddtree_speculation_allowed(int slot) const; + // Compare aggregate emitted yield (accepted children plus one replay + // bonus per request) against the cohort continuation floor. + static bool ddtree_cohort_should_suspend(uint64_t total_emitted, int active); + // Returns true exactly once, when this sample newly suspends the request. + bool record_ddtree_sample(int slot, bool suspend_cohort); + // One-token compatibility wrapper used by ordinary autoregressive decode. StepAppend append_token(int slot, int32_t fed_token); diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index c640d778f..cb3cf7d0a 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -548,6 +548,13 @@ bool Qwen35Backend::init() { max_concurrent_prefills, mixed_prefill_tokens, long_mixed_prefill_tokens, long_prefill_threshold, idle_prefill_tokens, prefill_quantum); + if (concurrent_local_ddtree) { + const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); + std::fprintf(stderr, + "[parallel-ddtree] enabled budget=%d width=%d mode=packed-verify-replay adaptive=%s\n", + cfg_.ddtree_budget, tree_width, + adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); + } std::printf("[parallel] %d decode slots, up to %d packed prefills " "(mixed short/long %d/%d at >=%d tokens, " "idle %d, quantum %d), " diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index c34c71fb6..045786cb2 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -41,6 +41,7 @@ struct SchedSlot { uint64_t engine_request_id = 0; uint64_t ddtree_steps = 0; uint64_t ddtree_accepted_tokens = 0; + uint64_t ddtree_suspensions = 0; uint64_t target_forwards = 0; uint64_t kvflash_page_ins = 0; uint64_t kvflash_page_outs = 0; @@ -326,6 +327,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { {"pflash_output_tokens", prompt_tokens}, {"ddtree_steps", s.ddtree_steps}, {"ddtree_accepted_tokens", s.ddtree_accepted_tokens}, + {"ddtree_suspensions", s.ddtree_suspensions}, {"target_forwards", s.target_forwards}, {"kvflash_page_ins", s.kvflash_page_ins}, {"kvflash_page_outs", s.kvflash_page_outs}, @@ -752,6 +754,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } s.ddtree_steps += out.ddtree_steps; s.ddtree_accepted_tokens += out.ddtree_accepted_tokens; + s.ddtree_suspensions += out.ddtree_suspensions; s.target_forwards += out.target_forwards; s.kvflash_page_ins += out.kvflash_page_ins; s.kvflash_page_outs += out.kvflash_page_outs; diff --git a/server/test/test_ddtree_path.cpp b/server/test/test_ddtree_path.cpp index e7d2bbf3a..a4118ce1e 100644 --- a/server/test/test_ddtree_path.cpp +++ b/server/test/test_ddtree_path.cpp @@ -1,6 +1,6 @@ #include "common/ddtree.h" +#include "host_check.h" -#include #include #include #include @@ -9,6 +9,8 @@ using dflash::common::DDTree; using dflash::common::follow_verified_tree; using dflash::common::truncate_verified_path; +static int g_checks = 0; + int main() { DDTree tree; tree.n_nodes = 2; @@ -23,23 +25,23 @@ int main() { int pending = -1; std::vector accepted = follow_verified_tree(tree, posterior, pending); - assert((accepted == std::vector{0, 1, 2})); - assert(pending == 33); + CHECK((accepted == std::vector{0, 1, 2})); + CHECK(pending == 33); // Truncating after node 1 means node 2's token becomes pending. Keeping // the old value (33) would skip token 22 and describe uncommitted state. - assert(truncate_verified_path(accepted, 2, posterior, pending)); - assert((accepted == std::vector{0, 1})); - assert(pending == 22); + CHECK(truncate_verified_path(accepted, 2, posterior, pending)); + CHECK((accepted == std::vector{0, 1})); + CHECK(pending == 22); // An unchanged path preserves the already-computed pending token. - assert(!truncate_verified_path(accepted, 2, posterior, pending)); - assert(pending == 22); + CHECK(!truncate_verified_path(accepted, 2, posterior, pending)); + CHECK(pending == 22); // No headroom is represented explicitly and never dereferences a tip. - assert(truncate_verified_path(accepted, 0, posterior, pending)); - assert(accepted.empty()); - assert(pending == -1); + CHECK(truncate_verified_path(accepted, 0, posterior, pending)); + CHECK(accepted.empty()); + CHECK(pending == -1); std::puts("ddtree path tests passed"); return 0; diff --git a/server/test/test_paged_kv_residency.cpp b/server/test/test_paged_kv_residency.cpp index d39a5a0a8..274c7f1e2 100644 --- a/server/test/test_paged_kv_residency.cpp +++ b/server/test/test_paged_kv_residency.cpp @@ -60,7 +60,7 @@ struct MockTransfers { pending.push_back({[this, physical, host, bytes] { std::memcpy(host, &device[(size_t)physical * block_bytes], bytes); }}); - return true; + return !fail_copy_out_after_queue; }, [this](PagedKvSequenceHandle, uint32_t, uint32_t physical, const void * host, size_t bytes) { @@ -68,7 +68,7 @@ struct MockTransfers { pending.push_back({[this, physical, host, bytes] { std::memcpy(&device[(size_t)physical * block_bytes], host, bytes); }}); - return true; + return !fail_copy_in_after_queue; }, [this] { syncs++; @@ -99,6 +99,8 @@ struct MockTransfers { bool fail_alloc = false; bool fail_copy_out = false; bool fail_copy_in = false; + bool fail_copy_out_after_queue = false; + bool fail_copy_in_after_queue = false; bool fail_sync = false; }; @@ -253,6 +255,79 @@ TEST_CASE(PagedKvResidencyFixture, allocation_and_copy_failures_are_explicit) { CHECK(pager.is_resident(handle, 0)); } +TEST_CASE(PagedKvResidencyFixture, + failed_copy_out_prefix_is_barriered_before_retry) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + + io.fail_copy_out_after_queue = true; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + CHECK(io.pending.empty()); + CHECK(pager.is_resident(handle, 0)); + CHECK(pager.stats().page_outs == 0); + + io.fail_copy_out_after_queue = false; + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + CHECK(!pager.is_resident(handle, 0)); +} + +TEST_CASE(PagedKvResidencyFixture, + rejected_append_remap_is_rolled_back_after_barrier) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 2)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + + io.fail_copy_in = true; + const auto append = pool.append(handle, 1); + CHECK(append.status == PagedKvStatus::Ok); + CHECK(append.remapped_cold_blocks.size() == 1); + CHECK(pager.observe_append(handle, append) == + PagedKvResidencyStatus::TransferFailed); + CHECK(snapshot(pool, handle).block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 1); + CHECK(pager.stats().page_ins == 0); +} + +TEST_CASE(PagedKvResidencyFixture, + failed_copy_in_prefix_quarantines_mapping_until_barrier) { + PagedKvPool pool(/*physical blocks=*/1, /*sequences=*/1, /*block size=*/4); + MockTransfers io(1, 16); + PagedKvResidencyManager pager(pool, config(16, 1), io.callbacks()); + const auto handle = acquire(pool, 1); + CHECK(pager.register_sequence(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(handle, 4)); + CHECK(pager.commit_pending_writes(handle) == PagedKvResidencyStatus::Ok); + CHECK(pager.evict_block(handle, 0, /*allow_protected=*/true) == + PagedKvResidencyStatus::Ok); + + io.fail_copy_in_after_queue = true; + io.fail_sync = true; + CHECK(pager.ensure_resident(handle, {0}) == + PagedKvResidencyStatus::TransferFailed); + CHECK(snapshot(pool, handle).block_table[0] != PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 0); + CHECK(io.pending.size() == 1); + + io.fail_sync = false; + CHECK(pager.synchronize_before_read() == PagedKvResidencyStatus::Ok); + CHECK(snapshot(pool, handle).block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(pool.free_block_count() == 1); + CHECK(pager.stats().page_ins == 0); +} + TEST_CASE(PagedKvResidencyFixture, forget_frees_host_backing_and_stale_is_rejected) { PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/1, /*block size=*/4); MockTransfers io(2, 16); @@ -534,6 +609,32 @@ TEST_CASE(PagedKvResidencyFixture, CHECK(snapshot(pool, second).block_table[0] == PAGED_KV_COLD_BLOCK); } +TEST_CASE(PagedKvResidencyFixture, + resident_partial_head_stays_within_budget_during_growth) { + PagedKvPool pool(/*physical blocks=*/2, /*sequences=*/2, /*block size=*/4); + MockTransfers io(2, 16); + PagedKvResidencyManager pager( + pool, config(16, /*budget=*/2, /*sink=*/0, /*tail=*/0), + io.callbacks()); + const auto first = acquire(pool, 1); + const auto second = acquire(pool, 2); + CHECK(pager.register_sequence(first) == PagedKvResidencyStatus::Ok); + CHECK(pager.register_sequence(second) == PagedKvResidencyStatus::Ok); + CHECK(pager.append(first, 2)); + CHECK(pager.commit_pending_writes(first) == PagedKvResidencyStatus::Ok); + const uint32_t append_head = snapshot(pool, first).block_table[0]; + CHECK(pager.append(second, 4)); + CHECK(pager.commit_pending_writes(second) == PagedKvResidencyStatus::Ok); + + const auto grown = pager.append(first, 3); + CHECK(grown); + const auto table = snapshot(pool, first).block_table; + CHECK(table.size() == 2); + CHECK(table[0] == append_head); + CHECK(table[1] != PAGED_KV_COLD_BLOCK); + CHECK(snapshot(pool, second).block_table[0] == PAGED_KV_COLD_BLOCK); + CHECK(pager.stats().resident_blocks == 2); +} TEST_CASE(PagedKvResidencyFixture, requested_restore_set_is_protected_as_one_batch) { PagedKvPool pool(/*physical blocks=*/3, /*sequences=*/2, /*block size=*/4); diff --git a/server/test/test_seq_batch_plan.cpp b/server/test/test_seq_batch_plan.cpp index a771d160b..e1e13065e 100644 --- a/server/test/test_seq_batch_plan.cpp +++ b/server/test/test_seq_batch_plan.cpp @@ -131,9 +131,22 @@ int main() { burst.decode[0].committed_tokens = {8, 9, 10}; burst.decode[0].ddtree_steps = 1; burst.decode[0].ddtree_accepted_tokens = 3; + burst.decode[0].ddtree_suspensions = 1; burst.decode[0].target_forwards = 1; CHECK(validate_step_result(work, burst, 2).empty()); + SeqEngine::StepResult orphan_suspension = good; + orphan_suspension.decode[0].ddtree_suspensions = 1; + CHECK(!validate_step_result(work, orphan_suspension, 2).empty()); + + SeqEngine::StepResult failed_suspension = good; + failed_suspension.decode[0].failed = true; + failed_suspension.decode[0].token = -1; + failed_suspension.decode[0].error = "decode failed"; + failed_suspension.decode[0].ddtree_steps = 1; + failed_suspension.decode[0].ddtree_suspensions = 1; + CHECK(!validate_step_result(work, failed_suspension, 2).empty()); + // A scheduler stop in the committed prefix must hide the remaining burst // and final pending token. The backend state is discarded at retirement. std::vector delivered; diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index 46a5287d7..7988733b2 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -7,10 +7,13 @@ #include "qwen35/concurrency/qwen35_slot_manager.h" #include "host_check.h" +#include "scoped_env.h" #include #include #include +#include +#include using namespace dflash::common; @@ -489,6 +492,117 @@ int main() { CHECK(!mgr.has_prefill_prompt_at_least(768)); } + // Adaptive DDTree judges aggregate cohort yield rather than allowing one + // unlucky request to suppress higher-yield peers. Equality at the + // six-token average continues; a low average suspends all participants. + { + const luce_test::ScopedEnvVar adaptive("DFLASH_DDTREE_ADAPTIVE", nullptr); + CHECK(Qwen35SlotManager::ddtree_cohort_should_suspend(5, 1)); + CHECK(!Qwen35SlotManager::ddtree_cohort_should_suspend(6, 1)); + // Mixed synthetic cohort: one child-less row plus an eleven-token row + // exactly meets the continuation floor; one fewer token does not. + CHECK(!Qwen35SlotManager::ddtree_cohort_should_suspend(12, 2)); + CHECK(Qwen35SlotManager::ddtree_cohort_should_suspend(11, 2)); + + PagedKvPool pool(8, 2, /*block_size=*/16); + Qwen35SlotManager mgr(pool, 64); + auto a = admit(mgr, 101, prompt_tokens(4), greedy_sampler()); + auto b = admit(mgr, 102, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(a)); + CHECK(is_admitted(b)); + CHECK(mgr.append_prefill(a.slot, 4).ok); + CHECK(mgr.append_prefill(b.slot, 4).ok); + mgr.commit_prefill(a.slot); + mgr.commit_prefill(b.slot); + CHECK(mgr.slot(a.slot).request_id == 101); + CHECK(mgr.slot(b.slot).request_id == 102); + CHECK(mgr.ddtree_speculation_allowed(a.slot)); + CHECK(mgr.ddtree_speculation_allowed(b.slot)); + + // The keep decision records a real sample without latching either + // participant, including the low-yield member. + CHECK(!mgr.record_ddtree_sample(a.slot, false)); + CHECK(!mgr.record_ddtree_sample(b.slot, false)); + CHECK(mgr.ddtree_speculation_allowed(a.slot)); + CHECK(mgr.ddtree_speculation_allowed(b.slot)); + CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 1); + CHECK(mgr.slot(b.slot).ddtree_sampled_steps == 1); + + // A low cohort decision is applied identically and atomically to both. + CHECK(mgr.record_ddtree_sample(a.slot, true)); + CHECK(mgr.record_ddtree_sample(b.slot, true)); + CHECK(!mgr.ddtree_speculation_allowed(a.slot)); + CHECK(!mgr.ddtree_speculation_allowed(b.slot)); + CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 2); + CHECK(mgr.slot(b.slot).ddtree_sampled_steps == 2); + + // Suspended requests cannot pay for another bad probe. + CHECK(!mgr.record_ddtree_sample(a.slot, true)); + CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 2); + + mgr.retire(a.slot); + auto reused = admit(mgr, 103, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(reused) && reused.slot == a.slot); + CHECK(mgr.slot(reused.slot).request_id == 103); + CHECK(mgr.ddtree_speculation_allowed(reused.slot)); + CHECK(mgr.slot(reused.slot).ddtree_sampled_steps == 0); + } + + // A failed residency barrier quarantines retirement ownership, and a + // later admission retries it before considering the slot reusable. + { + PagedKvPool pool(/*physical_block_count=*/1, + /*max_sequences=*/1, /*block_size=*/4); + bool fail_sync = false; + PagedKvResidencyTransferOps transfers{ + [](size_t bytes) -> void * { return new uint8_t[bytes]; }, + [](void * ptr) { delete[] static_cast(ptr); }, + [](PagedKvSequenceHandle, uint32_t, uint32_t, + void *, size_t) { return true; }, + [](PagedKvSequenceHandle, uint32_t, uint32_t, + const void *, size_t) { return true; }, + [&fail_sync] { return !fail_sync; }, + }; + PagedKvResidencyConfig config; + config.block_bytes = 16; + config.resident_budget_blocks = 1; + config.sink_blocks = 0; + config.tail_blocks = 0; + PagedKvResidencyManager residency(pool, config, std::move(transfers)); + Qwen35SlotManager mgr(pool, /*max_ctx=*/16, + /*speculative_headroom=*/1, &residency); + + auto first = admit(mgr, 301, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(first)); + CHECK(mgr.append_prefill(first.slot, 4).ok); + CHECK(residency.commit_pending_writes(mgr.slot(first.slot).handle) == + PagedKvResidencyStatus::Ok); + mgr.commit_prefill(first.slot); + + fail_sync = true; + CHECK(residency.evict_block( + mgr.slot(first.slot).handle, 0, + /*allow_protected=*/true) == + PagedKvResidencyStatus::TransferFailed); + mgr.retire(first.slot); + CHECK(!mgr.is_active(first.slot)); + CHECK(mgr.slot(first.slot).retiring()); + CHECK(pool.active_sequence_count() == 1); + + auto blocked = admit(mgr, 302, prompt_tokens(4), greedy_sampler()); + CHECK(!is_admitted(blocked)); + CHECK(is_busy(blocked)); + CHECK(mgr.slot(first.slot).retiring()); + CHECK(pool.active_sequence_count() == 1); + + fail_sync = false; + auto recovered = admit(mgr, 303, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(recovered)); + CHECK(recovered.slot == first.slot); + CHECK(mgr.is_prefilling(recovered.slot)); + CHECK(pool.active_sequence_count() == 1); + } + std::printf("OK test_seq_slot_manager (%d checks)\n", g_checks); return 0; } From ea5b737545ae310d11b871a53de976b8f51520fc Mon Sep 17 00:00:00 2001 From: Graffioh Date: Fri, 14 Aug 2026 11:16:14 +0000 Subject: [PATCH 05/18] perf(qwen35): adapt packed prefill budgeting --- .../concurrency/generate_ragged_prompts.py | 1 + .../concurrency/run_qwen36_concurrency.sh | 38 ++++++++++++++----- .../concurrency/qwen35_slot_manager.cpp | 9 +++++ server/src/qwen35/graph_builders.cpp | 6 +++ server/src/qwen35/qwen35_backend.cpp | 2 +- server/test/test_seq_slot_manager.cpp | 17 +++++++++ 6 files changed, 62 insertions(+), 11 deletions(-) diff --git a/harness/benchmarks/concurrency/generate_ragged_prompts.py b/harness/benchmarks/concurrency/generate_ragged_prompts.py index ca0b180a7..c04d4c434 100755 --- a/harness/benchmarks/concurrency/generate_ragged_prompts.py +++ b/harness/benchmarks/concurrency/generate_ragged_prompts.py @@ -10,6 +10,7 @@ PROFILES = { + "tiny": (64, 96, 128, 160), "short": (250, 350, 450, 550), "medium": (650, 850, 1150, 1350), "long": (2000, 2600, 3400, 4000), diff --git a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh index 62c1e513e..5583bacf4 100755 --- a/harness/benchmarks/concurrency/run_qwen36_concurrency.sh +++ b/harness/benchmarks/concurrency/run_qwen36_concurrency.sh @@ -40,7 +40,9 @@ if [[ $# -ne 0 ]]; then usage >&2; exit 2; fi for cmd in python3 curl sha256sum; do command -v "$cmd" >/dev/null || { echo "missing $cmd" >&2; exit 2; }; done [[ -r "$MODEL" ]] || { echo "set MODEL to a readable GGUF" >&2; exit 2; } [[ -x "$LUCE_SERVER_BIN" ]] || { echo "missing Lucebox server: $LUCE_SERVER_BIN" >&2; exit 2; } -[[ -x "$LLAMA_SERVER_BIN" ]] || { echo "missing llama.cpp server: $LLAMA_SERVER_BIN" >&2; exit 2; } +if [[ ",$VARIANTS," == *,llama,* ]]; then + [[ -x "$LLAMA_SERVER_BIN" ]] || { echo "missing llama.cpp server: $LLAMA_SERVER_BIN" >&2; exit 2; } +fi [[ "$REPEATS" =~ ^[1-9][0-9]*$ ]] || { echo "REPEATS must be positive" >&2; exit 2; } [[ ! -e "$OUT" ]] || { echo "refusing to overwrite $OUT" >&2; exit 2; } ambient_tuning="$(env | grep -E '^(GGML_|DFLASH_|LUCE_|HIP_|ROCR_|HSA_|LD_PRELOAD=|LD_LIBRARY_PATH=)' \ @@ -74,7 +76,9 @@ for c in "${client_list[@]}"; do [[ -n "${prompt_offsets[$c]+yes}" ]] || { echo "supported CLIENTS are 1,4,8,16" >&2; exit 2; } done for v in "${variant_list[@]}"; do - [[ "$v" == luce-k8 || "$v" == luce-k1 || "$v" == llama ]] || { echo "unknown variant $v" >&2; exit 2; } + [[ "$v" == luce-k8 || "$v" == luce-k1 || "$v" == luce-k16-b2 || + "$v" == luce-k16-b4 || "$v" == luce-k16-dyn || + "$v" == llama ]] || { echo "unknown variant $v" >&2; exit 2; } done mkdir -p "$OUT/prompts" @@ -95,7 +99,8 @@ stop_server() { fi server_pid="" } -trap stop_server EXIT INT TERM +trap stop_server EXIT +trap 'exit 130' INT TERM served_model_matches() { python3 - "$PORT" "$1" <<'PY' @@ -147,9 +152,9 @@ PY } write_metadata() { - local path="$1" variant="$2" workload="$3" clients="$4" repeat="$5" binary="$6" max_prefills="$7" command_file="$8" + local path="$1" variant="$2" workload="$3" clients="$4" repeat="$5" binary="$6" max_prefills="$7" mixed_budget="$8" idle_budget="$9" quantum="${10}" command_file="${11}" python3 -c 'import hashlib,json,pathlib,subprocess,sys -p,variant,workload,clients,repeat,binary,max_prefills,cmd_file,model_sha,prompts,repo=sys.argv[1:] +p,variant,workload,clients,repeat,binary,max_prefills,mixed,idle,quantum,cmd_file,model_sha,prompts,repo=sys.argv[1:] digest=lambda x: hashlib.sha256(pathlib.Path(x).read_bytes()).hexdigest() libs={} for line in subprocess.run(["ldd",binary],text=True,capture_output=True).stdout.splitlines(): @@ -165,18 +170,20 @@ if variant == "llama": raise RuntimeError(f"cannot identify llama.cpp source version from {binary} --version") obj={"variant":variant,"workload":workload,"clients":int(clients),"repeat":int(repeat), "max_concurrent_prefills":int(max_prefills),"server_binary":str(pathlib.Path(binary).resolve()), + "mixed_prefill_tokens":int(mixed),"idle_prefill_tokens":int(idle), + "prefill_allocation_quantum":int(quantum), "server_binary_sha256":digest(binary),"model_sha256":model_sha, "prompt_file_sha256":digest(prompts),"server_command":pathlib.Path(cmd_file).read_text().strip(), "resolved_shared_library_sha256":libs, "lucebox_git_head":lucebox_git_head if variant != "llama" else None, "server_version":server_version} pathlib.Path(p).write_text(json.dumps(obj,indent=2,sort_keys=True)+"\n")' \ - "$path" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$command_file" "$MODEL_SHA256" "$OUT/prompts/$workload.jsonl" "$REPO" + "$path" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$mixed_budget" "$idle_budget" "$quantum" "$command_file" "$MODEL_SHA256" "$OUT/prompts/$workload.jsonl" "$REPO" } run_case() { local repeat="$1" workload="$2" clients="$3" variant="$4" - local max_ctx timeout capacity max_prefills binary model_id + local max_ctx timeout capacity max_prefills mixed_budget idle_budget quantum binary model_id if [[ "$workload" == long ]]; then max_ctx=8192; timeout=1800 else @@ -188,20 +195,30 @@ run_case() { local -a command launch_env if [[ "$variant" == llama ]]; then binary="$LLAMA_SERVER_BIN"; model_id=qwen36-llama; max_prefills=0 + mixed_budget=0; idle_budget=0; quantum=0 command=("$binary" -m "$MODEL" -ngl all --parallel "$SLOTS" -c "$capacity" -b 2048 -ub 512 --cont-batching --no-context-shift --no-mmap -fa on -ctk q4_0 -ctv q4_0 --no-cache-prompt --host 127.0.0.1 --port "$PORT" --alias "$model_id") launch_env=() else binary="$LUCE_SERVER_BIN"; model_id=qwen36-luce - [[ "$variant" == luce-k8 ]] && max_prefills=8 || max_prefills=1 + case "$variant" in + luce-k8) max_prefills=8; mixed_budget=2048; idle_budget=4096; quantum=512 ;; + luce-k16-b2) max_prefills=16; mixed_budget=2048; idle_budget=2048; quantum=128 ;; + luce-k16-b4) max_prefills=16; mixed_budget=4096; idle_budget=4096; quantum=256 ;; + luce-k16-dyn) max_prefills=16; mixed_budget=2048; idle_budget=4096; quantum=256 ;; + *) max_prefills=1; mixed_budget=2048; idle_budget=4096; quantum=512 ;; + esac command=("$binary" "$MODEL" --target-device hip:0 --paged-attention --max-concurrency "$SLOTS" --kv-pool-tokens "$capacity" --max-ctx "$max_ctx" --cache-type-k q4_0 --cache-type-v q4_0 --fa-window 0 --prefix-cache-slots 0 --prefill-cache-slots 0 --admission-coalesce-ms 5 --host 127.0.0.1 --port "$PORT" --model-name "$model_id") launch_env=("DFLASH_MIN_TOKENS=$WARMUP_TOKENS" - "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills") + "DFLASH_MAX_CONCURRENT_PREFILLS=$max_prefills" + "DFLASH_MIXED_PREFILL_TOKENS=$mixed_budget" + "DFLASH_IDLE_PREFILL_TOKENS=$idle_budget" + "DFLASH_PREFILL_ALLOCATION_QUANTUM=$quantum") fi if ((${#launch_env[@]})); then printf 'env ' > "$case_dir/server-command.txt" @@ -210,7 +227,7 @@ run_case() { printf '%q ' "${command[@]}" > "$case_dir/server-command.txt" fi printf '\n' >> "$case_dir/server-command.txt" - write_metadata "$case_dir/server-metadata.json" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$case_dir/server-command.txt" + write_metadata "$case_dir/server-metadata.json" "$variant" "$workload" "$clients" "$repeat" "$binary" "$max_prefills" "$mixed_budget" "$idle_budget" "$quantum" "$case_dir/server-command.txt" echo "[run] $workload C=$clients repeat=$repeat variant=$variant" port_is_available || return 1 @@ -228,6 +245,7 @@ run_case() { --require-distinct-prompts --max-tokens "$WARMUP_TOKENS" --temperature 0 \ --ignore-eos --timeout "$timeout" --cooldown 0 --out "$case_dir/warmup.json" \ --label "$variant $workload C=$clients warmup" > "$case_dir/warmup.txt" + sleep 1 python3 "$CLIENT" --base-url "http://127.0.0.1:$PORT/v1" --model "$model_id" \ --clients "$clients" --prompt-file "$prompts" --prompt-offset "$offset" \ --require-distinct-prompts --max-tokens "$MAX_TOKENS" --temperature 0 \ diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index e0e64170d..e6ef4ac39 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -88,6 +88,15 @@ bool Qwen35SlotManager::is_active(int slot) const { bool Qwen35SlotManager::is_prefilling(int slot) const { return is_active(slot) && slots_[(size_t)slot].prefilling(); } + +bool Qwen35SlotManager::has_prefill_prompt_at_least(int tokens) const { + if (tokens <= 0) return true; + return std::any_of(slots_.begin(), slots_.end(), + [tokens](const Qwen35Slot & slot) { + return slot.prefilling() && slot.prompt_len >= tokens; + }); +} + void Qwen35SlotManager::accumulate_residency_delta( Qwen35Slot & slot, const PagedKvResidencyStats & before) { if (!residency_) return; diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index f96349af3..1becd302c 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -417,6 +417,12 @@ bool build_target_step( n_prefill_segments, graph_capacity)) { return false; } + // Experimental adaptive prefill: node count depends on aggregate width, + // ragged layout, and fused decode rows. Use the already-supported maximum + // capacity so the benchmark does not rely on an incomplete shape proxy. + if (n_prefill_tokens > 0) { + graph_capacity = std::max(graph_capacity, 131072); + } // Persistent thread_local arena: rebuilt step graphs land at identical // addresses, keeping the ggml-cuda CUDA-graph cache key (nodes[0]) and diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index cb3cf7d0a..efc7e617b 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -390,7 +390,7 @@ bool Qwen35Backend::init() { const int max_concurrent_prefills = n_slots > 1 ? std::clamp( env_int_or_default("DFLASH_MAX_CONCURRENT_PREFILLS", 8), - 1, std::min(n_slots, 8)) + 1, n_slots) : 1; const int mixed_prefill_tokens = std::max( 1, env_int_or_default("DFLASH_MIXED_PREFILL_TOKENS", 2048)); diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index 7988733b2..7a2d7a26c 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -603,6 +603,23 @@ int main() { CHECK(pool.active_sequence_count() == 1); } + // Long-prefill policy follows active request length and clears on retire. + { + PagedKvPool pool(128, 2, /*block_size=*/16); + Qwen35SlotManager mgr(pool, 2048); + CHECK(!mgr.has_prefill_prompt_at_least(768)); + auto short_req = + admit(mgr, 201, prompt_tokens(512), greedy_sampler()); + CHECK(is_admitted(short_req)); + CHECK(!mgr.has_prefill_prompt_at_least(768)); + auto long_req = + admit(mgr, 202, prompt_tokens(800), greedy_sampler()); + CHECK(is_admitted(long_req)); + CHECK(mgr.has_prefill_prompt_at_least(768)); + mgr.retire(long_req.slot); + CHECK(!mgr.has_prefill_prompt_at_least(768)); + } + std::printf("OK test_seq_slot_manager (%d checks)\n", g_checks); return 0; } From e3cee3222d14ea04ba8ebf002c7853f7ddde1a33 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 06:20:09 +0000 Subject: [PATCH 06/18] fix(qwen35): deduplicate prefill policy helper --- server/src/qwen35/concurrency/qwen35_slot_manager.cpp | 8 -------- server/test/test_recurrent_snapshot.cpp | 1 + 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index e6ef4ac39..e559a224e 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -89,14 +89,6 @@ bool Qwen35SlotManager::is_prefilling(int slot) const { return is_active(slot) && slots_[(size_t)slot].prefilling(); } -bool Qwen35SlotManager::has_prefill_prompt_at_least(int tokens) const { - if (tokens <= 0) return true; - return std::any_of(slots_.begin(), slots_.end(), - [tokens](const Qwen35Slot & slot) { - return slot.prefilling() && slot.prompt_len >= tokens; - }); -} - void Qwen35SlotManager::accumulate_residency_delta( Qwen35Slot & slot, const PagedKvResidencyStats & before) { if (!residency_) return; diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index b0bb0475b..41b38aea8 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -158,6 +158,7 @@ TEST_CASE(RecurrentSnapshotFixture, snapshot_and_restore_recurrent_state) { // C16 x width-23 selects a 32K graph. Prove that graph traversal and // gallocr can cross the old 16K hard ceiling without asserting. { + size_t graph_capacity = 0; CHECK(dflash::common::detail::target_paged_tree_graph_capacity( 23, 16, graph_capacity)); ggml_init_params graph_params{}; From 687f566ebe75598af864288efa603410e3e84401 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 17 Aug 2026 07:51:03 +0000 Subject: [PATCH 07/18] feat(concurrency): add speculation goodput controller --- server/CMakeLists.txt | 9 + .../concurrency/adaptive_verification.h | 264 ++++++++++++++++++ .../common/concurrency/speculation_goodput.h | 206 ++++++++++++++ server/test/test_speculation_goodput.cpp | 241 ++++++++++++++++ 4 files changed, 720 insertions(+) create mode 100644 server/src/common/concurrency/adaptive_verification.h create mode 100644 server/src/common/concurrency/speculation_goodput.h create mode 100644 server/test/test_speculation_goodput.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index ad6ef874f..eb32e698f 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1425,6 +1425,15 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/src) list(APPEND _raw_unit_test_targets test_seq_slot_manager) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_speculation_goodput.cpp") + # Model-neutral online AR/speculation policy: no model or GPU. + add_executable(test_speculation_goodput + test/test_speculation_goodput.cpp) + target_include_directories(test_speculation_goodput PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_speculation_goodput) + 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) diff --git a/server/src/common/concurrency/adaptive_verification.h b/server/src/common/concurrency/adaptive_verification.h new file mode 100644 index 000000000..386792cce --- /dev/null +++ b/server/src/common/concurrency/adaptive_verification.h @@ -0,0 +1,264 @@ +#pragma once + +// Model-neutral, hardware-aware request admission for adaptive verification. +// +// DSpark's scheduler ranks candidate tokens by cumulative survival probability +// and grows the verification batch while expected throughput improves. This +// helper applies the same policy at request granularity. The concrete +// speculator supplies expected useful tokens; the engine supplies observed AR +// and speculative subbatch costs. DDTree can learn value from accepted paths, +// while DSpark can use its calibrated confidence head directly. + +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct AdaptiveVerificationCandidate { + int request = -1; + // Includes the one token that ordinary AR would emit. + double expected_tokens = 1.0; + double maximum_tokens = 1.0; + bool calibrated = false; +}; + +struct AdaptiveVerificationDecision { + std::vector requests; + // The adapter may cheaply calibrate one request (for DDTree, from draft + // entropy; for DSpark, from its confidence head) before applying requests. + int calibration_request = -1; + bool exploring = false; + double predicted_gain = 1.0; +}; + +struct AdaptiveVerificationConfig { + // Preserve margin for timing noise and route-switch overhead. + double minimum_gain = 1.05; + double cost_ewma_alpha = 0.35; +}; + +// Convert conditional survival confidence into expected useful tokens: +// 1 AR token plus the probability of reaching every speculative prefix. +// DSpark supplies calibrated confidence-head values directly. +inline double expected_tokens_from_conditional_confidence( + const float * confidence, int count) { + double expected = 1.0; + double survival = 1.0; + for (int i = 0; confidence && i < count; ++i) { + const double conditional = std::clamp( + std::isfinite(confidence[i]) + ? static_cast(confidence[i]) : 0.0, + 0.0, 1.0); + survival *= conditional; + expected += survival; + } + return expected; +} + +class AdaptiveVerificationRanker { +public: + AdaptiveVerificationRanker() = default; + explicit AdaptiveVerificationRanker(AdaptiveVerificationConfig config) + : config_(sanitize(config)) {} + + void reset() { + ar_cost_us_.clear(); + spec_cost_us_.clear(); + ar_cost_known_.clear(); + spec_cost_known_.clear(); + } + + void observe_autoregressive(int batch_size, double elapsed_us) { + observe_cost(ar_cost_us_, ar_cost_known_, batch_size, elapsed_us); + } + + void observe_speculation(int batch_size, double elapsed_us) { + observe_cost(spec_cost_us_, spec_cost_known_, batch_size, elapsed_us); + } + + bool has_autoregressive_cost(int batch_size) const { + return has_cost(ar_cost_known_, batch_size); + } + + bool has_speculation_cost(int batch_size) const { + return has_cost(spec_cost_known_, batch_size); + } + + double autoregressive_cost_us(int batch_size) const { + return cost(ar_cost_us_, ar_cost_known_, batch_size); + } + + double speculation_cost_us(int batch_size) const { + return cost(spec_cost_us_, spec_cost_known_, batch_size); + } + + AdaptiveVerificationDecision select( + int active_requests, + const std::vector & candidates) + const { + AdaptiveVerificationDecision out; + if (active_requests <= 0 || candidates.empty() || + !has_autoregressive_cost(active_requests)) { + // First observe the exact all-AR baseline for this occupancy. + return out; + } + + std::vector known; + std::vector unknown; + known.reserve(candidates.size()); + unknown.reserve(candidates.size()); + for (AdaptiveVerificationCandidate candidate : candidates) { + if (candidate.request < 0) continue; + candidate.maximum_tokens = std::max( + 1.0, std::isfinite(candidate.maximum_tokens) + ? candidate.maximum_tokens : 1.0); + candidate.expected_tokens = std::clamp( + std::isfinite(candidate.expected_tokens) + ? candidate.expected_tokens : 1.0, + 1.0, candidate.maximum_tokens); + (candidate.calibrated ? known : unknown).push_back(candidate); + } + if (known.empty() && unknown.empty()) return out; + + auto by_value = [](const AdaptiveVerificationCandidate & a, + const AdaptiveVerificationCandidate & b) { + if (a.expected_tokens != b.expected_tokens) { + return a.expected_tokens > b.expected_tokens; + } + return a.request < b.request; + }; + std::stable_sort(known.begin(), known.end(), by_value); + std::stable_sort(unknown.begin(), unknown.end(), + [](const AdaptiveVerificationCandidate & a, + const AdaptiveVerificationCandidate & b) { + if (a.maximum_tokens != b.maximum_tokens) { + return a.maximum_tokens > b.maximum_tokens; + } + return a.request < b.request; + }); + if (!unknown.empty()) { + out.calibration_request = unknown.front().request; + } + + const double baseline = + static_cast(active_requests) / + autoregressive_cost_us(active_requests); + const double required = baseline * config_.minimum_gain; + + int admitted_prefix = 0; + int missing_cost_prefix = 0; + double best = baseline; + double admitted_goodput = baseline; + double expected_total = static_cast(active_requests); + for (int k = 1; k <= static_cast(known.size()) && + k <= active_requests; ++k) { + expected_total += known[(size_t)k - 1].expected_tokens - 1.0; + double route_us = 0.0; + if (!combined_cost(active_requests, k, route_us)) { + missing_cost_prefix = k; + break; + } + const double throughput = expected_total / route_us; + // DSpark's greedy policy stops at the first non-improving + // candidate because candidates are already ranked by survival. + if (throughput <= best) break; + best = throughput; + if (throughput >= required) { + admitted_prefix = k; + admitted_goodput = throughput; + } + } + + // Uncalibrated requests are never sent through an expensive target + // verification merely to discover their value. calibration_request + // asks the concrete adapter for its cheap confidence signal instead. + + // A calibrated route shape without a hardware sample gets one bounded + // probe. Larger prefixes are not explored after a measured decline. + if (missing_cost_prefix > 0) { + out.requests.reserve((size_t)missing_cost_prefix); + for (int i = 0; i < missing_cost_prefix; ++i) { + out.requests.push_back(known[(size_t)i].request); + } + out.exploring = true; + return out; + } + + if (admitted_prefix > 0) { + out.requests.reserve((size_t)admitted_prefix); + for (int i = 0; i < admitted_prefix; ++i) { + out.requests.push_back(known[(size_t)i].request); + } + out.predicted_gain = admitted_goodput / baseline; + } + return out; + } + +private: + static AdaptiveVerificationConfig sanitize( + AdaptiveVerificationConfig config) { + config.minimum_gain = std::max(1.0, config.minimum_gain); + config.cost_ewma_alpha = + std::clamp(config.cost_ewma_alpha, 0.0, 1.0); + return config; + } + + void observe_cost(std::vector & costs, + std::vector & known, + int batch_size, double elapsed_us) { + if (batch_size <= 0 || !std::isfinite(elapsed_us) || + elapsed_us <= 0.0) { + return; + } + const size_t needed = static_cast(batch_size) + 1; + if (costs.size() < needed) costs.resize(needed, 0.0); + if (known.size() < needed) known.resize(needed, false); + if (!known[(size_t)batch_size]) { + costs[(size_t)batch_size] = elapsed_us; + known[(size_t)batch_size] = true; + return; + } + costs[(size_t)batch_size] = + config_.cost_ewma_alpha * elapsed_us + + (1.0 - config_.cost_ewma_alpha) * + costs[(size_t)batch_size]; + } + + static bool has_cost(const std::vector & known, int batch_size) { + return batch_size == 0 || + (batch_size > 0 && static_cast(batch_size) < known.size() && + known[(size_t)batch_size]); + } + + static double cost(const std::vector & costs, + const std::vector & known, int batch_size) { + if (batch_size == 0) return 0.0; + return has_cost(known, batch_size) + ? costs[(size_t)batch_size] + : std::numeric_limits::infinity(); + } + + bool combined_cost(int active_requests, int speculative_requests, + double & elapsed_us) const { + const int ar_requests = active_requests - speculative_requests; + if (speculative_requests <= 0 || ar_requests < 0 || + !has_speculation_cost(speculative_requests) || + !has_autoregressive_cost(ar_requests)) { + return false; + } + elapsed_us = speculation_cost_us(speculative_requests) + + autoregressive_cost_us(ar_requests); + return std::isfinite(elapsed_us) && elapsed_us > 0.0; + } + + AdaptiveVerificationConfig config_; + std::vector ar_cost_us_; + std::vector spec_cost_us_; + std::vector ar_cost_known_; + std::vector spec_cost_known_; +}; + +} // namespace dflash::common diff --git a/server/src/common/concurrency/speculation_goodput.h b/server/src/common/concurrency/speculation_goodput.h new file mode 100644 index 000000000..259243e83 --- /dev/null +++ b/server/src/common/concurrency/speculation_goodput.h @@ -0,0 +1,206 @@ +#pragma once + +// Model-neutral online controller for deciding whether a request should use +// speculative decoding. The controller deliberately consumes only outcomes: +// useful tokens emitted and wall time for an AR or speculative step. A +// concrete speculator (DDTree today, DSpark later) remains responsible for +// producing and verifying candidates. + +#include +#include + +namespace dflash::common { + +enum class SpeculationGoodputTransition { + none, + enabled, + disabled, +}; + +struct SpeculationGoodputConfig { + // Require a measured advantage large enough to survive timing noise. + double minimum_gain = 1.05; + // Half-life is intentionally short: generation can move between + // predictable code/structure and high-entropy prose within one request. + double ewma_alpha = 0.5; + // Do not abandon a profitable route after one unlucky verification. + int bad_speculation_steps = 2; + // AR requests periodically earn one new speculative probe so a response + // can recover when its continuation becomes predictable again. Zero + // disables re-probing. + int ar_reprobe_steps = 16; +}; + +class SpeculationGoodputController { +public: + SpeculationGoodputController() = default; + explicit SpeculationGoodputController( + SpeculationGoodputConfig config) + : config_(sanitize(config)) {} + + void reset(bool adaptive = true) { + adaptive_ = adaptive; + phase_ = adaptive ? Phase::initial_spec_probe : Phase::speculate; + speculative_goodput_ = 0.0; + ar_goodput_ = 0.0; + expected_emitted_tokens_ = 1.0; + has_speculative_goodput_ = false; + has_ar_goodput_ = false; + has_expected_emitted_tokens_ = false; + bad_speculation_steps_ = 0; + ar_steps_since_probe_ = 0; + } + + bool wants_speculation() const { + return !adaptive_ || phase_ == Phase::initial_spec_probe || + phase_ == Phase::speculate || + phase_ == Phase::reprobe_speculation; + } + + // Seed request value from a cheap drafter-side confidence estimate without + // changing the measured AR/speculation route state. + void observe_expected_tokens(double expected_tokens) { + if (!adaptive_ || !std::isfinite(expected_tokens) || + expected_tokens <= 0.0) { + return; + } + update_ewma(expected_emitted_tokens_, + has_expected_emitted_tokens_, expected_tokens); + } + + SpeculationGoodputTransition observe_speculation( + double emitted_tokens, double elapsed_us) { + if (!adaptive_ || !valid_observation(emitted_tokens, elapsed_us)) { + return SpeculationGoodputTransition::none; + } + update_ewma(speculative_goodput_, has_speculative_goodput_, + emitted_tokens / elapsed_us); + update_ewma(expected_emitted_tokens_, + has_expected_emitted_tokens_, emitted_tokens); + + if (phase_ == Phase::initial_spec_probe) { + phase_ = Phase::initial_ar_probe; + return SpeculationGoodputTransition::none; + } + if (phase_ == Phase::reprobe_speculation) { + if (speculation_profitable()) { + phase_ = Phase::speculate; + bad_speculation_steps_ = 0; + return SpeculationGoodputTransition::enabled; + } + phase_ = Phase::autoregressive; + ar_steps_since_probe_ = 0; + return SpeculationGoodputTransition::none; + } + if (phase_ != Phase::speculate || !has_ar_goodput_) { + return SpeculationGoodputTransition::none; + } + + if (speculation_profitable()) { + bad_speculation_steps_ = 0; + return SpeculationGoodputTransition::none; + } + if (++bad_speculation_steps_ < config_.bad_speculation_steps) { + return SpeculationGoodputTransition::none; + } + phase_ = Phase::autoregressive; + bad_speculation_steps_ = 0; + ar_steps_since_probe_ = 0; + return SpeculationGoodputTransition::disabled; + } + + SpeculationGoodputTransition observe_autoregressive(double elapsed_us) { + if (!adaptive_ || !valid_observation(1.0, elapsed_us)) { + return SpeculationGoodputTransition::none; + } + update_ewma(ar_goodput_, has_ar_goodput_, 1.0 / elapsed_us); + + if (phase_ == Phase::initial_ar_probe) { + if (speculation_profitable()) { + phase_ = Phase::speculate; + bad_speculation_steps_ = 0; + return SpeculationGoodputTransition::none; + } + phase_ = Phase::autoregressive; + ar_steps_since_probe_ = 0; + return SpeculationGoodputTransition::disabled; + } + if (phase_ != Phase::autoregressive || + config_.ar_reprobe_steps <= 0) { + return SpeculationGoodputTransition::none; + } + if (++ar_steps_since_probe_ >= config_.ar_reprobe_steps) { + phase_ = Phase::reprobe_speculation; + ar_steps_since_probe_ = 0; + } + return SpeculationGoodputTransition::none; + } + + bool adaptive() const { return adaptive_; } + bool has_speculative_goodput() const { + return has_speculative_goodput_; + } + bool has_ar_goodput() const { return has_ar_goodput_; } + double speculative_goodput() const { return speculative_goodput_; } + double ar_goodput() const { return ar_goodput_; } + double expected_emitted_tokens() const { + return expected_emitted_tokens_; + } + bool has_expected_emitted_tokens() const { + return has_expected_emitted_tokens_; + } + +private: + enum class Phase { + initial_spec_probe, + initial_ar_probe, + speculate, + autoregressive, + reprobe_speculation, + }; + + static SpeculationGoodputConfig sanitize( + SpeculationGoodputConfig config) { + config.minimum_gain = std::max(1.0, config.minimum_gain); + config.ewma_alpha = std::clamp(config.ewma_alpha, 0.0, 1.0); + config.bad_speculation_steps = + std::max(1, config.bad_speculation_steps); + config.ar_reprobe_steps = std::max(0, config.ar_reprobe_steps); + return config; + } + + static bool valid_observation(double emitted_tokens, double elapsed_us) { + return std::isfinite(emitted_tokens) && emitted_tokens > 0.0 && + std::isfinite(elapsed_us) && elapsed_us > 0.0; + } + + void update_ewma(double & value, bool & initialized, double sample) { + if (!initialized) { + value = sample; + initialized = true; + return; + } + value = config_.ewma_alpha * sample + + (1.0 - config_.ewma_alpha) * value; + } + + bool speculation_profitable() const { + return has_speculative_goodput_ && has_ar_goodput_ && + speculative_goodput_ >= + config_.minimum_gain * ar_goodput_; + } + + SpeculationGoodputConfig config_; + bool adaptive_ = true; + Phase phase_ = Phase::initial_spec_probe; + double speculative_goodput_ = 0.0; + double ar_goodput_ = 0.0; + double expected_emitted_tokens_ = 1.0; + bool has_speculative_goodput_ = false; + bool has_ar_goodput_ = false; + bool has_expected_emitted_tokens_ = false; + int bad_speculation_steps_ = 0; + int ar_steps_since_probe_ = 0; +}; + +} // namespace dflash::common diff --git a/server/test/test_speculation_goodput.cpp b/server/test/test_speculation_goodput.cpp new file mode 100644 index 000000000..db59d9964 --- /dev/null +++ b/server/test/test_speculation_goodput.cpp @@ -0,0 +1,241 @@ +#include "common/concurrency/adaptive_verification.h" +#include "common/concurrency/speculation_goodput.h" +#include "host_check.h" + +#include + +using namespace dflash::common; + +static int g_checks = 0; + +int main() { + // Cold start measures one real speculative step and one neighboring AR + // step, then keeps the route with higher useful-token goodput. + { + SpeculationGoodputController policy; + CHECK(policy.wants_speculation()); + CHECK(policy.observe_speculation(/*emitted_tokens=*/4, + /*elapsed_us=*/200.0) == + SpeculationGoodputTransition::none); + CHECK(!policy.wants_speculation()); + CHECK(policy.observe_autoregressive(/*elapsed_us=*/100.0) == + SpeculationGoodputTransition::none); + CHECK(policy.wants_speculation()); + CHECK(policy.has_speculative_goodput()); + CHECK(policy.has_ar_goodput()); + CHECK(policy.has_expected_emitted_tokens()); + CHECK(policy.expected_emitted_tokens() == 4.0); + } + + // A chat-like low-yield probe loses to AR and is disabled per request. + { + SpeculationGoodputController policy; + CHECK(policy.observe_speculation(/*emitted_tokens=*/1, + /*elapsed_us=*/200.0) == + SpeculationGoodputTransition::none); + CHECK(policy.observe_autoregressive(/*elapsed_us=*/100.0) == + SpeculationGoodputTransition::disabled); + CHECK(!policy.wants_speculation()); + } + + // A request already benefiting from speculation tolerates one bad step, + // then leaves the route after a second consecutive bad observation. + { + SpeculationGoodputConfig config; + config.ewma_alpha = 1.0; + SpeculationGoodputController policy(config); + CHECK(policy.observe_speculation(4, 200.0) == + SpeculationGoodputTransition::none); + CHECK(policy.observe_autoregressive(100.0) == + SpeculationGoodputTransition::none); + CHECK(policy.wants_speculation()); + CHECK(policy.observe_speculation(1, 200.0) == + SpeculationGoodputTransition::none); + CHECK(policy.wants_speculation()); + CHECK(policy.observe_speculation(1, 200.0) == + SpeculationGoodputTransition::disabled); + CHECK(!policy.wants_speculation()); + } + + // AR periodically schedules one speculative re-probe. A profitable probe + // re-enables speculation; a losing probe returns directly to AR. + { + SpeculationGoodputConfig config; + config.ewma_alpha = 1.0; + config.ar_reprobe_steps = 2; + SpeculationGoodputController policy(config); + CHECK(policy.observe_speculation(1, 200.0) == + SpeculationGoodputTransition::none); + CHECK(policy.observe_autoregressive(100.0) == + SpeculationGoodputTransition::disabled); + CHECK(!policy.wants_speculation()); + CHECK(policy.observe_autoregressive(100.0) == + SpeculationGoodputTransition::none); + CHECK(!policy.wants_speculation()); + CHECK(policy.observe_autoregressive(100.0) == + SpeculationGoodputTransition::none); + CHECK(policy.wants_speculation()); + CHECK(policy.observe_speculation(4, 200.0) == + SpeculationGoodputTransition::enabled); + CHECK(policy.wants_speculation()); + } + + { + SpeculationGoodputConfig config; + config.ewma_alpha = 1.0; + config.ar_reprobe_steps = 1; + SpeculationGoodputController policy(config); + CHECK(policy.observe_speculation(1, 200.0) == + SpeculationGoodputTransition::none); + CHECK(policy.observe_autoregressive(100.0) == + SpeculationGoodputTransition::disabled); + CHECK(policy.observe_autoregressive(100.0) == + SpeculationGoodputTransition::none); + CHECK(policy.wants_speculation()); + CHECK(policy.observe_speculation(1, 200.0) == + SpeculationGoodputTransition::none); + CHECK(!policy.wants_speculation()); + } + + // Cheap confidence calibration does not masquerade as a measured target + // verification or advance the request route state. + { + SpeculationGoodputController policy; + policy.observe_expected_tokens(3.5); + CHECK(policy.wants_speculation()); + CHECK(policy.has_expected_emitted_tokens()); + CHECK(policy.expected_emitted_tokens() == 3.5); + CHECK(!policy.has_speculative_goodput()); + } + + // DSpark conditional confidence maps directly to the shared expected-token + // value used for request ranking. + { + const float confidence[] = {0.8f, 0.5f}; + const double expected = + expected_tokens_from_conditional_confidence(confidence, 2); + CHECK(expected > 2.19 && expected < 2.21); + } + + // The ranker first measures the all-AR hardware baseline at this exact + // occupancy. It then explores only one uncalibrated request. + { + AdaptiveVerificationRanker ranker; + std::vector candidates = { + {7, 1.0, 8.0, false}, + {3, 1.0, 8.0, false}, + }; + CHECK(ranker.select(5, candidates).requests.empty()); + ranker.observe_autoregressive(5, 100.0); + const AdaptiveVerificationDecision probe = + ranker.select(5, candidates); + CHECK(!probe.exploring); + CHECK(probe.requests.empty()); + CHECK(probe.calibration_request == 3); + } + + // A measured request is admitted at C=5 only when its expected useful + // tokens overcome both the speculative and remaining AR subbatch costs. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(5, 100.0); + ranker.observe_autoregressive(4, 80.0); + ranker.observe_speculation(1, 40.0); + std::vector candidates = { + {11, 4.0, 8.0, true}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(5, candidates); + CHECK(!decision.exploring); + CHECK(decision.requests.size() == 1); + CHECK(decision.requests[0] == 11); + CHECK(decision.predicted_gain > 1.05); + } + + // Candidates are ranked by expected value. Greedy growth stops at the + // first prefix that would reduce whole-batch throughput. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(5, 100.0); + ranker.observe_autoregressive(4, 80.0); + ranker.observe_autoregressive(3, 70.0); + ranker.observe_speculation(1, 40.0); + ranker.observe_speculation(2, 70.0); + std::vector candidates = { + {2, 1.5, 8.0, true}, + {9, 4.0, 8.0, true}, + {4, 1.2, 8.0, true}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(5, candidates); + CHECK(decision.requests.size() == 1); + CHECK(decision.requests[0] == 9); + } + + // Uncalibrated requests are returned to the concrete confidence adapter; + // they are never admitted to expensive target verification as a probe. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(16, 100.0); + ranker.observe_autoregressive(15, 95.0); + ranker.observe_speculation(1, 100.0); + std::vector candidates = { + {1, 1.0, 8.0, false}, + {2, 1.0, 8.0, false}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(16, candidates); + CHECK(decision.requests.empty()); + CHECK(decision.calibration_request == 1); + } + + // There is no concurrency cutoff: when one request pays for the route at + // C=16, that request alone remains speculative and all peers remain AR. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(16, 160.0); + ranker.observe_autoregressive(15, 150.0); + ranker.observe_autoregressive(14, 140.0); + ranker.observe_speculation(1, 20.0); + ranker.observe_speculation(2, 80.0); + std::vector candidates = { + {21, 8.0, 8.0, true}, + {22, 2.0, 8.0, true}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(16, candidates); + CHECK(decision.requests.size() == 1); + CHECK(decision.requests[0] == 21); + } + + // A promising calibrated prefix gets one bounded hardware-cost probe when + // that subbatch shape has not been observed yet. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(3, 90.0); + std::vector candidates = { + {31, 6.0, 8.0, true}, + }; + const AdaptiveVerificationDecision probe = + ranker.select(3, candidates); + CHECK(probe.exploring); + CHECK(probe.requests.size() == 1); + CHECK(probe.requests[0] == 31); + } + + // Kill-switch mode retains fixed speculation and ignores observations. + { + SpeculationGoodputController policy; + policy.reset(/*adaptive=*/false); + CHECK(policy.wants_speculation()); + CHECK(!policy.adaptive()); + CHECK(policy.observe_speculation(1, 1000.0) == + SpeculationGoodputTransition::none); + CHECK(policy.observe_autoregressive(1.0) == + SpeculationGoodputTransition::none); + CHECK(policy.wants_speculation()); + } + + std::printf("speculation goodput policy: %d checks passed\n", g_checks); + return 0; +} From ce70dd93dcf4c8d8ddc84017fdcb33a9770619bc Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 17 Aug 2026 08:04:45 +0000 Subject: [PATCH 08/18] feat(qwen35): route DDTree by measured goodput --- .../benchmarks/concurrency/FEATURE_MATRIX.md | 2 +- .../common/concurrency/speculation_goodput.h | 9 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 385 +++++++++++++++--- .../qwen35/concurrency/qwen35_seq_engine.h | 9 +- .../concurrency/qwen35_slot_manager.cpp | 35 +- .../qwen35/concurrency/qwen35_slot_manager.h | 22 +- server/src/qwen35/qwen35_backend.cpp | 4 +- server/test/test_seq_slot_manager.cpp | 107 ++--- 8 files changed, 433 insertions(+), 140 deletions(-) diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index 62fce6a46..19299a0b6 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -39,7 +39,7 @@ VARIANTS=ddtree,pflash,kvflash,full CLIENTS=1,4,8,16 REPEATS=5 \ harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh ``` -`DDTREE_ADAPTIVE=0` matches the blog's continuous DDTree probe policy; leave it at the default `1` when measuring the concurrent engine's adaptive fallback policy. The concurrent path now records a startup `[parallel-ddtree]` marker and per-request `ddtree_steps`; these are the proof that DDTree actually ran. +`DDTREE_ADAPTIVE=0` matches the blog's continuous DDTree probe policy; leave it at the default `1` when measuring adaptive verification. The default policy ranks requests independently at every concurrency from expected useful tokens and measured AR/speculation subbatch costs. DDTree calibrates cold requests from cumulative top-1 draft confidence; at C>3 these calibration passes are spaced by 64 decode steps by default (`DFLASH_ADAPTIVE_VERIFY_CALIBRATION_STEPS` overrides the interval). The speculator-neutral ranker also accepts calibrated DSpark confidence directly. The concurrent path records a startup `[parallel-ddtree]` marker and per-request `ddtree_steps`; these are the proof that DDTree actually ran. On a 128 GiB Strix Halo host, budget roughly 45–90 minutes for this smoke run; the long-context AR controls dominate and actual time depends on the build. diff --git a/server/src/common/concurrency/speculation_goodput.h b/server/src/common/concurrency/speculation_goodput.h index 259243e83..d2a41cba0 100644 --- a/server/src/common/concurrency/speculation_goodput.h +++ b/server/src/common/concurrency/speculation_goodput.h @@ -25,10 +25,11 @@ struct SpeculationGoodputConfig { double ewma_alpha = 0.5; // Do not abandon a profitable route after one unlucky verification. int bad_speculation_steps = 2; - // AR requests periodically earn one new speculative probe so a response - // can recover when its continuation becomes predictable again. Zero - // disables re-probing. - int ar_reprobe_steps = 16; + // Optional re-probing after the initial route decision. This is disabled + // by default because one speculative tree probe can be much more expensive + // than an AR step at high occupancy. Callers may opt in when the speculator + // has a sufficiently cheap proposal path. + int ar_reprobe_steps = 0; }; class SpeculationGoodputController { diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 52949acfc..7fdde2058 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -17,6 +17,8 @@ #include "internal.h" #include +#include +#include #include #include #include @@ -140,31 +142,105 @@ DraftKvState * Qwen35SeqEngine::ensure_slot_draft_kv(int slot) { return state.get(); } -bool Qwen35SeqEngine::ddtree_eligible(const StepPlan & plan) const { - if (tree_width_ <= 1 || !capture_features_ || !plan.prefills.empty() || - plan.decode.empty() || b_.dw_.block_size <= 1 || - b_.cfg_.ddtree_budget + 1 != tree_width_) { +bool Qwen35SeqEngine::ddtree_available(const StepPlan & plan) const { + return tree_width_ > 1 && capture_features_ && plan.prefills.empty() && + !plan.decode.empty() && b_.dw_.block_size > 1 && + b_.cfg_.ddtree_budget + 1 == tree_width_; +} + +bool Qwen35SeqEngine::ddtree_input_eligible(const StepInput & in) const { + if (!in.allow_speculation || in.slot < 0 || + in.slot >= slots_.slot_count() || + !slots_.slot(in.slot).decoding() || + slots_.slot(in.slot).sampler.needs_logit_processing() || + slots_.slot(in.slot).cur_pos < 1 || + slots_.slot(in.slot).cur_pos >= slots_.max_context()) { return false; } const int min_floor = []() { const char * value = std::getenv("DFLASH_MIN_TOKENS"); return value ? std::max(0, std::atoi(value)) : 0; }(); - for (const StepInput & in : plan.decode) { - if (!in.allow_speculation || in.slot < 0 || - in.slot >= slots_.slot_count() || - !slots_.slot(in.slot).decoding() || - !slots_.ddtree_speculation_allowed(in.slot) || - slots_.slot(in.slot).sampler.needs_logit_processing() || - slots_.slot(in.slot).cur_pos < 1 || - slots_.slot(in.slot).cur_pos >= slots_.max_context()) { - return false; - } - const Qwen35Slot & seq = slots_.slot(in.slot); - const int generated = seq.generated_tokens(); - if (generated < min_floor) return false; + return slots_.slot(in.slot).generated_tokens() >= min_floor; +} + +std::optional Qwen35SeqEngine::estimate_ddtree_expected_tokens( + const StepInput & in) { + const int q_len = b_.dw_.block_size; + const int hidden = b_.w_.n_embd; + if (q_len <= 1 || !build_lm_head_projection_step( + b_.proj_sg_, b_.w_, b_.target_backend_, q_len)) { + return std::nullopt; } - return true; + + DraftKvState * draft = ensure_slot_draft_kv(in.slot); + DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); + if (!draft || !mirror) return std::nullopt; + + auto fail = [&]() -> std::optional { + // begin_step advances only host cache bookkeeping, but a failed graph + // can leave the appended rows incomplete. Rebuild from committed + // target features on the next proposal. + draft_kv_reset(*draft); + return std::nullopt; + }; + + std::vector noise((size_t)q_len, b_.w_.mask_token_id); + noise[0] = in.token; + std::vector noise_embed((size_t)hidden * q_len); + if (!b_.w_.embedder.embed( + noise.data(), q_len, noise_embed.data()) || + !draft_kv_begin_step(*draft, b_.dw_, b_.draft_backend_, + *mirror, slots_.slot(in.slot).cur_pos)) { + return fail(); + } + ggml_backend_tensor_set( + draft->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + if (ggml_backend_graph_compute(b_.draft_backend_, draft->gf) != + GGML_STATUS_SUCCESS) { + return fail(); + } + ggml_backend_synchronize(b_.draft_backend_); + ggml_backend_tensor_copy( + draft->hidden_states, b_.proj_sg_.hidden_input); + if (ggml_backend_graph_compute( + b_.target_backend_, b_.proj_sg_.gf) != GGML_STATUS_SUCCESS) { + return fail(); + } + + std::vector top_lp((size_t)q_len); + std::vector top_ids((size_t)q_len); + bool topk_ready = false; +#ifdef DFLASH27B_HAVE_DRAFT_TOPK + topk_ready = geometric_extract_draft_topk_cuda( + b_.proj_sg_.logits->data, q_len, b_.w_.n_vocab, 1, + top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); +#endif + if (!topk_ready) { + std::vector logits( + (size_t)b_.w_.n_vocab * q_len); + ggml_backend_tensor_get( + b_.proj_sg_.logits, logits.data(), 0, + sizeof(float) * logits.size()); + extract_draft_topk( + logits.data(), q_len, b_.w_.n_vocab, 1, + top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); + } + + // Training-free SVIP-style confidence: the top-1 draft probability at + // each position estimates conditional survival, so the cumulative product + // estimates reaching that prefix. DSpark can replace this adapter with its + // calibrated confidence-head probabilities without changing the ranker. + std::vector confidence((size_t)q_len - 1); + for (int pos = 1; pos < q_len; ++pos) { + confidence[(size_t)pos - 1] = static_cast( + std::clamp(std::exp(static_cast( + top_lp[(size_t)pos])), + 0.0, 1.0)); + } + return expected_tokens_from_conditional_confidence( + confidence.data(), static_cast(confidence.size())); } std::optional Qwen35SeqEngine::step_ddtree( @@ -533,21 +609,6 @@ std::optional Qwen35SeqEngine::step_ddtree( } } - uint64_t cohort_emitted = 0; - for (const Proposal & p : proposals) { - // accepted contains the replay root plus accepted children. The - // output emits those children plus one separately computed pending - // scalar, so accepted.size() is this request's emitted yield. - cohort_emitted += (uint64_t)p.accepted.size(); - } - const bool suspend_cohort = - Qwen35SlotManager::ddtree_cohort_should_suspend( - cohort_emitted, active); - std::vector newly_suspended((size_t)slots_.slot_count(), false); - for (const Proposal & p : proposals) { - newly_suspended[(size_t)p.slot] = - slots_.record_ddtree_sample(p.slot, suspend_cohort); - } result.decode.reserve((size_t)active); for (Proposal & p : proposals) { @@ -563,20 +624,6 @@ std::optional Qwen35SeqEngine::step_ddtree( out.committed_tokens.push_back( p.tree.token_ids[(size_t)dfs - 1]); } - if (newly_suspended[(size_t)p.slot]) { - out.ddtree_suspensions = 1; - const Qwen35Slot & seq = slots_.slot(p.slot); - std::fprintf(stderr, - "[parallel-ddtree] adaptive suspend request=%llu slot=%d " - "sample=%llu emitted=%d accepted_children=%d " - "cohort_emitted=%llu cohort_size=%d target_forwards=2 " - "floor=%d\n", - (unsigned long long)seq.request_id, p.slot, - (unsigned long long)seq.ddtree_sampled_steps, - accepted_children + 1, accepted_children, - (unsigned long long)cohort_emitted, active, - Qwen35SlotManager::kDdtreeMinEmittedTokens); - } attach_residency_telemetry(out); result.decode.push_back(std::move(out)); } @@ -787,8 +834,6 @@ Qwen35SeqEngine::PrefillStage Qwen35SeqEngine::stage_prefill_chunk( SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { StepResult result; - std::vector & decode_outputs = result.decode; - std::vector & prefill_outputs = result.prefills; const std::vector & inputs = plan.decode; const int n_slots = slots_.slot_count(); @@ -835,13 +880,245 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (inputs.empty() && plan.prefills.empty()) return result; - if (ddtree_eligible(plan)) { - std::optional speculative = step_ddtree(plan); - if (speculative) return std::move(*speculative); - // Proposal setup failed before target/cache mutation. Preserve service - // by taking the existing packed AR path for this iteration. + if (!ddtree_available(plan)) return step_regular(plan); + + const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); + const bool adaptive_enabled = !(adaptive && std::atoi(adaptive) == 0); + + // Keep the scheduler independent of the concrete speculation algorithm. + // DDTree contributes either a cheap draft-confidence estimate or an EWMA + // of target-accepted useful tokens. A future DSpark adapter can contribute + // calibrated prefix survival from its confidence head. + std::vector candidates; + candidates.reserve(inputs.size()); + auto collect_candidates = [&]() { + candidates.clear(); + for (const StepInput & in : inputs) { + if (!ddtree_input_eligible(in) || + !slots_.ddtree_speculation_allowed(in.slot)) { + continue; + } + const SpeculationGoodputController & policy = + slots_.slot(in.slot).speculation; + candidates.push_back({ + in.slot, + policy.expected_emitted_tokens(), + static_cast(b_.dw_.block_size), + policy.has_expected_emitted_tokens(), + }); + } + }; + collect_candidates(); + + AdaptiveVerificationDecision decision; + if (adaptive_enabled) { + decision = adaptive_verification_.select( + static_cast(inputs.size()), candidates); + if (inputs.size() <= 3) { + adaptive_calibration_cooldown_ = 0; + } else if (adaptive_calibration_cooldown_ > 0) { + --adaptive_calibration_cooldown_; + } + if (decision.calibration_request >= 0 && + adaptive_calibration_cooldown_ == 0) { + const auto input = std::find_if( + inputs.begin(), inputs.end(), + [&](const StepInput & in) { + return in.slot == decision.calibration_request; + }); + if (input != inputs.end()) { + const std::optional expected = + estimate_ddtree_expected_tokens(*input); + if (expected) { + slots_.slot(input->slot).speculation + .observe_expected_tokens(*expected); + // Drafter confidence is cheap relative to target + // verification but not free on the current sequential + // per-slot DDTree path. At high occupancy, amortize cold + // calibration across decode steps; already-ranked requests + // remain independently selectable on every step. + if (inputs.size() > 3) { + static const int interval = []() { + const char * value = std::getenv( + "DFLASH_ADAPTIVE_VERIFY_CALIBRATION_STEPS"); + return value ? std::max(1, std::atoi(value)) : 64; + }(); + adaptive_calibration_cooldown_ = interval; + } + std::fprintf(stderr, + "[parallel-ddtree] confidence request=%llu slot=%d " + "expected_tokens=%.3f\n", + (unsigned long long) + slots_.slot(input->slot).request_id, + input->slot, *expected); + collect_candidates(); + decision = adaptive_verification_.select( + static_cast(inputs.size()), candidates); + } + } + } + } else { + decision.requests.reserve(candidates.size()); + for (const AdaptiveVerificationCandidate & candidate : candidates) { + decision.requests.push_back(candidate.request); + } + } + + std::vector selected((size_t)n_slots, 0); + for (int slot : decision.requests) { + if (slot >= 0 && slot < n_slots) selected[(size_t)slot] = 1; + } + + StepPlan speculative_plan; + StepPlan ar_plan; + speculative_plan.decode.reserve(decision.requests.size()); + ar_plan.decode.reserve(inputs.size() - decision.requests.size()); + std::vector observe_ar((size_t)n_slots, 0); + for (const StepInput & in : inputs) { + const bool eligible = ddtree_input_eligible(in); + if (eligible && selected[(size_t)in.slot]) { + speculative_plan.decode.push_back(in); + } else { + ar_plan.decode.push_back(in); + if (eligible) observe_ar[(size_t)in.slot] = 1; + } + } + + using Clock = std::chrono::steady_clock; + StepResult speculative_result; + StepResult ar_result; + double speculative_us = 0.0; + double ar_us = 0.0; + + if (!speculative_plan.decode.empty()) { + const auto started = Clock::now(); + std::optional speculative = + step_ddtree(speculative_plan); + speculative_us = std::max( + 1.0, std::chrono::duration( + Clock::now() - started).count()); + if (!speculative) { + // Proposal setup failed before target/cache mutation. Preserve + // service with one ordinary packed step and retry speculation on a + // later iteration. + return step_regular(plan); + } + if (!speculative->ok()) return std::move(*speculative); + speculative_result = std::move(*speculative); + } + + if (!ar_plan.decode.empty()) { + const auto started = Clock::now(); + ar_result = step_regular(ar_plan); + ar_us = std::max( + 1.0, std::chrono::duration( + Clock::now() - started).count()); + if (!ar_result.ok()) return ar_result; + } + + if (adaptive_enabled) { + if (!speculative_plan.decode.empty()) { + adaptive_verification_.observe_speculation( + static_cast(speculative_plan.decode.size()), + speculative_us); + } + if (!ar_plan.decode.empty()) { + adaptive_verification_.observe_autoregressive( + static_cast(ar_plan.decode.size()), ar_us); + } + } + if (decision.exploring && !speculative_plan.decode.empty()) { + std::fprintf(stderr, + "[parallel-ddtree] adaptive verification probe active=%zu " + "speculative=%zu ar=%zu\n", + inputs.size(), speculative_plan.decode.size(), + ar_plan.decode.size()); + } + + auto log_transition = [&](int slot, + SpeculationGoodputTransition transition, + const char * observed_route, + double emitted_tokens, + double elapsed_us) { + if (transition == SpeculationGoodputTransition::none) return; + const Qwen35Slot & seq = slots_.slot(slot); + const SpeculationGoodputController & policy = seq.speculation; + const char * action = + transition == SpeculationGoodputTransition::enabled + ? "enable" : "disable"; + const double spec_tps = policy.speculative_goodput() * 1.0e6; + const double ar_tps = policy.ar_goodput() * 1.0e6; + std::fprintf(stderr, + "[parallel-ddtree] adaptive route request=%llu slot=%d " + "action=%s observed=%s sample=%llu emitted=%.0f " + "elapsed_us=%.0f spec_tok_s=%.2f ar_tok_s=%.2f\n", + (unsigned long long)seq.request_id, slot, action, observed_route, + (unsigned long long)seq.ddtree_sampled_steps, emitted_tokens, + elapsed_us, spec_tps, ar_tps); + }; + + for (DecodeOutput & out : speculative_result.decode) { + if (out.failed || out.slot < 0 || out.slot >= n_slots) continue; + const double emitted = + static_cast(out.ddtree_accepted_tokens + 1); + const SpeculationGoodputTransition transition = + slots_.record_speculation_sample( + out.slot, emitted, speculative_us); + if (transition == SpeculationGoodputTransition::disabled) { + out.ddtree_suspensions = 1; + } + log_transition( + out.slot, transition, "speculation", emitted, speculative_us); } + for (DecodeOutput & out : ar_result.decode) { + if (out.failed || out.slot < 0 || out.slot >= n_slots || + !observe_ar[(size_t)out.slot]) { + continue; + } + const SpeculationGoodputTransition transition = + slots_.record_ar_sample(out.slot, ar_us); + log_transition(out.slot, transition, "ar", 1.0, ar_us); + } + + std::vector by_slot((size_t)n_slots); + std::vector present((size_t)n_slots, 0); + auto collect = [&](std::vector & outputs) { + for (DecodeOutput & out : outputs) { + if (out.slot < 0 || out.slot >= n_slots || + present[(size_t)out.slot]) { + return false; + } + present[(size_t)out.slot] = 1; + by_slot[(size_t)out.slot] = std::move(out); + } + return true; + }; + if (!collect(speculative_result.decode) || !collect(ar_result.decode)) { + return fail_step("adaptive decode produced duplicate slot output"); + } + result.decode.reserve(inputs.size()); + for (const StepInput & in : inputs) { + if (!present[(size_t)in.slot]) { + return fail_step("adaptive decode omitted a live slot output"); + } + result.decode.push_back(std::move(by_slot[(size_t)in.slot])); + } + return result; +} +SeqEngine::StepResult Qwen35SeqEngine::step_regular(const StepPlan & plan) { + StepResult result; + std::vector & decode_outputs = result.decode; + std::vector & prefill_outputs = result.prefills; + 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); + }; const TargetWeights & w = b_.w_; StepGraph & sg = b_.sg_; const int hidden = w.n_embd; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 75275df39..9361d549c 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -21,6 +21,7 @@ #pragma once +#include "common/concurrency/adaptive_verification.h" #include "common/concurrency/seq_engine.h" #include "common/dflash_draft_kv.h" #include "common/dflash_feature_ring.h" @@ -118,7 +119,11 @@ class Qwen35SeqEngine final : public SeqEngine { std::vector * logits_scratch = nullptr); DraftFeatureMirror * slot_feature_mirror(int slot); DraftKvState * ensure_slot_draft_kv(int slot); - bool ddtree_eligible(const StepPlan & plan) const; + bool ddtree_available(const StepPlan & plan) const; + bool ddtree_input_eligible(const StepInput & input) const; + std::optional estimate_ddtree_expected_tokens( + const StepInput & input); + StepResult step_regular(const StepPlan & plan); // nullopt means proposal setup failed before target/cache mutation and the // caller may safely use the ordinary packed AR path for this iteration. std::optional step_ddtree(const StepPlan & plan); @@ -130,6 +135,8 @@ class Qwen35SeqEngine final : public SeqEngine { int tree_scratch_base_ = 0; int tree_scratch_stride_ = 0; bool capture_features_ = false; + AdaptiveVerificationRanker adaptive_verification_; + int adaptive_calibration_cooldown_ = 0; ggml_context * feature_view_ctx_ = nullptr; std::vector slot_feature_mirrors_; std::vector> slot_draft_kv_; diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index e559a224e..755402943 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -13,6 +13,8 @@ Qwen35SlotManager::Qwen35SlotManager( headroom_tokens_(std::max(pool.block_size(), speculative_headroom)), residency_(residency) { slots_.assign(pool.max_sequences(), Qwen35Slot{}); + const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); + speculation_adaptive_ = !(adaptive && std::atoi(adaptive) == 0); } int Qwen35SlotManager::decoding_count() const { @@ -205,8 +207,8 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( Qwen35Slot & s = slots_[(size_t)slot]; s.phase = Qwen35SlotPhase::prefill; s.request_id = request_id; - s.ddtree_suspended = false; s.ddtree_sampled_steps = 0; + s.speculation.reset(speculation_adaptive_); s.handle = handle; s.cur_pos = 0; s.prompt_len = prompt_len; @@ -226,30 +228,23 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( } bool Qwen35SlotManager::ddtree_speculation_allowed(int slot) const { - return is_active(slot) && !slots_[(size_t)slot].ddtree_suspended; + return is_active(slot) && + slots_[(size_t)slot].speculation.wants_speculation(); } -bool Qwen35SlotManager::ddtree_cohort_should_suspend( - uint64_t total_emitted, int active) { - const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); - if (adaptive && std::atoi(adaptive) == 0) { - return false; - } - if (active <= 0) return false; - return total_emitted < - (uint64_t)active * (uint64_t)kDdtreeMinEmittedTokens; +SpeculationGoodputTransition Qwen35SlotManager::record_speculation_sample( + int slot, double emitted_tokens, double elapsed_us) { + if (!is_active(slot)) return SpeculationGoodputTransition::none; + Qwen35Slot & s = slots_[(size_t)slot]; + ++s.ddtree_sampled_steps; + return s.speculation.observe_speculation(emitted_tokens, elapsed_us); } -bool Qwen35SlotManager::record_ddtree_sample( - int slot, bool suspend_cohort) { - if (!is_active(slot)) return false; +SpeculationGoodputTransition Qwen35SlotManager::record_ar_sample( + int slot, double elapsed_us) { + if (!is_active(slot)) return SpeculationGoodputTransition::none; Qwen35Slot & s = slots_[(size_t)slot]; - // A suspended request must never pay for another probe. - if (s.ddtree_suspended) return false; - ++s.ddtree_sampled_steps; - if (!suspend_cohort) return false; - s.ddtree_suspended = true; - return true; + return s.speculation.observe_autoregressive(elapsed_us); } Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index 9a84fb3be..d73ea5824 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -18,6 +18,7 @@ #include "common/concurrency/paged_kv_pool.h" #include "common/concurrency/paged_kv_residency.h" +#include "common/concurrency/speculation_goodput.h" #include "common/sampler.h" #include "common/concurrency/seq_engine.h" @@ -66,10 +67,10 @@ struct Qwen35Slot { : 0; } - // Real packed-tree samples are counted per request. A low aggregate-yield - // cohort sample suspends every participating request; ordinary AR keeps - // every target cache and feature-ring row current. - bool ddtree_suspended = false; + // Route each request from measured useful-token goodput, independently of + // the concrete speculator. DDTree records observations today; DSpark can + // use the same controller later. + SpeculationGoodputController speculation; uint64_t ddtree_sampled_steps = 0; bool active() const { @@ -83,10 +84,6 @@ struct Qwen35Slot { class Qwen35SlotManager { public: - // Packed DDTree pays for verify + accepted-path replay. Requiring six - // emitted tokens makes continuation earn at least three tokens per target - // forward before accounting for its additional draft/tree work. - static constexpr int kDdtreeMinEmittedTokens = 6; // `max_ctx` is the per-sequence logical bound; slot count comes from the // pool's max_sequences. The pool must outlive the manager. @@ -157,11 +154,9 @@ class Qwen35SlotManager { void take_residency_telemetry(int slot, SeqEngine::DecodeOutput & out); bool ddtree_speculation_allowed(int slot) const; - // Compare aggregate emitted yield (accepted children plus one replay - // bonus per request) against the cohort continuation floor. - static bool ddtree_cohort_should_suspend(uint64_t total_emitted, int active); - // Returns true exactly once, when this sample newly suspends the request. - bool record_ddtree_sample(int slot, bool suspend_cohort); + SpeculationGoodputTransition record_speculation_sample( + int slot, double emitted_tokens, double elapsed_us); + SpeculationGoodputTransition record_ar_sample(int slot, double elapsed_us); // One-token compatibility wrapper used by ordinary autoregressive decode. StepAppend append_token(int slot, int32_t fed_token); @@ -198,6 +193,7 @@ class Qwen35SlotManager { int max_ctx_ = 0; int headroom_tokens_; PagedKvResidencyManager * residency_ = nullptr; + bool speculation_adaptive_ = true; std::vector slots_; }; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index efc7e617b..19a01b731 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -551,7 +551,9 @@ bool Qwen35Backend::init() { if (concurrent_local_ddtree) { const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); std::fprintf(stderr, - "[parallel-ddtree] enabled budget=%d width=%d mode=packed-verify-replay adaptive=%s\n", + "[parallel-ddtree] enabled budget=%d width=%d " + "mode=packed-verify-replay adaptive=%s " + "policy=ranked-goodput scope=all-concurrency\n", cfg_.ddtree_budget, tree_width, adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); } diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index 7a2d7a26c..67458afb4 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -492,62 +492,77 @@ int main() { CHECK(!mgr.has_prefill_prompt_at_least(768)); } - // Adaptive DDTree judges aggregate cohort yield rather than allowing one - // unlucky request to suppress higher-yield peers. Equality at the - // six-token average continues; a low average suspends all participants. + // Requests in one concurrent cohort make independent routing decisions + // from their measured useful-token goodput. A predictable request can keep + // DDTree while a low-yield chat-like peer switches to packed AR. { - const luce_test::ScopedEnvVar adaptive("DFLASH_DDTREE_ADAPTIVE", nullptr); - CHECK(Qwen35SlotManager::ddtree_cohort_should_suspend(5, 1)); - CHECK(!Qwen35SlotManager::ddtree_cohort_should_suspend(6, 1)); - // Mixed synthetic cohort: one child-less row plus an eleven-token row - // exactly meets the continuation floor; one fewer token does not. - CHECK(!Qwen35SlotManager::ddtree_cohort_should_suspend(12, 2)); - CHECK(Qwen35SlotManager::ddtree_cohort_should_suspend(11, 2)); - + const luce_test::ScopedEnvVar adaptive( + "DFLASH_DDTREE_ADAPTIVE", nullptr); PagedKvPool pool(8, 2, /*block_size=*/16); Qwen35SlotManager mgr(pool, 64); - auto a = admit(mgr, 101, prompt_tokens(4), greedy_sampler()); - auto b = admit(mgr, 102, prompt_tokens(4), greedy_sampler()); - CHECK(is_admitted(a)); - CHECK(is_admitted(b)); - CHECK(mgr.append_prefill(a.slot, 4).ok); - CHECK(mgr.append_prefill(b.slot, 4).ok); - mgr.commit_prefill(a.slot); - mgr.commit_prefill(b.slot); - CHECK(mgr.slot(a.slot).request_id == 101); - CHECK(mgr.slot(b.slot).request_id == 102); - CHECK(mgr.ddtree_speculation_allowed(a.slot)); - CHECK(mgr.ddtree_speculation_allowed(b.slot)); - - // The keep decision records a real sample without latching either - // participant, including the low-yield member. - CHECK(!mgr.record_ddtree_sample(a.slot, false)); - CHECK(!mgr.record_ddtree_sample(b.slot, false)); - CHECK(mgr.ddtree_speculation_allowed(a.slot)); - CHECK(mgr.ddtree_speculation_allowed(b.slot)); - CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 1); - CHECK(mgr.slot(b.slot).ddtree_sampled_steps == 1); - - // A low cohort decision is applied identically and atomically to both. - CHECK(mgr.record_ddtree_sample(a.slot, true)); - CHECK(mgr.record_ddtree_sample(b.slot, true)); - CHECK(!mgr.ddtree_speculation_allowed(a.slot)); - CHECK(!mgr.ddtree_speculation_allowed(b.slot)); - CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 2); - CHECK(mgr.slot(b.slot).ddtree_sampled_steps == 2); - - // Suspended requests cannot pay for another bad probe. - CHECK(!mgr.record_ddtree_sample(a.slot, true)); - CHECK(mgr.slot(a.slot).ddtree_sampled_steps == 2); + auto code = admit(mgr, 101, prompt_tokens(4), greedy_sampler()); + auto chat = admit(mgr, 102, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(code)); + CHECK(is_admitted(chat)); + CHECK(mgr.append_prefill(code.slot, 4).ok); + CHECK(mgr.append_prefill(chat.slot, 4).ok); + mgr.commit_prefill(code.slot); + mgr.commit_prefill(chat.slot); + CHECK(mgr.ddtree_speculation_allowed(code.slot)); + CHECK(mgr.ddtree_speculation_allowed(chat.slot)); + + CHECK(mgr.record_speculation_sample( + code.slot, /*emitted_tokens=*/8, /*elapsed_us=*/2000.0) == + SpeculationGoodputTransition::none); + CHECK(mgr.record_speculation_sample( + chat.slot, /*emitted_tokens=*/1, /*elapsed_us=*/2000.0) == + SpeculationGoodputTransition::none); + // Both requests take one neighboring AR calibration step. + CHECK(!mgr.ddtree_speculation_allowed(code.slot)); + CHECK(!mgr.ddtree_speculation_allowed(chat.slot)); + CHECK(mgr.record_ar_sample(code.slot, /*elapsed_us=*/1000.0) == + SpeculationGoodputTransition::none); + CHECK(mgr.record_ar_sample(chat.slot, /*elapsed_us=*/1000.0) == + SpeculationGoodputTransition::disabled); + CHECK(mgr.ddtree_speculation_allowed(code.slot)); + CHECK(!mgr.ddtree_speculation_allowed(chat.slot)); + CHECK(mgr.slot(code.slot).ddtree_sampled_steps == 1); + CHECK(mgr.slot(chat.slot).ddtree_sampled_steps == 1); + + // The production default makes one bounded decision per request; an + // expensive speculator does not periodically interrupt a winning AR + // route. Re-probing remains available through controller config. + for (int i = 0; i < 32; ++i) { + CHECK(mgr.record_ar_sample(chat.slot, 1000.0) == + SpeculationGoodputTransition::none); + CHECK(!mgr.ddtree_speculation_allowed(chat.slot)); + } - mgr.retire(a.slot); + mgr.retire(code.slot); auto reused = admit(mgr, 103, prompt_tokens(4), greedy_sampler()); - CHECK(is_admitted(reused) && reused.slot == a.slot); + CHECK(is_admitted(reused) && reused.slot == code.slot); CHECK(mgr.slot(reused.slot).request_id == 103); CHECK(mgr.ddtree_speculation_allowed(reused.slot)); CHECK(mgr.slot(reused.slot).ddtree_sampled_steps == 0); } + // The existing burn-in switch preserves fixed speculation. + { + const luce_test::ScopedEnvVar adaptive("DFLASH_DDTREE_ADAPTIVE", "0"); + PagedKvPool pool(4, 1, /*block_size=*/16); + Qwen35SlotManager mgr(pool, 64); + auto admitted = admit(mgr, 201, prompt_tokens(4), greedy_sampler()); + CHECK(is_admitted(admitted)); + CHECK(mgr.append_prefill(admitted.slot, 4).ok); + mgr.commit_prefill(admitted.slot); + CHECK(mgr.ddtree_speculation_allowed(admitted.slot)); + CHECK(mgr.record_speculation_sample(admitted.slot, 1, 2000.0) == + SpeculationGoodputTransition::none); + CHECK(mgr.record_ar_sample(admitted.slot, 1000.0) == + SpeculationGoodputTransition::none); + CHECK(mgr.ddtree_speculation_allowed(admitted.slot)); + } + // A failed residency barrier quarantines retirement ownership, and a // later admission retries it before considering the slot reusable. { From 1103701531d4228f91cb956e5f2760c5dbf53e17 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 17 Aug 2026 10:15:36 +0000 Subject: [PATCH 09/18] perf(qwen35): fuse mixed speculation replay --- .../benchmarks/concurrency/FEATURE_MATRIX.md | 2 +- .../concurrency/adaptive_verification.h | 133 +++++++------ .../qwen35/concurrency/qwen35_seq_engine.cpp | 179 +++++++++++------- .../qwen35/concurrency/qwen35_seq_engine.h | 3 +- server/src/qwen35/qwen35_backend.cpp | 2 +- server/test/test_speculation_goodput.cpp | 19 +- 6 files changed, 182 insertions(+), 156 deletions(-) diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index 19299a0b6..a8fc7ffd9 100644 --- a/harness/benchmarks/concurrency/FEATURE_MATRIX.md +++ b/harness/benchmarks/concurrency/FEATURE_MATRIX.md @@ -39,7 +39,7 @@ VARIANTS=ddtree,pflash,kvflash,full CLIENTS=1,4,8,16 REPEATS=5 \ harness/benchmarks/concurrency/run_qwen36_feature_matrix.sh ``` -`DDTREE_ADAPTIVE=0` matches the blog's continuous DDTree probe policy; leave it at the default `1` when measuring adaptive verification. The default policy ranks requests independently at every concurrency from expected useful tokens and measured AR/speculation subbatch costs. DDTree calibrates cold requests from cumulative top-1 draft confidence; at C>3 these calibration passes are spaced by 64 decode steps by default (`DFLASH_ADAPTIVE_VERIFY_CALIBRATION_STEPS` overrides the interval). The speculator-neutral ranker also accepts calibrated DSpark confidence directly. The concurrent path records a startup `[parallel-ddtree]` marker and per-request `ddtree_steps`; these are the proof that DDTree actually ran. +`DDTREE_ADAPTIVE=0` matches the blog-aligned continuous DDTree policy; leave it at the default `1` when measuring adaptive verification. The default policy ranks requests independently at every concurrency from expected useful tokens and the measured cost of the exact `(active requests, speculative requests)` route shape. DDTree calibrates cold requests from cumulative top-1 draft confidence; at C>3 these calibration passes are spaced by 64 decode steps by default (`DFLASH_ADAPTIVE_VERIFY_CALIBRATION_STEPS` overrides the interval). Selected DDTree paths and unselected AR roots share one ragged durable replay, so mixed routing adds a tree-verification pass but not a second AR target pass. The speculator-neutral ranker also accepts calibrated DSpark confidence directly. The concurrent path records a startup `[parallel-ddtree]` marker and per-request `ddtree_steps`; these are the proof that DDTree actually ran. On a 128 GiB Strix Halo host, budget roughly 45–90 minutes for this smoke run; the long-context AR controls dominate and actual time depends on the build. diff --git a/server/src/common/concurrency/adaptive_verification.h b/server/src/common/concurrency/adaptive_verification.h index 386792cce..f0c68f3e7 100644 --- a/server/src/common/concurrency/adaptive_verification.h +++ b/server/src/common/concurrency/adaptive_verification.h @@ -5,9 +5,10 @@ // DSpark's scheduler ranks candidate tokens by cumulative survival probability // and grows the verification batch while expected throughput improves. This // helper applies the same policy at request granularity. The concrete -// speculator supplies expected useful tokens; the engine supplies observed AR -// and speculative subbatch costs. DDTree can learn value from accepted paths, -// while DSpark can use its calibrated confidence head directly. +// speculator supplies expected useful tokens; the engine supplies the observed +// cost of each mixed route shape (active requests, speculative requests). +// DDTree can learn value from accepted paths, while DSpark can use its +// calibrated confidence head directly. #include #include @@ -65,34 +66,82 @@ class AdaptiveVerificationRanker { : config_(sanitize(config)) {} void reset() { - ar_cost_us_.clear(); - spec_cost_us_.clear(); - ar_cost_known_.clear(); - spec_cost_known_.clear(); + route_cost_us_.clear(); + route_cost_known_.clear(); } void observe_autoregressive(int batch_size, double elapsed_us) { - observe_cost(ar_cost_us_, ar_cost_known_, batch_size, elapsed_us); + observe_route(batch_size, /*speculative_requests=*/0, elapsed_us); } + // Compatibility helper for engines whose whole batch takes one + // speculative route. Mixed executors should call observe_route directly. void observe_speculation(int batch_size, double elapsed_us) { - observe_cost(spec_cost_us_, spec_cost_known_, batch_size, elapsed_us); + observe_route(batch_size, batch_size, elapsed_us); + } + + void observe_route(int active_requests, int speculative_requests, + double elapsed_us) { + if (active_requests <= 0 || speculative_requests < 0 || + speculative_requests > active_requests || + !std::isfinite(elapsed_us) || elapsed_us <= 0.0) { + return; + } + const size_t rows = static_cast(active_requests) + 1; + if (route_cost_us_.size() < rows) route_cost_us_.resize(rows); + if (route_cost_known_.size() < rows) route_cost_known_.resize(rows); + std::vector & costs = + route_cost_us_[(size_t)active_requests]; + std::vector & known = + route_cost_known_[(size_t)active_requests]; + const size_t cols = static_cast(speculative_requests) + 1; + if (costs.size() < cols) costs.resize(cols, 0.0); + if (known.size() < cols) known.resize(cols, false); + if (!known[(size_t)speculative_requests]) { + costs[(size_t)speculative_requests] = elapsed_us; + known[(size_t)speculative_requests] = true; + return; + } + costs[(size_t)speculative_requests] = + config_.cost_ewma_alpha * elapsed_us + + (1.0 - config_.cost_ewma_alpha) * + costs[(size_t)speculative_requests]; } bool has_autoregressive_cost(int batch_size) const { - return has_cost(ar_cost_known_, batch_size); + return has_route_cost(batch_size, /*speculative_requests=*/0); } bool has_speculation_cost(int batch_size) const { - return has_cost(spec_cost_known_, batch_size); + return has_route_cost(batch_size, batch_size); } double autoregressive_cost_us(int batch_size) const { - return cost(ar_cost_us_, ar_cost_known_, batch_size); + return route_cost_us(batch_size, /*speculative_requests=*/0); } double speculation_cost_us(int batch_size) const { - return cost(spec_cost_us_, spec_cost_known_, batch_size); + return route_cost_us(batch_size, batch_size); + } + + bool has_route_cost(int active_requests, + int speculative_requests) const { + return active_requests > 0 && speculative_requests >= 0 && + speculative_requests <= active_requests && + static_cast(active_requests) < + route_cost_known_.size() && + static_cast(speculative_requests) < + route_cost_known_[(size_t)active_requests].size() && + route_cost_known_[(size_t)active_requests] + [(size_t)speculative_requests]; + } + + double route_cost_us(int active_requests, + int speculative_requests) const { + return has_route_cost(active_requests, speculative_requests) + ? route_cost_us_[(size_t)active_requests] + [(size_t)speculative_requests] + : std::numeric_limits::infinity(); } AdaptiveVerificationDecision select( @@ -156,11 +205,11 @@ class AdaptiveVerificationRanker { for (int k = 1; k <= static_cast(known.size()) && k <= active_requests; ++k) { expected_total += known[(size_t)k - 1].expected_tokens - 1.0; - double route_us = 0.0; - if (!combined_cost(active_requests, k, route_us)) { + if (!has_route_cost(active_requests, k)) { missing_cost_prefix = k; break; } + const double route_us = route_cost_us(active_requests, k); const double throughput = expected_total / route_us; // DSpark's greedy policy stops at the first non-improving // candidate because candidates are already ranked by survival. @@ -206,59 +255,9 @@ class AdaptiveVerificationRanker { return config; } - void observe_cost(std::vector & costs, - std::vector & known, - int batch_size, double elapsed_us) { - if (batch_size <= 0 || !std::isfinite(elapsed_us) || - elapsed_us <= 0.0) { - return; - } - const size_t needed = static_cast(batch_size) + 1; - if (costs.size() < needed) costs.resize(needed, 0.0); - if (known.size() < needed) known.resize(needed, false); - if (!known[(size_t)batch_size]) { - costs[(size_t)batch_size] = elapsed_us; - known[(size_t)batch_size] = true; - return; - } - costs[(size_t)batch_size] = - config_.cost_ewma_alpha * elapsed_us + - (1.0 - config_.cost_ewma_alpha) * - costs[(size_t)batch_size]; - } - - static bool has_cost(const std::vector & known, int batch_size) { - return batch_size == 0 || - (batch_size > 0 && static_cast(batch_size) < known.size() && - known[(size_t)batch_size]); - } - - static double cost(const std::vector & costs, - const std::vector & known, int batch_size) { - if (batch_size == 0) return 0.0; - return has_cost(known, batch_size) - ? costs[(size_t)batch_size] - : std::numeric_limits::infinity(); - } - - bool combined_cost(int active_requests, int speculative_requests, - double & elapsed_us) const { - const int ar_requests = active_requests - speculative_requests; - if (speculative_requests <= 0 || ar_requests < 0 || - !has_speculation_cost(speculative_requests) || - !has_autoregressive_cost(ar_requests)) { - return false; - } - elapsed_us = speculation_cost_us(speculative_requests) + - autoregressive_cost_us(ar_requests); - return std::isfinite(elapsed_us) && elapsed_us > 0.0; - } - AdaptiveVerificationConfig config_; - std::vector ar_cost_us_; - std::vector spec_cost_us_; - std::vector ar_cost_known_; - std::vector spec_cost_known_; + std::vector> route_cost_us_; + std::vector> route_cost_known_; }; } // namespace dflash::common diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 7fdde2058..327067f24 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -244,9 +244,10 @@ std::optional Qwen35SeqEngine::estimate_ddtree_expected_tokens( } std::optional Qwen35SeqEngine::step_ddtree( - const StepPlan & plan) { + const StepPlan & speculative_plan, const StepPlan & ar_plan) { StepResult result; - const int active = (int)plan.decode.size(); + const int active = (int)speculative_plan.decode.size(); + const int total_active = active + (int)ar_plan.decode.size(); const int bucket = decode_bucket_width(active); const int T = tree_width_; const int q_len = b_.dw_.block_size; @@ -276,7 +277,7 @@ std::optional Qwen35SeqEngine::step_ddtree( // A failed proposal graph may therefore leave only part of that // cache valid. Reset all participating draft rings so a later // speculative round rebuilds them from committed target features. - for (const StepInput & in : plan.decode) { + for (const StepInput & in : speculative_plan.decode) { if (in.slot >= 0 && in.slot < (int)slot_draft_kv_.size() && slot_draft_kv_[(size_t)in.slot]) { draft_kv_reset(*slot_draft_kv_[(size_t)in.slot]); @@ -292,7 +293,7 @@ std::optional Qwen35SeqEngine::step_ddtree( // Proposal is sequential by slot: immutable draft weights are shared, // while each slot owns an independent persistent context-KV ring. - for (const StepInput & in : plan.decode) { + for (const StepInput & in : speculative_plan.decode) { DraftKvState * draft = ensure_slot_draft_kv(in.slot); DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); if (!draft || !mirror) return proposal_fallback(); @@ -463,6 +464,7 @@ std::optional Qwen35SeqEngine::step_ddtree( } replay_total += (int)p.accepted.size(); } + replay_total += (int)ar_plan.decode.size(); std::vector replay_segments; std::vector replay_tokens; @@ -470,12 +472,12 @@ std::optional Qwen35SeqEngine::step_ddtree( std::vector replay_positions; std::vector replay_rows; std::vector replay_logits_rows; - replay_segments.reserve((size_t)active); + replay_segments.reserve((size_t)total_active); replay_tokens.reserve((size_t)replay_total); replay_slots.reserve((size_t)replay_total); replay_positions.reserve((size_t)replay_total); replay_rows.assign((size_t)replay_total * n_head_kv, scratch_row_); - replay_logits_rows.reserve((size_t)active); + replay_logits_rows.reserve((size_t)total_active); seq_lens_.assign((size_t)n_slots, 0); int replay_offset = 0; @@ -513,6 +515,36 @@ std::optional Qwen35SeqEngine::step_ddtree( replay_logits_rows.push_back(replay_offset - 1); seq_lens_[(size_t)p.slot] = app.position + (int)path.size(); } + for (const StepInput & in : ar_plan.decode) { + const Qwen35SlotManager::StepAppend app = slots_.append_tokens( + in.slot, &in.token, 1); + const bool table_ok = slots_.residency_active() || + upload_block_table_delta(in.slot, app.first_new_block, + app.new_blocks.data(), app.new_blocks.size()); + if (!app.ok || app.physical_rows.size() != 1 || !table_ok) { + result.error = app.busy + ? "paged KV pool exhausted during mixed DDTree/AR replay" + : "mixed DDTree/AR replay K/V append failed"; + return result; + } + replay_segments.push_back({replay_offset, 1, in.slot}); + replay_tokens.push_back(in.token); + replay_slots.push_back(in.slot); + replay_positions.push_back(app.position); + for (int h = 0; h < n_head_kv; ++h) { + replay_rows[(size_t)h * replay_total + replay_offset] = + app.physical_rows[0]; + } + ++replay_offset; + replay_logits_rows.push_back(replay_offset - 1); + seq_lens_[(size_t)in.slot] = app.position + 1; + } + if (replay_offset != replay_total || + (int)replay_segments.size() != total_active || + (int)replay_logits_rows.size() != total_active) { + result.error = "mixed DDTree/AR replay staging mismatch"; + return result; + } if (!upload_all_active_block_tables()) { result.error = "DDTree replay block-table refresh failed"; @@ -526,7 +558,7 @@ std::optional Qwen35SeqEngine::step_ddtree( b_.cfg_.kq_stride_pad, false, false, false, true, 1, 0, *std::max_element(seq_lens_.begin(), seq_lens_.end()), replay_total, replay_segments.data(), - (int)replay_segments.size(), active, false) || + (int)replay_segments.size(), total_active, false) || !replay_sg.target_feat_rows || !replay_sg.paged_query_seq_ids || !replay_sg.paged_query_positions || !replay_sg.logits_row_indices || !replay_sg.argmax_tokens) { @@ -576,18 +608,20 @@ std::optional Qwen35SeqEngine::step_ddtree( return result; } std::vector replay_write_slots; - replay_write_slots.reserve(proposals.size()); + replay_write_slots.reserve((size_t)total_active); for (const Proposal & p : proposals) replay_write_slots.push_back(p.slot); + for (const StepInput & in : ar_plan.decode) { + replay_write_slots.push_back(in.slot); + } if (!commit_residency_writes(replay_write_slots)) { - result.error = "DDTree replay KV write commit failed"; + result.error = "mixed DDTree/AR replay KV write commit failed"; return result; } + for (int slot : replay_write_slots) slots_.commit_step(slot); - // The replay is the durable target forward: its recurrent/KV/feature - // state is what the next step consumes. Use its posterior rather than - // the tree-verify posterior so the pending scalar remains exact even if - // the two graph shapes differ numerically. - std::vector replay_next((size_t)active, -1); + // The replay is the one durable target forward for both routes. Its + // recurrent/KV/feature state is what the next step consumes. + std::vector replay_next((size_t)total_active, -1); ggml_backend_tensor_get( replay_sg.argmax_tokens, replay_next.data(), 0, sizeof(int32_t) * replay_next.size()); @@ -598,19 +632,28 @@ std::optional Qwen35SeqEngine::step_ddtree( } proposals[(size_t)s].bonus = replay_next[(size_t)s]; } + std::vector ar_next(ar_plan.decode.size(), -1); + for (size_t i = 0; i < ar_plan.decode.size(); ++i) { + const int logits_row = active + (int)i; + ar_next[i] = sample_graph_row( + ar_plan.decode[i].slot, logits_row, + &replay_next[(size_t)logits_row], &logits_buf_); + if (ar_next[i] < 0) { + result.error = "mixed replay produced an invalid AR token"; + return result; + } + } - for (Proposal & p : proposals) { - slots_.commit_step(p.slot); + for (int slot : replay_write_slots) { std::string reselect_error; - if (!maybe_reselect_residency(p.slot, reselect_error)) { + if (!maybe_reselect_residency(slot, reselect_error)) { result.error = reselect_error.empty() ? "KVFlash reselect failed" : reselect_error; return result; } } - - result.decode.reserve((size_t)active); + result.decode.reserve((size_t)total_active); for (Proposal & p : proposals) { DecodeOutput out; out.slot = p.slot; @@ -627,6 +670,14 @@ std::optional Qwen35SeqEngine::step_ddtree( attach_residency_telemetry(out); result.decode.push_back(std::move(out)); } + for (size_t i = 0; i < ar_plan.decode.size(); ++i) { + DecodeOutput out; + out.slot = ar_plan.decode[i].slot; + out.token = ar_next[i]; + out.target_forwards = 1; + attach_residency_telemetry(out); + result.decode.push_back(std::move(out)); + } return result; } @@ -985,47 +1036,32 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } using Clock = std::chrono::steady_clock; - StepResult speculative_result; - StepResult ar_result; - double speculative_us = 0.0; - double ar_us = 0.0; - - if (!speculative_plan.decode.empty()) { - const auto started = Clock::now(); - std::optional speculative = - step_ddtree(speculative_plan); - speculative_us = std::max( - 1.0, std::chrono::duration( - Clock::now() - started).count()); - if (!speculative) { + StepResult routed_result; + const int speculative_count = + static_cast(speculative_plan.decode.size()); + const auto started = Clock::now(); + if (speculative_count > 0) { + std::optional mixed = + step_ddtree(speculative_plan, ar_plan); + if (!mixed) { // Proposal setup failed before target/cache mutation. Preserve // service with one ordinary packed step and retry speculation on a - // later iteration. + // later iteration. Do not learn from this contaminated timing. return step_regular(plan); } - if (!speculative->ok()) return std::move(*speculative); - speculative_result = std::move(*speculative); - } - - if (!ar_plan.decode.empty()) { - const auto started = Clock::now(); - ar_result = step_regular(ar_plan); - ar_us = std::max( - 1.0, std::chrono::duration( - Clock::now() - started).count()); - if (!ar_result.ok()) return ar_result; + if (!mixed->ok()) return std::move(*mixed); + routed_result = std::move(*mixed); + } else { + routed_result = step_regular(ar_plan); + if (!routed_result.ok()) return routed_result; } + const double route_us = std::max( + 1.0, std::chrono::duration( + Clock::now() - started).count()); if (adaptive_enabled) { - if (!speculative_plan.decode.empty()) { - adaptive_verification_.observe_speculation( - static_cast(speculative_plan.decode.size()), - speculative_us); - } - if (!ar_plan.decode.empty()) { - adaptive_verification_.observe_autoregressive( - static_cast(ar_plan.decode.size()), ar_us); - } + adaptive_verification_.observe_route( + static_cast(inputs.size()), speculative_count, route_us); } if (decision.exploring && !speculative_plan.decode.empty()) { std::fprintf(stderr, @@ -1057,27 +1093,24 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { elapsed_us, spec_tps, ar_tps); }; - for (DecodeOutput & out : speculative_result.decode) { + for (DecodeOutput & out : routed_result.decode) { if (out.failed || out.slot < 0 || out.slot >= n_slots) continue; - const double emitted = - static_cast(out.ddtree_accepted_tokens + 1); - const SpeculationGoodputTransition transition = - slots_.record_speculation_sample( - out.slot, emitted, speculative_us); - if (transition == SpeculationGoodputTransition::disabled) { - out.ddtree_suspensions = 1; - } - log_transition( - out.slot, transition, "speculation", emitted, speculative_us); - } - for (DecodeOutput & out : ar_result.decode) { - if (out.failed || out.slot < 0 || out.slot >= n_slots || - !observe_ar[(size_t)out.slot]) { - continue; + if (selected[(size_t)out.slot]) { + const double emitted = + static_cast(out.ddtree_accepted_tokens + 1); + const SpeculationGoodputTransition transition = + slots_.record_speculation_sample( + out.slot, emitted, route_us); + if (transition == SpeculationGoodputTransition::disabled) { + out.ddtree_suspensions = 1; + } + log_transition( + out.slot, transition, "speculation", emitted, route_us); + } else if (observe_ar[(size_t)out.slot]) { + const SpeculationGoodputTransition transition = + slots_.record_ar_sample(out.slot, route_us); + log_transition(out.slot, transition, "ar", 1.0, route_us); } - const SpeculationGoodputTransition transition = - slots_.record_ar_sample(out.slot, ar_us); - log_transition(out.slot, transition, "ar", 1.0, ar_us); } std::vector by_slot((size_t)n_slots); @@ -1093,7 +1126,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } return true; }; - if (!collect(speculative_result.decode) || !collect(ar_result.decode)) { + if (!collect(routed_result.decode)) { return fail_step("adaptive decode produced duplicate slot output"); } result.decode.reserve(inputs.size()); diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 9361d549c..5bc4d3400 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -126,7 +126,8 @@ class Qwen35SeqEngine final : public SeqEngine { StepResult step_regular(const StepPlan & plan); // nullopt means proposal setup failed before target/cache mutation and the // caller may safely use the ordinary packed AR path for this iteration. - std::optional step_ddtree(const StepPlan & plan); + std::optional step_ddtree( + const StepPlan & speculative_plan, const StepPlan & ar_plan); Qwen35Backend & b_; Qwen35SlotManager slots_; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 19a01b731..da5cf0402 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -552,7 +552,7 @@ bool Qwen35Backend::init() { const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); std::fprintf(stderr, "[parallel-ddtree] enabled budget=%d width=%d " - "mode=packed-verify-replay adaptive=%s " + "mode=packed-verify-mixed-replay adaptive=%s " "policy=ranked-goodput scope=all-concurrency\n", cfg_.ddtree_budget, tree_width, adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); diff --git a/server/test/test_speculation_goodput.cpp b/server/test/test_speculation_goodput.cpp index db59d9964..853af8e0c 100644 --- a/server/test/test_speculation_goodput.cpp +++ b/server/test/test_speculation_goodput.cpp @@ -135,12 +135,11 @@ int main() { } // A measured request is admitted at C=5 only when its expected useful - // tokens overcome both the speculative and remaining AR subbatch costs. + // tokens overcome the measured cost of that exact fused route shape. { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(5, 100.0); - ranker.observe_autoregressive(4, 80.0); - ranker.observe_speculation(1, 40.0); + ranker.observe_route(5, 1, 80.0); std::vector candidates = { {11, 4.0, 8.0, true}, }; @@ -157,10 +156,8 @@ int main() { { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(5, 100.0); - ranker.observe_autoregressive(4, 80.0); - ranker.observe_autoregressive(3, 70.0); - ranker.observe_speculation(1, 40.0); - ranker.observe_speculation(2, 70.0); + ranker.observe_route(5, 1, 80.0); + ranker.observe_route(5, 2, 95.0); std::vector candidates = { {2, 1.5, 8.0, true}, {9, 4.0, 8.0, true}, @@ -177,8 +174,6 @@ int main() { { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(16, 100.0); - ranker.observe_autoregressive(15, 95.0); - ranker.observe_speculation(1, 100.0); std::vector candidates = { {1, 1.0, 8.0, false}, {2, 1.0, 8.0, false}, @@ -194,10 +189,8 @@ int main() { { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(16, 160.0); - ranker.observe_autoregressive(15, 150.0); - ranker.observe_autoregressive(14, 140.0); - ranker.observe_speculation(1, 20.0); - ranker.observe_speculation(2, 80.0); + ranker.observe_route(16, 1, 150.0); + ranker.observe_route(16, 2, 220.0); std::vector candidates = { {21, 8.0, 8.0, true}, {22, 2.0, 8.0, true}, From aff52605d0b3f426bb6ccc813eca994a5cf56b23 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 17 Aug 2026 11:26:30 +0000 Subject: [PATCH 10/18] perf(qwen35): commit concurrent DDTree paths directly --- server/src/internal.h | 24 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 281 +++++++++++++++++- server/src/qwen35/graph_builders.cpp | 27 +- server/src/qwen35/graph_builders.h | 3 +- server/src/qwen35/qwen35_backend.cpp | 33 +- server/src/qwen35/qwen35_target_graph.cpp | 76 +++-- 6 files changed, 403 insertions(+), 41 deletions(-) diff --git a/server/src/internal.h b/server/src/internal.h index 9cb2a03b6..45334fd44 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -413,6 +413,14 @@ struct TargetCache { std::vector ssm_intermediate; // size = n_delta (48) std::vector conv_input_cache; // size = n_delta (48) + // Bounded concurrent-tree checkpoint domain. A packed tree row is + // sequence-major, so recurrent checkpoint t for compact tree lane s lives + // at s*tree_capture_width+t. Keeping this lane count independent from the + // physical serving-slot count lets an all-C adaptive route directly commit + // a small profitable subset without reserving T*C recurrent states. + int tree_capture_width = 0; + int tree_capture_lanes = 0; + // Rolling target layer features captured during target forward passes. // Single-sequence shape: [5 * hidden, target_feat_cap] bf16. A multi-slot // cache owns one ring per physical sequence slot and one final dead row: @@ -422,6 +430,10 @@ struct TargetCache { // target_feat_cap remains the per-sequence ring width. ggml_tensor * target_feat = nullptr; int target_feat_cap = 0; + // Extra rows after the per-slot rings and dead row. Direct concurrent-tree + // verification writes candidate features here, then copies only the + // accepted spine into the owning slot's durable ring. + int target_feat_tree_scratch_base = 0; // KVFlash target-QK scorer: last token's post-RoPE (and post-FWHT when // kv_k_rotated) query per full-attention layer, written by the graph @@ -559,9 +571,9 @@ bool restore_target_cache_chain(const PrefixSnapshot * thick, // n_seq_slots > 1 the attention K/V tensors are sized by ctx_alloc (the shared // pool capacity plus one scratch block) rather than one sequence's max_ctx. // `concurrent_tree` declares that a paged multi-slot caller will build packed -// DDTree verification graphs. Those graphs are deliberately side-effect-free -// for recurrent state and commit accepted paths through a later replay, so no -// rollback snapshots/intermediates are allocated. +// DDTree verification graphs. `concurrent_tree_capture_lanes` reserves a +// bounded number of compact F16 recurrent checkpoint lanes for direct accepted +// path commit; routes wider than that retain the replay fallback. bool create_target_cache(const TargetWeights & w, int max_ctx, int max_verify_tokens, @@ -571,7 +583,8 @@ bool create_target_cache(const TargetWeights & w, int ctx_alloc = 0, bool paged_attention = false, int n_seq_slots = 1, - bool concurrent_tree = false); + bool concurrent_tree = false, + int concurrent_tree_capture_lanes = 0); // `f32_ssm_intermediates` enables exact per-token checkpoints for the opt-in // layer-split fast rollback path. The default preserves the established Q8_0 @@ -589,7 +602,8 @@ bool create_target_cache_partial(const TargetWeights & w, bool f32_ssm_intermediates = false, bool paged_attention = false, int n_seq_slots = 1, - bool concurrent_tree = false); + bool concurrent_tree = false, + int concurrent_tree_capture_lanes = 0); void free_target_cache(TargetCache & c); diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 327067f24..09de57004 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -15,6 +15,8 @@ #include "common/ddtree.h" #include "common/geometric_draft_topk_cuda.h" #include "internal.h" +#include "common/gpu_runtime_compat.h" +#include "ggml-backend-impl.h" #include #include @@ -23,6 +25,9 @@ #include #include #include +using to_fp32_cuda_t = void (*)(const void *, float *, int64_t, cudaStream_t); +extern "C++" to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type); + namespace dflash::common { @@ -355,6 +360,15 @@ std::optional Qwen35SeqEngine::step_ddtree( proposals.push_back(std::move(p)); } + const bool target_is_meta = b_.cache_.ssm_state.empty() || + !b_.cache_.ssm_state.front() || + ggml_backend_buft_is_meta(ggml_backend_buffer_get_type( + b_.cache_.ssm_state.front()->buffer)); + const bool direct_commit = + !target_is_meta && bucket <= b_.cache_.tree_capture_lanes && + T == b_.cache_.tree_capture_width && + b_.cache_.target_feat_tree_scratch_base > 0; + StepGraph & tree_sg = b_.sg_; int max_prefix = 1; for (const Proposal & p : proposals) { @@ -363,12 +377,20 @@ std::optional Qwen35SeqEngine::step_ddtree( if (!build_target_step_paged_tree( tree_sg, b_.w_, b_.cache_, b_.target_backend_, T, bucket, max_prefix, tree_scratch_base_, tree_scratch_stride_, - b_.cfg_.kq_stride_pad)) { + b_.cfg_.kq_stride_pad, direct_commit)) { result.error = "packed DDTree verify graph build failed"; return result; } const int total_tree = T * bucket; + std::vector tree_feature_rows; + if (direct_commit) { + tree_feature_rows.resize((size_t)total_tree); + for (int row = 0; row < total_tree; ++row) { + tree_feature_rows[(size_t)row] = + b_.cache_.target_feat_tree_scratch_base + row; + } + } std::vector flat_tokens((size_t)total_tree, 0); std::vector parents((size_t)total_tree, -1); std::vector sizes((size_t)bucket, 0); @@ -438,6 +460,11 @@ std::optional Qwen35SeqEngine::step_ddtree( sizeof(int32_t) * query_slots.size()); ggml_backend_tensor_set(tree_sg.kv_write_rows, tree_rows.data(), 0, sizeof(int64_t) * tree_rows.size()); + if (direct_commit) { + ggml_backend_tensor_set( + tree_sg.target_feat_rows, tree_feature_rows.data(), 0, + sizeof(int32_t) * tree_feature_rows.size()); + } ggml_backend_tensor_set(b_.cache_.paged_kv_seq_lens, seq_lens_.data(), 0, sizeof(int32_t) * seq_lens_.size()); if (ggml_backend_graph_compute(b_.target_backend_, tree_sg.gf) != @@ -464,6 +491,258 @@ std::optional Qwen35SeqEngine::step_ddtree( } replay_total += (int)p.accepted.size(); } + if (direct_commit) { + struct DirectAppend { + Qwen35SlotManager::StepAppend append; + std::vector tokens; + }; + std::vector appends((size_t)active); + std::vector write_slots; + write_slots.reserve((size_t)active); + + for (int s = 0; s < active; ++s) { + Proposal & p = proposals[(size_t)s]; + DirectAppend & staged = appends[(size_t)s]; + staged.tokens.reserve(p.accepted.size()); + for (int dfs : p.accepted) { + staged.tokens.push_back(dfs == 0 ? p.root : + p.tree.token_ids[(size_t)dfs - 1]); + } + staged.append = slots_.append_tokens( + p.slot, staged.tokens.data(), (int)staged.tokens.size()); + const bool table_ok = staged.append.ok && + (slots_.residency_active() || + upload_block_table_delta( + p.slot, staged.append.first_new_block, + staged.append.new_blocks.data(), + staged.append.new_blocks.size())); + if (!table_ok || + staged.append.physical_rows.size() != staged.tokens.size()) { + result.error = staged.append.busy + ? "paged KV pool exhausted during DDTree direct commit" + : "DDTree direct commit K/V append failed"; + return result; + } + write_slots.push_back(p.slot); + } + + const size_t n_delta = b_.cache_.ssm_state.size(); + if (tree_sg.delta_captures.size() != n_delta || + b_.cache_.conv_state.size() != n_delta) { + result.error = "DDTree direct commit capture count mismatch"; + return result; + } + const auto to_fp32 = ggml_get_to_fp32_cuda(GGML_TYPE_F16); + if (!to_fp32) { + result.error = "DDTree direct commit has no F16 state converter"; + return result; + } + + cudaStream_t stream = nullptr; + auto copy_2d = [&](void * dst, size_t dst_pitch, + const void * src, size_t src_pitch, + size_t width, size_t height) { + return cudaMemcpy2DAsync( + dst, dst_pitch, src, src_pitch, width, height, + cudaMemcpyDeviceToDevice, stream) == cudaSuccess; + }; + auto copy_row = [&](ggml_tensor * tensor, int64_t src_row, + int64_t dst_row) { + if (!tensor || src_row < 0 || dst_row < 0 || + src_row >= tensor->ne[1] || dst_row >= tensor->ne[1]) { + return false; + } + const size_t row_bytes = tensor->nb[1]; + for (int h = 0; h < (int)tensor->ne[2]; ++h) { + const char * src = (const char *)tensor->data + + (size_t)h * tensor->nb[2] + + (size_t)src_row * tensor->nb[1]; + char * dst = (char *)tensor->data + + (size_t)h * tensor->nb[2] + + (size_t)dst_row * tensor->nb[1]; + if (cudaMemcpyAsync( + dst, src, row_bytes, cudaMemcpyDeviceToDevice, + stream) != cudaSuccess) { + return false; + } + } + return true; + }; + + const int conv_kernel = b_.w_.ssm_d_conv; + if (conv_kernel < 2) { + result.error = "DDTree direct commit has invalid conv kernel"; + return result; + } + for (size_t il = 0; il < n_delta; ++il) { + const DeltaNetCapture & cap = tree_sg.delta_captures[il]; + ggml_tensor * state = b_.cache_.ssm_state[il]; + ggml_tensor * conv = b_.cache_.conv_state[il]; + if (!cap.ssm_intermediate_states || !cap.conv_input || + !state || !conv || + cap.ssm_intermediate_states->type != GGML_TYPE_F16 || + cap.ssm_intermediate_states->ne[3] < T * bucket || + cap.conv_input->ne[2] < bucket) { + result.error = "DDTree direct commit capture layout mismatch"; + return result; + } + const int64_t state_elems = + state->ne[0] * state->ne[1] * state->ne[2]; + for (int s = 0; s < active; ++s) { + const Proposal & p = proposals[(size_t)s]; + const int deepest = p.accepted.back(); + const int capture_row = s * T + deepest; + const char * state_src = + (const char *)cap.ssm_intermediate_states->data + + (size_t)capture_row * + cap.ssm_intermediate_states->nb[3]; + float * state_dst = (float *)((char *)state->data + + (size_t)p.slot * state->nb[3]); + to_fp32(state_src, state_dst, state_elems, stream); + if (cudaPeekAtLastError() != cudaSuccess) { + result.error = "DDTree direct commit SSM conversion failed"; + return result; + } + + std::vector ancestry((size_t)conv_kernel - 1); + ancestry.back() = deepest; + for (int k = conv_kernel - 3; k >= 0; --k) { + const int next = ancestry[(size_t)k + 1]; + ancestry[(size_t)k] = + next >= 0 ? p.tree.parents[(size_t)next] : next - 1; + } + for (int k = 0; k < conv_kernel - 1; ++k) { + const int source_col = + conv_kernel - 1 + ancestry[(size_t)k]; + if (source_col < 0 || + source_col >= cap.conv_input->ne[0]) { + result.error = + "DDTree direct commit conv ancestry is invalid"; + return result; + } + const char * conv_src = + (const char *)cap.conv_input->data + + (size_t)s * cap.conv_input->nb[2] + + (size_t)source_col * + ggml_element_size(cap.conv_input); + char * conv_dst = (char *)conv->data + + (size_t)p.slot * conv->nb[2] + + (size_t)k * ggml_element_size(conv); + if (!copy_2d( + conv_dst, conv->nb[1], conv_src, + cap.conv_input->nb[1], + ggml_element_size(conv), conv->ne[1])) { + result.error = + "DDTree direct commit conv state copy failed"; + return result; + } + } + } + } + + for (int s = 0; s < active; ++s) { + const Proposal & p = proposals[(size_t)s]; + const DirectAppend & staged = appends[(size_t)s]; + for (size_t d = 0; d < p.accepted.size(); ++d) { + const int dfs = p.accepted[d]; + const int64_t src_row = + (int64_t)tree_scratch_base_ + + (int64_t)p.slot * tree_scratch_stride_ + dfs; + const int64_t dst_row = staged.append.physical_rows[d]; + for (size_t il = 0; il < b_.cache_.attn_k.size(); ++il) { + if (!copy_row(b_.cache_.attn_k[il], src_row, dst_row) || + !copy_row(b_.cache_.attn_v[il], src_row, dst_row)) { + result.error = + "DDTree direct commit paged K/V copy failed"; + return result; + } + } + + ggml_tensor * feat = b_.cache_.target_feat; + const int src_feat = + b_.cache_.target_feat_tree_scratch_base + s * T + dfs; + const int dst_feat = p.slot * b_.cache_.target_feat_cap + + (staged.append.position + (int)d) % + b_.cache_.target_feat_cap; + if (!feat || src_feat < 0 || dst_feat < 0 || + src_feat >= feat->ne[1] || dst_feat >= feat->ne[1] || + cudaMemcpyAsync( + (char *)feat->data + + (size_t)dst_feat * feat->nb[1], + (const char *)feat->data + + (size_t)src_feat * feat->nb[1], + feat->nb[1], cudaMemcpyDeviceToDevice, + stream) != cudaSuccess) { + result.error = + "DDTree direct commit target feature copy failed"; + return result; + } + } + } + + if (cudaStreamSynchronize(stream) != cudaSuccess) { + result.error = "DDTree direct commit synchronization failed"; + return result; + } + if (!commit_residency_writes(write_slots)) { + result.error = "DDTree direct commit residency write failed"; + return result; + } + for (int slot : write_slots) slots_.commit_step(slot); + if (!upload_all_active_block_tables()) { + result.error = "DDTree direct commit block-table refresh failed"; + return result; + } + + StepResult ar_result; + if (!ar_plan.decode.empty()) { + ar_result = step_regular(ar_plan); + if (!ar_result.ok()) return ar_result; + } + + for (int slot : write_slots) { + std::string reselect_error; + if (!maybe_reselect_residency(slot, reselect_error)) { + result.error = reselect_error.empty() + ? "KVFlash reselect failed" : reselect_error; + return result; + } + } + + result.decode.reserve((size_t)total_active); + for (Proposal & p : proposals) { + if (p.bonus < 0 || p.bonus >= b_.w_.n_vocab) { + result.error = + "DDTree direct commit produced an invalid pending token"; + return result; + } + DecodeOutput out; + out.slot = p.slot; + out.token = p.bonus; + out.ddtree_steps = 1; + out.ddtree_accepted_tokens = + (uint64_t)((int)p.accepted.size() - 1); + out.target_forwards = 1; + for (size_t i = 1; i < p.accepted.size(); ++i) { + const int dfs = p.accepted[i]; + out.committed_tokens.push_back( + p.tree.token_ids[(size_t)dfs - 1]); + } + attach_residency_telemetry(out); + result.decode.push_back(std::move(out)); + } + for (DecodeOutput & out : ar_result.decode) { + result.decode.push_back(std::move(out)); + } + static const bool direct_commit_diag = + std::getenv("DFLASH_DDTREE_DIRECT_COMMIT_DIAG") != nullptr; + if (direct_commit_diag) { + std::fprintf(stderr, + "[parallel-ddtree] mode=direct-tree-commit speculative=%d ar=%zu\n", + active, ar_plan.decode.size()); + } + return result; + } replay_total += (int)ar_plan.decode.size(); std::vector replay_segments; diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 1becd302c..8cb390a31 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -740,7 +740,8 @@ bool build_target_step_paged_tree( int paged_max_kv_len, int tree_scratch_base, int tree_scratch_stride, - int kq_stride_pad) { + int kq_stride_pad, + bool capture_direct_commit) { (void)kq_stride_pad; step_graph_free(sg); @@ -749,6 +750,16 @@ bool build_target_step_paged_tree( tree_scratch_base, tree_scratch_stride)) { return false; } + if (capture_direct_commit && + (cache.tree_capture_width != tree_width || + cache.tree_capture_lanes < n_tree_seqs || + cache.target_feat_tree_scratch_base <= 0 || + cache.ssm_intermediate.empty() || + cache.conv_input_cache.empty() || + !cache.ssm_intermediate.front() || + !cache.conv_input_cache.front())) { + return false; + } size_t graph_capacity = 0; if (!detail::target_paged_tree_graph_capacity( tree_width, n_tree_seqs, graph_capacity)) { @@ -787,6 +798,10 @@ bool build_target_step_paged_tree( ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); sg.kv_write_rows = ggml_new_tensor_2d( sg.ctx, GGML_TYPE_I64, n_tokens, w.n_head_kv); + if (capture_direct_commit) { + sg.target_feat_rows = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + } const struct NamedInput { ggml_tensor * tensor; @@ -805,6 +820,10 @@ bool build_target_step_paged_tree( ggml_set_name(input.tensor, input.name); ggml_set_input(input.tensor); } + if (sg.target_feat_rows) { + ggml_set_name(sg.target_feat_rows, "target_feat_rows"); + ggml_set_input(sg.target_feat_rows); + } sg.gf = ggml_new_graph_custom(sg.ctx, graph_capacity, false); QwenGraphInputs gi{}; @@ -812,8 +831,8 @@ bool build_target_step_paged_tree( gi.positions = sg.positions; gi.n_tokens = n_tokens; gi.kv_start = 0; - gi.capture_layers = false; - gi.capture_delta_intermediate = false; + gi.capture_layers = capture_direct_commit; + gi.capture_delta_intermediate = capture_direct_commit; gi.parent_ids = sg.parent_ids; gi.tree_sizes = sg.tree_sizes; gi.kv_write_rows = sg.kv_write_rows; @@ -823,6 +842,7 @@ bool build_target_step_paged_tree( gi.state_slot_ids = sg.state_slot_ids; gi.paged_query_seq_ids = sg.paged_query_seq_ids; gi.n_seqs = n_tree_seqs; + gi.target_feat_rows = sg.target_feat_rows; gi.paged_max_kv_len = paged_max_kv_len; gi.tree_width = tree_width; gi.tree_scratch_base = tree_scratch_base; @@ -832,6 +852,7 @@ bool build_target_step_paged_tree( if (!go.logits) return false; sg.logits = go.logits; ggml_set_output(sg.logits); + sg.delta_captures = std::move(go.delta_captures); sg.argmax_tokens = ggml_argmax(sg.ctx, sg.logits); ggml_set_name(sg.argmax_tokens, "paged_tree_verify_argmax"); ggml_set_output(sg.argmax_tokens); diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index cbe58b786..e209d1b77 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -220,7 +220,8 @@ bool build_target_step_paged_tree( int paged_max_kv_len, int tree_scratch_base, int tree_scratch_stride, - int kq_stride_pad = KQ_MASK_PAD); + int kq_stride_pad = KQ_MASK_PAD, + bool capture_direct_commit = false); // LM-head projection: project draft hidden states through the target output matrix. bool build_lm_head_projection_step( diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index da5cf0402..309e84442 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -173,7 +173,8 @@ static FILE * open_dflash_floor_log() { // staging K/V or staging recurrent slab to reserve. static int64_t concurrent_fixed_cache_bytes( const TargetWeights & w, int max_ctx, int n_slots, - int64_t kv_bytes_per_token, int64_t scratch_tokens) { + int64_t kv_bytes_per_token, int64_t scratch_tokens, + int tree_width, int tree_capture_lanes) { const int64_t n_full_attn = w.n_layer / w.full_attention_interval; const int64_t n_delta = w.n_layer - n_full_attn; @@ -186,9 +187,18 @@ static int64_t concurrent_fixed_cache_bytes( (int64_t)sizeof(float); const int64_t recurrent = state_per_layer * n_delta * (int64_t)n_slots; + const int64_t tree_checkpoint = + tree_capture_lanes > 0 + ? n_delta * tree_capture_lanes * + ((head_v_dim * head_v_dim * w.ssm_dt_rank * + (int64_t)tree_width * (int64_t)sizeof(uint16_t)) + + (((int64_t)w.ssm_d_conv - 1 + tree_width) * conv_ch * + (int64_t)sizeof(float))) + : 0; const int64_t target_feat = (int64_t)w.n_capture_layers * w.n_embd * - ((int64_t)std::min(max_ctx, 4096) * n_slots + 1) * + ((int64_t)std::min(max_ctx, 4096) * n_slots + 1 + + (int64_t)tree_width * tree_capture_lanes) * (int64_t)sizeof(uint16_t); const int64_t q_capture = (int64_t)w.n_embd_head_k * w.n_head * n_full_attn * @@ -198,7 +208,7 @@ static int64_t concurrent_fixed_cache_bytes( (int64_t)sizeof(int32_t); const int64_t scratch = kv_bytes_per_token * scratch_tokens; - return recurrent + target_feat + q_capture + + return recurrent + tree_checkpoint + target_feat + q_capture + paged_metadata + scratch; } } // namespace @@ -412,6 +422,15 @@ bool Qwen35Backend::init() { target_backend_ == draft_backend_; const int tree_width = concurrent_local_ddtree ? cfg_.ddtree_budget + 1 : 0; + const int tree_capture_lanes = concurrent_local_ddtree + ? std::clamp( + env_int_or_default( + "DFLASH_DDTREE_DIRECT_COMMIT_MAX_REQUESTS", 3), + 0, n_slots) + : 0; + if (tree_capture_lanes > 0) { + std::fprintf(stderr, "[parallel-ddtree] direct commit lanes=%d\n", tree_capture_lanes); + } const int tree_stride = concurrent_local_ddtree ? paged_token_capacity(tree_width) : 0; const int64_t concurrent_scratch_tokens = @@ -446,7 +465,8 @@ bool Qwen35Backend::init() { budget.reserve_bytes = kvf_budget.reserve_bytes; budget.fixed_cache_bytes = concurrent_fixed_cache_bytes( w_, cfg_.device.max_ctx, n_slots, budget.bytes_per_token, - concurrent_scratch_tokens); + concurrent_scratch_tokens, tree_width, + tree_capture_lanes); pool_tokens = paged_kv_auto_pool_tokens( cfg_.device.max_ctx, n_slots, budget); const int64_t one_context = @@ -480,7 +500,8 @@ bool Qwen35Backend::init() { : kvflash_tokens_); if (!create_target_cache(w_, cfg_.device.max_ctx, max_verify_tokens, target_backend_, cache_, /*prefill_only=*/true, ctx_alloc, - cfg_.paged_attention, n_slots, concurrent_local_ddtree)) { + cfg_.paged_attention, n_slots, concurrent_local_ddtree, + tree_capture_lanes)) { std::fprintf(stderr, "cache: %s\n", dflash27b_last_error()); return false; } @@ -552,7 +573,7 @@ bool Qwen35Backend::init() { const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); std::fprintf(stderr, "[parallel-ddtree] enabled budget=%d width=%d " - "mode=packed-verify-mixed-replay adaptive=%s " + "mode=bounded-direct-commit+mixed-replay adaptive=%s " "policy=ranked-goodput scope=all-concurrency\n", cfg_.ddtree_budget, tree_width, adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 5f10dcdc9..ca7c3a3b2 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -81,13 +81,15 @@ bool create_target_cache(const TargetWeights & w, int ctx_alloc, bool paged_attention, int n_seq_slots, - bool concurrent_tree) { + bool concurrent_tree, + int concurrent_tree_capture_lanes) { return create_target_cache_partial(w, max_ctx, max_verify_tokens, backend, out, prefill_only, 0, w.n_layer, true, ctx_alloc, /*f32_ssm_intermediates=*/false, paged_attention, n_seq_slots, - concurrent_tree); + concurrent_tree, + concurrent_tree_capture_lanes); } // concurrent_fixed_cache_bytes() in qwen35_backend.cpp mirrors this @@ -106,7 +108,8 @@ bool create_target_cache_partial(const TargetWeights & w, bool f32_ssm_intermediates, bool paged_attention, int n_seq_slots, - bool concurrent_tree) { + bool concurrent_tree, + int concurrent_tree_capture_lanes) { if (layer_begin < 0) layer_begin = 0; if (layer_end < 0 || layer_end > w.n_layer) layer_end = w.n_layer; if (layer_begin > layer_end) { @@ -123,6 +126,9 @@ bool create_target_cache_partial(const TargetWeights & w, "concurrent tree cache requires paged multi-slot serving"); return false; } + if (!concurrent_tree) concurrent_tree_capture_lanes = 0; + concurrent_tree_capture_lanes = std::clamp( + concurrent_tree_capture_lanes, 0, n_seq_slots); out.backend = backend; out.max_ctx = max_ctx; out.cur_pos = 0; @@ -131,6 +137,10 @@ bool create_target_cache_partial(const TargetWeights & w, max_verify_tokens = DFLASH27B_DRAFT_BLOCK_SIZE; } + out.tree_capture_width = concurrent_tree_capture_lanes > 0 + ? max_verify_tokens : 0; + out.tree_capture_lanes = concurrent_tree_capture_lanes; + const int n_full_attn = w.n_layer / w.full_attention_interval; // 16 const int n_delta = w.n_layer - n_full_attn; // 48 const int head_dim = w.n_embd_head_k; @@ -236,11 +246,16 @@ bool create_target_cache_partial(const TargetWeights & w, if (allocate_target_feat) { const int fc_in = w.n_capture_layers * w.n_embd; // Concurrent slots own disjoint feature rings. The final row is - // dead scratch for padded bucket rows because set_rows does not - // accept negative destination indices. - const int feat_rows = multi_slot + // dead scratch for padded bucket rows. Direct-tree candidates use + // a separate compact scratch domain after the durable rings. + const int durable_feat_rows = multi_slot ? out.target_feat_cap * n_seq_slots + 1 : out.target_feat_cap; + const int tree_scratch_rows = + out.tree_capture_width * out.tree_capture_lanes; + const int feat_rows = durable_feat_rows + tree_scratch_rows; + out.target_feat_tree_scratch_base = tree_scratch_rows > 0 + ? durable_feat_rows : 0; out.target_feat = ggml_new_tensor_2d( out.base_ctx, GGML_TYPE_BF16, fc_in, feat_rows); ggml_set_name(out.target_feat, "target_feat"); @@ -281,12 +296,13 @@ bool create_target_cache_partial(const TargetWeights & w, } // ── Rollback context: snapshots + intermediates ─────────────────── - // Multi-slot caches skip these entirely. Packed tree verification gathers - // the selected slots' base state without mutating it, then a bounded - // replay commits accepted paths. T*S recurrent captures would be tens of - // GiB at useful Strix concurrency and are intentionally not allocated. - if (!prefill_only && !multi_slot) { - const int rb_tensors = 4 * n_delta; + // Concurrent trees reserve only a bounded compact set of F16 captures; + // physical serving slots that are not selected consume no checkpoint + // memory, and routes wider than the bound use accepted-path replay. + const bool concurrent_capture = multi_slot && + out.tree_capture_lanes > 0 && out.tree_capture_width > 0; + if ((!prefill_only && !multi_slot) || concurrent_capture) { + const int rb_tensors = (multi_slot ? 2 : 4) * n_delta; ggml_init_params ip{}; ip.mem_size = (size_t)(rb_tensors + 16) * ggml_tensor_overhead(); ip.mem_buffer = nullptr; @@ -299,26 +315,32 @@ bool create_target_cache_partial(const TargetWeights & w, if (((il + 1) % w.full_attention_interval) != 0) { const bool owns_layer = il >= layer_begin && il < layer_end; if (!owns_layer) { dn_idx++; continue; } - ggml_tensor * Sn = ggml_new_tensor_3d(out.rollback_ctx, GGML_TYPE_F32, - head_v_dim, head_v_dim, w.ssm_dt_rank); - ggml_tensor * Cn = ggml_new_tensor_2d(out.rollback_ctx, GGML_TYPE_F32, - w.ssm_d_conv - 1, conv_ch); + ggml_tensor * Sn = multi_slot ? nullptr : + ggml_new_tensor_3d(out.rollback_ctx, GGML_TYPE_F32, + head_v_dim, head_v_dim, w.ssm_dt_rank); + ggml_tensor * Cn = multi_slot ? nullptr : + ggml_new_tensor_2d(out.rollback_ctx, GGML_TYPE_F32, + w.ssm_d_conv - 1, conv_ch); // I0 domain: ne[3] is the root-inclusive flat verify-token // domain. Tree capture writes t=0 synthetic root through the // final/padded flat slot directly into slot t. - const ggml_type ssm_intermediate_type = f32_ssm_intermediates - ? GGML_TYPE_F32 : GGML_TYPE_Q8_0; + const ggml_type ssm_intermediate_type = multi_slot + ? GGML_TYPE_F16 + : (f32_ssm_intermediates ? GGML_TYPE_F32 : GGML_TYPE_Q8_0); ggml_tensor * Si = ggml_new_tensor_4d(out.rollback_ctx, ssm_intermediate_type, head_v_dim, head_v_dim, - w.ssm_dt_rank, max_verify_tokens); + w.ssm_dt_rank, + max_verify_tokens * + (multi_slot ? out.tree_capture_lanes : 1)); // I0 domain: ne[0] is [K_conv-1 prefix rows | // root-inclusive verify rows]. ggml_tensor * Ci = ggml_new_tensor_3d(out.rollback_ctx, GGML_TYPE_F32, (w.ssm_d_conv - 1) + max_verify_tokens, - conv_ch, 1); + conv_ch, + multi_slot ? out.tree_capture_lanes : 1); char name[64]; - std::snprintf(name, sizeof(name), "ssm_state_snap_%d", il); ggml_set_name(Sn, name); - std::snprintf(name, sizeof(name), "conv_state_snap_%d", il); ggml_set_name(Cn, name); + if (Sn) { std::snprintf(name, sizeof(name), "ssm_state_snap_%d", il); ggml_set_name(Sn, name); } + if (Cn) { std::snprintf(name, sizeof(name), "conv_state_snap_%d", il); ggml_set_name(Cn, name); } std::snprintf(name, sizeof(name), "ssm_intermediate_%d", il); ggml_set_name(Si, name); std::snprintf(name, sizeof(name), "conv_input_cache_%d", il); ggml_set_name(Ci, name); out.ssm_state_snap[dn_idx] = Sn; @@ -401,6 +423,9 @@ void free_target_cache(TargetCache & c) { c.ssm_intermediate.clear(); c.conv_input_cache.clear(); c.target_feat = nullptr; + c.tree_capture_width = 0; + c.tree_capture_lanes = 0; + c.target_feat_tree_scratch_base = 0; c.q_cap = nullptr; c.cur_pos = 0; } @@ -1095,7 +1120,7 @@ static ggml_tensor * build_delta_net_block( } GGML_ASSERT((active_slot_ids == nullptr) == (state_slot_ids == nullptr)); const bool mapped_tree = active_slot_ids && parent_ids; - GGML_ASSERT(!active_slot_ids || !cap); + GGML_ASSERT(!cap || !active_slot_ids || mapped_tree); GGML_ASSERT(!active_slot_ids || (mapped_tree ? (!ragged && prefill_total == 0 && @@ -1239,11 +1264,12 @@ static ggml_tensor * build_delta_net_block( // Copy into a matching-sized view of the cache destination. const int64_t ci_len = conv_input->ne[0]; ggml_tensor * dst; - if (ci_len == cap->conv_input->ne[0]) { + if (ci_len == cap->conv_input->ne[0] && + seg_seqs == cap->conv_input->ne[2]) { dst = cap->conv_input; } else { dst = ggml_view_3d(ctx, cap->conv_input, - ci_len, cap->conv_input->ne[1], cap->conv_input->ne[2], + ci_len, cap->conv_input->ne[1], seg_seqs, cap->conv_input->nb[1], cap->conv_input->nb[2], 0); } GGML_ASSERT(ggml_nelements(conv_input) == ggml_nelements(dst)); From 9c42275c4318a1aa85dfd18d74ba73e3e400c315 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 17 Aug 2026 12:08:41 +0000 Subject: [PATCH 11/18] perf(qwen35): verify mixed DDTree and AR in one pass --- .../concurrency/adaptive_verification.h | 9 +- server/src/common/step_graph.h | 8 + server/src/internal.h | 9 + .../qwen35/concurrency/qwen35_seq_engine.cpp | 137 ++++++++++++--- server/src/qwen35/graph_builders.cpp | 35 +++- server/src/qwen35/graph_builders.h | 25 +-- server/src/qwen35/qwen35_backend.cpp | 2 +- server/src/qwen35/qwen35_target_graph.cpp | 157 ++++++++++++------ server/test/test_recurrent_snapshot.cpp | 6 + server/test/test_speculation_goodput.cpp | 21 +++ 10 files changed, 319 insertions(+), 90 deletions(-) diff --git a/server/src/common/concurrency/adaptive_verification.h b/server/src/common/concurrency/adaptive_verification.h index f0c68f3e7..0de13f04a 100644 --- a/server/src/common/concurrency/adaptive_verification.h +++ b/server/src/common/concurrency/adaptive_verification.h @@ -146,8 +146,9 @@ class AdaptiveVerificationRanker { AdaptiveVerificationDecision select( int active_requests, - const std::vector & candidates) - const { + const std::vector & candidates, + int max_speculative_requests = + std::numeric_limits::max()) const { AdaptiveVerificationDecision out; if (active_requests <= 0 || candidates.empty() || !has_autoregressive_cost(active_requests)) { @@ -202,8 +203,10 @@ class AdaptiveVerificationRanker { double best = baseline; double admitted_goodput = baseline; double expected_total = static_cast(active_requests); + const int prefix_limit = std::max(0, std::min( + active_requests, max_speculative_requests)); for (int k = 1; k <= static_cast(known.size()) && - k <= active_requests; ++k) { + k <= prefix_limit; ++k) { expected_total += known[(size_t)k - 1].expected_tokens - 1.0; if (!has_route_cost(active_requests, k)) { missing_cost_prefix = k; diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h index 5e1814a72..69b5ff624 100644 --- a/server/src/common/step_graph.h +++ b/server/src/common/step_graph.h @@ -59,6 +59,12 @@ struct StepGraph { // Recurrent gather rows. Unlike active/paged IDs, padding must name a // valid harmless slot (normally 0): ggml_get_rows does not mask -1. ggml_tensor * state_slot_ids = nullptr; + // Mixed DDTree+AR target graph: compact AR rows use their own recurrent + // mapping after the tree rows. Keeping these separate lets the graph + // share projections/FFN while the small conv/GDN cores retain their + // proven uniform [timesteps,sequences] layouts. + ggml_tensor * ar_active_slot_ids = nullptr; + ggml_tensor * ar_state_slot_ids = nullptr; // Ragged paged read (concurrent prefill): per-row block-table column and // inclusive causal position, [n_tokens] i32 each. Padding rows carry -1. ggml_tensor * paged_query_seq_ids = nullptr; @@ -102,6 +108,8 @@ inline void step_graph_free(StepGraph & sg) { sg.kv_write_rows = nullptr; sg.active_slot_ids = nullptr; sg.state_slot_ids = nullptr; + sg.ar_active_slot_ids = nullptr; + sg.ar_state_slot_ids = nullptr; sg.paged_query_seq_ids = nullptr; sg.paged_query_positions = nullptr; sg.target_feat_rows = nullptr; diff --git a/server/src/internal.h b/server/src/internal.h index 45334fd44..0a4d0cc17 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -698,6 +698,12 @@ struct QwenGraphInputs { // each row's KV extent to position+1, which IS the causal mask. -1 on // padding rows. ggml_tensor * paged_query_positions = nullptr; + // Mixed packed-tree verification may append compact one-token AR rows + // after the sequence-major tree rows. These mappings name the physical + // recurrent slabs for that AR suffix; tree rows keep using + // active_slot_ids/state_slot_ids above. + ggml_tensor * ar_active_slot_ids = nullptr; + ggml_tensor * ar_state_slot_ids = nullptr; // Optional [n_rows] i32 gather of final-norm rows before the LM head: // multi-prompt steps sample scattered rows (each committing segment's // last row plus the decode rows), which a tail view cannot express. @@ -738,6 +744,9 @@ struct QwenGraphInputs { // Packed steps use logits_row_indices for scattered committing rows and // compact decode rows; logits_tail_rows remains the dense-path fallback. int n_seqs = 1; + // Number of compact one-token AR rows appended after a packed tree. + // Zero preserves the established pure-tree/prefill/decode layouts. + int n_ar_seqs = 0; int seq_slot = 0; int paged_max_kv_len = 0; int n_prefill_tokens = 0; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 09de57004..dd6afe3d9 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -368,30 +368,64 @@ std::optional Qwen35SeqEngine::step_ddtree( !target_is_meta && bucket <= b_.cache_.tree_capture_lanes && T == b_.cache_.tree_capture_width && b_.cache_.target_feat_tree_scratch_base > 0; + struct DirectArStage { + Qwen35SlotManager::StepAppend append; + }; + std::vector direct_ar; + if (direct_commit) { + direct_ar.resize(ar_plan.decode.size()); + for (size_t i = 0; i < ar_plan.decode.size(); ++i) { + const StepInput & in = ar_plan.decode[i]; + DirectArStage & staged = direct_ar[i]; + staged.append = slots_.append_token(in.slot, in.token); + const bool table_ok = staged.append.ok && + (slots_.residency_active() || + upload_block_table_delta( + in.slot, staged.append.first_new_block, + staged.append.new_blocks.data(), + staged.append.new_blocks.size())); + if (!table_ok || staged.append.physical_rows.size() != 1) { + result.error = staged.append.busy + ? "paged KV pool exhausted during mixed DDTree/AR step" + : "mixed DDTree/AR append failed"; + return result; + } + } + if (!direct_ar.empty() && !upload_all_active_block_tables()) { + result.error = + "mixed DDTree/AR block-table refresh failed"; + return result; + } + } StepGraph & tree_sg = b_.sg_; int max_prefix = 1; for (const Proposal & p : proposals) { max_prefix = std::max(max_prefix, slots_.slot(p.slot).cur_pos); } + for (const DirectArStage & staged : direct_ar) { + max_prefix = std::max(max_prefix, staged.append.position + 1); + } + const int n_ar = (int)direct_ar.size(); if (!build_target_step_paged_tree( tree_sg, b_.w_, b_.cache_, b_.target_backend_, T, bucket, max_prefix, tree_scratch_base_, tree_scratch_stride_, - b_.cfg_.kq_stride_pad, direct_commit)) { + b_.cfg_.kq_stride_pad, direct_commit, n_ar)) { result.error = "packed DDTree verify graph build failed"; return result; } const int total_tree = T * bucket; + const int total_packed = total_tree + n_ar; std::vector tree_feature_rows; if (direct_commit) { - tree_feature_rows.resize((size_t)total_tree); + tree_feature_rows.resize((size_t)total_packed); for (int row = 0; row < total_tree; ++row) { tree_feature_rows[(size_t)row] = b_.cache_.target_feat_tree_scratch_base + row; } } - std::vector flat_tokens((size_t)total_tree, 0); + std::vector flat_tokens((size_t)total_packed, 0); std::vector parents((size_t)total_tree, -1); std::vector sizes((size_t)bucket, 0); // Negative active IDs identify bucket padding. Recurrent gathers cannot @@ -400,11 +434,14 @@ std::optional Qwen35SeqEngine::step_ddtree( // their attention/output rows inactive and tree mode never persists it. std::vector tree_slots((size_t)bucket, -1); std::vector tree_state_slots((size_t)bucket, 0); - std::vector query_slots((size_t)total_tree, -1); + std::vector ar_slots((size_t)n_ar, -1); + std::vector ar_state_slots((size_t)n_ar, 0); + std::vector query_slots((size_t)total_packed, -1); + std::vector query_positions((size_t)total_packed, -1); std::vector tree_rows( - (size_t)total_tree * n_head_kv, scratch_row_); - std::vector tree_pos((size_t)4 * total_tree, 0); - std::vector tree_embed((size_t)hidden * total_tree, 0.0f); + (size_t)total_packed * n_head_kv, scratch_row_); + std::vector tree_pos((size_t)4 * total_packed, 0); + std::vector tree_embed((size_t)hidden * total_packed, 0.0f); seq_lens_.assign((size_t)n_slots, 0); for (int s = 0; s < active; ++s) { @@ -423,18 +460,40 @@ std::optional Qwen35SeqEngine::step_ddtree( const int depth = node == 0 ? 0 : p.tree.depths[(size_t)node - 1]; const int pos = slots_.slot(p.slot).cur_pos + depth; - tree_pos[(size_t)0 * total_tree + row] = pos; - tree_pos[(size_t)1 * total_tree + row] = pos; - tree_pos[(size_t)2 * total_tree + row] = pos; + tree_pos[(size_t)0 * total_packed + row] = pos; + tree_pos[(size_t)1 * total_packed + row] = pos; + tree_pos[(size_t)2 * total_packed + row] = pos; for (int h = 0; h < n_head_kv; ++h) { - tree_rows[(size_t)h * total_tree + row] = + tree_rows[(size_t)h * total_packed + row] = (int64_t)tree_scratch_base_ + (int64_t)p.slot * tree_scratch_stride_ + node; } } } + const int feature_cap = b_.cache_.target_feat_cap; + for (int i = 0; i < n_ar; ++i) { + const StepInput & in = ar_plan.decode[(size_t)i]; + const Qwen35SlotManager::StepAppend & app = + direct_ar[(size_t)i].append; + const int row = total_tree + i; + flat_tokens[(size_t)row] = in.token; + ar_slots[(size_t)i] = in.slot; + ar_state_slots[(size_t)i] = in.slot; + query_slots[(size_t)row] = in.slot; + query_positions[(size_t)row] = app.position; + seq_lens_[(size_t)in.slot] = app.position + 1; + tree_pos[(size_t)0 * total_packed + row] = app.position; + tree_pos[(size_t)1 * total_packed + row] = app.position; + tree_pos[(size_t)2 * total_packed + row] = app.position; + tree_feature_rows[(size_t)row] = + in.slot * feature_cap + app.position % feature_cap; + for (int h = 0; h < n_head_kv; ++h) { + tree_rows[(size_t)h * total_packed + row] = + app.physical_rows[0]; + } + } if (!b_.w_.embedder.embed( - flat_tokens.data(), total_tree, tree_embed.data())) { + flat_tokens.data(), total_packed, tree_embed.data())) { result.error = "packed DDTree embedding failed"; return result; } @@ -456,6 +515,17 @@ std::optional Qwen35SeqEngine::step_ddtree( } ggml_backend_tensor_set(tree_sg.state_slot_ids, tree_state_slots.data(), 0, sizeof(int32_t) * tree_state_slots.size()); + if (n_ar > 0) { + ggml_backend_tensor_set( + tree_sg.ar_active_slot_ids, ar_slots.data(), 0, + sizeof(int32_t) * ar_slots.size()); + ggml_backend_tensor_set( + tree_sg.ar_state_slot_ids, ar_state_slots.data(), 0, + sizeof(int32_t) * ar_state_slots.size()); + ggml_backend_tensor_set( + tree_sg.paged_query_positions, query_positions.data(), 0, + sizeof(int32_t) * query_positions.size()); + } ggml_backend_tensor_set(tree_sg.paged_query_seq_ids, query_slots.data(), 0, sizeof(int32_t) * query_slots.size()); ggml_backend_tensor_set(tree_sg.kv_write_rows, tree_rows.data(), 0, @@ -472,9 +542,21 @@ std::optional Qwen35SeqEngine::step_ddtree( result.error = "packed DDTree verify compute failed"; return result; } - std::vector posterior((size_t)total_tree, -1); + std::vector posterior((size_t)total_packed, -1); ggml_backend_tensor_get(tree_sg.argmax_tokens, posterior.data(), 0, sizeof(int32_t) * posterior.size()); + std::vector direct_ar_next((size_t)n_ar, -1); + for (int i = 0; i < n_ar; ++i) { + const int row = total_tree + i; + direct_ar_next[(size_t)i] = sample_graph_row( + ar_plan.decode[(size_t)i].slot, row, + &posterior[(size_t)row], &logits_buf_); + if (direct_ar_next[(size_t)i] < 0) { + result.error = + "mixed DDTree/AR graph produced an invalid AR token"; + return result; + } + } int replay_total = 0; for (int s = 0; s < active; ++s) { @@ -498,7 +580,10 @@ std::optional Qwen35SeqEngine::step_ddtree( }; std::vector appends((size_t)active); std::vector write_slots; - write_slots.reserve((size_t)active); + write_slots.reserve((size_t)total_active); + for (const StepInput & in : ar_plan.decode) { + write_slots.push_back(in.slot); + } for (int s = 0; s < active; ++s) { Proposal & p = proposals[(size_t)s]; @@ -694,11 +779,6 @@ std::optional Qwen35SeqEngine::step_ddtree( return result; } - StepResult ar_result; - if (!ar_plan.decode.empty()) { - ar_result = step_regular(ar_plan); - if (!ar_result.ok()) return ar_result; - } for (int slot : write_slots) { std::string reselect_error; @@ -731,14 +811,19 @@ std::optional Qwen35SeqEngine::step_ddtree( attach_residency_telemetry(out); result.decode.push_back(std::move(out)); } - for (DecodeOutput & out : ar_result.decode) { + for (size_t i = 0; i < ar_plan.decode.size(); ++i) { + DecodeOutput out; + out.slot = ar_plan.decode[i].slot; + out.token = direct_ar_next[i]; + out.target_forwards = 1; + attach_residency_telemetry(out); result.decode.push_back(std::move(out)); } static const bool direct_commit_diag = std::getenv("DFLASH_DDTREE_DIRECT_COMMIT_DIAG") != nullptr; if (direct_commit_diag) { std::fprintf(stderr, - "[parallel-ddtree] mode=direct-tree-commit speculative=%d ar=%zu\n", + "[parallel-ddtree] mode=one-pass-tree-ar speculative=%d ar=%zu\n", active, ar_plan.decode.size()); } return result; @@ -1241,9 +1326,14 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { collect_candidates(); AdaptiveVerificationDecision decision; + const int adaptive_speculation_limit = + b_.cache_.tree_capture_lanes > 0 + ? b_.cache_.tree_capture_lanes + : static_cast(inputs.size()); if (adaptive_enabled) { decision = adaptive_verification_.select( - static_cast(inputs.size()), candidates); + static_cast(inputs.size()), candidates, + adaptive_speculation_limit); if (inputs.size() <= 3) { adaptive_calibration_cooldown_ = 0; } else if (adaptive_calibration_cooldown_ > 0) { @@ -1283,7 +1373,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { input->slot, *expected); collect_candidates(); decision = adaptive_verification_.select( - static_cast(inputs.size()), candidates); + static_cast(inputs.size()), candidates, + adaptive_speculation_limit); } } } diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 8cb390a31..dc2796770 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -741,13 +741,15 @@ bool build_target_step_paged_tree( int tree_scratch_base, int tree_scratch_stride, int kq_stride_pad, - bool capture_direct_commit) { + bool capture_direct_commit, + int n_ar_seqs) { (void)kq_stride_pad; step_graph_free(sg); if (!detail::validate_target_paged_tree_layout( cache, tree_width, n_tree_seqs, paged_max_kv_len, - tree_scratch_base, tree_scratch_stride)) { + tree_scratch_base, tree_scratch_stride) || + n_ar_seqs < 0 || n_ar_seqs > cache.n_seq_slots) { return false; } if (capture_direct_commit && @@ -762,10 +764,12 @@ bool build_target_step_paged_tree( } size_t graph_capacity = 0; if (!detail::target_paged_tree_graph_capacity( - tree_width, n_tree_seqs, graph_capacity)) { + tree_width, n_tree_seqs, graph_capacity) || + !detail::target_graph_capacity_for_parallel_segments( + n_tree_seqs + n_ar_seqs, graph_capacity)) { return false; } - const int n_tokens = tree_width * n_tree_seqs; + const int n_tokens = tree_width * n_tree_seqs + n_ar_seqs; ggml_init_params ip{}; ip.mem_size = 512 * 1024 * 1024; @@ -778,7 +782,7 @@ bool build_target_step_paged_tree( // Salt graph addresses by the stable bucket shape so captured graphs for // different T/S buckets never alias in ggml-cuda's topology cache. - for (int i = 0; i < tree_width + n_tree_seqs; ++i) { + for (int i = 0; i < tree_width + n_tree_seqs + n_ar_seqs; ++i) { (void)ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, 1); } @@ -794,6 +798,14 @@ bool build_target_step_paged_tree( ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tree_seqs); sg.state_slot_ids = ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tree_seqs); + if (n_ar_seqs > 0) { + sg.ar_active_slot_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_ar_seqs); + sg.ar_state_slot_ids = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_ar_seqs); + sg.paged_query_positions = + ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); + } sg.paged_query_seq_ids = ggml_new_tensor_1d(sg.ctx, GGML_TYPE_I32, n_tokens); sg.kv_write_rows = ggml_new_tensor_2d( @@ -820,6 +832,15 @@ bool build_target_step_paged_tree( ggml_set_name(input.tensor, input.name); ggml_set_input(input.tensor); } + if (n_ar_seqs > 0) { + ggml_set_name(sg.ar_active_slot_ids, "ar_active_slot_ids"); + ggml_set_input(sg.ar_active_slot_ids); + ggml_set_name(sg.ar_state_slot_ids, "ar_state_slot_ids"); + ggml_set_input(sg.ar_state_slot_ids); + ggml_set_name( + sg.paged_query_positions, "paged_query_positions"); + ggml_set_input(sg.paged_query_positions); + } if (sg.target_feat_rows) { ggml_set_name(sg.target_feat_rows, "target_feat_rows"); ggml_set_input(sg.target_feat_rows); @@ -840,7 +861,11 @@ bool build_target_step_paged_tree( gi.paged_kv_seq_lens = cache.paged_kv_seq_lens; gi.active_slot_ids = sg.active_slot_ids; gi.state_slot_ids = sg.state_slot_ids; + gi.ar_active_slot_ids = sg.ar_active_slot_ids; + gi.ar_state_slot_ids = sg.ar_state_slot_ids; + gi.n_ar_seqs = n_ar_seqs; gi.paged_query_seq_ids = sg.paged_query_seq_ids; + gi.paged_query_positions = sg.paged_query_positions; gi.n_seqs = n_tree_seqs; gi.target_feat_rows = sg.target_feat_rows; gi.paged_max_kv_len = paged_max_kv_len; diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index e209d1b77..65018570c 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -60,7 +60,11 @@ inline bool target_paged_tree_uploads_ready(const StepGraph & sg) { const auto allocated = [](const ggml_tensor * tensor) { return tensor && tensor->buffer; }; - return sg.active_slot_ids && + const bool ar_ready = !sg.ar_active_slot_ids || + (allocated(sg.ar_active_slot_ids) && + allocated(sg.ar_state_slot_ids) && + allocated(sg.paged_query_positions)); + return ar_ready && sg.active_slot_ids && allocated(sg.inp_embed) && allocated(sg.positions) && allocated(sg.parent_ids) && allocated(sg.tree_sizes) && allocated(sg.state_slot_ids) && @@ -202,14 +206,14 @@ bool build_target_step_tree( int fa_window = 0, int kq_stride_pad = KQ_MASK_PAD); -// Packed concurrent DDTree verify over a paged multi-slot cache. Tokens are -// flattened sequence-major as [tree_width*n_tree_seqs]. n_tree_seqs is a -// stable graph-bucket width; inactive trees use tree_size=0 and dead/safe row -// mappings. In particular state_slot_ids padding must map to a valid harmless -// slot (normally 0), while active/paged sequence IDs may use -1. The graph -// writes candidate K/V into per-slot scratch slabs but -// does not mutate persistent recurrent state or target features. Accepted -// paths are committed by a later row-indexed replay through build_target_step. +// Packed concurrent DDTree verify over a paged multi-slot cache. Tree tokens +// are flattened sequence-major as [tree_width*n_tree_seqs]. An optional +// compact [n_ar_seqs] suffix advances ordinary AR requests in the same model +// graph. Inactive trees use tree_size=0 and dead/safe row mappings. Tree rows +// write candidate K/V and features into scratch and do not mutate recurrent +// state; AR rows update their mapped recurrent/KV/feature state in place. +// Accepted tree paths are copied from the captured scratch rows after the +// shared forward completes. bool build_target_step_paged_tree( StepGraph & sg, const TargetWeights & w, @@ -221,7 +225,8 @@ bool build_target_step_paged_tree( int tree_scratch_base, int tree_scratch_stride, int kq_stride_pad = KQ_MASK_PAD, - bool capture_direct_commit = false); + bool capture_direct_commit = false, + int n_ar_seqs = 0); // LM-head projection: project draft hidden states through the target output matrix. bool build_lm_head_projection_step( diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 309e84442..36a30e511 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -573,7 +573,7 @@ bool Qwen35Backend::init() { const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); std::fprintf(stderr, "[parallel-ddtree] enabled budget=%d width=%d " - "mode=bounded-direct-commit+mixed-replay adaptive=%s " + "mode=one-pass-tree-ar+bounded-replay adaptive=%s " "policy=ranked-goodput scope=all-concurrency\n", cfg_.ddtree_budget, tree_width, adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index ca7c3a3b2..5065ea225 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -864,10 +864,15 @@ static ggml_tensor * build_full_attn_block( GGML_ASSERT((paged_tree_parent_ids == nullptr) == (paged_tree_sizes == nullptr)); const bool ragged = paged_query_seq_ids != nullptr; + const int tree_rows = paged_tree + ? tree_width * (int)paged_tree_sizes->ne[0] : 0; + const bool mixed_tree_ar = paged_tree && tree_rows < n_tokens; GGML_ASSERT(!ragged || (paged_block_table && kv_write_rows)); GGML_ASSERT(!ragged || paged_tree || paged_query_positions); GGML_ASSERT(!paged_tree || - (ragged && !paged_query_positions && tree_width > 0)); + (ragged && tree_width > 0 && tree_rows > 0 && + tree_rows <= n_tokens && + (!mixed_tree_ar || paged_query_positions))); if (kv_write_rows) { // Step-invariant: the destination tensor stays fixed while the input // indices carry contiguous, KVFlash, or paged physical rows. @@ -932,7 +937,8 @@ static ggml_tensor * build_full_attn_block( auto paged_read = [&](ggml_tensor * q, int launch_kv_len, ggml_tensor * row_seq_ids, ggml_tensor * row_positions, - bool dense_token_layout) { + bool dense_token_layout, + bool tree_read) { // max_kv_seq_len sizes the logical partition grid. In bounded // KVFlash mode the block table maps that logical range onto a much // smaller physical K/V pool, so cache_k->ne[1] is not a valid clamp. @@ -956,8 +962,11 @@ static ggml_tensor * build_full_attn_block( ctx, q, cache_k, cache_v, paged_block_table, paged_kv_seq_lens, row_seq_ids, row_positions, kq_scale, PAGED_BLOCK_SIZE, launch_len, - paged_tree_parent_ids, paged_tree_sizes, - tree_width, tree_scratch_base, tree_scratch_stride); + tree_read ? paged_tree_parent_ids : nullptr, + tree_read ? paged_tree_sizes : nullptr, + tree_read ? tree_width : 0, + tree_read ? tree_scratch_base : 0, + tree_read ? tree_scratch_stride : 0); if (dense_token_layout) { out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3)); } @@ -966,17 +975,42 @@ static ggml_tensor * build_full_attn_block( ggml_tensor * attn = nullptr; if (paged_tree) { - // ── Packed concurrent tree verify. Every query row selects its - // physical sequence/scratch slab. The paged kernel combines the - // committed block-table prefix with only this node's ancestor chain; - // query_positions is intentionally absent in tree mode. - ggml_tensor * Qfa = q_segment(0, n_tokens); - if (q_fa_out) *q_fa_out = Qfa; + // Tree rows keep ancestor-chain visibility, while an optional compact + // AR suffix uses ordinary per-row causal positions. Both share QKV and + // output projections; only the small paged-attention calls split. const int launch_kv_len = paged_max_kv_len > 0 ? paged_max_kv_len : kv_start + n_tokens; - attn = paged_read(Qfa, launch_kv_len, - paged_query_seq_ids, /*row_positions=*/nullptr, - /*dense_token_layout=*/n_tokens > 1); + ggml_tensor * Qtree = q_segment(0, tree_rows); + ggml_tensor * tree_seq_ids = mixed_tree_ar + ? ggml_view_1d( + ctx, paged_query_seq_ids, tree_rows, 0) + : paged_query_seq_ids; + ggml_tensor * tree_attn = paged_read( + Qtree, launch_kv_len, tree_seq_ids, + /*row_positions=*/nullptr, + /*dense_token_layout=*/true, /*tree_read=*/true); + if (mixed_tree_ar) { + const int n_ar_rows = n_tokens - tree_rows; + ggml_tensor * Qar = q_segment(tree_rows, n_ar_rows); + ggml_tensor * ar_seq_ids = ggml_view_1d( + ctx, paged_query_seq_ids, n_ar_rows, + (size_t)tree_rows * + ggml_element_size(paged_query_seq_ids)); + ggml_tensor * ar_positions = ggml_view_1d( + ctx, paged_query_positions, n_ar_rows, + (size_t)tree_rows * + ggml_element_size(paged_query_positions)); + ggml_tensor * ar_attn = paged_read( + Qar, launch_kv_len, ar_seq_ids, ar_positions, + /*dense_token_layout=*/true, /*tree_read=*/false); + attn = ggml_concat(ctx, tree_attn, ar_attn, 2); + if (q_fa_out) { + *q_fa_out = ggml_concat(ctx, Qtree, Qar, 1); + } + } else { + attn = tree_attn; + if (q_fa_out) *q_fa_out = Qtree; + } } else if (ragged) { // ── Ragged concurrent step: prefill chunk rows and decode rows all // read the pool through one call, each row clamped to its own @@ -990,7 +1024,7 @@ static ggml_tensor * build_full_attn_block( : kv_start + n_tokens; attn = paged_read(Qfa, launch_kv_len, paged_query_seq_ids, paged_query_positions, - /*dense_token_layout=*/n_tokens > 1); + /*dense_token_layout=*/n_tokens > 1, /*tree_read=*/false); } else if (paged_block_table) { ggml_tensor * Qfa = q_segment(0, n_tokens); // Post-rotation Q matches the basis of the K rows in the cache, so a @@ -1012,7 +1046,7 @@ static ggml_tensor * build_full_attn_block( const int launch_kv_len = paged_max_kv_len > 0 ? paged_max_kv_len : kv_len; attn = paged_read( Qfa, launch_kv_len, active_slot_ids, /*row_positions=*/nullptr, - /*dense_token_layout=*/active_slot_ids && n_tokens > 1); + /*dense_token_layout=*/active_slot_ids && n_tokens > 1, /*tree_read=*/false); if (!active_slot_ids) { // The only non-mapped paged caller is classic single-token AR. GGML_ASSERT(n_tokens == 1); @@ -1102,6 +1136,9 @@ static ggml_tensor * build_delta_net_block( int n_prefill_segments = 0, ggml_tensor * active_slot_ids = nullptr, ggml_tensor * state_slot_ids = nullptr, + ggml_tensor * ar_active_slot_ids = nullptr, + ggml_tensor * ar_state_slot_ids = nullptr, + int n_ar_seqs = 0, bool allow_inplace_state = false ) { const int head_k_dim = w.ssm_d_state; @@ -1111,6 +1148,9 @@ static ggml_tensor * build_delta_net_block( const int conv_channels = w.ssm_d_inner + 2 * w.ssm_n_group * w.ssm_d_state; const bool ragged = n_prefill_segments > 0; GGML_ASSERT(n_seqs >= 1); + GGML_ASSERT(n_ar_seqs >= 0); + GGML_ASSERT((ar_active_slot_ids == nullptr) == + (ar_state_slot_ids == nullptr)); GGML_ASSERT(n_prefill_segments == 0 || prefill_segments); int prefill_total = 0; for (int i = 0; i < n_prefill_segments; ++i) { @@ -1120,18 +1160,24 @@ static ggml_tensor * build_delta_net_block( } GGML_ASSERT((active_slot_ids == nullptr) == (state_slot_ids == nullptr)); const bool mapped_tree = active_slot_ids && parent_ids; + const int tree_total = mapped_tree + ? (int)ggml_nelements(parent_ids) : 0; + const bool mixed_tree_ar = mapped_tree && n_ar_seqs > 0; + GGML_ASSERT(!n_ar_seqs || (mapped_tree && ar_active_slot_ids)); GGML_ASSERT(!cap || !active_slot_ids || mapped_tree); GGML_ASSERT(!active_slot_ids || (mapped_tree - ? (!ragged && prefill_total == 0 && - n_tokens % n_seqs == 0) + ? (!ragged && prefill_total == 0 && tree_total > 0 && + tree_total % n_seqs == 0 && + n_tokens == tree_total + n_ar_seqs) : (prefill_total + n_seqs == n_tokens))); if (!active_slot_ids) { GGML_ASSERT(n_seqs == 1); GGML_ASSERT(prefill_total == 0 || prefill_total == n_tokens); } GGML_ASSERT(!ragged || (!cap && !parent_ids)); - const bool can_skip_gdn_intermediate = skip_gdn_intermediate && !parent_ids && !cap; + GGML_ASSERT(!mixed_tree_ar || + ar_active_slot_ids->ne[0] == n_ar_seqs); // ── Whole-batch projections ───────────────────────────────────── // qkv_mixed = wqkv @ cur [10240, n_tokens] @@ -1160,9 +1206,11 @@ static ggml_tensor * build_delta_net_block( bool tree; // mapped tree: gather-only, no persistence ggml_tensor * conv_st; ggml_tensor * ssm_st; + ggml_tensor * active_ids; + ggml_tensor * state_ids; }; std::vector segs; - segs.reserve((size_t)n_prefill_segments + 1); + segs.reserve((size_t)n_prefill_segments + 2); for (int i = 0; i < n_prefill_segments; ++i) { const QwenPrefillSegment & pf = prefill_segments[i]; GGML_ASSERT(pf.seq_slot >= 0 && @@ -1176,16 +1224,22 @@ static ggml_tensor * build_delta_net_block( ssm_state->nb[1], ssm_state->nb[2], ssm_state->nb[3], (size_t)pf.seq_slot * ssm_state->nb[3]); segs.push_back({pf.token_offset, pf.n_tokens, 1, - false, false, c, s}); + false, false, c, s, nullptr, nullptr}); } if (active_slot_ids) { - const int tree_tokens = mapped_tree ? n_tokens / n_seqs : 1; + const int tree_tokens = mapped_tree ? tree_total / n_seqs : 1; segs.push_back({prefill_total, tree_tokens, n_seqs, true, - mapped_tree, conv_state, ssm_state}); + mapped_tree, conv_state, ssm_state, + active_slot_ids, state_slot_ids}); + if (mixed_tree_ar) { + segs.push_back({tree_total, 1, n_ar_seqs, true, false, + conv_state, ssm_state, + ar_active_slot_ids, ar_state_slot_ids}); + } } else if (segs.empty()) { // No general [timesteps x sequences] mode: one multi-token sequence. segs.push_back({0, n_tokens, n_seqs, false, false, - conv_state, ssm_state}); + conv_state, ssm_state, nullptr, nullptr}); } const int n_segs = (int)segs.size(); @@ -1206,6 +1260,10 @@ static ggml_tensor * build_delta_net_block( const int seg_tokens = seg.T * seg.S; const bool seg_active = seg.active; const bool seg_tree = seg.tree; + DeltaNetCapture * seg_cap = (!mixed_tree_ar || seg_tree) ? cap : nullptr; + ggml_tensor * seg_parent_ids = seg_tree ? parent_ids : nullptr; + const bool seg_can_skip_gdn_intermediate = + skip_gdn_intermediate && !seg_parent_ids && !seg_cap; // Plain one-token decode has no in-graph consumer of the updated state: // the next graph evaluation is the first read. Write the final state // directly into its persistent slab and avoid materializing/copying a @@ -1213,7 +1271,7 @@ static ggml_tensor * build_delta_net_block( // mapped physical slab directly; only its negative bucket-padding rows // use the result tensor's retained scratch state region. const bool inplace_state = (seg_active && !seg_tree) || - (allow_inplace_state && can_skip_gdn_intermediate && + (allow_inplace_state && seg_can_skip_gdn_intermediate && !ragged && n_seq_tokens == 1); ggml_tensor * qkv_mixed = ggml_reshape_3d(ctx, @@ -1235,7 +1293,7 @@ static ggml_tensor * build_delta_net_block( ggml_tensor * all_conv = ggml_reshape_2d( ctx, seg.conv_st, slab, seg.conv_st->ne[2]); ggml_tensor * gathered = - ggml_get_rows(ctx, all_conv, state_slot_ids); + ggml_get_rows(ctx, all_conv, seg.state_ids); conv_states_r = ggml_reshape_3d( ctx, gathered, w.ssm_d_conv - 1, conv_channels, seg_seqs); } else { @@ -1258,19 +1316,19 @@ static ggml_tensor * build_delta_net_block( // a graph output (which would force the gallocr to preserve its memory // past graph_compute). After graph_compute, the cache buffer's data is // always valid; the rollback code slices it at commit_n. - if (cap && cap->conv_input) { + if (seg_cap && seg_cap->conv_input) { // conv_input may be shorter than the pre-allocated cache // (e.g. during prefill when n_tokens < max_verify_tokens). // Copy into a matching-sized view of the cache destination. const int64_t ci_len = conv_input->ne[0]; ggml_tensor * dst; - if (ci_len == cap->conv_input->ne[0] && - seg_seqs == cap->conv_input->ne[2]) { - dst = cap->conv_input; + if (ci_len == seg_cap->conv_input->ne[0] && + seg_seqs == seg_cap->conv_input->ne[2]) { + dst = seg_cap->conv_input; } else { - dst = ggml_view_3d(ctx, cap->conv_input, - ci_len, cap->conv_input->ne[1], seg_seqs, - cap->conv_input->nb[1], cap->conv_input->nb[2], 0); + dst = ggml_view_3d(ctx, seg_cap->conv_input, + ci_len, seg_cap->conv_input->ne[1], seg_seqs, + seg_cap->conv_input->nb[1], seg_cap->conv_input->nb[2], 0); } GGML_ASSERT(ggml_nelements(conv_input) == ggml_nelements(dst)); ggml_build_forward_expand(gf, ggml_cpy(ctx, conv_input, dst)); @@ -1290,7 +1348,7 @@ static ggml_tensor * build_delta_net_block( ctx, seg.conv_st, slab, seg.conv_st->ne[2]); ggml_build_forward_expand( gf, ggml_set_rows_masked( - ctx, all_conv, compact_last, active_slot_ids)); + ctx, all_conv, compact_last, seg.active_ids)); } else if (!seg_tree) { ggml_build_forward_expand(gf, ggml_cpy(ctx, last_conv, seg.conv_st)); } @@ -1300,8 +1358,8 @@ static ggml_tensor * build_delta_net_block( // their conv window from their actual tree parent instead of the DFS // predecessor. Without this, siblings get garbage logits (the conv // output would mix unrelated branches). - ggml_tensor * conv_out = parent_ids - ? ggml_ssm_conv_tree(ctx, conv_input, L.ssm_conv1d, parent_ids) + ggml_tensor * conv_out = seg_parent_ids + ? ggml_ssm_conv_tree(ctx, conv_input, L.ssm_conv1d, seg_parent_ids) : ggml_ssm_conv (ctx, conv_input, L.ssm_conv1d); conv_out = ggml_silu(ctx, conv_out); @@ -1354,7 +1412,7 @@ static ggml_tensor * build_delta_net_block( ggml_tensor * all_ssm = ggml_reshape_2d( ctx, seg.ssm_st, slab, seg.ssm_st->ne[3]); ggml_tensor * gathered = - ggml_get_rows(ctx, all_ssm, state_slot_ids); + ggml_get_rows(ctx, all_ssm, seg.state_ids); s = ggml_reshape_4d(ctx, gathered, head_v_dim, head_v_dim, num_v_heads, seg_seqs); } else { @@ -1383,10 +1441,10 @@ static ggml_tensor * build_delta_net_block( // path is never quantized. In tree mode, n_seq_tokens is root-inclusive and // flat slot t is persisted directly at ne[3] slot t. // Q8_0 intermediates fall through to the guarded legacy copy path below. - ggml_tensor * persist_inter = (cap && cap->ssm_intermediate_states - && (cap->ssm_intermediate_states->type == GGML_TYPE_F32 - || cap->ssm_intermediate_states->type == GGML_TYPE_F16)) - ? cap->ssm_intermediate_states + ggml_tensor * persist_inter = (seg_cap && seg_cap->ssm_intermediate_states + && (seg_cap->ssm_intermediate_states->type == GGML_TYPE_F32 + || seg_cap->ssm_intermediate_states->type == GGML_TYPE_F16)) + ? seg_cap->ssm_intermediate_states : nullptr; // Chunked delta-net path: chain-only (no parent_ids), no per-token @@ -1398,7 +1456,7 @@ static ggml_tensor * build_delta_net_block( // causing AL degradation and loopy output. Set DFLASH27B_CHUNKED=1 to // opt in for A/B testing while debugging. bool use_chunked = false; - if (can_skip_gdn_intermediate && n_seq_tokens > 1) { + if (seg_can_skip_gdn_intermediate && n_seq_tokens > 1) { if (const char * s_env = std::getenv("DFLASH27B_CHUNKED")) { use_chunked = (std::atoi(s_env) != 0); } @@ -1417,12 +1475,12 @@ static ggml_tensor * build_delta_net_block( ggml_tensor * result; if (seg_active && !seg_tree) { result = ggml_gated_delta_net_active_inplace( - ctx, q_c, k_c, v_c, g_tensor, beta, s, active_slot_ids); - } else if (parent_ids) { + ctx, q_c, k_c, v_c, g_tensor, beta, s, seg.active_ids); + } else if (seg_parent_ids) { // Tree verify: _tree_persist wires src[7] internally. result = persist_inter - ? ggml_gated_delta_net_tree_persist(ctx, q_c, k_c, v_c, g_tensor, beta, s, parent_ids, persist_inter) - : ggml_gated_delta_net_tree(ctx, q_c, k_c, v_c, g_tensor, beta, s, parent_ids); + ? ggml_gated_delta_net_tree_persist(ctx, q_c, k_c, v_c, g_tensor, beta, s, seg_parent_ids, persist_inter) + : ggml_gated_delta_net_tree(ctx, q_c, k_c, v_c, g_tensor, beta, s, seg_parent_ids); } else { // Non-tree (chain/prefill). When capture is requested, set src[7] so // the kernel writes per-token intermediates directly to the persistent @@ -1436,7 +1494,7 @@ static ggml_tensor * build_delta_net_block( result->src[7] = persist_inter; } } - if (can_skip_gdn_intermediate) { + if (seg_can_skip_gdn_intermediate) { ggml_gated_delta_net_set_skip_intermediate(result, true); } @@ -1477,7 +1535,7 @@ static ggml_tensor * build_delta_net_block( // buffer cap->ssm_intermediate_states. The gallocr is unaware of the // persistent cache, so verify_build stays cheap. Matches SGLang's // mamba_caches.intermediate_ssm pattern. - if (cap && cap->ssm_intermediate_states && !persist_inter) { + if (seg_cap && seg_cap->ssm_intermediate_states && !persist_inter) { // This path is only reachable when the intermediate buffer is a type // persist routing can't handle (persist requires F32/F16; the cache // allocates F16, so this is normally dead). If the result tensor has no @@ -1486,7 +1544,7 @@ static ggml_tensor * build_delta_net_block( GGML_ABORT( "non-tree GDN intermediate capture requires an F32/F16 persist buffer " "(got type %d); use F16 intermediates (the default) or the tree-verify path.", - (int)cap->ssm_intermediate_states->type); + (int)seg_cap->ssm_intermediate_states->type); } } @@ -1767,6 +1825,9 @@ QwenGraphOutputs build_qwen35_graph( in.n_prefill_segments, in.active_slot_ids, in.state_slot_ids, + in.ar_active_slot_ids, + in.ar_state_slot_ids, + in.n_ar_seqs, /*allow_inplace_state=*/ in.n_prefill_tokens == 0); dn_idx++; diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index 41b38aea8..75c1efdca 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -92,6 +92,12 @@ TEST_CASE(RecurrentSnapshotFixture, hardens_feature_smoke_paths) { ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 4); tree.kv_write_rows = ggml_new_tensor_2d(live_ctx, GGML_TYPE_I64, 4, 1); + tree.ar_active_slot_ids = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 1); + tree.ar_state_slot_ids = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 1); + tree.paged_query_positions = + ggml_new_tensor_1d(live_ctx, GGML_TYPE_I32, 5); ggml_backend_buffer_t live_buffer = ggml_backend_alloc_ctx_tensors(live_ctx, tree_backend); CHECK(live_buffer != nullptr); diff --git a/server/test/test_speculation_goodput.cpp b/server/test/test_speculation_goodput.cpp index 853af8e0c..f4bf3fe4f 100644 --- a/server/test/test_speculation_goodput.cpp +++ b/server/test/test_speculation_goodput.cpp @@ -201,6 +201,27 @@ int main() { CHECK(decision.requests[0] == 21); } + // Executor capacity bounds the ranked prefix. This keeps every adaptive + // route on the one-pass implementation even when more requests rank well. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(16, 160.0); + ranker.observe_route(16, 1, 145.0); + ranker.observe_route(16, 2, 150.0); + ranker.observe_route(16, 3, 155.0); + std::vector candidates = { + {41, 8.0, 8.0, true}, + {42, 7.0, 8.0, true}, + {43, 6.0, 8.0, true}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(16, candidates, /*max_speculative_requests=*/2); + CHECK(!decision.exploring); + CHECK(decision.requests.size() == 2); + CHECK(decision.requests[0] == 41); + CHECK(decision.requests[1] == 42); + } + // A promising calibrated prefix gets one bounded hardware-cost probe when // that subbatch shape has not been observed yet. { From 98792f9341b638f5ea3ca98280cf0e9073a682ad Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 17 Aug 2026 13:55:06 +0000 Subject: [PATCH 12/18] perf(qwen35): route compact speculation by occupancy --- .../concurrency/adaptive_verification.h | 127 ++++++++++-- .../concurrency/speculation_prompt_prior.h | 58 ++++++ server/src/common/sampler.h | 4 + server/src/internal.h | 5 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 187 +++++++++++++++--- .../qwen35/concurrency/qwen35_seq_engine.h | 7 +- server/src/qwen35/graph_builders.cpp | 2 +- server/src/qwen35/qwen35_backend.cpp | 2 +- server/src/server/scheduler.cpp | 32 ++- server/test/test_speculation_goodput.cpp | 90 +++++++++ 10 files changed, 467 insertions(+), 47 deletions(-) create mode 100644 server/src/common/concurrency/speculation_prompt_prior.h diff --git a/server/src/common/concurrency/adaptive_verification.h b/server/src/common/concurrency/adaptive_verification.h index 0de13f04a..f2ecb500e 100644 --- a/server/src/common/concurrency/adaptive_verification.h +++ b/server/src/common/concurrency/adaptive_verification.h @@ -24,6 +24,9 @@ struct AdaptiveVerificationCandidate { double expected_tokens = 1.0; double maximum_tokens = 1.0; bool calibrated = false; + // Cheap request prior used only to order otherwise-unknown candidates. + // Measured expected tokens replace it after the first verification. + double routing_prior = 0.0; }; struct AdaptiveVerificationDecision { @@ -148,7 +151,10 @@ class AdaptiveVerificationRanker { int active_requests, const std::vector & candidates, int max_speculative_requests = - std::numeric_limits::max()) const { + std::numeric_limits::max(), + bool probe_uncalibrated_with_verifier = false, + bool project_cost_from_higher_occupancy = false, + bool evaluate_nonconvex_prefixes = false) const { AdaptiveVerificationDecision out; if (active_requests <= 0 || candidates.empty() || !has_autoregressive_cost(active_requests)) { @@ -184,6 +190,9 @@ class AdaptiveVerificationRanker { std::stable_sort(unknown.begin(), unknown.end(), [](const AdaptiveVerificationCandidate & a, const AdaptiveVerificationCandidate & b) { + if (a.routing_prior != b.routing_prior) { + return a.routing_prior > b.routing_prior; + } if (a.maximum_tokens != b.maximum_tokens) { return a.maximum_tokens > b.maximum_tokens; } @@ -205,32 +214,52 @@ class AdaptiveVerificationRanker { double expected_total = static_cast(active_requests); const int prefix_limit = std::max(0, std::min( active_requests, max_speculative_requests)); + auto effective_route_cost = [&](int speculative_requests) { + if (has_route_cost(active_requests, speculative_requests)) { + return route_cost_us( + active_requests, speculative_requests); + } + return project_cost_from_higher_occupancy + ? projected_route_cost_us( + active_requests, speculative_requests) + : std::numeric_limits::infinity(); + }; + bool measured_route = false; for (int k = 1; k <= static_cast(known.size()) && k <= prefix_limit; ++k) { expected_total += known[(size_t)k - 1].expected_tokens - 1.0; - if (!has_route_cost(active_requests, k)) { - missing_cost_prefix = k; - break; + const double route_us = effective_route_cost(k); + if (!std::isfinite(route_us)) { + if (missing_cost_prefix == 0) missing_cost_prefix = k; + if (!evaluate_nonconvex_prefixes) break; + continue; } - const double route_us = route_cost_us(active_requests, k); + measured_route = true; const double throughput = expected_total / route_us; - // DSpark's greedy policy stops at the first non-improving - // candidate because candidates are already ranked by survival. - if (throughput <= best) break; - best = throughput; - if (throughput >= required) { - admitted_prefix = k; - admitted_goodput = throughput; + // Request-level route costs can be non-convex: on wide GPUs a + // k=3 compact bundle may win even when k=1 and k=2 do not. + if (!evaluate_nonconvex_prefixes && throughput <= best) break; + if (throughput > best) { + best = throughput; + if (throughput >= required) { + admitted_prefix = k; + admitted_goodput = throughput; + } } } - // Uncalibrated requests are never sent through an expensive target - // verification merely to discover their value. calibration_request - // asks the concrete adapter for its cheap confidence signal instead. + // By default, calibration_request asks the concrete adapter for a + // cheap confidence signal. Verifier-side calibration is opt-in below + // for adapters whose separate confidence pass would duplicate work. // A calibrated route shape without a hardware sample gets one bounded - // probe. Larger prefixes are not explored after a measured decline. - if (missing_cost_prefix > 0) { + // probe. Grow a measured winning prefix by one, but do not walk every + // smaller shape after another measured width has already lost. + const bool probe_missing = missing_cost_prefix > 0 && + (!measured_route || + (admitted_prefix > 0 && + missing_cost_prefix == admitted_prefix + 1)); + if (probe_missing) { out.requests.reserve((size_t)missing_cost_prefix); for (int i = 0; i < missing_cost_prefix; ++i) { out.requests.push_back(known[(size_t)i].request); @@ -246,10 +275,74 @@ class AdaptiveVerificationRanker { } out.predicted_gain = admitted_goodput / baseline; } + // A speculator without a cheap calibrated confidence head can turn a + // bounded verification bundle into useful decoding and calibration. + // Prefer the executor's widest compact bundle because request-level + // route costs can be non-convex. Once sampled, its observed accepted + // yield gates later unknown bundles; perfect-acceptance optimism is + // used only before a yield sample exists. DDTree uses this at high + // occupancy, while DSpark can keep using its confidence head. + if (probe_uncalibrated_with_verifier && !unknown.empty() && + !out.exploring && + static_cast(out.requests.size()) < prefix_limit) { + const int probe_count = std::min( + static_cast(unknown.size()), + prefix_limit - static_cast(out.requests.size())); + const int probe_prefix = + static_cast(out.requests.size()) + probe_count; + double optimistic_total = + static_cast(active_requests); + for (int request : out.requests) { + const auto candidate = std::find_if( + known.begin(), known.end(), + [request](const AdaptiveVerificationCandidate & item) { + return item.request == request; + }); + if (candidate != known.end()) { + optimistic_total += candidate->expected_tokens - 1.0; + } + } + for (int i = 0; i < probe_count; ++i) { + optimistic_total += + unknown[(size_t)i].maximum_tokens - 1.0; + } + const double probe_route_us = + effective_route_cost(probe_prefix); + const bool unseen_shape = !std::isfinite(probe_route_us); + const bool can_win = unseen_shape || + optimistic_total / probe_route_us >= required; + if (can_win) { + for (int i = 0; i < probe_count; ++i) { + out.requests.push_back(unknown[(size_t)i].request); + } + out.calibration_request = -1; + out.exploring = true; + } + } return out; } private: + double projected_route_cost_us( + int active_requests, int speculative_requests) const { + if (!has_autoregressive_cost(active_requests)) { + return std::numeric_limits::infinity(); + } + for (int higher = active_requests + 1; + static_cast(higher) < route_cost_known_.size(); + ++higher) { + if (!has_autoregressive_cost(higher) || + !has_route_cost(higher, speculative_requests)) { + continue; + } + const double relative_cost = + route_cost_us(higher, speculative_requests) / + autoregressive_cost_us(higher); + return relative_cost * autoregressive_cost_us(active_requests); + } + return std::numeric_limits::infinity(); + } + static AdaptiveVerificationConfig sanitize( AdaptiveVerificationConfig config) { config.minimum_gain = std::max(1.0, config.minimum_gain); diff --git a/server/src/common/concurrency/speculation_prompt_prior.h b/server/src/common/concurrency/speculation_prompt_prior.h new file mode 100644 index 000000000..0c112860f --- /dev/null +++ b/server/src/common/concurrency/speculation_prompt_prior.h @@ -0,0 +1,58 @@ +#pragma once + +// Training-free request prior for cold-start speculative routing. +// +// This is deliberately a weak prior, not the final decision: obvious +// structured/code prompts are ranked ahead of neutral prompts, obvious +// conversational/creative prompts stay on AR, and measured verifier goodput +// remains authoritative after admission. The policy is model-neutral and can +// be replaced by a learned prompt ranker without changing the verifier. + +#include +#include +#include +#include + +namespace dflash::common { + +inline int speculation_prompt_hint(std::string_view prompt) { + constexpr size_t kMaxPromptChars = 16 * 1024; + std::string text(prompt.substr(0, kMaxPromptChars)); + std::transform(text.begin(), text.end(), text.begin(), + [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + auto has = [&](std::string_view cue) { + return text.find(cue) != std::string::npos; + }; + + // Explicit exclusions override incidental structured words such as + // "avoid code" in a creative-writing request. + if (has("avoid code") || has("chatting casually") || + has("casual conversation") || has("small talk") || + has("keep it conversational") || has("write a story") || + has("invent a story") || has("write a poem") || + has("roleplay")) { + return -1; + } + + int score = 0; + score += has("```") ? 4 : 0; + score += has("\ndef ") || has("\nclass ") || has("#include") ? 4 : 0; + score += has("public static") || has("fn ") || has("function ") ? 3 : 0; + score += has("implement") || has("debug") || has("unit test") ? 3 : 0; + score += has("algorithm") || has("sql query") || + has("regular expression") || has("json schema") ? 2 : 0; + score += has("python") || has("javascript") || has("typescript") || + has("rust") || has("c++") ? 2 : 0; + score += has("code") ? 1 : 0; + + if (score >= 3) return 1; + if (has("story") || has("conversational") || has("brainstorm") || + has("opinion") || has("friendly chat")) { + return -1; + } + return 0; +} + +} // namespace dflash::common diff --git a/server/src/common/sampler.h b/server/src/common/sampler.h index ff55ed5f8..8aab2718d 100644 --- a/server/src/common/sampler.h +++ b/server/src/common/sampler.h @@ -25,6 +25,10 @@ struct SamplerCfg { float rep_pen = 1.0f; // multiplicative repetition penalty (HF-style) int rep_window = 256; uint64_t seed = 0; + // Cold-start speculation prior: -1 conversational/creative, 0 neutral, + // +1 structured/code. It affects admission only; measured goodput remains + // authoritative. Non-HTTP callers naturally retain the neutral default. + int8_t speculation_prompt_hint = 0; // OpenAI-style additive penalties (applied per-token to logits before softmax). // frequency_penalty: subtract freq_pen * count(token_in_history) from logit. diff --git a/server/src/internal.h b/server/src/internal.h index 0a4d0cc17..23b095808 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -414,8 +414,9 @@ struct TargetCache { std::vector conv_input_cache; // size = n_delta (48) // Bounded concurrent-tree checkpoint domain. A packed tree row is - // sequence-major, so recurrent checkpoint t for compact tree lane s lives - // at s*tree_capture_width+t. Keeping this lane count independent from the + // sequence-major. The allocation reserves tree_capture_width rows per + // lane, while a runtime width T <= tree_capture_width writes the active + // packed prefix at s*T+t. Keeping this lane count independent from the // physical serving-slot count lets an all-C adaptive route directly commit // a small profitable subset without reserving T*C recurrent states. int tree_capture_width = 0; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index dd6afe3d9..dcdfd6ac6 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -42,6 +42,29 @@ int decode_bucket_width(int live_count) { return 64; } +struct AdaptiveVerificationOracle { + int speculative_requests = -1; + int tree_budget = -1; + + bool forces_selection() const { return speculative_requests >= 0; } +}; + +const AdaptiveVerificationOracle & adaptive_verification_oracle() { + static const AdaptiveVerificationOracle oracle = []() { + AdaptiveVerificationOracle out; + if (const char * value = std::getenv( + "DFLASH_ADAPTIVE_VERIFY_FORCE_SPECULATIVE_REQUESTS")) { + out.speculative_requests = std::max(0, std::atoi(value)); + } + if (const char * value = std::getenv( + "DFLASH_ADAPTIVE_VERIFY_FORCE_TREE_BUDGET")) { + out.tree_budget = std::max(1, std::atoi(value)); + } + return out; + }(); + return oracle; +} + } // namespace Qwen35SeqEngine::Qwen35SeqEngine( @@ -64,6 +87,7 @@ Qwen35SeqEngine::Qwen35SeqEngine( tree_scratch_stride_(tree_scratch_stride) { const int n_slots = slots_.slot_count(); slot_draft_kv_.resize((size_t)n_slots); + compact_tree_cohort_.resize((size_t)n_slots, 0); // The concurrent DDTree stack is gated to a local same-device drafter. // Build metadata-only BF16 views over each slot's disjoint target feature @@ -170,7 +194,7 @@ bool Qwen35SeqEngine::ddtree_input_eligible(const StepInput & in) const { } std::optional Qwen35SeqEngine::estimate_ddtree_expected_tokens( - const StepInput & in) { + const StepInput & in, int tree_budget) { const int q_len = b_.dw_.block_size; const int hidden = b_.w_.n_embd; if (q_len <= 1 || !build_lm_head_projection_step( @@ -237,8 +261,9 @@ std::optional Qwen35SeqEngine::estimate_ddtree_expected_tokens( // each position estimates conditional survival, so the cumulative product // estimates reaching that prefix. DSpark can replace this adapter with its // calibrated confidence-head probabilities without changing the ranker. - std::vector confidence((size_t)q_len - 1); - for (int pos = 1; pos < q_len; ++pos) { + const int confidence_tokens = std::min(tree_budget, q_len - 1); + std::vector confidence((size_t)confidence_tokens); + for (int pos = 1; pos <= confidence_tokens; ++pos) { confidence[(size_t)pos - 1] = static_cast( std::clamp(std::exp(static_cast( top_lp[(size_t)pos])), @@ -249,17 +274,18 @@ std::optional Qwen35SeqEngine::estimate_ddtree_expected_tokens( } std::optional Qwen35SeqEngine::step_ddtree( - const StepPlan & speculative_plan, const StepPlan & ar_plan) { + const StepPlan & speculative_plan, const StepPlan & ar_plan, + int tree_budget) { StepResult result; const int active = (int)speculative_plan.decode.size(); const int total_active = active + (int)ar_plan.decode.size(); const int bucket = decode_bucket_width(active); - const int T = tree_width_; + const int T = tree_budget + 1; const int q_len = b_.dw_.block_size; const int hidden = b_.w_.n_embd; const int n_head_kv = b_.w_.n_head_kv; const int n_slots = slots_.slot_count(); - const int K = b_.cfg_.ddtree_budget > q_len - 1 ? 8 : 1; + const int K = tree_budget > q_len - 1 ? 8 : 1; struct Proposal { int slot = -1; @@ -350,7 +376,7 @@ std::optional Qwen35SeqEngine::step_ddtree( p.root = in.token; p.tree = build_ddtree( top_lp.data() + K, top_ids.data() + K, - q_len - 1, K, b_.cfg_.ddtree_budget, + q_len - 1, K, tree_budget, b_.cfg_.ddtree_chain_seed); p.flat.assign((size_t)T, 0); p.flat[0] = in.token; @@ -366,7 +392,7 @@ std::optional Qwen35SeqEngine::step_ddtree( b_.cache_.ssm_state.front()->buffer)); const bool direct_commit = !target_is_meta && bucket <= b_.cache_.tree_capture_lanes && - T == b_.cache_.tree_capture_width && + T <= b_.cache_.tree_capture_width && b_.cache_.target_feat_tree_scratch_base > 0; struct DirectArStage { Qwen35SlotManager::StepAppend append; @@ -1055,6 +1081,7 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( const SamplerCfg & sampler) { AdmitResult result = slots_.admit(request_id, prompt, sampler); if (result.status == AdmitResult::Status::admitted) { + compact_tree_cohort_[(size_t)result.slot] = 0; reset_recurrent_slot(b_.cache_, result.slot); if (slots_.residency_active()) { slots_.slot(result.slot).kvflash_last_reselect_generated = @@ -1299,6 +1326,37 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); const bool adaptive_enabled = !(adaptive && std::atoi(adaptive) == 0); + const AdaptiveVerificationOracle & oracle = + adaptive_verification_oracle(); + const bool inherited_compact_tree = std::any_of( + inputs.begin(), inputs.end(), [&](const StepInput & in) { + return compact_tree_cohort_[(size_t)in.slot] != 0; + }); + // Full DDTree depth wins in the low-occupancy regime. Once target AR is + // well batched, a compact proposal lowers the marginal verification cost + // enough for a small profitable request subset. Keep that compact shape + // through the cohort's low-occupancy tail; a newly admitted C<=4 cohort + // still gets the established full-depth path. + constexpr int kFullTreeMaxConcurrency = 4; + constexpr int kCompactTreeBudget = 8; + constexpr uint64_t kVerifierCalibrationSamples = 16; + const bool compact_tree_route = + inputs.size() > kFullTreeMaxConcurrency || inherited_compact_tree; + const bool structured_compact_cohort = + !compact_tree_route || std::all_of( + inputs.begin(), inputs.end(), [&](const StepInput & in) { + return slots_.slot(in.slot).sampler.speculation_prompt_hint > 0; + }); + const int adaptive_tree_budget = + compact_tree_route + ? std::min(b_.cfg_.ddtree_budget, kCompactTreeBudget) + : b_.cfg_.ddtree_budget; + const int tree_budget = std::clamp( + oracle.tree_budget > 0 ? oracle.tree_budget + : adaptive_tree_budget, + 1, b_.cfg_.ddtree_budget); + AdaptiveVerificationRanker & route_ranker = compact_tree_route + ? compact_adaptive_verification_ : adaptive_verification_; // Keep the scheduler independent of the concrete speculation algorithm. // DDTree contributes either a cheap draft-confidence estimate or an EWMA @@ -1309,17 +1367,32 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { auto collect_candidates = [&]() { candidates.clear(); for (const StepInput & in : inputs) { - if (!ddtree_input_eligible(in) || - !slots_.ddtree_speculation_allowed(in.slot)) { + if (!ddtree_input_eligible(in)) continue; + const Qwen35Slot & seq = slots_.slot(in.slot); + const int prompt_hint = seq.sampler.speculation_prompt_hint; + const bool structured_warmup = + adaptive_enabled && !oracle.forces_selection() && + compact_tree_route && structured_compact_cohort && + prompt_hint > 0 && + seq.ddtree_sampled_steps < kVerifierCalibrationSamples; + if (!oracle.forces_selection() && + !slots_.ddtree_speculation_allowed(in.slot) && + !structured_warmup) { + continue; + } + if (adaptive_enabled && !oracle.forces_selection() && + compact_tree_route && prompt_hint <= 0) { continue; } const SpeculationGoodputController & policy = - slots_.slot(in.slot).speculation; + seq.speculation; candidates.push_back({ in.slot, policy.expected_emitted_tokens(), - static_cast(b_.dw_.block_size), + static_cast(compact_tree_route + ? tree_budget + 1 : b_.dw_.block_size), policy.has_expected_emitted_tokens(), + static_cast(prompt_hint), }); } }; @@ -1330,16 +1403,54 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { b_.cache_.tree_capture_lanes > 0 ? b_.cache_.tree_capture_lanes : static_cast(inputs.size()); - if (adaptive_enabled) { - decision = adaptive_verification_.select( + if (oracle.forces_selection()) { + const int force_limit = std::min( + adaptive_speculation_limit, + oracle.speculative_requests); + decision.requests.reserve((size_t)force_limit); + for (const AdaptiveVerificationCandidate & candidate : candidates) { + if ((int)decision.requests.size() >= force_limit) break; + decision.requests.push_back(candidate.request); + } + static bool logged = false; + if (!logged) { + std::fprintf(stderr, + "[parallel-ddtree] oracle force_speculative=%d " + "tree_budget=%d direct_limit=%d\n", + oracle.speculative_requests, tree_budget, + adaptive_speculation_limit); + logged = true; + } + } else if (adaptive_enabled) { + const bool drafter_side_calibration = !compact_tree_route; + // The compact verifier has three directly captured lanes on Strix + // Halo. Oracle sweeps show that a bundle remains worthwhile while + // those lanes cover at least one third of the active cohort. + constexpr int kCompactActiveRequestsPerLane = 3; + const int compact_calibration_limit = + kCompactActiveRequestsPerLane * adaptive_speculation_limit; + const bool verifier_side_calibration = + compact_tree_route && structured_compact_cohort && + static_cast(inputs.size()) <= compact_calibration_limit; + const int ranked_speculation_limit = + !compact_tree_route + ? adaptive_speculation_limit + : (structured_compact_cohort && + static_cast(inputs.size()) <= + compact_calibration_limit + ? adaptive_speculation_limit : 0); + decision = route_ranker.select( static_cast(inputs.size()), candidates, - adaptive_speculation_limit); + ranked_speculation_limit, verifier_side_calibration, + verifier_side_calibration, + /*evaluate_nonconvex_prefixes=*/compact_tree_route); if (inputs.size() <= 3) { adaptive_calibration_cooldown_ = 0; } else if (adaptive_calibration_cooldown_ > 0) { --adaptive_calibration_cooldown_; } - if (decision.calibration_request >= 0 && + if (drafter_side_calibration && + decision.calibration_request >= 0 && adaptive_calibration_cooldown_ == 0) { const auto input = std::find_if( inputs.begin(), inputs.end(), @@ -1348,7 +1459,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { }); if (input != inputs.end()) { const std::optional expected = - estimate_ddtree_expected_tokens(*input); + estimate_ddtree_expected_tokens(*input, tree_budget); if (expected) { slots_.slot(input->slot).speculation .observe_expected_tokens(*expected); @@ -1372,7 +1483,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { slots_.slot(input->slot).request_id, input->slot, *expected); collect_candidates(); - decision = adaptive_verification_.select( + decision = route_ranker.select( static_cast(inputs.size()), candidates, adaptive_speculation_limit); } @@ -1412,7 +1523,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const auto started = Clock::now(); if (speculative_count > 0) { std::optional mixed = - step_ddtree(speculative_plan, ar_plan); + step_ddtree(speculative_plan, ar_plan, tree_budget); if (!mixed) { // Proposal setup failed before target/cache mutation. Preserve // service with one ordinary packed step and retry speculation on a @@ -1421,6 +1532,12 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (!mixed->ok()) return std::move(*mixed); routed_result = std::move(*mixed); + if (adaptive_enabled && !oracle.forces_selection() && + compact_tree_route) { + for (const StepInput & in : inputs) { + compact_tree_cohort_[(size_t)in.slot] = 1; + } + } } else { routed_result = step_regular(ar_plan); if (!routed_result.ok()) return routed_result; @@ -1429,8 +1546,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { 1.0, std::chrono::duration( Clock::now() - started).count()); - if (adaptive_enabled) { - adaptive_verification_.observe_route( + if (adaptive_enabled && !oracle.forces_selection()) { + route_ranker.observe_route( static_cast(inputs.size()), speculative_count, route_us); } if (decision.exploring && !speculative_plan.decode.empty()) { @@ -1441,6 +1558,14 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { ar_plan.decode.size()); } + const bool use_route_ar_baseline = + adaptive_enabled && !oracle.forces_selection() && + compact_tree_route && decision.exploring && + route_ranker.has_autoregressive_cost( + static_cast(inputs.size())); + const double route_ar_us = use_route_ar_baseline + ? route_ranker.autoregressive_cost_us( + static_cast(inputs.size())) : 0.0; auto log_transition = [&](int slot, SpeculationGoodputTransition transition, const char * observed_route, @@ -1464,18 +1589,33 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { }; for (DecodeOutput & out : routed_result.decode) { + if (oracle.forces_selection()) break; if (out.failed || out.slot < 0 || out.slot >= n_slots) continue; if (selected[(size_t)out.slot]) { const double emitted = static_cast(out.ddtree_accepted_tokens + 1); - const SpeculationGoodputTransition transition = + SpeculationGoodputTransition transition = slots_.record_speculation_sample( out.slot, emitted, route_us); + bool seeded_route_ar = false; + const Qwen35Slot & observed_seq = slots_.slot(out.slot); + const bool warmup_complete = + observed_seq.sampler.speculation_prompt_hint <= 0 || + observed_seq.ddtree_sampled_steps >= + kVerifierCalibrationSamples; + if (use_route_ar_baseline && warmup_complete && + !observed_seq.speculation.has_ar_goodput()) { + transition = slots_.record_ar_sample( + out.slot, route_ar_us); + seeded_route_ar = true; + } if (transition == SpeculationGoodputTransition::disabled) { out.ddtree_suspensions = 1; } log_transition( - out.slot, transition, "speculation", emitted, route_us); + out.slot, transition, + seeded_route_ar ? "speculation-vs-route-ar" : "speculation", + emitted, route_us); } else if (observe_ar[(size_t)out.slot]) { const SpeculationGoodputTransition transition = slots_.record_ar_sample(out.slot, route_us); @@ -1903,6 +2043,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step_regular(const StepPlan & plan) { void Qwen35SeqEngine::retire(int slot) { if (!slots_.is_active(slot)) return; slots_.retire(slot); + compact_tree_cohort_[(size_t)slot] = 0; } } // namespace dflash::common diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 5bc4d3400..510c1ac5f 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -122,12 +122,13 @@ class Qwen35SeqEngine final : public SeqEngine { bool ddtree_available(const StepPlan & plan) const; bool ddtree_input_eligible(const StepInput & input) const; std::optional estimate_ddtree_expected_tokens( - const StepInput & input); + const StepInput & input, int tree_budget); StepResult step_regular(const StepPlan & plan); // nullopt means proposal setup failed before target/cache mutation and the // caller may safely use the ordinary packed AR path for this iteration. std::optional step_ddtree( - const StepPlan & speculative_plan, const StepPlan & ar_plan); + const StepPlan & speculative_plan, const StepPlan & ar_plan, + int tree_budget); Qwen35Backend & b_; Qwen35SlotManager slots_; @@ -137,7 +138,9 @@ class Qwen35SeqEngine final : public SeqEngine { int tree_scratch_stride_ = 0; bool capture_features_ = false; AdaptiveVerificationRanker adaptive_verification_; + AdaptiveVerificationRanker compact_adaptive_verification_; int adaptive_calibration_cooldown_ = 0; + std::vector compact_tree_cohort_; ggml_context * feature_view_ctx_ = nullptr; std::vector slot_feature_mirrors_; std::vector> slot_draft_kv_; diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index dc2796770..3bd4b067b 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -753,7 +753,7 @@ bool build_target_step_paged_tree( return false; } if (capture_direct_commit && - (cache.tree_capture_width != tree_width || + (cache.tree_capture_width < tree_width || cache.tree_capture_lanes < n_tree_seqs || cache.target_feat_tree_scratch_base <= 0 || cache.ssm_intermediate.empty() || diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 36a30e511..15541e7d8 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -574,7 +574,7 @@ bool Qwen35Backend::init() { std::fprintf(stderr, "[parallel-ddtree] enabled budget=%d width=%d " "mode=one-pass-tree-ar+bounded-replay adaptive=%s " - "policy=ranked-goodput scope=all-concurrency\n", + "policy=ranked-goodput scope=occupancy-aware\n", cfg_.ddtree_budget, tree_width, adaptive && std::atoi(adaptive) == 0 ? "off" : "on"); } diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 045786cb2..c2fb270cf 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -10,6 +10,7 @@ #include "http_server.h" #include "common/concurrency/seq_engine.h" +#include "common/concurrency/speculation_prompt_prior.h" #include #include @@ -21,6 +22,32 @@ namespace dflash::common { namespace { +std::string last_user_prompt_text(const json & messages) { + if (!messages.is_array()) return {}; + for (int i = static_cast(messages.size()) - 1; i >= 0; --i) { + const json & message = messages[(size_t)i]; + if (!message.is_object() || + message.value("role", "") != "user" || + !message.contains("content")) { + continue; + } + const json & content = message["content"]; + if (content.is_string()) return content.get(); + if (!content.is_array()) return {}; + std::string text; + for (const json & part : content) { + if (!part.is_object() || !part.contains("text") || + !part["text"].is_string()) { + continue; + } + if (!text.empty()) text.push_back('\n'); + text += part["text"].get(); + } + return text; + } + return {}; +} + // 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 @@ -523,8 +550,11 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // Admission only claims the slot and queues the prompt. Prefill // advances one chunk per engine step alongside live decode. const uint64_t engine_request_id = next_request_id; + SamplerCfg sampler = req.sampler; + sampler.speculation_prompt_hint = + speculation_prompt_hint(last_user_prompt_text(req.messages)); auto ar = engine.admit(engine_request_id, effective_prompt, - req.sampler); + sampler); if (ar.status == SeqEngine::AdmitResult::Status::busy) return AdmissionDisposition::Deferred; if (ar.status != SeqEngine::AdmitResult::Status::admitted) { diff --git a/server/test/test_speculation_goodput.cpp b/server/test/test_speculation_goodput.cpp index f4bf3fe4f..980609144 100644 --- a/server/test/test_speculation_goodput.cpp +++ b/server/test/test_speculation_goodput.cpp @@ -1,5 +1,6 @@ #include "common/concurrency/adaptive_verification.h" #include "common/concurrency/speculation_goodput.h" +#include "common/concurrency/speculation_prompt_prior.h" #include "host_check.h" #include @@ -9,6 +10,17 @@ using namespace dflash::common; static int g_checks = 0; int main() { + // The cold-start prior separates obvious structured/code requests from + // conversational writing while leaving ambiguous requests neutral. + { + CHECK(speculation_prompt_hint( + "Complete this Python function:\n\ndef solve(values):") == 1); + CHECK(speculation_prompt_hint( + "Chatting casually, write a story and avoid code.") == -1); + CHECK(speculation_prompt_hint( + "What happened during the Apollo 11 mission?") == 0); + } + // Cold start measures one real speculative step and one neighboring AR // step, then keeps the route with higher useful-token goodput. { @@ -184,6 +196,84 @@ int main() { CHECK(decision.calibration_request == 1); } + // A concrete adapter may calibrate a bounded compact bundle inside useful + // verification. The first unseen route probes up to executor capacity, + // never the whole unknown cohort. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(8, 100.0); + std::vector candidates = { + {5, 1.0, 9.0, false, 1.0}, + {2, 1.0, 9.0, false, -1.0}, + }; + const AdaptiveVerificationDecision probe = + ranker.select(8, candidates, 3, + /*probe_uncalibrated_with_verifier=*/true); + CHECK(probe.exploring); + CHECK(probe.requests.size() == 2); + CHECK(probe.requests[0] == 5); + CHECK(probe.requests[1] == 2); + CHECK(probe.calibration_request == -1); + } + + // A losing route ratio suppresses repeated probes as occupancy falls. + // Even perfect acceptance cannot make this projected route beat AR. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(16, 100.0); + ranker.observe_route(16, 1, 200.0); + ranker.observe_autoregressive(15, 95.0); + std::vector candidates = { + {7, 1.0, 9.0, false}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(15, candidates, 3, + /*probe_uncalibrated_with_verifier=*/true, + /*project_cost_from_higher_occupancy=*/true); + CHECK(!decision.exploring); + CHECK(decision.requests.empty()); + CHECK(decision.calibration_request == 7); + } + // A profitable measured route ratio can rank calibrated requests after + // occupancy drops, without spending a fresh hardware probe. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(8, 100.0); + ranker.observe_route(8, 1, 80.0); + ranker.observe_autoregressive(7, 90.0); + std::vector candidates = { + {12, 4.0, 9.0, true}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(7, candidates, 3, + /*probe_uncalibrated_with_verifier=*/false, + /*project_cost_from_higher_occupancy=*/true); + CHECK(!decision.exploring); + CHECK(decision.requests.size() == 1); + CHECK(decision.requests[0] == 12); + } + + // Prefix widths are compared independently because hardware occupancy can + // make k=3 profitable even when the measured k=1 route loses. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(8, 100.0); + ranker.observe_route(8, 1, 120.0); + ranker.observe_route(8, 3, 80.0); + std::vector candidates = { + {1, 4.0, 9.0, true}, + {2, 4.0, 9.0, true}, + {3, 4.0, 9.0, true}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(8, candidates, 3, + /*probe_uncalibrated_with_verifier=*/false, + /*project_cost_from_higher_occupancy=*/false, + /*evaluate_nonconvex_prefixes=*/true); + CHECK(!decision.exploring); + CHECK(decision.requests.size() == 3); + } + // There is no concurrency cutoff: when one request pays for the route at // C=16, that request alone remains speculative and all peers remain AR. { From a1b93a43f040937dc4b2d614ea2ddcef96b30ed2 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 17 Aug 2026 17:02:51 +0000 Subject: [PATCH 13/18] perf(qwen35): route speculation dynamically across concurrency --- .../concurrency/adaptive_verification.h | 464 ++++++++++++++---- .../concurrency/speculation_prompt_prior.h | 11 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 270 +++++----- .../qwen35/concurrency/qwen35_seq_engine.h | 1 + server/test/test_seq_slot_manager.cpp | 11 +- server/test/test_speculation_goodput.cpp | 405 +++++++++++++-- 6 files changed, 870 insertions(+), 292 deletions(-) diff --git a/server/src/common/concurrency/adaptive_verification.h b/server/src/common/concurrency/adaptive_verification.h index f2ecb500e..7dd983a4d 100644 --- a/server/src/common/concurrency/adaptive_verification.h +++ b/server/src/common/concurrency/adaptive_verification.h @@ -6,14 +6,19 @@ // and grows the verification batch while expected throughput improves. This // helper applies the same policy at request granularity. The concrete // speculator supplies expected useful tokens; the engine supplies the observed -// cost of each mixed route shape (active requests, speculative requests). +// cost of each exact mixed route shape (active requests, speculative requests). // DDTree can learn value from accepted paths, while DSpark can use its -// calibrated confidence head directly. +// calibrated confidence head directly. One ranker instance represents one +// fixed speculator/proposal shape; adapters with ragged verification work must +// keep separate rankers for distinct work buckets. #include +#include #include #include #include +#include +#include #include namespace dflash::common { @@ -25,8 +30,19 @@ struct AdaptiveVerificationCandidate { double maximum_tokens = 1.0; bool calibrated = false; // Cheap request prior used only to order otherwise-unknown candidates. - // Measured expected tokens replace it after the first verification. + // Request-local measurements smoothly replace it during warmup. double routing_prior = 0.0; + // Number of request-local observations supporting the estimate. Shared + // cohort evidence must never relax AR-peer protection for a new request. + std::size_t evidence_samples = 0; + // Generated output already committed for this request. Proven-useful peers + // with less progress receive compact lanes first to avoid cohort stragglers. + int progress_tokens = 0; +}; + +struct AdaptiveVerificationYieldEstimate { + double expected_tokens = 1.0; + std::size_t evidence_samples = 0; }; struct AdaptiveVerificationDecision { @@ -41,6 +57,16 @@ struct AdaptiveVerificationDecision { struct AdaptiveVerificationConfig { // Preserve margin for timing noise and route-switch overhead. double minimum_gain = 1.05; + // Protect requests that remain AR in a mixed route. A wider slowdown is + // allowed only after every active request has demonstrated useful + // speculative yield, so a homogeneous cohort can amortize the work. + double maximum_ar_peer_slowdown = 1.10; + double homogeneous_minimum_expected_tokens = 1.5; + // A slow mixed route is safe only when every request benefits similarly; + // otherwise scarce verifier lanes create a low-yield AR tail. + double homogeneous_minimum_relative_yield = 0.60; + std::size_t homogeneous_minimum_samples = 4; + std::size_t routing_prior_minimum_samples = 4; double cost_ewma_alpha = 0.35; }; @@ -71,6 +97,124 @@ class AdaptiveVerificationRanker { void reset() { route_cost_us_.clear(); route_cost_known_.clear(); + request_yield_.clear(); + routing_prior_yield_.clear(); + } + + void observe_request_yield(std::uint64_t request, + double emitted_tokens) { + if (!valid_yield(emitted_tokens)) return; + update_yield(request_yield_[request], emitted_tokens); + } + + void forget_request(std::uint64_t request) { + request_yield_.erase(request); + } + + std::optional request_expected_tokens( + std::uint64_t request) const { + const auto it = request_yield_.find(request); + return it == request_yield_.end() + ? std::nullopt + : std::optional(it->second.expected_tokens); + } + + std::size_t request_yield_samples(std::uint64_t request) const { + const auto it = request_yield_.find(request); + return it == request_yield_.end() ? 0 : it->second.samples; + } + + bool has_stable_evidence(std::size_t samples) const { + return samples >= config_.homogeneous_minimum_samples; + } + + bool has_useful_yield(double expected_tokens) const { + return std::isfinite(expected_tokens) && + expected_tokens >= + config_.homogeneous_minimum_expected_tokens; + } + + bool forms_homogeneous_cohort( + const std::vector & candidates, + bool require_stable_local_evidence = true) const { + if (candidates.empty()) return false; + double minimum = std::numeric_limits::infinity(); + double maximum = 1.0; + for (const AdaptiveVerificationCandidate & candidate : candidates) { + if (!candidate.calibrated || + !has_useful_yield(candidate.expected_tokens) || + (require_stable_local_evidence && + !has_stable_evidence(candidate.evidence_samples))) { + return false; + } + minimum = std::min(minimum, candidate.expected_tokens); + maximum = std::max(maximum, candidate.expected_tokens); + } + return minimum >= + config_.homogeneous_minimum_relative_yield * maximum; + } + + void observe_routing_prior_yield(double routing_prior, + double emitted_tokens) { + if (!std::isfinite(routing_prior) || !valid_yield(emitted_tokens)) { + return; + } + update_yield(routing_prior_yield_[routing_prior], emitted_tokens); + } + + std::optional routing_prior_expected_tokens( + double routing_prior) const { + if (!std::isfinite(routing_prior)) return std::nullopt; + const auto it = routing_prior_yield_.find(routing_prior); + return it == routing_prior_yield_.end() || + it->second.samples < config_.routing_prior_minimum_samples + ? std::nullopt : std::optional(it->second.expected_tokens); + } + + std::size_t routing_prior_yield_samples(double routing_prior) const { + if (!std::isfinite(routing_prior)) return 0; + const auto it = routing_prior_yield_.find(routing_prior); + return it == routing_prior_yield_.end() ? 0 : it->second.samples; + } + + std::optional estimate_request_yield( + std::uint64_t request, double routing_prior) const { + const auto request_it = request_yield_.find(request); + const auto prior_it = std::isfinite(routing_prior) + ? routing_prior_yield_.find(routing_prior) + : routing_prior_yield_.end(); + const bool has_request = request_it != request_yield_.end(); + const bool has_prior = prior_it != routing_prior_yield_.end() && + prior_it->second.samples >= config_.routing_prior_minimum_samples; + if (!has_request && !has_prior) return std::nullopt; + + AdaptiveVerificationYieldEstimate out; + if (!has_request) { + out.expected_tokens = prior_it->second.expected_tokens; + out.evidence_samples = 0; + return out; + } + + const YieldEstimate & request_estimate = request_it->second; + out.expected_tokens = request_estimate.expected_tokens; + out.evidence_samples = request_estimate.samples; + if (!has_prior) return out; + + const YieldEstimate & prior_estimate = prior_it->second; + if (request_estimate.samples < config_.homogeneous_minimum_samples) { + // Shrink the first few noisy request observations toward a stable + // cohort mean. Once request-local evidence is stable, its measured + // magnitude fully replaces the prior for goodput decisions. + const double local_weight = std::min( + 1.0, + static_cast(request_estimate.samples) / + static_cast( + config_.homogeneous_minimum_samples)); + out.expected_tokens = prior_estimate.expected_tokens + + local_weight * (request_estimate.expected_tokens - + prior_estimate.expected_tokens); + } + return out; } void observe_autoregressive(int batch_size, double elapsed_us) { @@ -147,14 +291,34 @@ class AdaptiveVerificationRanker { : std::numeric_limits::infinity(); } + bool has_exact_profile(int active_requests, + int max_speculative_requests) const { + if (!has_autoregressive_cost(active_requests)) return false; + const int limit = std::max(0, std::min( + active_requests, max_speculative_requests)); + for (int k = 1; k <= limit; ++k) { + if (!has_route_cost(active_requests, k)) return false; + } + return true; + } + + bool has_speculative_profile_sample( + int active_requests, int max_speculative_requests) const { + const int limit = std::max(0, std::min( + active_requests, max_speculative_requests)); + for (int k = 1; k <= limit; ++k) { + if (has_route_cost(active_requests, k)) return true; + } + return false; + } + AdaptiveVerificationDecision select( int active_requests, const std::vector & candidates, int max_speculative_requests = std::numeric_limits::max(), bool probe_uncalibrated_with_verifier = false, - bool project_cost_from_higher_occupancy = false, - bool evaluate_nonconvex_prefixes = false) const { + bool probe_missing_routes = true) const { AdaptiveVerificationDecision out; if (active_requests <= 0 || candidates.empty() || !has_autoregressive_cost(active_requests)) { @@ -175,18 +339,90 @@ class AdaptiveVerificationRanker { std::isfinite(candidate.expected_tokens) ? candidate.expected_tokens : 1.0, 1.0, candidate.maximum_tokens); + if (!std::isfinite(candidate.routing_prior)) { + candidate.routing_prior = 0.0; + } (candidate.calibrated ? known : unknown).push_back(candidate); } if (known.empty() && unknown.empty()) return out; - auto by_value = [](const AdaptiveVerificationCandidate & a, - const AdaptiveVerificationCandidate & b) { + auto has_stable_useful_yield = + [&](const AdaptiveVerificationCandidate & candidate) { + return candidate.expected_tokens >= + config_.homogeneous_minimum_expected_tokens && + candidate.evidence_samples >= + config_.homogeneous_minimum_samples; + }; + auto by_value = [&](const AdaptiveVerificationCandidate & a, + const AdaptiveVerificationCandidate & b) { if (a.expected_tokens != b.expected_tokens) { return a.expected_tokens > b.expected_tokens; } return a.request < b.request; }; + const bool homogeneous_speculative_cohort = + static_cast(known.size()) == active_requests && + forms_homogeneous_cohort( + known, /*require_stable_local_evidence=*/true); std::stable_sort(known.begin(), known.end(), by_value); + if (homogeneous_speculative_cohort) { + // Once every active request has proven useful, favor the lagging + // requests so scarce verifier lanes do not create a long AR tail. + std::stable_sort( + known.begin(), known.end(), + [](const AdaptiveVerificationCandidate & a, + const AdaptiveVerificationCandidate & b) { + if (a.progress_tokens != b.progress_tokens) { + return a.progress_tokens < b.progress_tokens; + } + if (a.expected_tokens != b.expected_tokens) { + return a.expected_tokens > b.expected_tokens; + } + return a.request < b.request; + }); + } else { + // For mixed cohorts, rotate only within contiguous half-token + // request-value buckets. This preserves ordering across materially + // different expected yields and never uses the raw prompt hint. + std::map> fair_group_positions; + for (std::size_t i = 0; i < known.size(); ++i) { + const int yield_bucket = static_cast( + std::floor(known[i].expected_tokens * 2.0)); + fair_group_positions[yield_bucket].push_back(i); + } + for (const auto & [yield_bucket, positions] : + fair_group_positions) { + (void)yield_bucket; + if (positions.size() < 2 || + !std::all_of( + positions.begin(), positions.end(), + [&](std::size_t position) { + return has_stable_useful_yield(known[position]); + })) { + continue; + } + std::vector members; + members.reserve(positions.size()); + for (std::size_t position : positions) { + members.push_back(known[position]); + } + std::stable_sort( + members.begin(), members.end(), + [](const AdaptiveVerificationCandidate & a, + const AdaptiveVerificationCandidate & b) { + if (a.progress_tokens != b.progress_tokens) { + return a.progress_tokens < b.progress_tokens; + } + if (a.expected_tokens != b.expected_tokens) { + return a.expected_tokens > b.expected_tokens; + } + return a.request < b.request; + }); + for (std::size_t i = 0; i < positions.size(); ++i) { + known[positions[i]] = members[i]; + } + } + } std::stable_sort(unknown.begin(), unknown.end(), [](const AdaptiveVerificationCandidate & a, const AdaptiveVerificationCandidate & b) { @@ -206,39 +442,25 @@ class AdaptiveVerificationRanker { static_cast(active_requests) / autoregressive_cost_us(active_requests); const double required = baseline * config_.minimum_gain; - int admitted_prefix = 0; - int missing_cost_prefix = 0; double best = baseline; double admitted_goodput = baseline; double expected_total = static_cast(active_requests); const int prefix_limit = std::max(0, std::min( active_requests, max_speculative_requests)); - auto effective_route_cost = [&](int speculative_requests) { - if (has_route_cost(active_requests, speculative_requests)) { - return route_cost_us( - active_requests, speculative_requests); - } - return project_cost_from_higher_occupancy - ? projected_route_cost_us( - active_requests, speculative_requests) - : std::numeric_limits::infinity(); - }; - bool measured_route = false; + // Evaluate every measured width independently. GPU occupancy makes + // these costs non-convex: k=3 may win even when k=1 and k=2 lose. for (int k = 1; k <= static_cast(known.size()) && k <= prefix_limit; ++k) { expected_total += known[(size_t)k - 1].expected_tokens - 1.0; - const double route_us = effective_route_cost(k); - if (!std::isfinite(route_us)) { - if (missing_cost_prefix == 0) missing_cost_prefix = k; - if (!evaluate_nonconvex_prefixes) break; - continue; - } - measured_route = true; + if (!has_route_cost(active_requests, k)) continue; + const double route_us = route_cost_us(active_requests, k); const double throughput = expected_total / route_us; - // Request-level route costs can be non-convex: on wide GPUs a - // k=3 compact bundle may win even when k=1 and k=2 do not. - if (!evaluate_nonconvex_prefixes && throughput <= best) break; + const bool protects_ar_peers = k == active_requests || + homogeneous_speculative_cohort || + route_us <= config_.maximum_ar_peer_slowdown * + autoregressive_cost_us(active_requests); + if (!protects_ar_peers) continue; if (throughput > best) { best = throughput; if (throughput >= required) { @@ -248,104 +470,105 @@ class AdaptiveVerificationRanker { } } - // By default, calibration_request asks the concrete adapter for a - // cheap confidence signal. Verifier-side calibration is opt-in below - // for adapters whose separate confidence pass would duplicate work. - - // A calibrated route shape without a hardware sample gets one bounded - // probe. Grow a measured winning prefix by one, but do not walk every - // smaller shape after another measured width has already lost. - const bool probe_missing = missing_cost_prefix > 0 && - (!measured_route || - (admitted_prefix > 0 && - missing_cost_prefix == admitted_prefix + 1)); - if (probe_missing) { - out.requests.reserve((size_t)missing_cost_prefix); - for (int i = 0; i < missing_cost_prefix; ++i) { + // Profile one missing exact width per useful decode step. Do not infer + // C from a neighboring occupancy or stop after a losing smaller width: + // both assumptions hide real GPU occupancy boundaries. Adapters with + // no cheap confidence signal may include unknown requests in this + // bounded prefix and learn their yield from accepted output. + const int probe_candidates = std::min( + prefix_limit, + static_cast(known.size()) + + (probe_uncalibrated_with_verifier + ? static_cast(unknown.size()) : 0)); + int missing_width = 0; + for (int k = 1; probe_missing_routes && + k <= probe_candidates; ++k) { + if (has_route_cost(active_requests, k)) continue; + missing_width = k; + break; + } + if (missing_width > 0) { + out.requests.reserve((size_t)missing_width); + const int known_count = std::min( + missing_width, static_cast(known.size())); + for (int i = 0; i < known_count; ++i) { out.requests.push_back(known[(size_t)i].request); } + for (int i = 0; + static_cast(out.requests.size()) < missing_width; + ++i) { + out.requests.push_back(unknown[(size_t)i].request); + } + out.calibration_request = -1; out.exploring = true; return out; } - if (admitted_prefix > 0) { - out.requests.reserve((size_t)admitted_prefix); - for (int i = 0; i < admitted_prefix; ++i) { - out.requests.push_back(known[(size_t)i].request); - } - out.predicted_gain = admitted_goodput / baseline; - } - // A speculator without a cheap calibrated confidence head can turn a - // bounded verification bundle into useful decoding and calibration. - // Prefer the executor's widest compact bundle because request-level - // route costs can be non-convex. Once sampled, its observed accepted - // yield gates later unknown bundles; perfect-acceptance optimism is - // used only before a yield sample exists. DDTree uses this at high - // occupancy, while DSpark can keep using its confidence head. + // Once the hardware widths are known, an adapter without a cheap + // confidence signal may calibrate one new request inside a useful + // measured route. The optimistic maximum-token bound prevents probes + // that cannot possibly clear the normal safety margin. if (probe_uncalibrated_with_verifier && !unknown.empty() && - !out.exploring && - static_cast(out.requests.size()) < prefix_limit) { - const int probe_count = std::min( - static_cast(unknown.size()), - prefix_limit - static_cast(out.requests.size())); - const int probe_prefix = - static_cast(out.requests.size()) + probe_count; - double optimistic_total = - static_cast(active_requests); - for (int request : out.requests) { - const auto candidate = std::find_if( - known.begin(), known.end(), - [request](const AdaptiveVerificationCandidate & item) { - return item.request == request; - }); - if (candidate != known.end()) { - optimistic_total += candidate->expected_tokens - 1.0; + prefix_limit > 0) { + int calibration_width = 0; + double calibration_goodput = required; + const int calibration_limit = std::min( + prefix_limit, static_cast(known.size()) + 1); + double known_total = static_cast(active_requests); + for (int k = 1; k <= calibration_limit; ++k) { + if (k > 1) { + known_total += + known[(size_t)k - 2].expected_tokens - 1.0; + } + if (!has_route_cost(active_requests, k)) continue; + const double route_us = route_cost_us(active_requests, k); + // This is one bounded evidence-gathering step, not a steady + // route. It may cross the steady AR-peer guard so an otherwise + // deadlocked homogeneous cohort can prove (or reject) itself. + const double optimistic_total = known_total + + unknown.front().maximum_tokens - 1.0; + const double throughput = optimistic_total / route_us; + if (throughput > calibration_goodput) { + calibration_goodput = throughput; + calibration_width = k; } } - for (int i = 0; i < probe_count; ++i) { - optimistic_total += - unknown[(size_t)i].maximum_tokens - 1.0; - } - const double probe_route_us = - effective_route_cost(probe_prefix); - const bool unseen_shape = !std::isfinite(probe_route_us); - const bool can_win = unseen_shape || - optimistic_total / probe_route_us >= required; - if (can_win) { - for (int i = 0; i < probe_count; ++i) { - out.requests.push_back(unknown[(size_t)i].request); + if (calibration_width > 0) { + out.requests.reserve((size_t)calibration_width); + for (int i = 0; i < calibration_width - 1; ++i) { + out.requests.push_back(known[(size_t)i].request); } + out.requests.push_back(unknown.front().request); out.calibration_request = -1; out.exploring = true; + return out; } } - return out; - } -private: - double projected_route_cost_us( - int active_requests, int speculative_requests) const { - if (!has_autoregressive_cost(active_requests)) { - return std::numeric_limits::infinity(); - } - for (int higher = active_requests + 1; - static_cast(higher) < route_cost_known_.size(); - ++higher) { - if (!has_autoregressive_cost(higher) || - !has_route_cost(higher, speculative_requests)) { - continue; + if (admitted_prefix > 0) { + out.requests.reserve((size_t)admitted_prefix); + for (int i = 0; i < admitted_prefix; ++i) { + out.requests.push_back(known[(size_t)i].request); } - const double relative_cost = - route_cost_us(higher, speculative_requests) / - autoregressive_cost_us(higher); - return relative_cost * autoregressive_cost_us(active_requests); + out.predicted_gain = admitted_goodput / baseline; } - return std::numeric_limits::infinity(); + return out; } +private: static AdaptiveVerificationConfig sanitize( AdaptiveVerificationConfig config) { config.minimum_gain = std::max(1.0, config.minimum_gain); + config.maximum_ar_peer_slowdown = + std::max(1.0, config.maximum_ar_peer_slowdown); + config.homogeneous_minimum_expected_tokens = + std::max(1.0, config.homogeneous_minimum_expected_tokens); + config.homogeneous_minimum_relative_yield = std::clamp( + config.homogeneous_minimum_relative_yield, 0.0, 1.0); + config.homogeneous_minimum_samples = + std::max(1, config.homogeneous_minimum_samples); + config.routing_prior_minimum_samples = + std::max(1, config.routing_prior_minimum_samples); config.cost_ewma_alpha = std::clamp(config.cost_ewma_alpha, 0.0, 1.0); return config; @@ -354,6 +577,35 @@ class AdaptiveVerificationRanker { AdaptiveVerificationConfig config_; std::vector> route_cost_us_; std::vector> route_cost_known_; + struct YieldEstimate { + double expected_tokens = 1.0; + std::size_t samples = 0; + }; + + static constexpr std::size_t kMaximumYieldSamples = 64; + + static bool valid_yield(double emitted_tokens) { + return std::isfinite(emitted_tokens) && emitted_tokens >= 1.0; + } + + static void update_yield(YieldEstimate & estimate, + double emitted_tokens) { + const std::size_t next_samples = + std::min(estimate.samples + 1, kMaximumYieldSamples); + const double alpha = 1.0 / static_cast(next_samples); + if (estimate.samples == 0) { + estimate.expected_tokens = emitted_tokens; + } else { + estimate.expected_tokens = alpha * emitted_tokens + + (1.0 - alpha) * estimate.expected_tokens; + } + estimate.samples = next_samples; + } + + // Capped running means represent the complete request/profile while still + // adapting gradually if a model or speculator changes behavior. + std::map request_yield_; + std::map routing_prior_yield_; }; } // namespace dflash::common diff --git a/server/src/common/concurrency/speculation_prompt_prior.h b/server/src/common/concurrency/speculation_prompt_prior.h index 0c112860f..66c90f1aa 100644 --- a/server/src/common/concurrency/speculation_prompt_prior.h +++ b/server/src/common/concurrency/speculation_prompt_prior.h @@ -2,11 +2,12 @@ // Training-free request prior for cold-start speculative routing. // -// This is deliberately a weak prior, not the final decision: obvious -// structured/code prompts are ranked ahead of neutral prompts, obvious -// conversational/creative prompts stay on AR, and measured verifier goodput -// remains authoritative after admission. The policy is model-neutral and can -// be replaced by a learned prompt ranker without changing the verifier. +// This is deliberately a weak prior, not an eligibility decision: obvious +// structured/code prompts are sampled before neutral or conversational ones +// during cold start. A bucket is reused only after several measured verifier +// outcomes, and request-local evidence takes over as it stabilizes. The policy +// is model-neutral and can be replaced by a learned prompt ranker without +// changing the verifier. #include #include diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index dcdfd6ac6..256258921 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -65,6 +65,15 @@ const AdaptiveVerificationOracle & adaptive_verification_oracle() { return oracle; } +int adaptive_calibration_interval() { + static const int interval = []() { + const char * value = std::getenv( + "DFLASH_ADAPTIVE_VERIFY_CALIBRATION_STEPS"); + return value ? std::max(1, std::atoi(value)) : 64; + }(); + return interval; +} + } // namespace Qwen35SeqEngine::Qwen35SeqEngine( @@ -1081,6 +1090,8 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( const SamplerCfg & sampler) { AdmitResult result = slots_.admit(request_id, prompt, sampler); if (result.status == AdmitResult::Status::admitted) { + adaptive_verification_.forget_request(request_id); + compact_adaptive_verification_.forget_request(request_id); compact_tree_cohort_[(size_t)result.slot] = 0; reset_recurrent_slot(b_.cache_, result.slot); if (slots_.residency_active()) { @@ -1328,10 +1339,16 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const bool adaptive_enabled = !(adaptive && std::atoi(adaptive) == 0); const AdaptiveVerificationOracle & oracle = adaptive_verification_oracle(); - const bool inherited_compact_tree = std::any_of( + const bool has_compact_cohort_member = std::any_of( inputs.begin(), inputs.end(), [&](const StepInput & in) { return compact_tree_cohort_[(size_t)in.slot] != 0; }); + const bool inherited_compact_tree = + !inputs.empty() && std::all_of( + inputs.begin(), inputs.end(), [&](const StepInput & in) { + return compact_tree_cohort_[(size_t)in.slot] != 0; + }); + const bool starts_compact_cohort = !has_compact_cohort_member; // Full DDTree depth wins in the low-occupancy regime. Once target AR is // well batched, a compact proposal lowers the marginal verification cost // enough for a small profitable request subset. Keep that compact shape @@ -1339,14 +1356,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // still gets the established full-depth path. constexpr int kFullTreeMaxConcurrency = 4; constexpr int kCompactTreeBudget = 8; - constexpr uint64_t kVerifierCalibrationSamples = 16; const bool compact_tree_route = inputs.size() > kFullTreeMaxConcurrency || inherited_compact_tree; - const bool structured_compact_cohort = - !compact_tree_route || std::all_of( - inputs.begin(), inputs.end(), [&](const StepInput & in) { - return slots_.slot(in.slot).sampler.speculation_prompt_hint > 0; - }); const int adaptive_tree_budget = compact_tree_route ? std::min(b_.cfg_.ddtree_budget, kCompactTreeBudget) @@ -1359,9 +1370,10 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { ? compact_adaptive_verification_ : adaptive_verification_; // Keep the scheduler independent of the concrete speculation algorithm. - // DDTree contributes either a cheap draft-confidence estimate or an EWMA - // of target-accepted useful tokens. A future DSpark adapter can contribute - // calibrated prefix survival from its confidence head. + // DDTree contributes either a cheap draft-confidence estimate or a + // profile-local mean of target-accepted useful tokens. A future DSpark + // adapter can contribute calibrated prefix survival from its confidence + // head while reusing the same selection policy. std::vector candidates; candidates.reserve(inputs.size()); auto collect_candidates = [&]() { @@ -1370,39 +1382,27 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (!ddtree_input_eligible(in)) continue; const Qwen35Slot & seq = slots_.slot(in.slot); const int prompt_hint = seq.sampler.speculation_prompt_hint; - const bool structured_warmup = - adaptive_enabled && !oracle.forces_selection() && - compact_tree_route && structured_compact_cohort && - prompt_hint > 0 && - seq.ddtree_sampled_steps < kVerifierCalibrationSamples; - if (!oracle.forces_selection() && - !slots_.ddtree_speculation_allowed(in.slot) && - !structured_warmup) { - continue; - } - if (adaptive_enabled && !oracle.forces_selection() && - compact_tree_route && prompt_hint <= 0) { - continue; - } - const SpeculationGoodputController & policy = - seq.speculation; + const std::optional estimate = + route_ranker.estimate_request_yield( + seq.request_id, static_cast(prompt_hint)); candidates.push_back({ in.slot, - policy.expected_emitted_tokens(), + estimate ? estimate->expected_tokens : 1.0, static_cast(compact_tree_route ? tree_budget + 1 : b_.dw_.block_size), - policy.has_expected_emitted_tokens(), + estimate.has_value(), static_cast(prompt_hint), + estimate ? estimate->evidence_samples : 0, + seq.generated_tokens(), }); } }; collect_candidates(); AdaptiveVerificationDecision decision; - const int adaptive_speculation_limit = - b_.cache_.tree_capture_lanes > 0 - ? b_.cache_.tree_capture_lanes - : static_cast(inputs.size()); + const int adaptive_speculation_limit = compact_tree_route + ? std::max(0, b_.cache_.tree_capture_lanes) + : static_cast(inputs.size()); if (oracle.forces_selection()) { const int force_limit = std::min( adaptive_speculation_limit, @@ -1422,36 +1422,85 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { logged = true; } } else if (adaptive_enabled) { - const bool drafter_side_calibration = !compact_tree_route; - // The compact verifier has three directly captured lanes on Strix - // Halo. Oracle sweeps show that a bundle remains worthwhile while - // those lanes cover at least one third of the active cohort. - constexpr int kCompactActiveRequestsPerLane = 3; - const int compact_calibration_limit = - kCompactActiveRequestsPerLane * adaptive_speculation_limit; + const int exact_profile_width = std::min( + adaptive_speculation_limit, + static_cast(candidates.size())); + const AdaptiveVerificationDecision steady_decision = + route_ranker.select( + static_cast(inputs.size()), candidates, + adaptive_speculation_limit, + /*probe_uncalibrated_with_verifier=*/false, + /*probe_missing_routes=*/false); + const bool request_frontier_ready = + static_cast(steady_decision.requests.size()) >= + exact_profile_width; + int & calibration_cooldown = compact_tree_route + ? compact_adaptive_calibration_cooldown_ + : adaptive_calibration_cooldown_; + const bool exact_profile_ready = route_ranker.has_exact_profile( + static_cast(inputs.size()), exact_profile_width); + const bool exact_profile_started = + route_ranker.has_speculative_profile_sample( + static_cast(inputs.size()), exact_profile_width); + // Finish a profile already started at this occupancy, but do not make + // a shrinking compact cohort pay for brand-new C-tail exploration. + // A fresh batch at that same C still profiles normally, and a prior + // exact profile remains reusable. + const bool probe_missing_routes = + !compact_tree_route || !inherited_compact_tree || + exact_profile_started; + // Full DDTree uses its drafter-side confidence adapter. Compact DDTree + // avoids a duplicate draft pass and calibrates one unknown request in + // useful verification instead. DSpark can take the former path with + // its calibrated confidence head. + const bool drafter_side_calibration = + !compact_tree_route && probe_missing_routes; + if (!exact_profile_ready || !request_frontier_ready || + inputs.size() <= 3) { + calibration_cooldown = 0; + } else if (calibration_cooldown > 0) { + --calibration_cooldown; + } const bool verifier_side_calibration = - compact_tree_route && structured_compact_cohort && - static_cast(inputs.size()) <= compact_calibration_limit; - const int ranked_speculation_limit = - !compact_tree_route - ? adaptive_speculation_limit - : (structured_compact_cohort && - static_cast(inputs.size()) <= - compact_calibration_limit - ? adaptive_speculation_limit : 0); + compact_tree_route && probe_missing_routes && + (!exact_profile_ready || + calibration_cooldown == 0); + std::vector verifier_candidates; + const std::vector * decision_candidates = + &candidates; + const bool potentially_homogeneous = + route_ranker.forms_homogeneous_cohort( + candidates, /*require_stable_local_evidence=*/false); + if (verifier_side_calibration) { + verifier_candidates = candidates; + for (AdaptiveVerificationCandidate & candidate : + verifier_candidates) { + if (candidate.calibrated && + potentially_homogeneous && + !route_ranker.has_stable_evidence( + candidate.evidence_samples) && + route_ranker.has_useful_yield( + candidate.expected_tokens)) { + // A slow route may earn the homogeneous exception only + // when every active request still looks useful. Stable + // low-yield peers make further local proof pure overhead. + candidate.calibrated = false; + } + } + decision_candidates = &verifier_candidates; + } decision = route_ranker.select( - static_cast(inputs.size()), candidates, - ranked_speculation_limit, verifier_side_calibration, - verifier_side_calibration, - /*evaluate_nonconvex_prefixes=*/compact_tree_route); - if (inputs.size() <= 3) { - adaptive_calibration_cooldown_ = 0; - } else if (adaptive_calibration_cooldown_ > 0) { - --adaptive_calibration_cooldown_; + static_cast(inputs.size()), *decision_candidates, + adaptive_speculation_limit, verifier_side_calibration, + probe_missing_routes); + if (compact_tree_route && exact_profile_ready && + verifier_side_calibration && decision.exploring) { + calibration_cooldown = + adaptive_calibration_interval(); } if (drafter_side_calibration && decision.calibration_request >= 0 && - adaptive_calibration_cooldown_ == 0) { + calibration_cooldown == 0) { const auto input = std::find_if( inputs.begin(), inputs.end(), [&](const StepInput & in) { @@ -1461,21 +1510,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const std::optional expected = estimate_ddtree_expected_tokens(*input, tree_budget); if (expected) { - slots_.slot(input->slot).speculation - .observe_expected_tokens(*expected); - // Drafter confidence is cheap relative to target - // verification but not free on the current sequential - // per-slot DDTree path. At high occupancy, amortize cold - // calibration across decode steps; already-ranked requests - // remain independently selectable on every step. - if (inputs.size() > 3) { - static const int interval = []() { - const char * value = std::getenv( - "DFLASH_ADAPTIVE_VERIFY_CALIBRATION_STEPS"); - return value ? std::max(1, std::atoi(value)) : 64; - }(); - adaptive_calibration_cooldown_ = interval; - } + route_ranker.observe_request_yield( + slots_.slot(input->slot).request_id, *expected); std::fprintf(stderr, "[parallel-ddtree] confidence request=%llu slot=%d " "expected_tokens=%.3f\n", @@ -1483,9 +1519,26 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { slots_.slot(input->slot).request_id, input->slot, *expected); collect_candidates(); + const AdaptiveVerificationDecision steady_after = + route_ranker.select( + static_cast(inputs.size()), candidates, + adaptive_speculation_limit, + /*probe_uncalibrated_with_verifier=*/false, + /*probe_missing_routes=*/false); + if (inputs.size() > 3 && exact_profile_ready && + static_cast(steady_after.requests.size()) >= + exact_profile_width) { + // Drafter confidence is not free on the sequential + // per-slot DDTree path. Fill the executor frontier + // immediately, then amortize lower-priority newcomers. + calibration_cooldown = + adaptive_calibration_interval(); + } decision = route_ranker.select( static_cast(inputs.size()), candidates, - adaptive_speculation_limit); + adaptive_speculation_limit, + verifier_side_calibration, + probe_missing_routes); } } } @@ -1505,14 +1558,12 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { StepPlan ar_plan; speculative_plan.decode.reserve(decision.requests.size()); ar_plan.decode.reserve(inputs.size() - decision.requests.size()); - std::vector observe_ar((size_t)n_slots, 0); for (const StepInput & in : inputs) { const bool eligible = ddtree_input_eligible(in); if (eligible && selected[(size_t)in.slot]) { speculative_plan.decode.push_back(in); } else { ar_plan.decode.push_back(in); - if (eligible) observe_ar[(size_t)in.slot] = 1; } } @@ -1533,7 +1584,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (!mixed->ok()) return std::move(*mixed); routed_result = std::move(*mixed); if (adaptive_enabled && !oracle.forces_selection() && - compact_tree_route) { + compact_tree_route && starts_compact_cohort) { for (const StepInput & in : inputs) { compact_tree_cohort_[(size_t)in.slot] = 1; } @@ -1558,68 +1609,20 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { ar_plan.decode.size()); } - const bool use_route_ar_baseline = - adaptive_enabled && !oracle.forces_selection() && - compact_tree_route && decision.exploring && - route_ranker.has_autoregressive_cost( - static_cast(inputs.size())); - const double route_ar_us = use_route_ar_baseline - ? route_ranker.autoregressive_cost_us( - static_cast(inputs.size())) : 0.0; - auto log_transition = [&](int slot, - SpeculationGoodputTransition transition, - const char * observed_route, - double emitted_tokens, - double elapsed_us) { - if (transition == SpeculationGoodputTransition::none) return; - const Qwen35Slot & seq = slots_.slot(slot); - const SpeculationGoodputController & policy = seq.speculation; - const char * action = - transition == SpeculationGoodputTransition::enabled - ? "enable" : "disable"; - const double spec_tps = policy.speculative_goodput() * 1.0e6; - const double ar_tps = policy.ar_goodput() * 1.0e6; - std::fprintf(stderr, - "[parallel-ddtree] adaptive route request=%llu slot=%d " - "action=%s observed=%s sample=%llu emitted=%.0f " - "elapsed_us=%.0f spec_tok_s=%.2f ar_tok_s=%.2f\n", - (unsigned long long)seq.request_id, slot, action, observed_route, - (unsigned long long)seq.ddtree_sampled_steps, emitted_tokens, - elapsed_us, spec_tps, ar_tps); - }; - for (DecodeOutput & out : routed_result.decode) { if (oracle.forces_selection()) break; if (out.failed || out.slot < 0 || out.slot >= n_slots) continue; if (selected[(size_t)out.slot]) { const double emitted = static_cast(out.ddtree_accepted_tokens + 1); - SpeculationGoodputTransition transition = - slots_.record_speculation_sample( - out.slot, emitted, route_us); - bool seeded_route_ar = false; - const Qwen35Slot & observed_seq = slots_.slot(out.slot); - const bool warmup_complete = - observed_seq.sampler.speculation_prompt_hint <= 0 || - observed_seq.ddtree_sampled_steps >= - kVerifierCalibrationSamples; - if (use_route_ar_baseline && warmup_complete && - !observed_seq.speculation.has_ar_goodput()) { - transition = slots_.record_ar_sample( - out.slot, route_ar_us); - seeded_route_ar = true; - } - if (transition == SpeculationGoodputTransition::disabled) { - out.ddtree_suspensions = 1; - } - log_transition( - out.slot, transition, - seeded_route_ar ? "speculation-vs-route-ar" : "speculation", - emitted, route_us); - } else if (observe_ar[(size_t)out.slot]) { - const SpeculationGoodputTransition transition = - slots_.record_ar_sample(out.slot, route_us); - log_transition(out.slot, transition, "ar", 1.0, route_us); + Qwen35Slot & observed_seq = slots_.slot(out.slot); + ++observed_seq.ddtree_sampled_steps; + route_ranker.observe_request_yield( + observed_seq.request_id, emitted); + route_ranker.observe_routing_prior_yield( + static_cast( + observed_seq.sampler.speculation_prompt_hint), + emitted); } } @@ -2042,6 +2045,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step_regular(const StepPlan & plan) { void Qwen35SeqEngine::retire(int slot) { if (!slots_.is_active(slot)) return; + const std::uint64_t request_id = slots_.slot(slot).request_id; + adaptive_verification_.forget_request(request_id); + compact_adaptive_verification_.forget_request(request_id); slots_.retire(slot); compact_tree_cohort_[(size_t)slot] = 0; } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 510c1ac5f..b2cb972c1 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -140,6 +140,7 @@ class Qwen35SeqEngine final : public SeqEngine { AdaptiveVerificationRanker adaptive_verification_; AdaptiveVerificationRanker compact_adaptive_verification_; int adaptive_calibration_cooldown_ = 0; + int compact_adaptive_calibration_cooldown_ = 0; std::vector compact_tree_cohort_; ggml_context * feature_view_ctx_ = nullptr; std::vector slot_feature_mirrors_; diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index 67458afb4..49cff096d 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -492,9 +492,9 @@ int main() { CHECK(!mgr.has_prefill_prompt_at_least(768)); } - // Requests in one concurrent cohort make independent routing decisions - // from their measured useful-token goodput. A predictable request can keep - // DDTree while a low-yield chat-like peer switches to packed AR. + // The legacy per-slot controller remains a tested primitive. Production + // concurrency routing uses the shared exact-C ranker so one authority owns + // request selection and mixed-route cost. { const luce_test::ScopedEnvVar adaptive( "DFLASH_DDTREE_ADAPTIVE", nullptr); @@ -529,9 +529,8 @@ int main() { CHECK(mgr.slot(code.slot).ddtree_sampled_steps == 1); CHECK(mgr.slot(chat.slot).ddtree_sampled_steps == 1); - // The production default makes one bounded decision per request; an - // expensive speculator does not periodically interrupt a winning AR - // route. Re-probing remains available through controller config. + // Its default still makes one bounded decision per request; explicit + // controller users may opt into periodic re-probing. for (int i = 0; i < 32; ++i) { CHECK(mgr.record_ar_sample(chat.slot, 1000.0) == SpeculationGoodputTransition::none); diff --git a/server/test/test_speculation_goodput.cpp b/server/test/test_speculation_goodput.cpp index 980609144..5ce384533 100644 --- a/server/test/test_speculation_goodput.cpp +++ b/server/test/test_speculation_goodput.cpp @@ -163,8 +163,8 @@ int main() { CHECK(decision.predicted_gain > 1.05); } - // Candidates are ranked by expected value. Greedy growth stops at the - // first prefix that would reduce whole-batch throughput. + // Candidates are ranked by expected value and every measured width is + // compared against the whole-batch AR baseline. { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(5, 100.0); @@ -176,7 +176,7 @@ int main() { {4, 1.2, 8.0, true}, }; const AdaptiveVerificationDecision decision = - ranker.select(5, candidates); + ranker.select(5, candidates, /*max_speculative_requests=*/2); CHECK(decision.requests.size() == 1); CHECK(decision.requests[0] == 9); } @@ -196,82 +196,295 @@ int main() { CHECK(decision.calibration_request == 1); } - // A concrete adapter may calibrate a bounded compact bundle inside useful - // verification. The first unseen route probes up to executor capacity, - // never the whole unknown cohort. + // Compact verifier calibration discovers every exact route width in order, + // even when k=1 and k=2 lose. The prompt prior orders unknown requests but + // never filters them, and the globally best non-convex k=3 route wins. { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(8, 100.0); std::vector candidates = { {5, 1.0, 9.0, false, 1.0}, {2, 1.0, 9.0, false, -1.0}, + {9, 1.0, 9.0, false, 0.0}, }; - const AdaptiveVerificationDecision probe = + AdaptiveVerificationDecision probe = ranker.select(8, candidates, 3, /*probe_uncalibrated_with_verifier=*/true); CHECK(probe.exploring); - CHECK(probe.requests.size() == 2); + CHECK(probe.requests.size() == 1); CHECK(probe.requests[0] == 5); - CHECK(probe.requests[1] == 2); CHECK(probe.calibration_request == -1); + ranker.observe_route(8, 1, 160.0); + candidates[0] = {5, 4.0, 9.0, true, 1.0}; + + probe = ranker.select(8, candidates, 3, + /*probe_uncalibrated_with_verifier=*/true); + CHECK(probe.exploring); + CHECK(probe.requests.size() == 2); + CHECK(probe.requests[0] == 5); + CHECK(probe.requests[1] == 9); + ranker.observe_route(8, 2, 170.0); + candidates[2] = {9, 4.0, 9.0, true, 0.0}; + + probe = ranker.select(8, candidates, 3, + /*probe_uncalibrated_with_verifier=*/true); + CHECK(probe.exploring); + CHECK(probe.requests.size() == 3); + CHECK(probe.requests[0] == 5); + CHECK(probe.requests[1] == 9); + CHECK(probe.requests[2] == 2); + ranker.observe_route(8, 3, 105.0); + candidates[1] = {2, 4.0, 9.0, true, -1.0}; + + const AdaptiveVerificationDecision decision = + ranker.select(8, candidates, 3, + /*probe_uncalibrated_with_verifier=*/true); + CHECK(!decision.exploring); + CHECK(decision.requests.size() == 3); + CHECK(ranker.has_exact_profile(8, 3)); } - // A losing route ratio suppresses repeated probes as occupancy falls. - // Even perfect acceptance cannot make this projected route beat AR. + // A profitable token-count estimate cannot strand slow AR peers. The same + // route is allowed only after every active request has enough supporting + // evidence; one cache hit or lucky accepted path is insufficient. { AdaptiveVerificationRanker ranker; - ranker.observe_autoregressive(16, 100.0); - ranker.observe_route(16, 1, 200.0); - ranker.observe_autoregressive(15, 95.0); - std::vector candidates = { - {7, 1.0, 9.0, false}, + ranker.observe_autoregressive(5, 100.0); + ranker.observe_route(5, 1, 150.0); + ranker.observe_route(5, 2, 150.0); + std::vector mixed = { + {81, 8.0, 8.0, true}, + {82, 8.0, 8.0, true}, + {83, 1.0, 8.0, true}, + {84, 1.0, 8.0, true}, + {85, 1.0, 8.0, true}, }; - const AdaptiveVerificationDecision decision = - ranker.select(15, candidates, 3, - /*probe_uncalibrated_with_verifier=*/true, - /*project_cost_from_higher_occupancy=*/true); - CHECK(!decision.exploring); - CHECK(decision.requests.empty()); - CHECK(decision.calibration_request == 7); + CHECK(ranker.select(5, mixed, 2).requests.empty()); + + for (AdaptiveVerificationCandidate & candidate : mixed) { + candidate.expected_tokens = 8.0; + } + CHECK(ranker.select(5, mixed, 2).requests.empty()); + + // Even individually useful requests cannot waive the peer guard when + // their measured yields are too dissimilar; that creates a slow tail. + for (AdaptiveVerificationCandidate & candidate : mixed) { + candidate.expected_tokens = 3.0; + candidate.evidence_samples = 4; + } + mixed[0].expected_tokens = 8.0; + mixed[1].expected_tokens = 8.0; + CHECK(ranker.select(5, mixed, 2).requests.empty()); + + // A genuinely proven, comparable high-yield cohort may amortize the + // slower mixed step because its remaining requests take verifier lanes + // as the first requests retire. + for (AdaptiveVerificationCandidate & candidate : mixed) { + candidate.expected_tokens = 8.0; + } + CHECK(ranker.select(5, mixed, 2).requests.size() == 2); } - // A profitable measured route ratio can rank calibrated requests after - // occupancy drops, without spending a fresh hardware probe. + + // Timings from a neighboring occupancy never stand in for the exact-C + // baseline or route profile. { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(8, 100.0); ranker.observe_route(8, 1, 80.0); - ranker.observe_autoregressive(7, 90.0); std::vector candidates = { {12, 4.0, 9.0, true}, }; const AdaptiveVerificationDecision decision = - ranker.select(7, candidates, 3, - /*probe_uncalibrated_with_verifier=*/false, - /*project_cost_from_higher_occupancy=*/true); + ranker.select(7, candidates, 3); + CHECK(decision.requests.empty()); CHECK(!decision.exploring); + CHECK(!ranker.has_exact_profile(7, 1)); + } + + // Proven high-yield peers rotate scarce verifier lanes toward the request + // with less generated progress, preventing a homogeneous cohort tail. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(2, 100.0); + ranker.observe_route(2, 1, 150.0); + std::vector candidates = { + {91, 8.0, 8.0, true, 1.0, 4, 10}, + {92, 8.0, 8.0, true, 1.0, 4, 2}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(2, candidates, 1); CHECK(decision.requests.size() == 1); - CHECK(decision.requests[0] == 12); + CHECK(decision.requests[0] == 92); } - // Prefix widths are compared independently because hardware occupancy can - // make k=3 profitable even when the measured k=1 route loses. + // A losing higher-C route never suppresses a fresh exact lower-C + // measurement: GPU occupancy boundaries are non-convex. { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(8, 100.0); - ranker.observe_route(8, 1, 120.0); - ranker.observe_route(8, 3, 80.0); + ranker.observe_route(8, 1, 200.0); + ranker.observe_autoregressive(7, 95.0); + + const AdaptiveVerificationDecision decision = ranker.select( + 7, {{13, 1.0, 9.0, true}}, 1); + CHECK(decision.exploring); + CHECK(decision.requests.size() == 1); + CHECK(decision.requests[0] == 13); + } + + // A prior affects cold-start order only. Once calibrated, the high-yield + // conversational-prior request ranks ahead of a low-yield code prior. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(2, 100.0); + ranker.observe_route(2, 1, 100.0); + ranker.observe_route(2, 2, 1000.0); std::vector candidates = { - {1, 4.0, 9.0, true}, - {2, 4.0, 9.0, true}, - {3, 4.0, 9.0, true}, + {1, 1.0, 8.0, true, 1.0}, + {2, 8.0, 8.0, true, -1.0}, }; const AdaptiveVerificationDecision decision = - ranker.select(8, candidates, 3, - /*probe_uncalibrated_with_verifier=*/false, - /*project_cost_from_higher_occupancy=*/false, - /*evaluate_nonconvex_prefixes=*/true); - CHECK(!decision.exploring); - CHECK(decision.requests.size() == 3); + ranker.select(2, candidates, 2); + CHECK(decision.requests.size() == 1); + CHECK(decision.requests[0] == 2); + } + + // Verified yield can seed later requests in the same coarse routing-prior + // bucket. The cache belongs to this speculator/profile ranker and resets + // with it; no model-specific table is required. + { + AdaptiveVerificationRanker ranker; + CHECK(!ranker.routing_prior_expected_tokens(1.0).has_value()); + ranker.observe_routing_prior_yield(1.0, 4.0); + ranker.observe_routing_prior_yield(1.0, 6.0); + ranker.observe_routing_prior_yield(1.0, 4.0); + CHECK(!ranker.routing_prior_expected_tokens(1.0).has_value()); + CHECK(ranker.routing_prior_yield_samples(1.0) == 3); + ranker.observe_routing_prior_yield(1.0, 6.0); + CHECK(ranker.routing_prior_expected_tokens(1.0).value() == 5.0); + CHECK(ranker.routing_prior_yield_samples(1.0) == 4); + CHECK(!ranker.routing_prior_expected_tokens(-1.0).has_value()); + + // Early request observations are shrunk toward the stable cohort, but + // four local samples make the request's own magnitude authoritative. + for (int sample = 1; sample <= 4; ++sample) { + ranker.observe_request_yield(98, 1.0); + const auto estimate = + ranker.estimate_request_yield(98, 1.0); + CHECK(estimate.has_value()); + CHECK(estimate->expected_tokens == 5.0 - sample); + CHECK(estimate->evidence_samples == + static_cast(sample)); + } + const auto local = ranker.estimate_request_yield(98, 1.0); + CHECK(local.has_value()); + CHECK(local->expected_tokens == 1.0); + + // Per-request evidence is local to one proposal-shape ranker and is + // explicitly forgotten at request retirement. + AdaptiveVerificationRanker compact; + ranker.observe_request_yield(99, 7.0); + CHECK(ranker.request_expected_tokens(99).value() == 7.0); + CHECK(ranker.request_yield_samples(99) == 1); + CHECK(!compact.request_expected_tokens(99).has_value()); + ranker.forget_request(99); + CHECK(!ranker.request_expected_tokens(99).has_value()); + ranker.reset(); + CHECK(!ranker.routing_prior_expected_tokens(1.0).has_value()); + } + + // Shared cohort evidence can rank a new request, but it cannot unlock the + // slow-route homogeneous exception. A bounded verifier probe may cross the + // steady peer guard to gather the missing request-local evidence. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(5, 100.0); + ranker.observe_route(5, 1, 150.0); + std::vector prior_backed; + for (int request = 1; request <= 5; ++request) { + prior_backed.push_back( + {request, 8.0, 8.0, true, 1.0, 0}); + } + CHECK(ranker.select(5, prior_backed, 1).requests.empty()); + + std::vector probe = prior_backed; + for (AdaptiveVerificationCandidate & candidate : probe) { + candidate.calibrated = false; + } + const AdaptiveVerificationDecision exploring = ranker.select( + 5, probe, 1, + /*probe_uncalibrated_with_verifier=*/true); + CHECK(exploring.exploring); + CHECK(exploring.requests.size() == 1); + + for (AdaptiveVerificationCandidate & candidate : prior_backed) { + candidate.evidence_samples = 4; + } + CHECK(ranker.select(5, prior_backed, 1).requests.size() == 1); + prior_backed.back().expected_tokens = 1.0; + CHECK(ranker.select(5, prior_backed, 1).requests.empty()); + } + + // With route costs already profiled, a no-confidence adapter calibrates + // exactly one unknown request in the best optimistic measured width. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(5, 100.0); + ranker.observe_route(5, 1, 80.0); + ranker.observe_route(5, 2, 95.0); + std::vector candidates = { + {14, 4.0, 8.0, true, 1.0}, + {15, 1.0, 8.0, false, -1.0}, + {16, 1.0, 8.0, false, 0.0}, + }; + const AdaptiveVerificationDecision decision = ranker.select( + 5, candidates, 2, + /*probe_uncalibrated_with_verifier=*/true); + CHECK(decision.exploring); + CHECK(decision.requests.size() == 2); + CHECK(decision.requests[0] == 14); + CHECK(decision.requests[1] == 16); + } + + // A promising newcomer can replace one incumbent even when every executor + // lane is occupied; calibration remains one bounded K-wide step. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(5, 100.0); + ranker.observe_route(5, 1, 80.0); + ranker.observe_route(5, 2, 90.0); + std::vector candidates = { + {17, 4.0, 8.0, true}, + {18, 4.0, 8.0, true}, + {19, 1.0, 8.0, false, 1.0}, + }; + const AdaptiveVerificationDecision decision = ranker.select( + 5, candidates, 2, + /*probe_uncalibrated_with_verifier=*/true); + CHECK(decision.exploring); + CHECK(decision.requests.size() == 2); + CHECK(decision.requests[0] == 17); + CHECK(decision.requests[1] == 19); + } + + // Selection is tied to request value, not to a lane count. When the + // profitable request retires, speculation does not migrate to chat. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(2, 100.0); + ranker.observe_route(2, 1, 100.0); + std::vector code_and_chat = { + {21, 8.0, 8.0, true, 1.0}, + {22, 1.0, 8.0, true, -1.0}, + }; + AdaptiveVerificationDecision decision = + ranker.select(2, code_and_chat, 1); + CHECK(decision.requests.size() == 1); + CHECK(decision.requests[0] == 21); + ranker.observe_autoregressive(1, 60.0); + ranker.observe_route(1, 1, 100.0); + decision = ranker.select(1, {code_and_chat[1]}, 1); + CHECK(decision.requests.empty()); } // There is no concurrency cutoff: when one request pays for the route at @@ -291,6 +504,60 @@ int main() { CHECK(decision.requests[0] == 21); } + // Executor capacity is independent of C, and only eligible candidates + // consume speculative lanes. Three non-candidates remain ordinary AR. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(5, 100.0); + ranker.observe_route(5, 1, 160.0); + ranker.observe_route(5, 2, 100.0); + std::vector candidates = { + {51, 4.0, 8.0, true}, + {52, 4.0, 8.0, true}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(5, candidates, 3); + CHECK(ranker.has_exact_profile(5, 2)); + CHECK(!ranker.has_exact_profile(5, 3)); + CHECK(decision.requests.size() == 2); + } + + // The safety margin rejects a noisy 4% estimate and admits a 6% gain. + { + std::vector candidates = { + {61, 1.0, 1.0, true}, + }; + AdaptiveVerificationRanker below_margin; + below_margin.observe_autoregressive(1, 100.0); + below_margin.observe_route(1, 1, 96.0); + CHECK(below_margin.select(1, candidates).requests.empty()); + + AdaptiveVerificationRanker above_margin; + above_margin.observe_autoregressive(1, 100.0); + above_margin.observe_route(1, 1, 94.0); + CHECK(above_margin.select(1, candidates).requests.size() == 1); + } + + // DDTree accepted-path yield and DSpark conditional survival use the same + // model-neutral expected-token contract and therefore make the same choice. + { + const float confidence[] = {0.8f, 0.5f}; + const double dspark_expected = + expected_tokens_from_conditional_confidence(confidence, 2); + AdaptiveVerificationRanker ddtree; + AdaptiveVerificationRanker dspark; + for (AdaptiveVerificationRanker * ranker : {&ddtree, &dspark}) { + ranker->observe_autoregressive(4, 100.0); + ranker->observe_route(4, 1, 80.0); + } + const auto ddtree_decision = ddtree.select( + 4, {{71, 2.2, 3.0, true}}, 1); + const auto dspark_decision = dspark.select( + 4, {{71, dspark_expected, 3.0, true}}, 1); + CHECK(ddtree_decision.requests == dspark_decision.requests); + CHECK(ddtree_decision.requests.size() == 1); + } + // Executor capacity bounds the ranked prefix. This keeps every adaptive // route on the one-pass implementation even when more requests rank well. { @@ -312,6 +579,33 @@ int main() { CHECK(decision.requests[1] == 42); } + // Exact profiling and executor bounds are independent of C. Every fresh + // occupancy from 1 through 16 discovers k=1..min(C, 3), including widths + // whose smaller neighbors lose. + { + for (int concurrency = 1; concurrency <= 16; ++concurrency) { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(concurrency, 100.0); + const int limit = std::min(concurrency, 3); + std::vector candidates; + for (int i = 0; i < limit; ++i) { + candidates.push_back({ + concurrency * 10 + i, 4.0, 8.0, true}); + } + for (int width = 1; width <= limit; ++width) { + const AdaptiveVerificationDecision probe = ranker.select( + concurrency, candidates, limit); + CHECK(probe.exploring); + CHECK(static_cast(probe.requests.size()) == width); + ranker.observe_route( + concurrency, width, 150.0 - 20.0 * width); + } + CHECK(ranker.has_exact_profile(concurrency, limit)); + CHECK(ranker.select(concurrency, candidates, 0) + .requests.empty()); + } + } + // A promising calibrated prefix gets one bounded hardware-cost probe when // that subbatch shape has not been observed yet. { @@ -327,6 +621,31 @@ int main() { CHECK(probe.requests[0] == 31); } + // A shrinking cohort may suppress a brand-new tail profile without + // disabling a route that was already measured at that exact occupancy. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(3, 90.0); + std::vector candidates = { + {32, 6.0, 8.0, true}, + }; + AdaptiveVerificationDecision decision = ranker.select( + 3, candidates, 1, + /*probe_uncalibrated_with_verifier=*/false, + /*probe_missing_routes=*/false); + CHECK(decision.requests.empty()); + CHECK(!decision.exploring); + CHECK(!ranker.has_speculative_profile_sample(3, 1)); + + ranker.observe_route(3, 1, 70.0); + decision = ranker.select( + 3, candidates, 1, + /*probe_uncalibrated_with_verifier=*/false, + /*probe_missing_routes=*/false); + CHECK(decision.requests.size() == 1); + CHECK(ranker.has_speculative_profile_sample(3, 1)); + } + // Kill-switch mode retains fixed speculation and ignores observations. { SpeculationGoodputController policy; From 36d3427c24a7a3f9aecfdbf1f59b2a16d48469c6 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 17 Aug 2026 19:19:56 +0000 Subject: [PATCH 14/18] perf(qwen35): make adaptive routing backlog aware --- .../concurrency/adaptive_verification.h | 36 +++- server/src/common/concurrency/seq_engine.h | 5 + .../qwen35/concurrency/qwen35_seq_engine.cpp | 158 +++++++++++------- .../qwen35/concurrency/qwen35_seq_engine.h | 4 + .../concurrency/qwen35_slot_manager.cpp | 25 --- .../qwen35/concurrency/qwen35_slot_manager.h | 13 -- server/src/qwen35/graph_builders.cpp | 31 +++- server/src/qwen35/graph_builders.h | 6 + server/src/server/http_server.cpp | 5 + server/src/server/http_server.h | 1 + server/src/server/scheduler.cpp | 2 + server/test/test_recurrent_snapshot.cpp | 15 ++ server/test/test_seq_slot_manager.cpp | 72 -------- server/test/test_speculation_goodput.cpp | 47 +++++- 14 files changed, 236 insertions(+), 184 deletions(-) diff --git a/server/src/common/concurrency/adaptive_verification.h b/server/src/common/concurrency/adaptive_verification.h index 7dd983a4d..411d7dd9d 100644 --- a/server/src/common/concurrency/adaptive_verification.h +++ b/server/src/common/concurrency/adaptive_verification.h @@ -66,10 +66,22 @@ struct AdaptiveVerificationConfig { // otherwise scarce verifier lanes create a low-yield AR tail. double homogeneous_minimum_relative_yield = 0.60; std::size_t homogeneous_minimum_samples = 4; - std::size_t routing_prior_minimum_samples = 4; + // A complete two-lane verifier profile observes k=1 then k=2, yielding + // three request outcomes before any steady route is selected. + std::size_t routing_prior_minimum_samples = 3; + // Shared evidence ranks a new request after the minimum above, but its + // yield magnitude is shrunk toward AR until this many outcomes exist. + std::size_t routing_prior_full_weight_samples = 8; double cost_ewma_alpha = 0.35; }; +inline bool adaptive_verification_can_relax_peer_guard( + int active_requests, int verifier_request_lanes) { + return active_requests > 0 && verifier_request_lanes > 0 && + static_cast(active_requests) <= + 2 * static_cast(verifier_request_lanes) + 1; +} + // Convert conditional survival confidence into expected useful tokens: // 1 AR token plus the probability of reaching every speculative prefix. // DSpark supplies calibrated confidence-head values directly. @@ -189,8 +201,14 @@ class AdaptiveVerificationRanker { if (!has_request && !has_prior) return std::nullopt; AdaptiveVerificationYieldEstimate out; + const double prior_weight = std::min( + 1.0, static_cast(prior_it->second.samples) / + static_cast( + config_.routing_prior_full_weight_samples)); + const double trusted_prior = 1.0 + prior_weight * + (prior_it->second.expected_tokens - 1.0); if (!has_request) { - out.expected_tokens = prior_it->second.expected_tokens; + out.expected_tokens = trusted_prior; out.evidence_samples = 0; return out; } @@ -200,7 +218,6 @@ class AdaptiveVerificationRanker { out.evidence_samples = request_estimate.samples; if (!has_prior) return out; - const YieldEstimate & prior_estimate = prior_it->second; if (request_estimate.samples < config_.homogeneous_minimum_samples) { // Shrink the first few noisy request observations toward a stable // cohort mean. Once request-local evidence is stable, its measured @@ -210,9 +227,9 @@ class AdaptiveVerificationRanker { static_cast(request_estimate.samples) / static_cast( config_.homogeneous_minimum_samples)); - out.expected_tokens = prior_estimate.expected_tokens + + out.expected_tokens = trusted_prior + local_weight * (request_estimate.expected_tokens - - prior_estimate.expected_tokens); + trusted_prior); } return out; } @@ -318,7 +335,8 @@ class AdaptiveVerificationRanker { int max_speculative_requests = std::numeric_limits::max(), bool probe_uncalibrated_with_verifier = false, - bool probe_missing_routes = true) const { + bool probe_missing_routes = true, + bool enforce_ar_peer_guard = true) const { AdaptiveVerificationDecision out; if (active_requests <= 0 || candidates.empty() || !has_autoregressive_cost(active_requests)) { @@ -456,7 +474,8 @@ class AdaptiveVerificationRanker { if (!has_route_cost(active_requests, k)) continue; const double route_us = route_cost_us(active_requests, k); const double throughput = expected_total / route_us; - const bool protects_ar_peers = k == active_requests || + const bool protects_ar_peers = !enforce_ar_peer_guard || + k == active_requests || homogeneous_speculative_cohort || route_us <= config_.maximum_ar_peer_slowdown * autoregressive_cost_us(active_requests); @@ -569,6 +588,9 @@ class AdaptiveVerificationRanker { std::max(1, config.homogeneous_minimum_samples); config.routing_prior_minimum_samples = std::max(1, config.routing_prior_minimum_samples); + config.routing_prior_full_weight_samples = + std::max(config.routing_prior_minimum_samples, + config.routing_prior_full_weight_samples); config.cost_ewma_alpha = std::clamp(config.cost_ewma_alpha, 0.0, 1.0); return config; diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 6371e2552..c723b68c7 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -234,6 +234,11 @@ class SeqEngine { struct StepPlan { std::vector decode; std::vector prefills; + // True when completing a live request can immediately admit queued + // work. Engines may then optimize sustained goodput rather than the + // makespan of the current finite cohort, using measured total output + // rate for any mixed route. + bool has_refill_backlog = false; }; struct StepResult { diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 256258921..adf51676f 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -33,15 +33,6 @@ namespace dflash::common { namespace { -// Denser than pure power-of-2: reduces padding waste at non-power-of-2 -// live counts (e.g. C=5 uses bucket=6 at 17% waste vs bucket=8 at 37.5%). -int decode_bucket_width(int live_count) { - static constexpr int buckets[] = {1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64}; - for (int b : buckets) - if (b >= live_count) return b; - return 64; -} - struct AdaptiveVerificationOracle { int speculative_requests = -1; int tree_budget = -1; @@ -65,13 +56,17 @@ const AdaptiveVerificationOracle & adaptive_verification_oracle() { return oracle; } -int adaptive_calibration_interval() { - static const int interval = []() { +int adaptive_calibration_interval(int active_requests) { + static const int configured_interval = []() { const char * value = std::getenv( "DFLASH_ADAPTIVE_VERIFY_CALIBRATION_STEPS"); - return value ? std::max(1, std::atoi(value)) : 64; + return value ? std::max(1, std::atoi(value)) : 0; }(); - return interval; + if (configured_interval > 0) return configured_interval; + // A verifier probe is increasingly expensive relative to packed AR as C + // grows. Keep adaptation responsive at low occupancy while amortizing a + // rejected high-C probe over proportionally more decode work. + return 64 * std::max(1, (active_requests + 3) / 4); } } // namespace @@ -288,7 +283,7 @@ std::optional Qwen35SeqEngine::step_ddtree( StepResult result; const int active = (int)speculative_plan.decode.size(); const int total_active = active + (int)ar_plan.decode.size(); - const int bucket = decode_bucket_width(active); + const int bucket = detail::target_paged_tree_bucket_width(active); const int T = tree_budget + 1; const int q_len = b_.dw_.block_size; const int hidden = b_.w_.n_embd; @@ -394,7 +389,6 @@ std::optional Qwen35SeqEngine::step_ddtree( } proposals.push_back(std::move(p)); } - const bool target_is_meta = b_.cache_.ssm_state.empty() || !b_.cache_.ssm_state.front() || ggml_backend_buft_is_meta(ggml_backend_buffer_get_type( @@ -449,7 +443,6 @@ std::optional Qwen35SeqEngine::step_ddtree( result.error = "packed DDTree verify graph build failed"; return result; } - const int total_tree = T * bucket; const int total_packed = total_tree + n_ar; std::vector tree_feature_rows; @@ -1091,6 +1084,7 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( AdmitResult result = slots_.admit(request_id, prompt, sampler); if (result.status == AdmitResult::Status::admitted) { adaptive_verification_.forget_request(request_id); + compact_short_adaptive_verification_.forget_request(request_id); compact_adaptive_verification_.forget_request(request_id); compact_tree_cohort_[(size_t)result.slot] = 0; reset_recurrent_slot(b_.cache_, result.slot); @@ -1339,35 +1333,58 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const bool adaptive_enabled = !(adaptive && std::atoi(adaptive) == 0); const AdaptiveVerificationOracle & oracle = adaptive_verification_oracle(); - const bool has_compact_cohort_member = std::any_of( - inputs.begin(), inputs.end(), [&](const StepInput & in) { - return compact_tree_cohort_[(size_t)in.slot] != 0; - }); + const int inherited_compact_budget = inputs.empty() + ? 0 : compact_tree_cohort_[(size_t)inputs.front().slot]; const bool inherited_compact_tree = - !inputs.empty() && std::all_of( + inherited_compact_budget > 0 && + std::all_of( inputs.begin(), inputs.end(), [&](const StepInput & in) { - return compact_tree_cohort_[(size_t)in.slot] != 0; + return compact_tree_cohort_[(size_t)in.slot] == + inherited_compact_budget; }); - const bool starts_compact_cohort = !has_compact_cohort_member; - // Full DDTree depth wins in the low-occupancy regime. Once target AR is - // well batched, a compact proposal lowers the marginal verification cost - // enough for a small profitable request subset. Keep that compact shape - // through the cohort's low-occupancy tail; a newly admitted C<=4 cohort - // still gets the established full-depth path. + const int positive_prompt_candidates = static_cast(std::count_if( + inputs.begin(), inputs.end(), [&](const StepInput & in) { + return ddtree_input_eligible(in) && + slots_.slot(in.slot).sampler.speculation_prompt_hint > 0; + })); + // Full DDTree depth wins for low-occupancy homogeneous code. Once target + // AR is well batched, or a queued C=4 cohort is genuinely mixed, a compact + // proposal lowers the marginal verification cost enough for a small + // profitable request subset. Keep that compact shape through the cohort's + // tail; all-code C<=4 still gets the established full-depth path. constexpr int kFullTreeMaxConcurrency = 4; + constexpr int kCompactShortTreeBudget = 4; constexpr int kCompactTreeBudget = 8; + const bool starts_short_backlog_shape = + plan.has_refill_backlog && inputs.size() == kFullTreeMaxConcurrency && + positive_prompt_candidates > 0 && + positive_prompt_candidates <= 2 && + positive_prompt_candidates < static_cast(inputs.size()); const bool compact_tree_route = - inputs.size() > kFullTreeMaxConcurrency || inherited_compact_tree; + inputs.size() > kFullTreeMaxConcurrency || inherited_compact_tree || + starts_short_backlog_shape; + const int selected_compact_budget = + inherited_compact_tree && !plan.has_refill_backlog + ? inherited_compact_budget + : (positive_prompt_candidates > 0 && + positive_prompt_candidates <= 2 + ? kCompactShortTreeBudget + : kCompactTreeBudget); const int adaptive_tree_budget = compact_tree_route - ? std::min(b_.cfg_.ddtree_budget, kCompactTreeBudget) + ? std::min(b_.cfg_.ddtree_budget, selected_compact_budget) : b_.cfg_.ddtree_budget; const int tree_budget = std::clamp( oracle.tree_budget > 0 ? oracle.tree_budget : adaptive_tree_budget, 1, b_.cfg_.ddtree_budget); + const bool compact_short_shape = + compact_tree_route && tree_budget <= kCompactShortTreeBudget; AdaptiveVerificationRanker & route_ranker = compact_tree_route - ? compact_adaptive_verification_ : adaptive_verification_; + ? (compact_short_shape + ? compact_short_adaptive_verification_ + : compact_adaptive_verification_) + : adaptive_verification_; // Keep the scheduler independent of the concrete speculation algorithm. // DDTree contributes either a cheap draft-confidence estimate or a @@ -1400,9 +1417,23 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { collect_candidates(); AdaptiveVerificationDecision decision; + const int direct_speculation_limit = + detail::target_paged_tree_direct_request_limit( + b_.cache_.tree_capture_lanes); const int adaptive_speculation_limit = compact_tree_route - ? std::max(0, b_.cache_.tree_capture_lanes) + ? std::min( + direct_speculation_limit, + compact_short_shape ? 2 : direct_speculation_limit) : static_cast(inputs.size()); + // Relax peer protection for continuous batching only when the one-pass + // verifier can cover roughly half the live cohort. This derives the + // boundary from executor capacity rather than a fixed C: a wider DSpark + // adapter automatically expands the eligible concurrency range. + const bool relax_ar_peer_guard = + plan.has_refill_backlog && compact_tree_route && + adaptive_verification_can_relax_peer_guard( + static_cast(inputs.size()), adaptive_speculation_limit); + const bool enforce_ar_peer_guard = !relax_ar_peer_guard; if (oracle.forces_selection()) { const int force_limit = std::min( adaptive_speculation_limit, @@ -1425,17 +1456,10 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const int exact_profile_width = std::min( adaptive_speculation_limit, static_cast(candidates.size())); - const AdaptiveVerificationDecision steady_decision = - route_ranker.select( - static_cast(inputs.size()), candidates, - adaptive_speculation_limit, - /*probe_uncalibrated_with_verifier=*/false, - /*probe_missing_routes=*/false); - const bool request_frontier_ready = - static_cast(steady_decision.requests.size()) >= - exact_profile_width; int & calibration_cooldown = compact_tree_route - ? compact_adaptive_calibration_cooldown_ + ? (compact_short_shape + ? compact_short_adaptive_calibration_cooldown_ + : compact_adaptive_calibration_cooldown_) : adaptive_calibration_cooldown_; const bool exact_profile_ready = route_ranker.has_exact_profile( static_cast(inputs.size()), exact_profile_width); @@ -1455,16 +1479,24 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // its calibrated confidence head. const bool drafter_side_calibration = !compact_tree_route && probe_missing_routes; - if (!exact_profile_ready || !request_frontier_ready || - inputs.size() <= 3) { + const bool prompt_prior_marks_mixed = + positive_prompt_candidates > 0 && + positive_prompt_candidates < static_cast(inputs.size()); + if (!exact_profile_ready || inputs.size() <= 3) { calibration_cooldown = 0; } else if (calibration_cooldown > 0) { --calibration_cooldown; } + // Once a closed mixed cohort has measured every exact width, probing + // each remaining low-priority request cannot remove its AR critical + // path. A refill backlog changes that objective, so keep calibrating + // newcomers there; homogeneous/unknown cohorts can still prove the + // existing all-speculative exception. const bool verifier_side_calibration = compact_tree_route && probe_missing_routes && (!exact_profile_ready || - calibration_cooldown == 0); + (calibration_cooldown == 0 && + (plan.has_refill_backlog || !prompt_prior_marks_mixed))); std::vector verifier_candidates; const std::vector * decision_candidates = &candidates; @@ -1492,11 +1524,13 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { decision = route_ranker.select( static_cast(inputs.size()), *decision_candidates, adaptive_speculation_limit, verifier_side_calibration, - probe_missing_routes); + probe_missing_routes, + enforce_ar_peer_guard); if (compact_tree_route && exact_profile_ready && verifier_side_calibration && decision.exploring) { calibration_cooldown = - adaptive_calibration_interval(); + adaptive_calibration_interval( + static_cast(inputs.size())); } if (drafter_side_calibration && decision.calibration_request >= 0 && @@ -1524,7 +1558,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { static_cast(inputs.size()), candidates, adaptive_speculation_limit, /*probe_uncalibrated_with_verifier=*/false, - /*probe_missing_routes=*/false); + /*probe_missing_routes=*/false, + enforce_ar_peer_guard); if (inputs.size() > 3 && exact_profile_ready && static_cast(steady_after.requests.size()) >= exact_profile_width) { @@ -1532,13 +1567,15 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // per-slot DDTree path. Fill the executor frontier // immediately, then amortize lower-priority newcomers. calibration_cooldown = - adaptive_calibration_interval(); + adaptive_calibration_interval( + static_cast(inputs.size())); } decision = route_ranker.select( static_cast(inputs.size()), candidates, adaptive_speculation_limit, verifier_side_calibration, - probe_missing_routes); + probe_missing_routes, + enforce_ar_peer_guard); } } } @@ -1566,6 +1603,16 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { ar_plan.decode.push_back(in); } } + if (adaptive_enabled && !oracle.forces_selection() && + compact_tree_route) { + // This is work-shape metadata, not a decision to speculate. Propagate + // it on AR steps too so newly refilled slots cannot make a draining + // cohort rediscover C-tail graph costs. + for (const StepInput & in : inputs) { + compact_tree_cohort_[(size_t)in.slot] = + static_cast(tree_budget); + } + } using Clock = std::chrono::steady_clock; StepResult routed_result; @@ -1583,12 +1630,6 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (!mixed->ok()) return std::move(*mixed); routed_result = std::move(*mixed); - if (adaptive_enabled && !oracle.forces_selection() && - compact_tree_route && starts_compact_cohort) { - for (const StepInput & in : inputs) { - compact_tree_cohort_[(size_t)in.slot] = 1; - } - } } else { routed_result = step_regular(ar_plan); if (!routed_result.ok()) return routed_result; @@ -1616,7 +1657,6 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const double emitted = static_cast(out.ddtree_accepted_tokens + 1); Qwen35Slot & observed_seq = slots_.slot(out.slot); - ++observed_seq.ddtree_sampled_steps; route_ranker.observe_request_yield( observed_seq.request_id, emitted); route_ranker.observe_routing_prior_yield( @@ -1743,7 +1783,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step_regular(const StepPlan & plan) { const int live_count = (int)live_tokens_.size(); const bool with_decode = live_count > 0; - const int decode_bucket = with_decode ? decode_bucket_width(live_count) : 0; + const int decode_bucket = with_decode + ? detail::target_paged_tree_bucket_width(live_count) : 0; dec_tokens_.assign((size_t)decode_bucket, 0); dec_rows_.assign((size_t)decode_bucket * n_head_kv, scratch_row_); @@ -2047,6 +2088,7 @@ void Qwen35SeqEngine::retire(int slot) { if (!slots_.is_active(slot)) return; const std::uint64_t request_id = slots_.slot(slot).request_id; adaptive_verification_.forget_request(request_id); + compact_short_adaptive_verification_.forget_request(request_id); compact_adaptive_verification_.forget_request(request_id); slots_.retire(slot); compact_tree_cohort_[(size_t)slot] = 0; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index b2cb972c1..2d91f126d 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -138,9 +138,13 @@ class Qwen35SeqEngine final : public SeqEngine { int tree_scratch_stride_ = 0; bool capture_features_ = false; AdaptiveVerificationRanker adaptive_verification_; + AdaptiveVerificationRanker compact_short_adaptive_verification_; AdaptiveVerificationRanker compact_adaptive_verification_; int adaptive_calibration_cooldown_ = 0; + int compact_short_adaptive_calibration_cooldown_ = 0; int compact_adaptive_calibration_cooldown_ = 0; + // Zero means no sticky compact cohort. Non-zero entries store the exact + // DDTree work-shape budget so short and wide observations never mix. std::vector compact_tree_cohort_; ggml_context * feature_view_ctx_ = nullptr; std::vector slot_feature_mirrors_; diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index 755402943..e689219f1 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -2,7 +2,6 @@ #include #include -#include namespace dflash::common { @@ -13,8 +12,6 @@ Qwen35SlotManager::Qwen35SlotManager( headroom_tokens_(std::max(pool.block_size(), speculative_headroom)), residency_(residency) { slots_.assign(pool.max_sequences(), Qwen35Slot{}); - const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); - speculation_adaptive_ = !(adaptive && std::atoi(adaptive) == 0); } int Qwen35SlotManager::decoding_count() const { @@ -207,8 +204,6 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( Qwen35Slot & s = slots_[(size_t)slot]; s.phase = Qwen35SlotPhase::prefill; s.request_id = request_id; - s.ddtree_sampled_steps = 0; - s.speculation.reset(speculation_adaptive_); s.handle = handle; s.cur_pos = 0; s.prompt_len = prompt_len; @@ -227,26 +222,6 @@ SeqEngine::AdmitResult Qwen35SlotManager::admit( return r; } -bool Qwen35SlotManager::ddtree_speculation_allowed(int slot) const { - return is_active(slot) && - slots_[(size_t)slot].speculation.wants_speculation(); -} - -SpeculationGoodputTransition Qwen35SlotManager::record_speculation_sample( - int slot, double emitted_tokens, double elapsed_us) { - if (!is_active(slot)) return SpeculationGoodputTransition::none; - Qwen35Slot & s = slots_[(size_t)slot]; - ++s.ddtree_sampled_steps; - return s.speculation.observe_speculation(emitted_tokens, elapsed_us); -} - -SpeculationGoodputTransition Qwen35SlotManager::record_ar_sample( - int slot, double elapsed_us) { - if (!is_active(slot)) return SpeculationGoodputTransition::none; - Qwen35Slot & s = slots_[(size_t)slot]; - return s.speculation.observe_autoregressive(elapsed_us); -} - Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( int slot, int n_tokens) { PrefillChunk out; diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index d73ea5824..a2c13f968 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -18,7 +18,6 @@ #include "common/concurrency/paged_kv_pool.h" #include "common/concurrency/paged_kv_residency.h" -#include "common/concurrency/speculation_goodput.h" #include "common/sampler.h" #include "common/concurrency/seq_engine.h" @@ -67,12 +66,6 @@ struct Qwen35Slot { : 0; } - // Route each request from measured useful-token goodput, independently of - // the concrete speculator. DDTree records observations today; DSpark can - // use the same controller later. - SpeculationGoodputController speculation; - uint64_t ddtree_sampled_steps = 0; - bool active() const { return phase == Qwen35SlotPhase::prefill || phase == Qwen35SlotPhase::decode; @@ -153,11 +146,6 @@ class Qwen35SlotManager { std::string * error = nullptr); void take_residency_telemetry(int slot, SeqEngine::DecodeOutput & out); - bool ddtree_speculation_allowed(int slot) const; - SpeculationGoodputTransition record_speculation_sample( - int slot, double emitted_tokens, double elapsed_us); - SpeculationGoodputTransition record_ar_sample(int slot, double elapsed_us); - // One-token compatibility wrapper used by ordinary autoregressive decode. StepAppend append_token(int slot, int32_t fed_token); @@ -193,7 +181,6 @@ class Qwen35SlotManager { int max_ctx_ = 0; int headroom_tokens_; PagedKvResidencyManager * residency_ = nullptr; - bool speculation_adaptive_ = true; std::vector slots_; }; diff --git a/server/src/qwen35/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 3bd4b067b..df7a79c2b 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -10,6 +10,31 @@ namespace dflash::common { +namespace { + +constexpr int kTreeSequenceBuckets[] = { + 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, +}; + +} // namespace + +int detail::target_paged_tree_bucket_width(int requested_sequences) { + if (requested_sequences <= 0) return 0; + for (int bucket : kTreeSequenceBuckets) { + if (bucket >= requested_sequences) return bucket; + } + return 0; +} + +int detail::target_paged_tree_direct_request_limit(int capture_lanes) { + for (int requested = std::min(capture_lanes, 64); + requested > 0; --requested) { + const int bucket = target_paged_tree_bucket_width(requested); + if (bucket > 0 && bucket <= capture_lanes) return requested; + } + return 0; +} + bool detail::target_graph_capacity_for_parallel_segments( int n_parallel_segments, size_t & capacity) { @@ -42,12 +67,8 @@ bool detail::target_paged_tree_graph_capacity( int tree_width, int n_tree_seqs, size_t & capacity) { - static constexpr int tree_buckets[] = { - 1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64, - }; if (tree_width < 1 || tree_width > 256 || - std::find(std::begin(tree_buckets), std::end(tree_buckets), - n_tree_seqs) == std::end(tree_buckets) || + target_paged_tree_bucket_width(n_tree_seqs) != n_tree_seqs || (int64_t)tree_width * n_tree_seqs > INT32_MAX) { return false; } diff --git a/server/src/qwen35/graph_builders.h b/server/src/qwen35/graph_builders.h index 65018570c..6759d2dde 100644 --- a/server/src/qwen35/graph_builders.h +++ b/server/src/qwen35/graph_builders.h @@ -25,6 +25,12 @@ namespace dflash::common { namespace detail { +// Stable tree-sequence buckets shared by graph validation and the concurrent +// engine. The direct-request limit rounds capture capacity down so every +// admitted width stays on the one-pass state-capture path. +int target_paged_tree_bucket_width(int requested_sequences); +int target_paged_tree_direct_request_limit(int capture_lanes); + // Qwen's recurrent graph duplicates one small subgraph per ragged sequence. // Return a graph capacity that covers every supported concurrent bucket while // keeping the legacy allocation for the common <= 8-sequence case. diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 52a4a845c..b25a69758 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -4006,6 +4006,11 @@ ServerJob * HttpServer::try_dequeue() { return job; } +bool HttpServer::has_queued_job() { + std::lock_guard lk(queue_mu_); + return queue_head_ != nullptr; +} + ServerJob * HttpServer::dequeue_for( std::chrono::steady_clock::duration timeout) { std::unique_lock lk(queue_mu_); diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 0fa0ef718..92e4f9a9e 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -479,6 +479,7 @@ class HttpServer { // Non-blocking dequeue used for admission polling between decode steps. ServerJob * try_dequeue(); + bool has_queued_job(); // Bounded wait used only during an idle-to-busy admission window. ServerJob * dequeue_for( std::chrono::steady_clock::duration timeout); diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index c2fb270cf..428d6b2a9 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -744,6 +744,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { engine.step_plan_limits((int)step_plan.decode.size()); step_plan.prefills = plan_prefill_slices( prefill_candidates, step_limits, prefill_round_robin_start); + step_plan.has_refill_backlog = + deferred != nullptr || has_queued_job(); if (!prefill_candidates.empty()) { ++prefill_round_robin_start; } diff --git a/server/test/test_recurrent_snapshot.cpp b/server/test/test_recurrent_snapshot.cpp index 75c1efdca..7b4cb28c4 100644 --- a/server/test/test_recurrent_snapshot.cpp +++ b/server/test/test_recurrent_snapshot.cpp @@ -55,6 +55,21 @@ TEST_CASE(RecurrentSnapshotFixture, hardens_feature_smoke_paths) { 23, 16, graph_capacity) && graph_capacity == 32768); CHECK(!dflash::common::detail::target_paged_tree_graph_capacity( 257, 16, graph_capacity)); + for (int lanes = 0; lanes <= 16; ++lanes) { + const int limit = + dflash::common::detail::target_paged_tree_direct_request_limit( + lanes); + CHECK(limit >= 0); + CHECK(limit <= lanes); + if (limit > 0) { + CHECK(dflash::common::detail::target_paged_tree_bucket_width( + limit) <= lanes); + } + if (limit < lanes) { + CHECK(dflash::common::detail::target_paged_tree_bucket_width( + limit + 1) > lanes); + } + } // Mapped-tree active_slot_ids is a topology marker and can legitimately // be left without gallocr storage. Required state/query/write metadata diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index 49cff096d..31b1cb232 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -7,11 +7,9 @@ #include "qwen35/concurrency/qwen35_slot_manager.h" #include "host_check.h" -#include "scoped_env.h" #include #include -#include #include #include @@ -492,76 +490,6 @@ int main() { CHECK(!mgr.has_prefill_prompt_at_least(768)); } - // The legacy per-slot controller remains a tested primitive. Production - // concurrency routing uses the shared exact-C ranker so one authority owns - // request selection and mixed-route cost. - { - const luce_test::ScopedEnvVar adaptive( - "DFLASH_DDTREE_ADAPTIVE", nullptr); - PagedKvPool pool(8, 2, /*block_size=*/16); - Qwen35SlotManager mgr(pool, 64); - auto code = admit(mgr, 101, prompt_tokens(4), greedy_sampler()); - auto chat = admit(mgr, 102, prompt_tokens(4), greedy_sampler()); - CHECK(is_admitted(code)); - CHECK(is_admitted(chat)); - CHECK(mgr.append_prefill(code.slot, 4).ok); - CHECK(mgr.append_prefill(chat.slot, 4).ok); - mgr.commit_prefill(code.slot); - mgr.commit_prefill(chat.slot); - CHECK(mgr.ddtree_speculation_allowed(code.slot)); - CHECK(mgr.ddtree_speculation_allowed(chat.slot)); - - CHECK(mgr.record_speculation_sample( - code.slot, /*emitted_tokens=*/8, /*elapsed_us=*/2000.0) == - SpeculationGoodputTransition::none); - CHECK(mgr.record_speculation_sample( - chat.slot, /*emitted_tokens=*/1, /*elapsed_us=*/2000.0) == - SpeculationGoodputTransition::none); - // Both requests take one neighboring AR calibration step. - CHECK(!mgr.ddtree_speculation_allowed(code.slot)); - CHECK(!mgr.ddtree_speculation_allowed(chat.slot)); - CHECK(mgr.record_ar_sample(code.slot, /*elapsed_us=*/1000.0) == - SpeculationGoodputTransition::none); - CHECK(mgr.record_ar_sample(chat.slot, /*elapsed_us=*/1000.0) == - SpeculationGoodputTransition::disabled); - CHECK(mgr.ddtree_speculation_allowed(code.slot)); - CHECK(!mgr.ddtree_speculation_allowed(chat.slot)); - CHECK(mgr.slot(code.slot).ddtree_sampled_steps == 1); - CHECK(mgr.slot(chat.slot).ddtree_sampled_steps == 1); - - // Its default still makes one bounded decision per request; explicit - // controller users may opt into periodic re-probing. - for (int i = 0; i < 32; ++i) { - CHECK(mgr.record_ar_sample(chat.slot, 1000.0) == - SpeculationGoodputTransition::none); - CHECK(!mgr.ddtree_speculation_allowed(chat.slot)); - } - - mgr.retire(code.slot); - auto reused = admit(mgr, 103, prompt_tokens(4), greedy_sampler()); - CHECK(is_admitted(reused) && reused.slot == code.slot); - CHECK(mgr.slot(reused.slot).request_id == 103); - CHECK(mgr.ddtree_speculation_allowed(reused.slot)); - CHECK(mgr.slot(reused.slot).ddtree_sampled_steps == 0); - } - - // The existing burn-in switch preserves fixed speculation. - { - const luce_test::ScopedEnvVar adaptive("DFLASH_DDTREE_ADAPTIVE", "0"); - PagedKvPool pool(4, 1, /*block_size=*/16); - Qwen35SlotManager mgr(pool, 64); - auto admitted = admit(mgr, 201, prompt_tokens(4), greedy_sampler()); - CHECK(is_admitted(admitted)); - CHECK(mgr.append_prefill(admitted.slot, 4).ok); - mgr.commit_prefill(admitted.slot); - CHECK(mgr.ddtree_speculation_allowed(admitted.slot)); - CHECK(mgr.record_speculation_sample(admitted.slot, 1, 2000.0) == - SpeculationGoodputTransition::none); - CHECK(mgr.record_ar_sample(admitted.slot, 1000.0) == - SpeculationGoodputTransition::none); - CHECK(mgr.ddtree_speculation_allowed(admitted.slot)); - } - // A failed residency barrier quarantines retirement ownership, and a // later admission retries it before considering the slot reusable. { diff --git a/server/test/test_speculation_goodput.cpp b/server/test/test_speculation_goodput.cpp index 5ce384533..f9ce4da13 100644 --- a/server/test/test_speculation_goodput.cpp +++ b/server/test/test_speculation_goodput.cpp @@ -285,6 +285,42 @@ int main() { CHECK(ranker.select(5, mixed, 2).requests.size() == 2); } + // A continuous-batching scheduler may optimize total goodput when a + // completed request immediately refills its slot. Closed cohorts keep the + // AR-peer guard; backlog mode may select the measured goodput winner. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(5, 100.0); + ranker.observe_route(5, 2, 140.0); + std::vector candidates = { + {86, 5.0, 5.0, true}, + {87, 5.0, 5.0, true}, + }; + CHECK(ranker.select( + 5, candidates, 2, + /*probe_uncalibrated_with_verifier=*/false, + /*probe_missing_routes=*/false) + .requests.empty()); + const AdaptiveVerificationDecision backlog = ranker.select( + 5, candidates, 2, + /*probe_uncalibrated_with_verifier=*/false, + /*probe_missing_routes=*/false, + /*enforce_ar_peer_guard=*/false); + CHECK(backlog.requests.size() == 2); + } + + // Peer-guard relaxation follows verifier capacity rather than a hard-coded + // concurrency cutoff. Two lanes cover C=5 but not C=6; a wider executor + // raises the boundary automatically. + { + CHECK(adaptive_verification_can_relax_peer_guard(5, 2)); + CHECK(!adaptive_verification_can_relax_peer_guard(6, 2)); + CHECK(adaptive_verification_can_relax_peer_guard(7, 3)); + CHECK(!adaptive_verification_can_relax_peer_guard(8, 3)); + CHECK(adaptive_verification_can_relax_peer_guard(16, 8)); + CHECK(!adaptive_verification_can_relax_peer_guard(1, 0)); + } + // Timings from a neighboring occupancy never stand in for the exact-C // baseline or route profile. { @@ -358,21 +394,24 @@ int main() { ranker.observe_routing_prior_yield(1.0, 4.0); ranker.observe_routing_prior_yield(1.0, 6.0); ranker.observe_routing_prior_yield(1.0, 4.0); - CHECK(!ranker.routing_prior_expected_tokens(1.0).has_value()); + CHECK(std::abs( + ranker.routing_prior_expected_tokens(1.0).value() - + 14.0 / 3.0) < 1e-9); CHECK(ranker.routing_prior_yield_samples(1.0) == 3); ranker.observe_routing_prior_yield(1.0, 6.0); CHECK(ranker.routing_prior_expected_tokens(1.0).value() == 5.0); CHECK(ranker.routing_prior_yield_samples(1.0) == 4); CHECK(!ranker.routing_prior_expected_tokens(-1.0).has_value()); - // Early request observations are shrunk toward the stable cohort, but - // four local samples make the request's own magnitude authoritative. + // Shared yield magnitude itself starts conservatively shrunk toward + // AR; four local samples still make the request authoritative. + const double expected[] = {2.5, 2.0, 1.5, 1.0}; for (int sample = 1; sample <= 4; ++sample) { ranker.observe_request_yield(98, 1.0); const auto estimate = ranker.estimate_request_yield(98, 1.0); CHECK(estimate.has_value()); - CHECK(estimate->expected_tokens == 5.0 - sample); + CHECK(estimate->expected_tokens == expected[sample - 1]); CHECK(estimate->evidence_samples == static_cast(sample)); } From 66c1578dc20398e3e0f02fdf599177bb24286f05 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 17 Aug 2026 20:59:36 +0000 Subject: [PATCH 15/18] perf(qwen35): warm profitable high-concurrency routes --- .../concurrency/adaptive_verification.h | 118 ++++++++++++++++-- .../qwen35/concurrency/qwen35_seq_engine.cpp | 78 +++++++++--- server/test/test_speculation_goodput.cpp | 56 +++++++++ 3 files changed, 225 insertions(+), 27 deletions(-) diff --git a/server/src/common/concurrency/adaptive_verification.h b/server/src/common/concurrency/adaptive_verification.h index 411d7dd9d..afa42d6b2 100644 --- a/server/src/common/concurrency/adaptive_verification.h +++ b/server/src/common/concurrency/adaptive_verification.h @@ -72,6 +72,11 @@ struct AdaptiveVerificationConfig { // Shared evidence ranks a new request after the minimum above, but its // yield magnitude is shrunk toward AR until this many outcomes exist. std::size_t routing_prior_full_weight_samples = 8; + // A continuous-backlog scheduler may optimize aggregate output after a + // routing-prior bucket has accumulated this many target-verified outcomes. + // This is deliberately stronger than cold ordering, but does not count as + // request-local proof for closed-cohort AR-peer protection. + std::size_t routing_prior_peer_guard_samples = 6; double cost_ewma_alpha = 0.35; }; @@ -82,6 +87,13 @@ inline bool adaptive_verification_can_relax_peer_guard( 2 * static_cast(verifier_request_lanes) + 1; } +inline bool adaptive_verification_can_extend_stable_cohort( + int active_requests, int verifier_request_lanes) { + return active_requests > 0 && verifier_request_lanes > 0 && + static_cast(active_requests) <= + 3 * static_cast(verifier_request_lanes); +} + // Convert conditional survival confidence into expected useful tokens: // 1 AR token plus the probability of reaching every speculative prefix. // DSpark supplies calibrated confidence-head values directly. @@ -109,6 +121,7 @@ class AdaptiveVerificationRanker { void reset() { route_cost_us_.clear(); route_cost_known_.clear(); + route_cost_samples_.clear(); request_yield_.clear(); routing_prior_yield_.clear(); } @@ -189,8 +202,36 @@ class AdaptiveVerificationRanker { return it == routing_prior_yield_.end() ? 0 : it->second.samples; } + bool has_stable_routing_prior_yield(double routing_prior) const { + if (!std::isfinite(routing_prior)) return false; + const auto it = routing_prior_yield_.find(routing_prior); + return it != routing_prior_yield_.end() && + it->second.samples >= + config_.routing_prior_peer_guard_samples && + has_useful_yield(it->second.expected_tokens); + } + + bool forms_stable_routing_prior_cohort( + const std::vector & candidates) + const { + if (candidates.empty()) return false; + double minimum = std::numeric_limits::infinity(); + double maximum = 1.0; + for (const AdaptiveVerificationCandidate & candidate : candidates) { + if (!has_stable_routing_prior_yield(candidate.routing_prior)) { + return false; + } + const auto it = routing_prior_yield_.find(candidate.routing_prior); + minimum = std::min(minimum, it->second.expected_tokens); + maximum = std::max(maximum, it->second.expected_tokens); + } + return minimum >= + config_.homogeneous_minimum_relative_yield * maximum; + } + std::optional estimate_request_yield( - std::uint64_t request, double routing_prior) const { + std::uint64_t request, double routing_prior, + bool trust_stable_routing_prior = false) const { const auto request_it = request_yield_.find(request); const auto prior_it = std::isfinite(routing_prior) ? routing_prior_yield_.find(routing_prior) @@ -201,10 +242,16 @@ class AdaptiveVerificationRanker { if (!has_request && !has_prior) return std::nullopt; AdaptiveVerificationYieldEstimate out; - const double prior_weight = std::min( - 1.0, static_cast(prior_it->second.samples) / - static_cast( - config_.routing_prior_full_weight_samples)); + const bool stable_backlog_prior = + trust_stable_routing_prior && + prior_it->second.samples >= + config_.routing_prior_peer_guard_samples; + const double prior_weight = stable_backlog_prior + ? 1.0 + : std::min( + 1.0, static_cast(prior_it->second.samples) / + static_cast( + config_.routing_prior_full_weight_samples)); const double trusted_prior = 1.0 + prior_weight * (prior_it->second.expected_tokens - 1.0); if (!has_request) { @@ -245,7 +292,8 @@ class AdaptiveVerificationRanker { } void observe_route(int active_requests, int speculative_requests, - double elapsed_us) { + double elapsed_us, + bool discard_first_sample = false) { if (active_requests <= 0 || speculative_requests < 0 || speculative_requests > active_requests || !std::isfinite(elapsed_us) || elapsed_us <= 0.0) { @@ -254,22 +302,41 @@ class AdaptiveVerificationRanker { const size_t rows = static_cast(active_requests) + 1; if (route_cost_us_.size() < rows) route_cost_us_.resize(rows); if (route_cost_known_.size() < rows) route_cost_known_.resize(rows); + if (route_cost_samples_.size() < rows) route_cost_samples_.resize(rows); std::vector & costs = route_cost_us_[(size_t)active_requests]; std::vector & known = route_cost_known_[(size_t)active_requests]; + std::vector & samples = + route_cost_samples_[(size_t)active_requests]; const size_t cols = static_cast(speculative_requests) + 1; if (costs.size() < cols) costs.resize(cols, 0.0); if (known.size() < cols) known.resize(cols, false); + if (samples.size() < cols) samples.resize(cols, 0); if (!known[(size_t)speculative_requests]) { costs[(size_t)speculative_requests] = elapsed_us; known[(size_t)speculative_requests] = true; + samples[(size_t)speculative_requests] = 1; + return; + } + if (discard_first_sample && + samples[(size_t)speculative_requests] == 1) { + // The first execution of a new graph shape includes capture and + // allocator warmup. Replace it with the first replay before using + // the normal EWMA, rather than permanently teaching the policy + // that a profitable steady route is cold-start slow. + costs[(size_t)speculative_requests] = elapsed_us; + samples[(size_t)speculative_requests] = 2; return; } costs[(size_t)speculative_requests] = config_.cost_ewma_alpha * elapsed_us + (1.0 - config_.cost_ewma_alpha) * costs[(size_t)speculative_requests]; + if (samples[(size_t)speculative_requests] < + std::numeric_limits::max()) { + ++samples[(size_t)speculative_requests]; + } } bool has_autoregressive_cost(int batch_size) const { @@ -308,13 +375,29 @@ class AdaptiveVerificationRanker { : std::numeric_limits::infinity(); } + std::size_t route_cost_samples(int active_requests, + int speculative_requests) const { + return has_route_cost(active_requests, speculative_requests) && + static_cast(active_requests) < + route_cost_samples_.size() && + static_cast(speculative_requests) < + route_cost_samples_[(size_t)active_requests].size() + ? route_cost_samples_[(size_t)active_requests] + [(size_t)speculative_requests] + : 0; + } + bool has_exact_profile(int active_requests, - int max_speculative_requests) const { + int max_speculative_requests, + std::size_t minimum_route_samples = 1) const { if (!has_autoregressive_cost(active_requests)) return false; + const std::size_t required_samples = + std::max(1, minimum_route_samples); const int limit = std::max(0, std::min( active_requests, max_speculative_requests)); for (int k = 1; k <= limit; ++k) { - if (!has_route_cost(active_requests, k)) return false; + if (route_cost_samples(active_requests, k) < + required_samples) return false; } return true; } @@ -336,7 +419,8 @@ class AdaptiveVerificationRanker { std::numeric_limits::max(), bool probe_uncalibrated_with_verifier = false, bool probe_missing_routes = true, - bool enforce_ar_peer_guard = true) const { + bool enforce_ar_peer_guard = true, + std::size_t minimum_speculative_route_samples = 1) const { AdaptiveVerificationDecision out; if (active_requests <= 0 || candidates.empty() || !has_autoregressive_cost(active_requests)) { @@ -460,6 +544,9 @@ class AdaptiveVerificationRanker { static_cast(active_requests) / autoregressive_cost_us(active_requests); const double required = baseline * config_.minimum_gain; + const std::size_t required_route_samples = + std::max( + 1, minimum_speculative_route_samples); int admitted_prefix = 0; double best = baseline; double admitted_goodput = baseline; @@ -471,7 +558,8 @@ class AdaptiveVerificationRanker { for (int k = 1; k <= static_cast(known.size()) && k <= prefix_limit; ++k) { expected_total += known[(size_t)k - 1].expected_tokens - 1.0; - if (!has_route_cost(active_requests, k)) continue; + if (route_cost_samples(active_requests, k) < + required_route_samples) continue; const double route_us = route_cost_us(active_requests, k); const double throughput = expected_total / route_us; const bool protects_ar_peers = !enforce_ar_peer_guard || @@ -502,7 +590,8 @@ class AdaptiveVerificationRanker { int missing_width = 0; for (int k = 1; probe_missing_routes && k <= probe_candidates; ++k) { - if (has_route_cost(active_requests, k)) continue; + if (route_cost_samples(active_requests, k) >= + required_route_samples) continue; missing_width = k; break; } @@ -539,7 +628,8 @@ class AdaptiveVerificationRanker { known_total += known[(size_t)k - 2].expected_tokens - 1.0; } - if (!has_route_cost(active_requests, k)) continue; + if (route_cost_samples(active_requests, k) < + required_route_samples) continue; const double route_us = route_cost_us(active_requests, k); // This is one bounded evidence-gathering step, not a steady // route. It may cross the steady AR-peer guard so an otherwise @@ -591,6 +681,9 @@ class AdaptiveVerificationRanker { config.routing_prior_full_weight_samples = std::max(config.routing_prior_minimum_samples, config.routing_prior_full_weight_samples); + config.routing_prior_peer_guard_samples = + std::max(config.routing_prior_minimum_samples, + config.routing_prior_peer_guard_samples); config.cost_ewma_alpha = std::clamp(config.cost_ewma_alpha, 0.0, 1.0); return config; @@ -599,6 +692,7 @@ class AdaptiveVerificationRanker { AdaptiveVerificationConfig config_; std::vector> route_cost_us_; std::vector> route_cost_known_; + std::vector> route_cost_samples_; struct YieldEstimate { double expected_tokens = 1.0; std::size_t samples = 0; diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index adf51676f..1a11f5300 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -1393,7 +1393,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // head while reusing the same selection policy. std::vector candidates; candidates.reserve(inputs.size()); - auto collect_candidates = [&]() { + auto collect_candidates = [&](bool trust_stable_routing_prior = false) { candidates.clear(); for (const StepInput & in : inputs) { if (!ddtree_input_eligible(in)) continue; @@ -1401,7 +1401,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const int prompt_hint = seq.sampler.speculation_prompt_hint; const std::optional estimate = route_ranker.estimate_request_yield( - seq.request_id, static_cast(prompt_hint)); + seq.request_id, static_cast(prompt_hint), + trust_stable_routing_prior); candidates.push_back({ in.slot, estimate ? estimate->expected_tokens : 1.0, @@ -1425,15 +1426,21 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { direct_speculation_limit, compact_short_shape ? 2 : direct_speculation_limit) : static_cast(inputs.size()); - // Relax peer protection for continuous batching only when the one-pass - // verifier can cover roughly half the live cohort. This derives the - // boundary from executor capacity rather than a fixed C: a wider DSpark - // adapter automatically expands the eligible concurrency range. - const bool relax_ar_peer_guard = - plan.has_refill_backlog && compact_tree_route && + const int active_requests = static_cast(inputs.size()); + const bool broad_verifier_coverage = adaptive_verification_can_relax_peer_guard( - static_cast(inputs.size()), adaptive_speculation_limit); - const bool enforce_ar_peer_guard = !relax_ar_peer_guard; + active_requests, adaptive_speculation_limit); + const bool bounded_stable_cohort_extension = + adaptive_verification_can_extend_stable_cohort( + active_requests, adaptive_speculation_limit); + const bool use_stable_cohort_extension = + plan.has_refill_backlog && compact_tree_route && + !broad_verifier_coverage && bounded_stable_cohort_extension; + if (use_stable_cohort_extension) { + collect_candidates(/*trust_stable_routing_prior=*/true); + } + bool enforce_ar_peer_guard = true; + std::size_t minimum_route_samples = 1; if (oracle.forces_selection()) { const int force_limit = std::min( adaptive_speculation_limit, @@ -1456,13 +1463,46 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const int exact_profile_width = std::min( adaptive_speculation_limit, static_cast(candidates.size())); + // Replay a cold graph shape only in the narrow capacity extension that + // can plausibly become steady. This avoids doubling rejected probes at + // occupancies such as Strix C=16 and leaves established low-C profiles + // unchanged. + minimum_route_samples = + use_stable_cohort_extension && + positive_prompt_candidates == active_requests + ? 2 : 1; int & calibration_cooldown = compact_tree_route ? (compact_short_shape ? compact_short_adaptive_calibration_cooldown_ : compact_adaptive_calibration_cooldown_) : adaptive_calibration_cooldown_; const bool exact_profile_ready = route_ranker.has_exact_profile( - static_cast(inputs.size()), exact_profile_width); + static_cast(inputs.size()), exact_profile_width, + minimum_route_samples); + // A closed cohort keeps request-local AR-peer protection. Continuous + // backlog may instead optimize aggregate output once either verifier + // coverage is broad enough or target-verified outcomes have made a + // useful routing-prior bucket stable. Exact route costs still decide + // whether speculation wins, so the same evidence can unlock C=8 while + // an unprofitable C=16 profile remains AR. A future speculator supplies + // its own work-bucket ranker and inherits the same rule. + const bool has_stable_backlog_cohort = + plan.has_refill_backlog && exact_profile_ready && + route_ranker.forms_stable_routing_prior_cohort(candidates); + const bool has_stable_backlog_candidate = + plan.has_refill_backlog && exact_profile_ready && + std::any_of( + candidates.begin(), candidates.end(), + [&](const AdaptiveVerificationCandidate & candidate) { + return route_ranker.has_stable_routing_prior_yield( + candidate.routing_prior); + }); + const bool relax_ar_peer_guard = + plan.has_refill_backlog && compact_tree_route && + (broad_verifier_coverage || + (bounded_stable_cohort_extension && + has_stable_backlog_cohort)); + enforce_ar_peer_guard = !relax_ar_peer_guard; const bool exact_profile_started = route_ranker.has_speculative_profile_sample( static_cast(inputs.size()), exact_profile_width); @@ -1482,6 +1522,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const bool prompt_prior_marks_mixed = positive_prompt_candidates > 0 && positive_prompt_candidates < static_cast(inputs.size()); + const bool suppress_uncovered_mixed_calibration = + exact_profile_ready && prompt_prior_marks_mixed && + !broad_verifier_coverage && has_stable_backlog_candidate; if (!exact_profile_ready || inputs.size() <= 3) { calibration_cooldown = 0; } else if (calibration_cooldown > 0) { @@ -1494,6 +1537,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // existing all-speculative exception. const bool verifier_side_calibration = compact_tree_route && probe_missing_routes && + !suppress_uncovered_mixed_calibration && (!exact_profile_ready || (calibration_cooldown == 0 && (plan.has_refill_backlog || !prompt_prior_marks_mixed))); @@ -1525,7 +1569,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { static_cast(inputs.size()), *decision_candidates, adaptive_speculation_limit, verifier_side_calibration, probe_missing_routes, - enforce_ar_peer_guard); + enforce_ar_peer_guard, minimum_route_samples); if (compact_tree_route && exact_profile_ready && verifier_side_calibration && decision.exploring) { calibration_cooldown = @@ -1559,7 +1603,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { adaptive_speculation_limit, /*probe_uncalibrated_with_verifier=*/false, /*probe_missing_routes=*/false, - enforce_ar_peer_guard); + enforce_ar_peer_guard, + minimum_route_samples); if (inputs.size() > 3 && exact_profile_ready && static_cast(steady_after.requests.size()) >= exact_profile_width) { @@ -1575,7 +1620,8 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { adaptive_speculation_limit, verifier_side_calibration, probe_missing_routes, - enforce_ar_peer_guard); + enforce_ar_peer_guard, + minimum_route_samples); } } } @@ -1640,7 +1686,9 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (adaptive_enabled && !oracle.forces_selection()) { route_ranker.observe_route( - static_cast(inputs.size()), speculative_count, route_us); + static_cast(inputs.size()), speculative_count, route_us, + /*discard_first_sample=*/ + speculative_count > 0 && minimum_route_samples > 1); } if (decision.exploring && !speculative_plan.decode.empty()) { std::fprintf(stderr, diff --git a/server/test/test_speculation_goodput.cpp b/server/test/test_speculation_goodput.cpp index f9ce4da13..f667680cc 100644 --- a/server/test/test_speculation_goodput.cpp +++ b/server/test/test_speculation_goodput.cpp @@ -319,6 +319,12 @@ int main() { CHECK(!adaptive_verification_can_relax_peer_guard(8, 3)); CHECK(adaptive_verification_can_relax_peer_guard(16, 8)); CHECK(!adaptive_verification_can_relax_peer_guard(1, 0)); + CHECK(adaptive_verification_can_extend_stable_cohort(8, 3)); + CHECK(adaptive_verification_can_extend_stable_cohort(9, 3)); + CHECK(!adaptive_verification_can_extend_stable_cohort(10, 3)); + CHECK(!adaptive_verification_can_extend_stable_cohort(16, 3)); + CHECK(adaptive_verification_can_extend_stable_cohort(16, 6)); + CHECK(!adaptive_verification_can_extend_stable_cohort(1, 0)); } // Timings from a neighboring occupancy never stand in for the exact-C @@ -337,6 +343,28 @@ int main() { CHECK(!ranker.has_exact_profile(7, 1)); } + // The first graph-shape execution is a capture sample. A backlog profile + // can require a second replay, which replaces that cold timing before the + // normal EWMA begins. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(8, 100.0); + ranker.observe_route(8, 1, 250.0); + CHECK(ranker.route_cost_samples(8, 1) == 1); + CHECK(ranker.route_cost_us(8, 1) == 250.0); + CHECK(ranker.has_exact_profile(8, 1)); + CHECK(!ranker.has_exact_profile(8, 2, 0)); + CHECK(!ranker.has_exact_profile(8, 1, 2)); + ranker.observe_route( + 8, 1, 100.0, /*discard_first_sample=*/true); + CHECK(ranker.route_cost_samples(8, 1) == 2); + CHECK(ranker.route_cost_us(8, 1) == 100.0); + CHECK(ranker.has_exact_profile(8, 1, 2)); + ranker.observe_route(8, 1, 200.0); + CHECK(ranker.route_cost_samples(8, 1) == 3); + CHECK(std::abs(ranker.route_cost_us(8, 1) - 135.0) < 1e-9); + } + // Proven high-yield peers rotate scarce verifier lanes toward the request // with less generated progress, preventing a homogeneous cohort tail. { @@ -401,7 +429,9 @@ int main() { ranker.observe_routing_prior_yield(1.0, 6.0); CHECK(ranker.routing_prior_expected_tokens(1.0).value() == 5.0); CHECK(ranker.routing_prior_yield_samples(1.0) == 4); + CHECK(!ranker.has_stable_routing_prior_yield(1.0)); CHECK(!ranker.routing_prior_expected_tokens(-1.0).has_value()); + CHECK(!ranker.has_stable_routing_prior_yield(-1.0)); // Shared yield magnitude itself starts conservatively shrunk toward // AR; four local samples still make the request authoritative. @@ -419,6 +449,32 @@ int main() { CHECK(local.has_value()); CHECK(local->expected_tokens == 1.0); + ranker.observe_routing_prior_yield(1.0, 4.0); + CHECK(!ranker.has_stable_routing_prior_yield(1.0)); + ranker.observe_routing_prior_yield(1.0, 6.0); + CHECK(ranker.has_stable_routing_prior_yield(1.0)); + const auto closed_prior = + ranker.estimate_request_yield(100, 1.0); + const auto backlog_prior = + ranker.estimate_request_yield( + 100, 1.0, /*trust_stable_routing_prior=*/true); + CHECK(closed_prior.has_value()); + CHECK(backlog_prior.has_value()); + CHECK(closed_prior->expected_tokens == 4.0); + CHECK(backlog_prior->expected_tokens == 5.0); + for (int i = 0; i < 6; ++i) { + ranker.observe_routing_prior_yield(-1.0, 1.25); + } + CHECK(!ranker.has_stable_routing_prior_yield(-1.0)); + CHECK(ranker.forms_stable_routing_prior_cohort({ + {101, 5.0, 9.0, true, 1.0}, + {102, 5.0, 9.0, true, 1.0}, + })); + CHECK(!ranker.forms_stable_routing_prior_cohort({ + {101, 5.0, 9.0, true, 1.0}, + {103, 1.25, 9.0, true, -1.0}, + })); + // Per-request evidence is local to one proposal-shape ranker and is // explicitly forgotten at request retirement. AdaptiveVerificationRanker compact; From 4935d3769e41601e9ff0417a067c3d7726e82ae3 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 07:17:10 +0000 Subject: [PATCH 16/18] feat(server): add per-request decode modes --- server/README.md | 24 +++++ .../concurrency/adaptive_verification.h | 92 ++++++++++++++----- server/src/common/concurrency/seq_engine.h | 10 +- server/src/common/model_backend.h | 4 + server/src/common/speculation_policy.h | 35 +++++++ .../qwen35/concurrency/qwen35_seq_engine.cpp | 4 +- server/src/server/http_server.cpp | 36 +++++++- server/src/server/http_server.h | 5 + server/src/server/scheduler.cpp | 5 +- server/src/server/server_main.cpp | 32 ++++++- server/test/test_server_unit.cpp | 18 ++++ server/test/test_speculation_goodput.cpp | 64 +++++++++++++ 12 files changed, 300 insertions(+), 29 deletions(-) create mode 100644 server/src/common/speculation_policy.h diff --git a/server/README.md b/server/README.md index 1e96b697d..6dc8878d8 100644 --- a/server/README.md +++ b/server/README.md @@ -170,6 +170,30 @@ Run it directly: --model-name luce-dflash ``` +With a decode drafter configured (`--draft`), the server default is adaptive +routing. The operator can choose a different default without unloading the +draft model; DDTree is shown here, but the control is not DDTree-specific: + +```bash +--ddtree --decode-mode adaptive # measured request-level routing +--ddtree --decode-mode speculation # force speculation when eligible +--ddtree --decode-mode ar # force autoregressive decode +``` + +Every generation endpoint also accepts a top-level per-request override: + +```json +{ + "model": "luce-dflash", + "messages": [{"role": "user", "content": "Implement a Python parser"}], + "decode_mode": "adaptive" +} +``` + +The accepted values are `adaptive`, `ar`, and `speculation`. `speculation` +requires the server to have speculative decode enabled. The mode is speculator-neutral: +DDTree consumes it today, and DSpark or a later verifier can use the same API. + ### Compression proxy mode `dflash_server` can run as a **PFlash compression proxy** in front of any diff --git a/server/src/common/concurrency/adaptive_verification.h b/server/src/common/concurrency/adaptive_verification.h index afa42d6b2..8730b073b 100644 --- a/server/src/common/concurrency/adaptive_verification.h +++ b/server/src/common/concurrency/adaptive_verification.h @@ -38,6 +38,10 @@ struct AdaptiveVerificationCandidate { // Generated output already committed for this request. Proven-useful peers // with less progress receive compact lanes first to avoid cohort stragglers. int progress_tokens = 0; + // A user-forced request is always present in the selected prefix. Its + // observed yield still contributes to deciding whether adaptive peers + // should share the same verifier pass. + bool required = false; }; struct AdaptiveVerificationYieldEstimate { @@ -422,14 +426,12 @@ class AdaptiveVerificationRanker { bool enforce_ar_peer_guard = true, std::size_t minimum_speculative_route_samples = 1) const { AdaptiveVerificationDecision out; - if (active_requests <= 0 || candidates.empty() || - !has_autoregressive_cost(active_requests)) { - // First observe the exact all-AR baseline for this occupancy. - return out; - } + if (active_requests <= 0 || candidates.empty()) return out; + std::vector required_candidates; std::vector known; std::vector unknown; + required_candidates.reserve(candidates.size()); known.reserve(candidates.size()); unknown.reserve(candidates.size()); for (AdaptiveVerificationCandidate candidate : candidates) { @@ -444,9 +446,37 @@ class AdaptiveVerificationRanker { if (!std::isfinite(candidate.routing_prior)) { candidate.routing_prior = 0.0; } - (candidate.calibrated ? known : unknown).push_back(candidate); + if (candidate.required) { + required_candidates.push_back(candidate); + } else { + (candidate.calibrated ? known : unknown).push_back(candidate); + } + } + if (required_candidates.empty() && known.empty() && unknown.empty()) { + return out; + } + + std::stable_sort( + required_candidates.begin(), required_candidates.end(), + [](const AdaptiveVerificationCandidate & a, + const AdaptiveVerificationCandidate & b) { + return a.request < b.request; + }); + auto select_required = [&]() { + out.requests.clear(); + out.requests.reserve(required_candidates.size()); + for (const AdaptiveVerificationCandidate & candidate : + required_candidates) { + out.requests.push_back(candidate.request); + } + }; + if (!has_autoregressive_cost(active_requests)) { + // Adaptive-only cohorts first observe the exact AR baseline. User + // overrides are stronger: an Always request must not be delayed + // by profiling, and its route timing is learned from this step. + select_required(); + return out; } - if (known.empty() && unknown.empty()) return out; auto has_stable_useful_yield = [&](const AdaptiveVerificationCandidate & candidate) { @@ -540,6 +570,16 @@ class AdaptiveVerificationRanker { out.calibration_request = unknown.front().request; } + std::vector ordered_known; + ordered_known.reserve(required_candidates.size() + known.size()); + ordered_known.insert( + ordered_known.end(), required_candidates.begin(), + required_candidates.end()); + ordered_known.insert( + ordered_known.end(), known.begin(), known.end()); + const int required_count = + static_cast(required_candidates.size()); + const double baseline = static_cast(active_requests) / autoregressive_cost_us(active_requests); @@ -547,17 +587,21 @@ class AdaptiveVerificationRanker { const std::size_t required_route_samples = std::max( 1, minimum_speculative_route_samples); - int admitted_prefix = 0; + int admitted_prefix = required_count; double best = baseline; double admitted_goodput = baseline; double expected_total = static_cast(active_requests); - const int prefix_limit = std::max(0, std::min( - active_requests, max_speculative_requests)); + const int prefix_limit = std::max( + required_count, + std::max(0, std::min( + active_requests, max_speculative_requests))); // Evaluate every measured width independently. GPU occupancy makes // these costs non-convex: k=3 may win even when k=1 and k=2 lose. - for (int k = 1; k <= static_cast(known.size()) && + for (int k = 1; k <= static_cast(ordered_known.size()) && k <= prefix_limit; ++k) { - expected_total += known[(size_t)k - 1].expected_tokens - 1.0; + expected_total += + ordered_known[(size_t)k - 1].expected_tokens - 1.0; + if (k < required_count) continue; if (route_cost_samples(active_requests, k) < required_route_samples) continue; const double route_us = route_cost_us(active_requests, k); @@ -584,11 +628,11 @@ class AdaptiveVerificationRanker { // bounded prefix and learn their yield from accepted output. const int probe_candidates = std::min( prefix_limit, - static_cast(known.size()) + + static_cast(ordered_known.size()) + (probe_uncalibrated_with_verifier ? static_cast(unknown.size()) : 0)); int missing_width = 0; - for (int k = 1; probe_missing_routes && + for (int k = std::max(1, required_count); probe_missing_routes && k <= probe_candidates; ++k) { if (route_cost_samples(active_requests, k) >= required_route_samples) continue; @@ -598,9 +642,10 @@ class AdaptiveVerificationRanker { if (missing_width > 0) { out.requests.reserve((size_t)missing_width); const int known_count = std::min( - missing_width, static_cast(known.size())); + missing_width, static_cast(ordered_known.size())); for (int i = 0; i < known_count; ++i) { - out.requests.push_back(known[(size_t)i].request); + out.requests.push_back( + ordered_known[(size_t)i].request); } for (int i = 0; static_cast(out.requests.size()) < missing_width; @@ -621,13 +666,14 @@ class AdaptiveVerificationRanker { int calibration_width = 0; double calibration_goodput = required; const int calibration_limit = std::min( - prefix_limit, static_cast(known.size()) + 1); + prefix_limit, static_cast(ordered_known.size()) + 1); double known_total = static_cast(active_requests); for (int k = 1; k <= calibration_limit; ++k) { if (k > 1) { known_total += - known[(size_t)k - 2].expected_tokens - 1.0; + ordered_known[(size_t)k - 2].expected_tokens - 1.0; } + if (k <= required_count) continue; if (route_cost_samples(active_requests, k) < required_route_samples) continue; const double route_us = route_cost_us(active_requests, k); @@ -645,7 +691,8 @@ class AdaptiveVerificationRanker { if (calibration_width > 0) { out.requests.reserve((size_t)calibration_width); for (int i = 0; i < calibration_width - 1; ++i) { - out.requests.push_back(known[(size_t)i].request); + out.requests.push_back( + ordered_known[(size_t)i].request); } out.requests.push_back(unknown.front().request); out.calibration_request = -1; @@ -657,9 +704,12 @@ class AdaptiveVerificationRanker { if (admitted_prefix > 0) { out.requests.reserve((size_t)admitted_prefix); for (int i = 0; i < admitted_prefix; ++i) { - out.requests.push_back(known[(size_t)i].request); + out.requests.push_back( + ordered_known[(size_t)i].request); + } + if (admitted_prefix > required_count) { + out.predicted_gain = admitted_goodput / baseline; } - out.predicted_gain = admitted_goodput / baseline; } return out; } diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index c723b68c7..38d5bfaf1 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -56,6 +56,7 @@ #include #include "common/sampler.h" +#include "common/speculation_policy.h" namespace dflash::common { @@ -182,6 +183,11 @@ class SeqEngine { struct StepInput { int slot = -1; int32_t token = -1; // token to commit at this slot's next position + // User policy for this request. `Always` is mandatory speculation, + // `Never` is AR, and `Adaptive` delegates to the engine's measured + // policy. The mechanism is deliberately independent of DDTree/DSpark. + SpeculationPolicy speculation_policy = + SpeculationPolicy::Adaptive; // False when scheduler-side policy may replace the sampled token // before it is committed (currently the thinking-budget close hook). bool allow_speculation = true; @@ -311,7 +317,9 @@ inline std::string validate_step_result( return "decode plan contains a duplicate slot"; decode_planned[(size_t)input.slot] = 1; speculation_allowed[(size_t)input.slot] = - input.allow_speculation ? 1 : 0; + input.allow_speculation && + input.speculation_policy != SpeculationPolicy::Never + ? 1 : 0; } for (const PrefillSlice & slice : plan.prefills) { if (slice.slot < 0 || slice.slot >= slot_count || diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 12445e6b2..bce1c97e9 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -170,6 +170,10 @@ struct GenerateRequest { std::vector prompt; int n_gen = 0; SamplerCfg sampler; + // Public decode policy. Concurrent engines consume it per request; + // sequential backends currently map Never to their established AR path. + SpeculationPolicy speculation_policy = + SpeculationPolicy::Adaptive; bool do_sample = false; bool stream = false; // emit tokens to stream_fd // Optional inline-snap: snapshot at this position after prefill. diff --git a/server/src/common/speculation_policy.h b/server/src/common/speculation_policy.h new file mode 100644 index 000000000..0cc8745f1 --- /dev/null +++ b/server/src/common/speculation_policy.h @@ -0,0 +1,35 @@ +#pragma once + +// User-selectable, speculator-neutral decode policy. Concrete backends map +// this to DDTree, DSpark, or any later speculative implementation. + +#include +#include + +namespace dflash::common { + +enum class SpeculationPolicy { + Adaptive, + Always, + Never, +}; + +constexpr std::string_view decode_mode_name( + SpeculationPolicy policy) { + switch (policy) { + case SpeculationPolicy::Adaptive: return "adaptive"; + case SpeculationPolicy::Always: return "speculation"; + case SpeculationPolicy::Never: return "ar"; + } + return "adaptive"; +} + +inline std::optional parse_decode_mode( + std::string_view value) { + if (value == "adaptive") return SpeculationPolicy::Adaptive; + if (value == "speculation") return SpeculationPolicy::Always; + if (value == "ar") return SpeculationPolicy::Never; + return std::nullopt; +} + +} // namespace dflash::common diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 1a11f5300..a7ea34f3b 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -182,7 +182,8 @@ bool Qwen35SeqEngine::ddtree_available(const StepPlan & plan) const { } bool Qwen35SeqEngine::ddtree_input_eligible(const StepInput & in) const { - if (!in.allow_speculation || in.slot < 0 || + if (!in.allow_speculation || + in.speculation_policy == SpeculationPolicy::Never || in.slot < 0 || in.slot >= slots_.slot_count() || !slots_.slot(in.slot).decoding() || slots_.slot(in.slot).sampler.needs_logit_processing() || @@ -1412,6 +1413,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { static_cast(prompt_hint), estimate ? estimate->evidence_samples : 0, seq.generated_tokens(), + in.speculation_policy == SpeculationPolicy::Always, }); } }; diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index b25a69758..6efd8589a 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -764,6 +764,7 @@ json build_props_body(const ServerConfig & config, {"build_info", std::string(kServerName) + " v" DFLASH_SERVER_VERSION " props_schema=" + std::to_string(kPropsSchema)}, {"speculative_mode", speculative_mode}, + {"decode_mode", std::string(decode_mode_name(config.decode_mode))}, {"server", server}, {"model", { {"arch", config.arch}, @@ -824,7 +825,9 @@ json build_props_body(const ServerConfig & config, }}, {"speculative", { {"enabled", config.speculative_enabled}, - {"ddtree_budget", config.speculative_enabled + {"decode_mode", std::string(decode_mode_name( + config.decode_mode))}, + {"ddtree_budget", config.ddtree_budget > 0 ? json(config.ddtree_budget) : json(nullptr)}, }}, {"sampling", { @@ -1661,6 +1664,29 @@ bool HttpServer::parse_common_request_fields( req.stream = body.value("stream", false); req.model = body.value("model", config_.model_name); req.disk_cache_policy = config_.disk_cache_policy; + req.decode_mode = config_.decode_mode; + + if (body.contains("decode_mode")) { + if (!body["decode_mode"].is_string()) { + send_error(fd, 400, + "decode_mode must be adaptive, ar, or speculation"); + return false; + } + const auto parsed = parse_decode_mode( + body["decode_mode"].get()); + if (!parsed) { + send_error(fd, 400, + "decode_mode must be adaptive, ar, or speculation"); + return false; + } + if (*parsed == SpeculationPolicy::Always && + !config_.speculative_enabled) { + send_error(fd, 400, + "decode_mode=speculation requires speculative decode to be enabled"); + return false; + } + req.decode_mode = *parsed; + } // Accept the output-token names used by each supported API dialect. // Default when the client omits all three: --default-max-tokens, so @@ -1975,13 +2001,16 @@ bool HttpServer::validate_request_context( void HttpServer::log_parsed_request(const ParsedRequest & req) const { std::fprintf(stderr, "[server] chat %s format=%s stream=%s msgs=%zu tools=%zu prompt_tokens=%zu " - "max_tokens=%d max_ctx=%d thinking=%s started_in_thinking=%s stops=%zu model=%s\n", + "max_tokens=%d max_ctx=%d thinking=%s started_in_thinking=%s " + "decode_mode=%.*s stops=%zu model=%s\n", req.response_id.c_str(), api_format_name(req.format), req.stream ? "true" : "false", json_array_size(req.messages), json_array_size(req.tools), req.prompt_tokens.size(), req.max_output, config_.max_ctx, req.thinking_enabled ? "true" : "false", req.started_in_thinking ? "true" : "false", + static_cast(decode_mode_name(req.decode_mode).size()), + decode_mode_name(req.decode_mode).data(), req.stop_sequences.size(), req.model.c_str()); } @@ -3475,6 +3504,9 @@ void HttpServer::prepare_generation_inputs( inputs.request.prompt = prepared.tokens; inputs.request.n_gen = inputs.generation_cap; inputs.request.sampler = req.sampler; + inputs.request.speculation_policy = req.decode_mode; + inputs.request.force_ar_decode = + req.decode_mode == SpeculationPolicy::Never; inputs.request.do_sample = req.sampler.needs_logit_processing(); // Tokens are delivered through DaemonIO so all API formats share the // same disconnect and streaming state machine. diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 92e4f9a9e..8a4739c82 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -176,6 +176,7 @@ struct ServerConfig { int fa_window = 0; int ddtree_budget = 0; bool speculative_enabled = false; + SpeculationPolicy decode_mode = SpeculationPolicy::Adaptive; bool target_sharding = false; // Prefill chunk size (bargs.chunk). Exposed at /props.runtime.chunk so // bench/snapshot tooling can capture the full server config — needed @@ -302,6 +303,10 @@ struct ParsedRequest { int max_output = 4096; bool stream = true; SamplerCfg sampler; + // Defaults to the server policy and may be overridden by the top-level + // `decode_mode` request field. + SpeculationPolicy decode_mode = + SpeculationPolicy::Adaptive; std::string model; // Tool definitions (stored as JSON for response formatting) json tools; diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 428d6b2a9..c9b860990 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -732,8 +732,11 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { SeqEngine::StepInput input; input.slot = i; input.token = slots[(size_t)i].pending_tok; + input.speculation_policy = + slots[(size_t)i].job->req.decode_mode; input.allow_speculation = - slots[(size_t)i].hook.close_token_ids.empty(); + slots[(size_t)i].hook.close_token_ids.empty() && + input.speculation_policy != SpeculationPolicy::Never; step_plan.decode.push_back(input); } else if (slots[(size_t)i].job) { prefill_candidates.push_back( diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 7373aace4..69f6ddaa3 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -122,6 +122,9 @@ static void print_usage(const char * prog) { " --no-fast-rollback Disable speculative fast rollback, even with --ddtree\n" " --ddtree Enable DDTree speculative decode\n" " --ddtree-budget DDTree budget (default: 22)\n" + " --decode-mode \n" + " Default decode policy; requests may override it\n" + " with the decode_mode field (default: adaptive)\n" " --verify-width laguna chain spec verify width (default: base 8,\n" " trimmed per step by drafter confidence; N = fixed base)\n" " --adaptive-experts [tau] MoE expert-count gating on verify batches\n" @@ -412,6 +415,14 @@ int main(int argc, char ** argv) { bargs.fast_rollback = true; } else if (std::strcmp(argv[i], "--ddtree-budget") == 0 && i + 1 < argc) { bargs.ddtree_budget = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--decode-mode") == 0 && i + 1 < argc) { + const auto policy = parse_decode_mode(argv[++i]); + if (!policy) { + std::fprintf(stderr, + "--decode-mode must be adaptive, ar, or speculation\n"); + return 2; + } + sconfig.decode_mode = *policy; } else if (std::strcmp(argv[i], "--adaptive-experts") == 0) { const char * tau = "0.80"; if (i + 1 < argc && argv[i + 1][0] != '-') { @@ -654,7 +665,6 @@ int main(int argc, char ** argv) { 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. @@ -681,6 +691,18 @@ int main(int argc, char ** argv) { ? 2 : 1; } + const bool speculative_decode_enabled = + bargs.draft_path != nullptr && + arch_supports_decode_draft( + backend_preparation.plan.arch(), + bargs.device.is_layer_split()); + if (sconfig.decode_mode == SpeculationPolicy::Always && + !speculative_decode_enabled) { + std::fprintf(stderr, + "[server] --decode-mode=speculation requires a supported " + "decode drafter (--draft)\n"); + return 2; + } // Options that parsed cleanly but do nothing on this model. Reported up // front so they are visible before the backend's own startup chatter. for (const std::string & warning : backend_preparation.warnings) { @@ -1108,6 +1130,10 @@ int main(int argc, char ** argv) { "[server] │ Use --fa-window 0 for tool-call workloads.\n"); } std::fprintf(stderr, "[server] │ ddtree = %s\n", bargs.ddtree_mode ? "ON" : "off"); + std::fprintf(stderr, "[server] │ decode_mode = %.*s\n", + static_cast(decode_mode_name( + sconfig.decode_mode).size()), + decode_mode_name(sconfig.decode_mode).data()); std::fprintf(stderr, "[server] │ fast_rollback = %s\n", bargs.fast_rollback ? "ON" : "off"); if (bargs.device.is_layer_split()) { std::fprintf(stderr, "[server] │ split_rollback = %s\n", @@ -1157,8 +1183,8 @@ int main(int argc, char ** argv) { sconfig.model_path = bargs.model_path ? bargs.model_path : ""; sconfig.draft_path = bargs.draft_path ? bargs.draft_path : ""; sconfig.fa_window = bargs.fa_window; - sconfig.ddtree_budget = bargs.ddtree_budget; - sconfig.speculative_enabled = bargs.ddtree_mode; + sconfig.ddtree_budget = bargs.ddtree_mode ? bargs.ddtree_budget : 0; + sconfig.speculative_enabled = speculative_decode_enabled; sconfig.target_sharding = bargs.device.is_layer_split(); // KV type: report the operator's choice if set, else the family default // the backend resolves (the tq3_0 auto policy was removed; laguna uses diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 2ff471a3d..09d788419 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -2557,6 +2557,24 @@ TEST_CASE(ServerUnitFixture, test_pflash_raw_body_preserved) { TEST_ASSERT(req.raw_body["temperature"].get() > 0.6f); } +TEST_CASE(ServerUnitFixture, test_decode_mode_defaults_and_props) { + ParsedRequest req; + TEST_ASSERT(req.decode_mode == SpeculationPolicy::Adaptive); + + ServerConfig cfg; + cfg.arch = "qwen35"; + cfg.speculative_enabled = true; + cfg.decode_mode = SpeculationPolicy::Always; + Tokenizer tok; + PrefixCache pc(0, tok); + ToolMemory tm; + const json body = build_props_body(cfg, pc, tm); + TEST_ASSERT(body["speculative"]["enabled"].get()); + TEST_ASSERT(body["decode_mode"].get() == "speculation"); + TEST_ASSERT(body["speculative"]["decode_mode"].get() == + "speculation"); +} + TEST_CASE(ServerUnitFixture, test_parse_request_sampler_applies_defaults_and_overrides) { SamplingDefaults defaults; defaults.has_temperature = true; diff --git a/server/test/test_speculation_goodput.cpp b/server/test/test_speculation_goodput.cpp index f667680cc..478ee8e26 100644 --- a/server/test/test_speculation_goodput.cpp +++ b/server/test/test_speculation_goodput.cpp @@ -1,6 +1,7 @@ #include "common/concurrency/adaptive_verification.h" #include "common/concurrency/speculation_goodput.h" #include "common/concurrency/speculation_prompt_prior.h" +#include "common/speculation_policy.h" #include "host_check.h" #include @@ -10,6 +11,19 @@ using namespace dflash::common; static int g_checks = 0; int main() { + // The public policy vocabulary is intentionally small and stable across + // DDTree, DSpark, and later speculators. + { + CHECK(parse_decode_mode("adaptive") == + SpeculationPolicy::Adaptive); + CHECK(parse_decode_mode("speculation") == + SpeculationPolicy::Always); + CHECK(parse_decode_mode("ar") == + SpeculationPolicy::Never); + CHECK(!parse_decode_mode("always")); + CHECK(decode_mode_name(SpeculationPolicy::Never) == "ar"); + } + // The cold-start prior separates obvious structured/code requests from // conversational writing while leaving ambiguous requests neutral. { @@ -163,6 +177,56 @@ int main() { CHECK(decision.predicted_gain > 1.05); } + // User-forced speculation does not wait for an AR baseline and remains + // selected even when its measured route is slower than AR. + { + AdaptiveVerificationRanker ranker; + std::vector candidates = { + {41, 1.0, 8.0, false, 0.0, 0, 0, true}, + }; + AdaptiveVerificationDecision decision = + ranker.select(2, candidates, 1); + CHECK(decision.requests.size() == 1); + CHECK(decision.requests[0] == 41); + + ranker.observe_autoregressive(2, 100.0); + ranker.observe_route(2, 1, 200.0); + decision = ranker.select(2, candidates, 1); + CHECK(decision.requests.size() == 1); + CHECK(decision.requests[0] == 41); + } + + // Mandatory requests form the prefix. Adaptive peers may join only when + // the exact combined width is profitable, and a user override can exceed + // the normal efficient-lane limit without being silently ignored. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(2, 100.0); + ranker.observe_route(2, 1, 200.0); + ranker.observe_route(2, 2, 100.0); + std::vector candidates = { + {51, 4.0, 8.0, true}, + {52, 1.0, 8.0, true, 0.0, 0, 0, true}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(2, candidates, 2); + CHECK(decision.requests.size() == 2); + CHECK(decision.requests[0] == 52); + CHECK(decision.requests[1] == 51); + } + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(4, 100.0); + std::vector candidates = { + {61, 1.0, 8.0, false, 0.0, 0, 0, true}, + {62, 1.0, 8.0, false, 0.0, 0, 0, true}, + }; + const AdaptiveVerificationDecision decision = + ranker.select(4, candidates, /*max_speculative_requests=*/1); + CHECK(decision.requests.size() == 2); + CHECK(decision.exploring); + } + // Candidates are ranked by expected value and every measured width is // compared against the whole-batch AR baseline. { From 1c49cd2ee0faed7b9afa5a7d8b5725a6dbb70315 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 08:08:49 +0000 Subject: [PATCH 17/18] feat(concurrency): route speculation from drafter confidence --- .../concurrency/adaptive_verification.h | 319 +++++++++++------- .../concurrency/speculation_confidence.h | 192 +++++++++++ .../concurrency/speculation_prompt_prior.h | 59 ---- server/src/common/sampler.h | 5 - .../src/deepseek4/deepseek4_dspark_spec.cpp | 11 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 296 ++++++++++++---- .../qwen35/concurrency/qwen35_seq_engine.h | 21 +- server/src/server/scheduler.cpp | 32 +- server/test/test_speculation_goodput.cpp | 309 +++++++++++------ 9 files changed, 864 insertions(+), 380 deletions(-) create mode 100644 server/src/common/concurrency/speculation_confidence.h delete mode 100644 server/src/common/concurrency/speculation_prompt_prior.h diff --git a/server/src/common/concurrency/adaptive_verification.h b/server/src/common/concurrency/adaptive_verification.h index 8730b073b..ece531281 100644 --- a/server/src/common/concurrency/adaptive_verification.h +++ b/server/src/common/concurrency/adaptive_verification.h @@ -7,11 +7,14 @@ // helper applies the same policy at request granularity. The concrete // speculator supplies expected useful tokens; the engine supplies the observed // cost of each exact mixed route shape (active requests, speculative requests). -// DDTree can learn value from accepted paths, while DSpark can use its -// calibrated confidence head directly. One ranker instance represents one +// DDTree and DSpark expose the same conditional-survival contract. Raw model +// confidence orders cold probes while target-verified accepted paths calibrate +// it online. One ranker instance represents one // fixed speculator/proposal shape; adapters with ragged verification work must // keep separate rankers for distinct work buckets. +#include "speculation_confidence.h" + #include #include #include @@ -28,10 +31,13 @@ struct AdaptiveVerificationCandidate { // Includes the one token that ordinary AR would emit. double expected_tokens = 1.0; double maximum_tokens = 1.0; + // True when either the concrete speculator or target outcomes provide an + // estimate. This does not mean the raw confidence was post-hoc calibrated. bool calibrated = false; - // Cheap request prior used only to order otherwise-unknown candidates. - // Request-local measurements smoothly replace it during warmup. - double routing_prior = 0.0; + // Raw speculator estimate used to find a target-calibrated confidence + // profile. Prompt semantics never enter this value. + double confidence_expected_tokens = + std::numeric_limits::quiet_NaN(); // Number of request-local observations supporting the estimate. Shared // cohort evidence must never relax AR-peer protection for a new request. std::size_t evidence_samples = 0; @@ -42,11 +48,15 @@ struct AdaptiveVerificationCandidate { // observed yield still contributes to deciding whether adaptive peers // should share the same verifier pass. bool required = false; + SpeculatorKind speculator = SpeculatorKind::DDTree; }; struct AdaptiveVerificationYieldEstimate { double expected_tokens = 1.0; std::size_t evidence_samples = 0; + double confidence_expected_tokens = + std::numeric_limits::quiet_NaN(); + SpeculatorKind speculator = SpeculatorKind::DDTree; }; struct AdaptiveVerificationDecision { @@ -72,15 +82,15 @@ struct AdaptiveVerificationConfig { std::size_t homogeneous_minimum_samples = 4; // A complete two-lane verifier profile observes k=1 then k=2, yielding // three request outcomes before any steady route is selected. - std::size_t routing_prior_minimum_samples = 3; + std::size_t confidence_profile_minimum_samples = 3; // Shared evidence ranks a new request after the minimum above, but its // yield magnitude is shrunk toward AR until this many outcomes exist. - std::size_t routing_prior_full_weight_samples = 8; + std::size_t confidence_profile_full_weight_samples = 8; // A continuous-backlog scheduler may optimize aggregate output after a - // routing-prior bucket has accumulated this many target-verified outcomes. + // confidence bucket has accumulated this many target-verified outcomes. // This is deliberately stronger than cold ordering, but does not count as // request-local proof for closed-cohort AR-peer protection. - std::size_t routing_prior_peer_guard_samples = 6; + std::size_t confidence_profile_peer_guard_samples = 6; double cost_ewma_alpha = 0.35; }; @@ -98,22 +108,12 @@ inline bool adaptive_verification_can_extend_stable_cohort( 3 * static_cast(verifier_request_lanes); } -// Convert conditional survival confidence into expected useful tokens: -// 1 AR token plus the probability of reaching every speculative prefix. -// DSpark supplies calibrated confidence-head values directly. -inline double expected_tokens_from_conditional_confidence( - const float * confidence, int count) { - double expected = 1.0; - double survival = 1.0; - for (int i = 0; confidence && i < count; ++i) { - const double conditional = std::clamp( - std::isfinite(confidence[i]) - ? static_cast(confidence[i]) : 0.0, - 0.0, 1.0); - survival *= conditional; - expected += survival; - } - return expected; +inline bool adaptive_verification_confidence_is_stale( + int current_progress, int observed_progress, + int refresh_interval) { + return refresh_interval > 0 && current_progress >= observed_progress && + static_cast(current_progress) - observed_progress >= + refresh_interval; } class AdaptiveVerificationRanker { @@ -127,17 +127,49 @@ class AdaptiveVerificationRanker { route_cost_known_.clear(); route_cost_samples_.clear(); request_yield_.clear(); - routing_prior_yield_.clear(); + request_confidence_.clear(); + confidence_yield_.clear(); + } + + void observe_request_estimate( + std::uint64_t request, + const SpeculationConfidenceEstimate & estimate) { + if (!estimate.available()) return; + const auto current = request_confidence_.find(request); + if (current != request_confidence_.end() && + current->second.available()) { + const SpeculationConfidenceProfile old_profile = + speculation_confidence_profile(current->second); + const SpeculationConfidenceProfile new_profile = + speculation_confidence_profile(estimate); + if (old_profile.speculator != new_profile.speculator || + std::abs(old_profile.expected_half_tokens - + new_profile.expected_half_tokens) >= 2) { + // A materially different drafter regime invalidates the local + // running target mean. Profile-level target outcomes stay + // isolated in their original buckets. + request_yield_.erase(request); + } + } + request_confidence_[request] = estimate; } void observe_request_yield(std::uint64_t request, double emitted_tokens) { if (!valid_yield(emitted_tokens)) return; update_yield(request_yield_[request], emitted_tokens); + const auto confidence = request_confidence_.find(request); + if (confidence != request_confidence_.end() && + confidence->second.available()) { + update_yield( + confidence_yield_[speculation_confidence_profile( + confidence->second)], emitted_tokens); + } } void forget_request(std::uint64_t request) { request_yield_.erase(request); + request_confidence_.erase(request); } std::optional request_expected_tokens( @@ -183,49 +215,59 @@ class AdaptiveVerificationRanker { config_.homogeneous_minimum_relative_yield * maximum; } - void observe_routing_prior_yield(double routing_prior, - double emitted_tokens) { - if (!std::isfinite(routing_prior) || !valid_yield(emitted_tokens)) { - return; - } - update_yield(routing_prior_yield_[routing_prior], emitted_tokens); - } - - std::optional routing_prior_expected_tokens( - double routing_prior) const { - if (!std::isfinite(routing_prior)) return std::nullopt; - const auto it = routing_prior_yield_.find(routing_prior); - return it == routing_prior_yield_.end() || - it->second.samples < config_.routing_prior_minimum_samples - ? std::nullopt : std::optional(it->second.expected_tokens); + std::optional confidence_profile_expected_tokens( + SpeculatorKind speculator, + double confidence_expected_tokens) const { + if (!std::isfinite(confidence_expected_tokens)) return std::nullopt; + const auto profile = speculation_confidence_profile( + speculator, confidence_expected_tokens); + const auto it = confidence_yield_.find(profile); + return it == confidence_yield_.end() || + it->second.samples < + config_.confidence_profile_minimum_samples + ? std::nullopt + : std::optional(it->second.expected_tokens); } - std::size_t routing_prior_yield_samples(double routing_prior) const { - if (!std::isfinite(routing_prior)) return 0; - const auto it = routing_prior_yield_.find(routing_prior); - return it == routing_prior_yield_.end() ? 0 : it->second.samples; + std::size_t confidence_profile_yield_samples( + SpeculatorKind speculator, + double confidence_expected_tokens) const { + if (!std::isfinite(confidence_expected_tokens)) return 0; + const auto it = confidence_yield_.find( + speculation_confidence_profile( + speculator, confidence_expected_tokens)); + return it == confidence_yield_.end() ? 0 : it->second.samples; } - bool has_stable_routing_prior_yield(double routing_prior) const { - if (!std::isfinite(routing_prior)) return false; - const auto it = routing_prior_yield_.find(routing_prior); - return it != routing_prior_yield_.end() && + bool has_stable_confidence_yield( + SpeculatorKind speculator, + double confidence_expected_tokens) const { + if (!std::isfinite(confidence_expected_tokens)) return false; + const auto it = confidence_yield_.find( + speculation_confidence_profile( + speculator, confidence_expected_tokens)); + return it != confidence_yield_.end() && it->second.samples >= - config_.routing_prior_peer_guard_samples && + config_.confidence_profile_peer_guard_samples && has_useful_yield(it->second.expected_tokens); } - bool forms_stable_routing_prior_cohort( + bool forms_stable_confidence_cohort( const std::vector & candidates) const { if (candidates.empty()) return false; double minimum = std::numeric_limits::infinity(); double maximum = 1.0; for (const AdaptiveVerificationCandidate & candidate : candidates) { - if (!has_stable_routing_prior_yield(candidate.routing_prior)) { + if (!has_stable_confidence_yield( + candidate.speculator, + candidate.confidence_expected_tokens)) { return false; } - const auto it = routing_prior_yield_.find(candidate.routing_prior); + const auto it = confidence_yield_.find( + speculation_confidence_profile( + candidate.speculator, + candidate.confidence_expected_tokens)); minimum = std::min(minimum, it->second.expected_tokens); maximum = std::max(maximum, it->second.expected_tokens); } @@ -234,32 +276,51 @@ class AdaptiveVerificationRanker { } std::optional estimate_request_yield( - std::uint64_t request, double routing_prior, - bool trust_stable_routing_prior = false) const { + std::uint64_t request, + bool trust_stable_confidence = false) const { const auto request_it = request_yield_.find(request); - const auto prior_it = std::isfinite(routing_prior) - ? routing_prior_yield_.find(routing_prior) - : routing_prior_yield_.end(); + const auto confidence_it = request_confidence_.find(request); const bool has_request = request_it != request_yield_.end(); - const bool has_prior = prior_it != routing_prior_yield_.end() && - prior_it->second.samples >= config_.routing_prior_minimum_samples; - if (!has_request && !has_prior) return std::nullopt; + const bool has_confidence = + confidence_it != request_confidence_.end() && + confidence_it->second.available(); + auto profile_it = confidence_yield_.end(); + if (has_confidence) { + profile_it = confidence_yield_.find( + speculation_confidence_profile(confidence_it->second)); + } + const bool has_profile = profile_it != confidence_yield_.end() && + profile_it->second.samples >= + config_.confidence_profile_minimum_samples; + if (!has_request && !has_confidence && !has_profile) { + return std::nullopt; + } AdaptiveVerificationYieldEstimate out; - const bool stable_backlog_prior = - trust_stable_routing_prior && - prior_it->second.samples >= - config_.routing_prior_peer_guard_samples; - const double prior_weight = stable_backlog_prior - ? 1.0 - : std::min( - 1.0, static_cast(prior_it->second.samples) / - static_cast( - config_.routing_prior_full_weight_samples)); - const double trusted_prior = 1.0 + prior_weight * - (prior_it->second.expected_tokens - 1.0); + if (has_confidence) { + out.confidence_expected_tokens = + confidence_it->second.expected_tokens(); + out.speculator = confidence_it->second.speculator; + } + double trusted_confidence = has_confidence + ? confidence_it->second.expected_tokens() : 1.0; + if (has_profile) { + const bool stable_backlog_confidence = + trust_stable_confidence && + profile_it->second.samples >= + config_.confidence_profile_peer_guard_samples; + const double profile_weight = stable_backlog_confidence + ? 1.0 + : std::min( + 1.0, + static_cast(profile_it->second.samples) / + static_cast( + config_.confidence_profile_full_weight_samples)); + trusted_confidence += profile_weight * + (profile_it->second.expected_tokens - trusted_confidence); + } if (!has_request) { - out.expected_tokens = trusted_prior; + out.expected_tokens = trusted_confidence; out.evidence_samples = 0; return out; } @@ -267,20 +328,20 @@ class AdaptiveVerificationRanker { const YieldEstimate & request_estimate = request_it->second; out.expected_tokens = request_estimate.expected_tokens; out.evidence_samples = request_estimate.samples; - if (!has_prior) return out; - - if (request_estimate.samples < config_.homogeneous_minimum_samples) { - // Shrink the first few noisy request observations toward a stable - // cohort mean. Once request-local evidence is stable, its measured - // magnitude fully replaces the prior for goodput decisions. + if (has_confidence && + request_estimate.samples < + config_.homogeneous_minimum_samples) { + // Blend the first few noisy target observations with the concrete + // speculator estimate (and any profile calibration). Once local + // evidence is stable, measured target yield fully replaces it. const double local_weight = std::min( 1.0, static_cast(request_estimate.samples) / static_cast( config_.homogeneous_minimum_samples)); - out.expected_tokens = trusted_prior + + out.expected_tokens = trusted_confidence + local_weight * (request_estimate.expected_tokens - - trusted_prior); + trusted_confidence); } return out; } @@ -393,13 +454,27 @@ class AdaptiveVerificationRanker { bool has_exact_profile(int active_requests, int max_speculative_requests, - std::size_t minimum_route_samples = 1) const { - if (!has_autoregressive_cost(active_requests)) return false; + std::size_t minimum_route_samples = 1, + int baseline_speculative_requests = 0) const { + if (baseline_speculative_requests < 0 || + baseline_speculative_requests > active_requests) { + return false; + } const std::size_t required_samples = std::max(1, minimum_route_samples); - const int limit = std::max(0, std::min( - active_requests, max_speculative_requests)); - for (int k = 1; k <= limit; ++k) { + const std::size_t baseline_samples = + baseline_speculative_requests == 0 ? 1 : required_samples; + if (route_cost_samples( + active_requests, baseline_speculative_requests) < + baseline_samples) { + return false; + } + const int limit = std::max( + baseline_speculative_requests, + std::max(0, std::min( + active_requests, max_speculative_requests))); + for (int k = baseline_speculative_requests + 1; + k <= limit; ++k) { if (route_cost_samples(active_requests, k) < required_samples) return false; } @@ -443,9 +518,6 @@ class AdaptiveVerificationRanker { std::isfinite(candidate.expected_tokens) ? candidate.expected_tokens : 1.0, 1.0, candidate.maximum_tokens); - if (!std::isfinite(candidate.routing_prior)) { - candidate.routing_prior = 0.0; - } if (candidate.required) { required_candidates.push_back(candidate); } else { @@ -470,10 +542,13 @@ class AdaptiveVerificationRanker { out.requests.push_back(candidate.request); } }; - if (!has_autoregressive_cost(active_requests)) { - // Adaptive-only cohorts first observe the exact AR baseline. User - // overrides are stronger: an Always request must not be delayed - // by profiling, and its route timing is learned from this step. + const int required_count = + static_cast(required_candidates.size()); + const int baseline_width = required_count; + if (!has_route_cost(active_requests, baseline_width)) { + // Adaptive-only cohorts first observe all-AR. With user-forced + // requests, all-AR is unattainable, so their mandatory mixed route + // is the baseline from which optional adaptive peers are judged. select_required(); return out; } @@ -515,7 +590,7 @@ class AdaptiveVerificationRanker { } else { // For mixed cohorts, rotate only within contiguous half-token // request-value buckets. This preserves ordering across materially - // different expected yields and never uses the raw prompt hint. + // different expected yields without inspecting prompt semantics. std::map> fair_group_positions; for (std::size_t i = 0; i < known.size(); ++i) { const int yield_bucket = static_cast( @@ -558,12 +633,12 @@ class AdaptiveVerificationRanker { std::stable_sort(unknown.begin(), unknown.end(), [](const AdaptiveVerificationCandidate & a, const AdaptiveVerificationCandidate & b) { - if (a.routing_prior != b.routing_prior) { - return a.routing_prior > b.routing_prior; - } if (a.maximum_tokens != b.maximum_tokens) { return a.maximum_tokens > b.maximum_tokens; } + if (a.progress_tokens != b.progress_tokens) { + return a.progress_tokens < b.progress_tokens; + } return a.request < b.request; }); if (!unknown.empty()) { @@ -577,16 +652,24 @@ class AdaptiveVerificationRanker { required_candidates.end()); ordered_known.insert( ordered_known.end(), known.begin(), known.end()); - const int required_count = - static_cast(required_candidates.size()); - - const double baseline = - static_cast(active_requests) / - autoregressive_cost_us(active_requests); - const double required = baseline * config_.minimum_gain; const std::size_t required_route_samples = std::max( 1, minimum_speculative_route_samples); + if (route_cost_samples(active_requests, baseline_width) < + required_route_samples) { + select_required(); + return out; + } + double baseline_expected_total = + static_cast(active_requests); + for (const AdaptiveVerificationCandidate & candidate : + required_candidates) { + baseline_expected_total += candidate.expected_tokens - 1.0; + } + const double baseline_cost = + route_cost_us(active_requests, baseline_width); + const double baseline = baseline_expected_total / baseline_cost; + const double required = baseline * config_.minimum_gain; int admitted_prefix = required_count; double best = baseline; double admitted_goodput = baseline; @@ -610,7 +693,7 @@ class AdaptiveVerificationRanker { k == active_requests || homogeneous_speculative_cohort || route_us <= config_.maximum_ar_peer_slowdown * - autoregressive_cost_us(active_requests); + baseline_cost; if (!protects_ar_peers) continue; if (throughput > best) { best = throughput; @@ -632,7 +715,8 @@ class AdaptiveVerificationRanker { (probe_uncalibrated_with_verifier ? static_cast(unknown.size()) : 0)); int missing_width = 0; - for (int k = std::max(1, required_count); probe_missing_routes && + for (int k = std::max(1, required_count + 1); + probe_missing_routes && k <= probe_candidates; ++k) { if (route_cost_samples(active_requests, k) >= required_route_samples) continue; @@ -726,14 +810,15 @@ class AdaptiveVerificationRanker { config.homogeneous_minimum_relative_yield, 0.0, 1.0); config.homogeneous_minimum_samples = std::max(1, config.homogeneous_minimum_samples); - config.routing_prior_minimum_samples = - std::max(1, config.routing_prior_minimum_samples); - config.routing_prior_full_weight_samples = - std::max(config.routing_prior_minimum_samples, - config.routing_prior_full_weight_samples); - config.routing_prior_peer_guard_samples = - std::max(config.routing_prior_minimum_samples, - config.routing_prior_peer_guard_samples); + config.confidence_profile_minimum_samples = + std::max( + 1, config.confidence_profile_minimum_samples); + config.confidence_profile_full_weight_samples = + std::max(config.confidence_profile_minimum_samples, + config.confidence_profile_full_weight_samples); + config.confidence_profile_peer_guard_samples = + std::max(config.confidence_profile_minimum_samples, + config.confidence_profile_peer_guard_samples); config.cost_ewma_alpha = std::clamp(config.cost_ewma_alpha, 0.0, 1.0); return config; @@ -771,7 +856,9 @@ class AdaptiveVerificationRanker { // Capped running means represent the complete request/profile while still // adapting gradually if a model or speculator changes behavior. std::map request_yield_; - std::map routing_prior_yield_; + std::map + request_confidence_; + std::map confidence_yield_; }; } // namespace dflash::common diff --git a/server/src/common/concurrency/speculation_confidence.h b/server/src/common/concurrency/speculation_confidence.h new file mode 100644 index 000000000..cd1932e13 --- /dev/null +++ b/server/src/common/concurrency/speculation_confidence.h @@ -0,0 +1,192 @@ +#pragma once + +// Model-neutral confidence value shared by adaptive speculative decoders. +// +// The scheduler never interprets prompt text. A concrete speculator exposes +// conditional prefix-survival confidence and the ranker converts it to useful +// tokens. DDTree obtains the values from draft top-1 probabilities; DSpark +// obtains them from its confidence head. Target-verified accepted tokens are +// deliberately tracked separately and calibrate these estimates online. + +#include +#include +#include +#include + +namespace dflash::common { + +enum class SpeculatorKind : std::uint8_t { + DDTree, + DSpark, +}; + +// Confidence is not necessarily available before routing. In particular, +// both current adapters must run their drafter; an already-selected proposal +// can expose it for free, while a pre-route DDTree probe costs an extra pass. +enum class SpeculationConfidenceCost : std::uint8_t { + Unavailable, + ExtraDraftPass, + PiggybacksOnProposal, +}; + +inline double conditional_prefix_survival( + const float * confidence, int count) { + double survival = 1.0; + for (int i = 0; confidence && i < count; ++i) { + const double conditional = std::clamp( + std::isfinite(confidence[i]) + ? static_cast(confidence[i]) : 0.0, + 0.0, 1.0); + survival *= conditional; + } + return survival; +} + +inline double expected_tokens_from_conditional_confidence( + const float * confidence, int count) { + double expected = 1.0; + double survival = 1.0; + for (int i = 0; confidence && i < count; ++i) { + const double conditional = std::clamp( + std::isfinite(confidence[i]) + ? static_cast(confidence[i]) : 0.0, + 0.0, 1.0); + survival *= conditional; + expected += survival; + } + return expected; +} + +// Non-owning view for a speculator hot path. DSpark can select its verifier +// width without allocating; engines that retain confidence across steps copy +// the same view into SpeculationConfidenceEstimate below. +struct SpeculationConfidenceView { + SpeculatorKind speculator = SpeculatorKind::DDTree; + SpeculationConfidenceCost cost = + SpeculationConfidenceCost::Unavailable; + bool posthoc_calibrated = false; + const float * conditional = nullptr; + int count = 0; + + bool available() const { + return cost != SpeculationConfidenceCost::Unavailable && + conditional && count > 0; + } + + double prefix_survival(int prefix_tokens) const { + return conditional_prefix_survival( + conditional, std::max(0, std::min(count, prefix_tokens))); + } + + double expected_tokens() const { + return expected_tokens_from_conditional_confidence( + conditional, std::max(0, count)); + } +}; + +inline SpeculationConfidenceView make_speculation_confidence_view( + SpeculatorKind speculator, const float * confidence, int count, + SpeculationConfidenceCost cost, + bool posthoc_calibrated = false) { + return { + speculator, + confidence && count > 0 + ? cost : SpeculationConfidenceCost::Unavailable, + posthoc_calibrated, + confidence, + std::max(0, count), + }; +} + +struct SpeculationConfidenceEstimate { + SpeculatorKind speculator = SpeculatorKind::DDTree; + SpeculationConfidenceCost cost = + SpeculationConfidenceCost::Unavailable; + // True only when the artifact guarantees probability calibration. Raw + // DDTree softmax and the current DSpark sigmoid head both leave this false; + // the ranker still uses them as ordering estimates and corrects them from + // target outcomes. + bool posthoc_calibrated = false; + std::vector conditional; + + bool available() const { + return cost != SpeculationConfidenceCost::Unavailable && + !conditional.empty(); + } + + double maximum_tokens() const { + return 1.0 + static_cast(conditional.size()); + } + + double expected_tokens() const { + return expected_tokens_from_conditional_confidence( + conditional.data(), static_cast(conditional.size())); + } + + SpeculationConfidenceEstimate limited_to( + int speculative_tokens) const { + SpeculationConfidenceEstimate out = *this; + const std::size_t limit = static_cast( + std::max(0, speculative_tokens)); + if (out.conditional.size() > limit) { + out.conditional.resize(limit); + } + return out; + } +}; + +inline SpeculationConfidenceEstimate make_speculation_confidence_estimate( + SpeculatorKind speculator, const float * confidence, int count, + SpeculationConfidenceCost cost, + bool posthoc_calibrated = false) { + const SpeculationConfidenceView view = make_speculation_confidence_view( + speculator, confidence, count, cost, posthoc_calibrated); + SpeculationConfidenceEstimate out; + out.speculator = view.speculator; + out.cost = view.cost; + out.posthoc_calibrated = view.posthoc_calibrated; + if (!view.available()) return out; + out.conditional.reserve(static_cast(view.count)); + for (int i = 0; i < view.count; ++i) { + out.conditional.push_back(static_cast(std::clamp( + std::isfinite(view.conditional[i]) + ? static_cast(view.conditional[i]) : 0.0, + 0.0, 1.0))); + } + return out; +} + +struct SpeculationConfidenceProfile { + SpeculatorKind speculator = SpeculatorKind::DDTree; + // Half-token bins are coarse enough to transfer target observations + // between similar requests without pretending raw confidence is exact. + int expected_half_tokens = 2; + + bool operator<(const SpeculationConfidenceProfile & other) const { + if (speculator != other.speculator) { + return static_cast(speculator) < + static_cast(other.speculator); + } + return expected_half_tokens < other.expected_half_tokens; + } +}; + +inline SpeculationConfidenceProfile speculation_confidence_profile( + const SpeculationConfidenceEstimate & estimate) { + return { + estimate.speculator, + static_cast(std::floor(estimate.expected_tokens() * 2.0)), + }; +} + +inline SpeculationConfidenceProfile speculation_confidence_profile( + SpeculatorKind speculator, double expected_tokens) { + const double sanitized = std::max( + 1.0, std::isfinite(expected_tokens) ? expected_tokens : 1.0); + return { + speculator, + static_cast(std::floor(sanitized * 2.0)), + }; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/speculation_prompt_prior.h b/server/src/common/concurrency/speculation_prompt_prior.h deleted file mode 100644 index 66c90f1aa..000000000 --- a/server/src/common/concurrency/speculation_prompt_prior.h +++ /dev/null @@ -1,59 +0,0 @@ -#pragma once - -// Training-free request prior for cold-start speculative routing. -// -// This is deliberately a weak prior, not an eligibility decision: obvious -// structured/code prompts are sampled before neutral or conversational ones -// during cold start. A bucket is reused only after several measured verifier -// outcomes, and request-local evidence takes over as it stabilizes. The policy -// is model-neutral and can be replaced by a learned prompt ranker without -// changing the verifier. - -#include -#include -#include -#include - -namespace dflash::common { - -inline int speculation_prompt_hint(std::string_view prompt) { - constexpr size_t kMaxPromptChars = 16 * 1024; - std::string text(prompt.substr(0, kMaxPromptChars)); - std::transform(text.begin(), text.end(), text.begin(), - [](unsigned char c) { - return static_cast(std::tolower(c)); - }); - auto has = [&](std::string_view cue) { - return text.find(cue) != std::string::npos; - }; - - // Explicit exclusions override incidental structured words such as - // "avoid code" in a creative-writing request. - if (has("avoid code") || has("chatting casually") || - has("casual conversation") || has("small talk") || - has("keep it conversational") || has("write a story") || - has("invent a story") || has("write a poem") || - has("roleplay")) { - return -1; - } - - int score = 0; - score += has("```") ? 4 : 0; - score += has("\ndef ") || has("\nclass ") || has("#include") ? 4 : 0; - score += has("public static") || has("fn ") || has("function ") ? 3 : 0; - score += has("implement") || has("debug") || has("unit test") ? 3 : 0; - score += has("algorithm") || has("sql query") || - has("regular expression") || has("json schema") ? 2 : 0; - score += has("python") || has("javascript") || has("typescript") || - has("rust") || has("c++") ? 2 : 0; - score += has("code") ? 1 : 0; - - if (score >= 3) return 1; - if (has("story") || has("conversational") || has("brainstorm") || - has("opinion") || has("friendly chat")) { - return -1; - } - return 0; -} - -} // namespace dflash::common diff --git a/server/src/common/sampler.h b/server/src/common/sampler.h index 8aab2718d..53a4c2ce9 100644 --- a/server/src/common/sampler.h +++ b/server/src/common/sampler.h @@ -25,11 +25,6 @@ struct SamplerCfg { float rep_pen = 1.0f; // multiplicative repetition penalty (HF-style) int rep_window = 256; uint64_t seed = 0; - // Cold-start speculation prior: -1 conversational/creative, 0 neutral, - // +1 structured/code. It affects admission only; measured goodput remains - // authoritative. Non-HTTP callers naturally retain the neutral default. - int8_t speculation_prompt_hint = 0; - // OpenAI-style additive penalties (applied per-token to logits before softmax). // frequency_penalty: subtract freq_pen * count(token_in_history) from logit. // presence_penalty: subtract pres_pen * 1(token_appeared_in_history) from logit. diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 1d687b586..c6ce1f3cf 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -25,6 +25,7 @@ #include "deepseek4_internal.h" #include "deepseek4_roctx.h" #include "internal.h" +#include "common/concurrency/speculation_confidence.h" #include "common/dspark_head.h" #include "ggml.h" @@ -930,10 +931,16 @@ bool run_deepseek4_dspark_spec_decode( // traces and keep q=4 for high-confidence prefixes while avoiding its // extra verify cost on low-acceptance prompts. if (use_confidence_width && draft_confidence.size() >= 2 && draft_tok.size() >= 3) { - const float confidence_p2 = draft_confidence[0] * draft_confidence[1]; + const SpeculationConfidenceView confidence = + make_speculation_confidence_view( + SpeculatorKind::DSpark, draft_confidence.data(), + static_cast(draft_confidence.size()), + SpeculationConfidenceCost::PiggybacksOnProposal, + /*posthoc_calibrated=*/false); + const double confidence_p2 = confidence.prefix_survival(2); int selected_q = confidence_p2 >= kConfidenceQ3Threshold ? 3 : 2; if (selected_q == 3 && draft_confidence.size() >= 3 && draft_tok.size() >= 4) { - const float confidence_p3 = confidence_p2 * draft_confidence[2]; + const double confidence_p3 = confidence.prefix_survival(3); if (confidence_p3 >= kConfidenceQ4Threshold) selected_q = 4; } if ((int) draft_tok.size() > selected_q) draft_tok.resize((size_t) selected_q); diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index a7ea34f3b..af02236f5 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -69,6 +69,12 @@ int adaptive_calibration_interval(int active_requests) { return 64 * std::max(1, (active_requests + 3) / 4); } +int adaptive_confidence_refresh_interval(int active_requests) { + // Revisit AR-routed requests without turning confidence estimation into a + // per-token tax. Wider batches amortize a rejected probe for longer. + return 64 * std::max(1, (active_requests + 3) / 4); +} + } // namespace Qwen35SeqEngine::Qwen35SeqEngine( @@ -198,8 +204,9 @@ bool Qwen35SeqEngine::ddtree_input_eligible(const StepInput & in) const { return slots_.slot(in.slot).generated_tokens() >= min_floor; } -std::optional Qwen35SeqEngine::estimate_ddtree_expected_tokens( - const StepInput & in, int tree_budget) { +std::optional +Qwen35SeqEngine::estimate_ddtree_confidence( + const StepInput & in) { const int q_len = b_.dw_.block_size; const int hidden = b_.w_.n_embd; if (q_len <= 1 || !build_lm_head_projection_step( @@ -211,7 +218,8 @@ std::optional Qwen35SeqEngine::estimate_ddtree_expected_tokens( DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); if (!draft || !mirror) return std::nullopt; - auto fail = [&]() -> std::optional { + auto fail = [&]() -> + std::optional { // begin_step advances only host cache bookkeeping, but a failed graph // can leave the appended rows incomplete. Rebuild from committed // target features on the next proposal. @@ -265,8 +273,8 @@ std::optional Qwen35SeqEngine::estimate_ddtree_expected_tokens( // Training-free SVIP-style confidence: the top-1 draft probability at // each position estimates conditional survival, so the cumulative product // estimates reaching that prefix. DSpark can replace this adapter with its - // calibrated confidence-head probabilities without changing the ranker. - const int confidence_tokens = std::min(tree_budget, q_len - 1); + // confidence-head probabilities without changing the ranker. + const int confidence_tokens = q_len - 1; std::vector confidence((size_t)confidence_tokens); for (int pos = 1; pos <= confidence_tokens; ++pos) { confidence[(size_t)pos - 1] = static_cast( @@ -274,8 +282,33 @@ std::optional Qwen35SeqEngine::estimate_ddtree_expected_tokens( top_lp[(size_t)pos])), 0.0, 1.0)); } - return expected_tokens_from_conditional_confidence( - confidence.data(), static_cast(confidence.size())); + return make_speculation_confidence_estimate( + SpeculatorKind::DDTree, confidence.data(), + static_cast(confidence.size()), + SpeculationConfidenceCost::ExtraDraftPass, + /*posthoc_calibrated=*/false); +} + +void Qwen35SeqEngine::remember_ddtree_confidence( + std::uint64_t request_id, + const SpeculationConfidenceEstimate & estimate, + int progress_tokens) { + if (!estimate.available()) return; + const auto current = ddtree_confidence_.find(request_id); + if (current != ddtree_confidence_.end() && + current->second.available()) { + const SpeculationConfidenceProfile old_profile = + speculation_confidence_profile(current->second); + const SpeculationConfidenceProfile new_profile = + speculation_confidence_profile(estimate); + if (old_profile.speculator != new_profile.speculator || + std::abs(old_profile.expected_half_tokens - + new_profile.expected_half_tokens) >= 2) { + ddtree_target_yield_.erase(request_id); + } + } + ddtree_confidence_[request_id] = estimate; + ddtree_confidence_progress_[request_id] = progress_tokens; } std::optional Qwen35SeqEngine::step_ddtree( @@ -299,6 +332,7 @@ std::optional Qwen35SeqEngine::step_ddtree( std::vector flat; std::vector accepted; int32_t bonus = -1; + SpeculationConfidenceEstimate confidence; }; std::vector proposals; proposals.reserve((size_t)active); @@ -379,6 +413,18 @@ std::optional Qwen35SeqEngine::step_ddtree( Proposal p; p.slot = in.slot; p.root = in.token; + const int confidence_tokens = q_len - 1; + std::vector confidence((size_t)confidence_tokens); + for (int pos = 1; pos <= confidence_tokens; ++pos) { + confidence[(size_t)pos - 1] = static_cast( + std::clamp(std::exp(static_cast( + top_lp[(size_t)pos * K])), + 0.0, 1.0)); + } + p.confidence = make_speculation_confidence_estimate( + SpeculatorKind::DDTree, confidence.data(), confidence_tokens, + SpeculationConfidenceCost::PiggybacksOnProposal, + /*posthoc_calibrated=*/false); p.tree = build_ddtree( top_lp.data() + K, top_ids.data() + K, q_len - 1, K, tree_budget, @@ -390,6 +436,13 @@ std::optional Qwen35SeqEngine::step_ddtree( } proposals.push_back(std::move(p)); } + for (const Proposal & proposal : proposals) { + const std::uint64_t request_id = + slots_.slot(proposal.slot).request_id; + remember_ddtree_confidence( + request_id, proposal.confidence, + slots_.slot(proposal.slot).generated_tokens()); + } const bool target_is_meta = b_.cache_.ssm_state.empty() || !b_.cache_.ssm_state.front() || ggml_backend_buft_is_meta(ggml_backend_buffer_get_type( @@ -1084,10 +1137,30 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( const SamplerCfg & sampler) { AdmitResult result = slots_.admit(request_id, prompt, sampler); if (result.status == AdmitResult::Status::admitted) { + std::uint8_t inherited_compact_budget = 0; + bool has_active_peer = false; + bool compact_peers_agree = true; + for (int slot = 0; slot < slots_.slot_count(); ++slot) { + if (slot == result.slot || !slots_.is_active(slot)) continue; + const std::uint8_t peer_budget = + compact_tree_cohort_[(size_t)slot]; + if (!has_active_peer) { + inherited_compact_budget = peer_budget; + has_active_peer = true; + } else if (peer_budget != inherited_compact_budget) { + compact_peers_agree = false; + } + if (peer_budget == 0) compact_peers_agree = false; + } adaptive_verification_.forget_request(request_id); compact_short_adaptive_verification_.forget_request(request_id); compact_adaptive_verification_.forget_request(request_id); - compact_tree_cohort_[(size_t)result.slot] = 0; + ddtree_confidence_.erase(request_id); + ddtree_confidence_progress_.erase(request_id); + ddtree_target_yield_.erase(request_id); + compact_tree_cohort_[(size_t)result.slot] = + has_active_peer && compact_peers_agree + ? inherited_compact_budget : 0; reset_recurrent_slot(b_.cache_, result.slot); if (slots_.residency_active()) { slots_.slot(result.slot).kvflash_last_reselect_generated = @@ -1334,6 +1407,31 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const bool adaptive_enabled = !(adaptive && std::atoi(adaptive) == 0); const AdaptiveVerificationOracle & oracle = adaptive_verification_oracle(); + if (adaptive_enabled && !oracle.forces_selection()) { + const int refresh_interval = adaptive_confidence_refresh_interval( + static_cast(inputs.size())); + for (const StepInput & in : inputs) { + const Qwen35Slot & seq = slots_.slot(in.slot); + const auto observed = + ddtree_confidence_progress_.find(seq.request_id); + if (observed == ddtree_confidence_progress_.end() || + !adaptive_verification_confidence_is_stale( + seq.generated_tokens(), observed->second, + refresh_interval)) { + continue; + } + // Confidence is a decode-regime signal, not a permanent prompt + // label. Periodically forget it so an AR-routed request can be + // re-probed after its continuation changes character. + adaptive_verification_.forget_request(seq.request_id); + compact_short_adaptive_verification_.forget_request( + seq.request_id); + compact_adaptive_verification_.forget_request(seq.request_id); + ddtree_confidence_.erase(seq.request_id); + ddtree_confidence_progress_.erase(observed); + ddtree_target_yield_.erase(seq.request_id); + } + } const int inherited_compact_budget = inputs.empty() ? 0 : compact_tree_cohort_[(size_t)inputs.front().slot]; const bool inherited_compact_tree = @@ -1343,32 +1441,63 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { return compact_tree_cohort_[(size_t)in.slot] == inherited_compact_budget; }); - const int positive_prompt_candidates = static_cast(std::count_if( - inputs.begin(), inputs.end(), [&](const StepInput & in) { - return ddtree_input_eligible(in) && - slots_.slot(in.slot).sampler.speculation_prompt_hint > 0; - })); - // Full DDTree depth wins for low-occupancy homogeneous code. Once target - // AR is well batched, or a queued C=4 cohort is genuinely mixed, a compact - // proposal lowers the marginal verification cost enough for a small - // profitable request subset. Keep that compact shape through the cohort's - // tail; all-code C<=4 still gets the established full-depth path. + int confidence_candidates = 0; + int useful_confidence_candidates = 0; + int eligible_candidates = 0; + for (const StepInput & in : inputs) { + if (!ddtree_input_eligible(in)) continue; + ++eligible_candidates; + const Qwen35Slot & seq = slots_.slot(in.slot); + const auto estimate = ddtree_confidence_.find(seq.request_id); + if (estimate == ddtree_confidence_.end() || + !estimate->second.available()) { + continue; + } + ++confidence_candidates; + double expected_tokens = estimate->second.expected_tokens(); + const auto target = ddtree_target_yield_.find(seq.request_id); + if (target != ddtree_target_yield_.end() && target->second.samples > 0) { + // Raw confidence orders cold requests. Target-accepted output + // progressively corrects that estimate, independent of which + // proposal shape produced the observation. + const double target_weight = std::min( + 1.0, static_cast(target->second.samples) / 4.0); + expected_tokens += target_weight * + (target->second.expected_tokens - expected_tokens); + } + if (expected_tokens >= 1.5) { + ++useful_confidence_candidates; + } + } + const bool confidence_cohort_ready = eligible_candidates > 0 && + confidence_candidates == eligible_candidates; + const bool mixed_confidence_cohort = confidence_cohort_ready && + useful_confidence_candidates > 0 && + useful_confidence_candidates < static_cast(inputs.size()); + // Full DDTree depth remains the neutral low-occupancy shape. Once the + // concrete speculator has estimated the whole eligible cohort, a mixed C=4 + // backlog can use a compact proposal for its small useful subset. No prompt + // category or model-specific task table participates in this decision. constexpr int kFullTreeMaxConcurrency = 4; constexpr int kCompactShortTreeBudget = 4; constexpr int kCompactTreeBudget = 8; const bool starts_short_backlog_shape = plan.has_refill_backlog && inputs.size() == kFullTreeMaxConcurrency && - positive_prompt_candidates > 0 && - positive_prompt_candidates <= 2 && - positive_prompt_candidates < static_cast(inputs.size()); + confidence_cohort_ready && useful_confidence_candidates > 0 && + useful_confidence_candidates <= 2 && mixed_confidence_cohort; + const bool retains_compact_mixed_cohort = + inherited_compact_tree && + (!confidence_cohort_ready || mixed_confidence_cohort); const bool compact_tree_route = - inputs.size() > kFullTreeMaxConcurrency || inherited_compact_tree || - starts_short_backlog_shape; + inputs.size() > kFullTreeMaxConcurrency || + retains_compact_mixed_cohort || starts_short_backlog_shape; const int selected_compact_budget = - inherited_compact_tree && !plan.has_refill_backlog + inherited_compact_tree && + (!plan.has_refill_backlog || !confidence_cohort_ready) ? inherited_compact_budget - : (positive_prompt_candidates > 0 && - positive_prompt_candidates <= 2 + : (confidence_cohort_ready && + useful_confidence_candidates > 0 && + useful_confidence_candidates <= 2 ? kCompactShortTreeBudget : kCompactTreeBudget); const int adaptive_tree_budget = @@ -1394,26 +1523,33 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // head while reusing the same selection policy. std::vector candidates; candidates.reserve(inputs.size()); - auto collect_candidates = [&](bool trust_stable_routing_prior = false) { + auto collect_candidates = [&](bool trust_stable_confidence = false) { candidates.clear(); for (const StepInput & in : inputs) { if (!ddtree_input_eligible(in)) continue; const Qwen35Slot & seq = slots_.slot(in.slot); - const int prompt_hint = seq.sampler.speculation_prompt_hint; + const auto confidence = ddtree_confidence_.find(seq.request_id); + if (confidence != ddtree_confidence_.end() && + confidence->second.available()) { + route_ranker.observe_request_estimate( + seq.request_id, + confidence->second.limited_to(tree_budget)); + } const std::optional estimate = route_ranker.estimate_request_yield( - seq.request_id, static_cast(prompt_hint), - trust_stable_routing_prior); + seq.request_id, trust_stable_confidence); candidates.push_back({ in.slot, estimate ? estimate->expected_tokens : 1.0, static_cast(compact_tree_route ? tree_budget + 1 : b_.dw_.block_size), estimate.has_value(), - static_cast(prompt_hint), + estimate ? estimate->confidence_expected_tokens + : std::numeric_limits::quiet_NaN(), estimate ? estimate->evidence_samples : 0, seq.generated_tokens(), in.speculation_policy == SpeculationPolicy::Always, + SpeculatorKind::DDTree, }); } }; @@ -1439,7 +1575,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { plan.has_refill_backlog && compact_tree_route && !broad_verifier_coverage && bounded_stable_cohort_extension; if (use_stable_cohort_extension) { - collect_candidates(/*trust_stable_routing_prior=*/true); + collect_candidates(/*trust_stable_confidence=*/true); } bool enforce_ar_peer_guard = true; std::size_t minimum_route_samples = 1; @@ -1462,16 +1598,24 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { logged = true; } } else if (adaptive_enabled) { - const int exact_profile_width = std::min( - adaptive_speculation_limit, - static_cast(candidates.size())); + const int mandatory_speculative_requests = + static_cast(std::count_if( + candidates.begin(), candidates.end(), + [](const AdaptiveVerificationCandidate & candidate) { + return candidate.required; + })); + const int exact_profile_width = std::max( + mandatory_speculative_requests, + std::min(adaptive_speculation_limit, + static_cast(candidates.size()))); // Replay a cold graph shape only in the narrow capacity extension that // can plausibly become steady. This avoids doubling rejected probes at // occupancies such as Strix C=16 and leaves established low-C profiles // unchanged. minimum_route_samples = use_stable_cohort_extension && - positive_prompt_candidates == active_requests + confidence_cohort_ready && + useful_confidence_candidates == active_requests ? 2 : 1; int & calibration_cooldown = compact_tree_route ? (compact_short_shape @@ -1480,24 +1624,25 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { : adaptive_calibration_cooldown_; const bool exact_profile_ready = route_ranker.has_exact_profile( static_cast(inputs.size()), exact_profile_width, - minimum_route_samples); + minimum_route_samples, mandatory_speculative_requests); // A closed cohort keeps request-local AR-peer protection. Continuous // backlog may instead optimize aggregate output once either verifier // coverage is broad enough or target-verified outcomes have made a - // useful routing-prior bucket stable. Exact route costs still decide + // useful confidence bucket stable. Exact route costs still decide // whether speculation wins, so the same evidence can unlock C=8 while // an unprofitable C=16 profile remains AR. A future speculator supplies // its own work-bucket ranker and inherits the same rule. const bool has_stable_backlog_cohort = plan.has_refill_backlog && exact_profile_ready && - route_ranker.forms_stable_routing_prior_cohort(candidates); + route_ranker.forms_stable_confidence_cohort(candidates); const bool has_stable_backlog_candidate = plan.has_refill_backlog && exact_profile_ready && std::any_of( candidates.begin(), candidates.end(), [&](const AdaptiveVerificationCandidate & candidate) { - return route_ranker.has_stable_routing_prior_yield( - candidate.routing_prior); + return route_ranker.has_stable_confidence_yield( + candidate.speculator, + candidate.confidence_expected_tokens); }); const bool relax_ar_peer_guard = plan.has_refill_backlog && compact_tree_route && @@ -1518,14 +1663,11 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // Full DDTree uses its drafter-side confidence adapter. Compact DDTree // avoids a duplicate draft pass and calibrates one unknown request in // useful verification instead. DSpark can take the former path with - // its calibrated confidence head. + // its confidence head; target outcomes calibrate both adapters. const bool drafter_side_calibration = !compact_tree_route && probe_missing_routes; - const bool prompt_prior_marks_mixed = - positive_prompt_candidates > 0 && - positive_prompt_candidates < static_cast(inputs.size()); const bool suppress_uncovered_mixed_calibration = - exact_profile_ready && prompt_prior_marks_mixed && + exact_profile_ready && mixed_confidence_cohort && !broad_verifier_coverage && has_stable_backlog_candidate; if (!exact_profile_ready || inputs.size() <= 3) { calibration_cooldown = 0; @@ -1542,7 +1684,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { !suppress_uncovered_mixed_calibration && (!exact_profile_ready || (calibration_cooldown == 0 && - (plan.has_refill_backlog || !prompt_prior_marks_mixed))); + (plan.has_refill_backlog || !mixed_confidence_cohort))); std::vector verifier_candidates; const std::vector * decision_candidates = &candidates; @@ -1587,17 +1729,22 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { return in.slot == decision.calibration_request; }); if (input != inputs.end()) { - const std::optional expected = - estimate_ddtree_expected_tokens(*input, tree_budget); - if (expected) { - route_ranker.observe_request_yield( - slots_.slot(input->slot).request_id, *expected); + const std::optional + confidence = + estimate_ddtree_confidence(*input); + if (confidence) { + const std::uint64_t request_id = + slots_.slot(input->slot).request_id; + remember_ddtree_confidence( + request_id, *confidence, + slots_.slot(input->slot).generated_tokens()); + route_ranker.observe_request_estimate( + request_id, confidence->limited_to(tree_budget)); std::fprintf(stderr, "[parallel-ddtree] confidence request=%llu slot=%d " "expected_tokens=%.3f\n", - (unsigned long long) - slots_.slot(input->slot).request_id, - input->slot, *expected); + (unsigned long long)request_id, + input->slot, confidence->expected_tokens()); collect_candidates(); const AdaptiveVerificationDecision steady_after = route_ranker.select( @@ -1691,6 +1838,20 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { static_cast(inputs.size()), speculative_count, route_us, /*discard_first_sample=*/ speculative_count > 0 && minimum_route_samples > 1); + if (speculative_count == 0) { + // AR has the same work shape for every proposal profile. Sharing + // this exact (C,0) timing avoids re-running the baseline merely + // because confidence moves a cohort from full to compact DDTree. + for (AdaptiveVerificationRanker * ranker : { + &adaptive_verification_, + &compact_short_adaptive_verification_, + &compact_adaptive_verification_}) { + if (ranker != &route_ranker) { + ranker->observe_route( + static_cast(inputs.size()), 0, route_us); + } + } + } } if (decision.exploring && !speculative_plan.decode.empty()) { std::fprintf(stderr, @@ -1707,12 +1868,26 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { const double emitted = static_cast(out.ddtree_accepted_tokens + 1); Qwen35Slot & observed_seq = slots_.slot(out.slot); + const auto confidence = + ddtree_confidence_.find(observed_seq.request_id); + if (confidence != ddtree_confidence_.end()) { + route_ranker.observe_request_estimate( + observed_seq.request_id, + confidence->second.limited_to(tree_budget)); + } route_ranker.observe_request_yield( observed_seq.request_id, emitted); - route_ranker.observe_routing_prior_yield( - static_cast( - observed_seq.sampler.speculation_prompt_hint), - emitted); + RequestTargetYield & target = + ddtree_target_yield_[observed_seq.request_id]; + if (target.samples == 0) { + target.expected_tokens = emitted; + } else { + constexpr double kTargetYieldAlpha = 0.35; + target.expected_tokens += kTargetYieldAlpha * + (emitted - target.expected_tokens); + } + target.samples = std::min( + target.samples + 1, 64); } } @@ -2140,6 +2315,9 @@ void Qwen35SeqEngine::retire(int slot) { adaptive_verification_.forget_request(request_id); compact_short_adaptive_verification_.forget_request(request_id); compact_adaptive_verification_.forget_request(request_id); + ddtree_confidence_.erase(request_id); + ddtree_confidence_progress_.erase(request_id); + ddtree_target_yield_.erase(request_id); slots_.retire(slot); compact_tree_cohort_[(size_t)slot] = 0; } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 2d91f126d..3149c6e65 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -30,6 +30,7 @@ #include #include +#include #include #include #include @@ -121,8 +122,12 @@ class Qwen35SeqEngine final : public SeqEngine { DraftKvState * ensure_slot_draft_kv(int slot); bool ddtree_available(const StepPlan & plan) const; bool ddtree_input_eligible(const StepInput & input) const; - std::optional estimate_ddtree_expected_tokens( - const StepInput & input, int tree_budget); + std::optional estimate_ddtree_confidence( + const StepInput & input); + void remember_ddtree_confidence( + std::uint64_t request_id, + const SpeculationConfidenceEstimate & estimate, + int progress_tokens); StepResult step_regular(const StepPlan & plan); // nullopt means proposal setup failed before target/cache mutation and the // caller may safely use the ordinary packed AR path for this iteration. @@ -140,6 +145,18 @@ class Qwen35SeqEngine final : public SeqEngine { AdaptiveVerificationRanker adaptive_verification_; AdaptiveVerificationRanker compact_short_adaptive_verification_; AdaptiveVerificationRanker compact_adaptive_verification_; + // Latest drafter confidence is request-owned and can be projected onto + // the current full/compact work budget without looking at prompt text. + std::map + ddtree_confidence_; + std::map ddtree_confidence_progress_; + struct RequestTargetYield { + double expected_tokens = 1.0; + std::size_t samples = 0; + }; + // Shape choice uses target outcomes when available; this map deliberately + // spans full/compact rankers while their hardware route costs stay split. + std::map ddtree_target_yield_; int adaptive_calibration_cooldown_ = 0; int compact_short_adaptive_calibration_cooldown_ = 0; int compact_adaptive_calibration_cooldown_ = 0; diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index c9b860990..b4d279474 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -10,7 +10,6 @@ #include "http_server.h" #include "common/concurrency/seq_engine.h" -#include "common/concurrency/speculation_prompt_prior.h" #include #include @@ -22,32 +21,6 @@ namespace dflash::common { namespace { -std::string last_user_prompt_text(const json & messages) { - if (!messages.is_array()) return {}; - for (int i = static_cast(messages.size()) - 1; i >= 0; --i) { - const json & message = messages[(size_t)i]; - if (!message.is_object() || - message.value("role", "") != "user" || - !message.contains("content")) { - continue; - } - const json & content = message["content"]; - if (content.is_string()) return content.get(); - if (!content.is_array()) return {}; - std::string text; - for (const json & part : content) { - if (!part.is_object() || !part.contains("text") || - !part["text"].is_string()) { - continue; - } - if (!text.empty()) text.push_back('\n'); - text += part["text"].get(); - } - return text; - } - return {}; -} - // 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 @@ -550,11 +523,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // Admission only claims the slot and queues the prompt. Prefill // advances one chunk per engine step alongside live decode. const uint64_t engine_request_id = next_request_id; - SamplerCfg sampler = req.sampler; - sampler.speculation_prompt_hint = - speculation_prompt_hint(last_user_prompt_text(req.messages)); auto ar = engine.admit(engine_request_id, effective_prompt, - sampler); + req.sampler); if (ar.status == SeqEngine::AdmitResult::Status::busy) return AdmissionDisposition::Deferred; if (ar.status != SeqEngine::AdmitResult::Status::admitted) { diff --git a/server/test/test_speculation_goodput.cpp b/server/test/test_speculation_goodput.cpp index 478ee8e26..2a77b2547 100644 --- a/server/test/test_speculation_goodput.cpp +++ b/server/test/test_speculation_goodput.cpp @@ -1,6 +1,5 @@ #include "common/concurrency/adaptive_verification.h" #include "common/concurrency/speculation_goodput.h" -#include "common/concurrency/speculation_prompt_prior.h" #include "common/speculation_policy.h" #include "host_check.h" @@ -24,17 +23,6 @@ int main() { CHECK(decode_mode_name(SpeculationPolicy::Never) == "ar"); } - // The cold-start prior separates obvious structured/code requests from - // conversational writing while leaving ambiguous requests neutral. - { - CHECK(speculation_prompt_hint( - "Complete this Python function:\n\ndef solve(values):") == 1); - CHECK(speculation_prompt_hint( - "Chatting casually, write a story and avoid code.") == -1); - CHECK(speculation_prompt_hint( - "What happened during the Apollo 11 mission?") == 0); - } - // Cold start measures one real speculative step and one neighboring AR // step, then keeps the route with higher useful-token goodput. { @@ -53,7 +41,7 @@ int main() { CHECK(policy.expected_emitted_tokens() == 4.0); } - // A chat-like low-yield probe loses to AR and is disabled per request. + // A low-yield probe loses to AR and is disabled per request. { SpeculationGoodputController policy; CHECK(policy.observe_speculation(/*emitted_tokens=*/1, @@ -224,7 +212,37 @@ int main() { const AdaptiveVerificationDecision decision = ranker.select(4, candidates, /*max_speculative_requests=*/1); CHECK(decision.requests.size() == 2); + CHECK(!decision.exploring); + } + + // A forced request makes (C,r), not the impossible all-AR route, the cold + // baseline. Optional adaptive peers can therefore be profiled and joined. + { + AdaptiveVerificationRanker ranker; + std::vector candidates = { + {71, 4.0, 8.0, true, 4.0, 0, 0, true}, + {72, 4.0, 8.0, true, 4.0}, + }; + AdaptiveVerificationDecision decision = + ranker.select(2, candidates, 2); + CHECK(decision.requests == std::vector{71}); + CHECK(!ranker.has_exact_profile( + 2, 2, /*minimum_route_samples=*/1, + /*baseline_speculative_requests=*/1)); + ranker.observe_route(2, 1, 100.0); + decision = ranker.select(2, candidates, 2); CHECK(decision.exploring); + CHECK(decision.requests == std::vector({71, 72})); + CHECK(!ranker.has_exact_profile( + 2, 2, /*minimum_route_samples=*/1, + /*baseline_speculative_requests=*/1)); + ranker.observe_route(2, 2, 80.0); + decision = ranker.select(2, candidates, 2); + CHECK(!decision.exploring); + CHECK(decision.requests == std::vector({71, 72})); + CHECK(ranker.has_exact_profile( + 2, 2, /*minimum_route_samples=*/1, + /*baseline_speculative_requests=*/1)); } // Candidates are ranked by expected value and every measured width is @@ -261,44 +279,44 @@ int main() { } // Compact verifier calibration discovers every exact route width in order, - // even when k=1 and k=2 lose. The prompt prior orders unknown requests but - // never filters them, and the globally best non-convex k=3 route wins. + // even when k=1 and k=2 lose. Unknown requests use deterministic neutral + // ordering, and the globally best non-convex k=3 route wins. { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(8, 100.0); std::vector candidates = { - {5, 1.0, 9.0, false, 1.0}, - {2, 1.0, 9.0, false, -1.0}, - {9, 1.0, 9.0, false, 0.0}, + {5, 1.0, 9.0, false}, + {2, 1.0, 9.0, false}, + {9, 1.0, 9.0, false}, }; AdaptiveVerificationDecision probe = ranker.select(8, candidates, 3, /*probe_uncalibrated_with_verifier=*/true); CHECK(probe.exploring); CHECK(probe.requests.size() == 1); - CHECK(probe.requests[0] == 5); + CHECK(probe.requests[0] == 2); CHECK(probe.calibration_request == -1); ranker.observe_route(8, 1, 160.0); - candidates[0] = {5, 4.0, 9.0, true, 1.0}; + candidates[1] = {2, 4.0, 9.0, true}; probe = ranker.select(8, candidates, 3, /*probe_uncalibrated_with_verifier=*/true); CHECK(probe.exploring); CHECK(probe.requests.size() == 2); - CHECK(probe.requests[0] == 5); - CHECK(probe.requests[1] == 9); + CHECK(probe.requests[0] == 2); + CHECK(probe.requests[1] == 5); ranker.observe_route(8, 2, 170.0); - candidates[2] = {9, 4.0, 9.0, true, 0.0}; + candidates[0] = {5, 4.0, 9.0, true}; probe = ranker.select(8, candidates, 3, /*probe_uncalibrated_with_verifier=*/true); CHECK(probe.exploring); CHECK(probe.requests.size() == 3); - CHECK(probe.requests[0] == 5); - CHECK(probe.requests[1] == 9); - CHECK(probe.requests[2] == 2); + CHECK(probe.requests[0] == 2); + CHECK(probe.requests[1] == 5); + CHECK(probe.requests[2] == 9); ranker.observe_route(8, 3, 105.0); - candidates[1] = {2, 4.0, 9.0, true, -1.0}; + candidates[2] = {9, 4.0, 9.0, true}; const AdaptiveVerificationDecision decision = ranker.select(8, candidates, 3, @@ -389,6 +407,12 @@ int main() { CHECK(!adaptive_verification_can_extend_stable_cohort(16, 3)); CHECK(adaptive_verification_can_extend_stable_cohort(16, 6)); CHECK(!adaptive_verification_can_extend_stable_cohort(1, 0)); + + CHECK(!adaptive_verification_confidence_is_stale(63, 0, 64)); + CHECK(adaptive_verification_confidence_is_stale(64, 0, 64)); + CHECK(adaptive_verification_confidence_is_stale(192, 64, 128)); + CHECK(!adaptive_verification_confidence_is_stale(63, 64, 64)); + CHECK(!adaptive_verification_confidence_is_stale(128, 0, 0)); } // Timings from a neighboring occupancy never stand in for the exact-C @@ -460,8 +484,7 @@ int main() { CHECK(decision.requests[0] == 13); } - // A prior affects cold-start order only. Once calibrated, the high-yield - // conversational-prior request ranks ahead of a low-yield code prior. + // Expected target yield, rather than a prompt category, determines order. { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(2, 100.0); @@ -477,79 +500,126 @@ int main() { CHECK(decision.requests[0] == 2); } - // Verified yield can seed later requests in the same coarse routing-prior - // bucket. The cache belongs to this speculator/profile ranker and resets - // with it; no model-specific table is required. + // DDTree softmax and DSpark confidence-head output use one sanitized + // conditional-survival contract. Raw estimates are explicitly distinct + // from target-verified calibration evidence. { + const float confidence[] = {0.8f, 0.5f}; + const SpeculationConfidenceEstimate ddtree_confidence = + make_speculation_confidence_estimate( + SpeculatorKind::DDTree, confidence, 2, + SpeculationConfidenceCost::ExtraDraftPass); + const SpeculationConfidenceEstimate dspark_confidence = + make_speculation_confidence_estimate( + SpeculatorKind::DSpark, confidence, 2, + SpeculationConfidenceCost::PiggybacksOnProposal); + CHECK(ddtree_confidence.available()); + CHECK(dspark_confidence.available()); + CHECK(!ddtree_confidence.posthoc_calibrated); + CHECK(!dspark_confidence.posthoc_calibrated); + CHECK(std::abs(ddtree_confidence.expected_tokens() - 2.2) < 1e-6); + CHECK(std::abs(dspark_confidence.expected_tokens() - 2.2) < 1e-6); + CHECK(std::abs(conditional_prefix_survival(confidence, 2) - 0.4) < + 1e-6); + const SpeculationConfidenceEstimate limited = + ddtree_confidence.limited_to(1); + CHECK(limited.maximum_tokens() == 2.0); + CHECK(ddtree_confidence.maximum_tokens() == 3.0); + CHECK(std::abs(ddtree_confidence.expected_tokens() - 2.2) < 1e-6); + + const float invalid[] = { + 2.0f, -1.0f, std::numeric_limits::quiet_NaN()}; + const auto sanitized = make_speculation_confidence_estimate( + SpeculatorKind::DDTree, invalid, 3, + SpeculationConfidenceCost::PiggybacksOnProposal); + CHECK(sanitized.expected_tokens() == 2.0); + const auto unavailable = make_speculation_confidence_estimate( + SpeculatorKind::DSpark, nullptr, 0, + SpeculationConfidenceCost::PiggybacksOnProposal); + CHECK(!unavailable.available()); + CHECK(unavailable.expected_tokens() == 1.0); + AdaptiveVerificationRanker ranker; - CHECK(!ranker.routing_prior_expected_tokens(1.0).has_value()); - ranker.observe_routing_prior_yield(1.0, 4.0); - ranker.observe_routing_prior_yield(1.0, 6.0); - ranker.observe_routing_prior_yield(1.0, 4.0); - CHECK(std::abs( - ranker.routing_prior_expected_tokens(1.0).value() - - 14.0 / 3.0) < 1e-9); - CHECK(ranker.routing_prior_yield_samples(1.0) == 3); - ranker.observe_routing_prior_yield(1.0, 6.0); - CHECK(ranker.routing_prior_expected_tokens(1.0).value() == 5.0); - CHECK(ranker.routing_prior_yield_samples(1.0) == 4); - CHECK(!ranker.has_stable_routing_prior_yield(1.0)); - CHECK(!ranker.routing_prior_expected_tokens(-1.0).has_value()); - CHECK(!ranker.has_stable_routing_prior_yield(-1.0)); - - // Shared yield magnitude itself starts conservatively shrunk toward - // AR; four local samples still make the request authoritative. - const double expected[] = {2.5, 2.0, 1.5, 1.0}; - for (int sample = 1; sample <= 4; ++sample) { + ranker.observe_request_estimate(98, ddtree_confidence); + auto estimate = ranker.estimate_request_yield(98); + CHECK(estimate.has_value()); + CHECK(std::abs(estimate->expected_tokens - 2.2) < 1e-6); + CHECK(estimate->evidence_samples == 0); + CHECK(ranker.request_yield_samples(98) == 0); + CHECK(!ranker.confidence_profile_expected_tokens( + SpeculatorKind::DDTree, 2.2).has_value()); + + // Early target outcomes smoothly correct an optimistic raw estimate; + // drafter confidence never counts as manufactured target evidence. + ranker.observe_request_yield(98, 1.0); + estimate = ranker.estimate_request_yield(98); + CHECK(estimate.has_value()); + CHECK(std::abs(estimate->expected_tokens - 1.9) < 1e-6); + CHECK(estimate->evidence_samples == 1); + CHECK(ranker.confidence_profile_yield_samples( + SpeculatorKind::DDTree, 2.2) == 1); + for (int sample = 1; sample < 4; ++sample) { ranker.observe_request_yield(98, 1.0); - const auto estimate = - ranker.estimate_request_yield(98, 1.0); - CHECK(estimate.has_value()); - CHECK(estimate->expected_tokens == expected[sample - 1]); - CHECK(estimate->evidence_samples == - static_cast(sample)); } - const auto local = ranker.estimate_request_yield(98, 1.0); - CHECK(local.has_value()); - CHECK(local->expected_tokens == 1.0); - - ranker.observe_routing_prior_yield(1.0, 4.0); - CHECK(!ranker.has_stable_routing_prior_yield(1.0)); - ranker.observe_routing_prior_yield(1.0, 6.0); - CHECK(ranker.has_stable_routing_prior_yield(1.0)); - const auto closed_prior = - ranker.estimate_request_yield(100, 1.0); - const auto backlog_prior = - ranker.estimate_request_yield( - 100, 1.0, /*trust_stable_routing_prior=*/true); - CHECK(closed_prior.has_value()); - CHECK(backlog_prior.has_value()); - CHECK(closed_prior->expected_tokens == 4.0); - CHECK(backlog_prior->expected_tokens == 5.0); - for (int i = 0; i < 6; ++i) { - ranker.observe_routing_prior_yield(-1.0, 1.25); + estimate = ranker.estimate_request_yield(98); + CHECK(estimate.has_value()); + CHECK(std::abs(estimate->expected_tokens - 1.0) < 1e-6); + CHECK(estimate->evidence_samples == 4); + + // A material confidence-regime change invalidates local target history + // rather than treating the request as a permanent semantic category. + const float low_confidence_values[] = {0.1f, 0.1f}; + const SpeculationConfidenceEstimate low_confidence = + make_speculation_confidence_estimate( + SpeculatorKind::DDTree, low_confidence_values, 2, + SpeculationConfidenceCost::PiggybacksOnProposal); + AdaptiveVerificationRanker changing; + changing.observe_request_estimate(197, ddtree_confidence); + for (int sample = 0; sample < 4; ++sample) { + changing.observe_request_yield(197, 4.0); + } + CHECK(changing.request_yield_samples(197) == 4); + changing.observe_request_estimate(197, low_confidence); + CHECK(changing.request_yield_samples(197) == 0); + const auto changed = changing.estimate_request_yield(197); + CHECK(changed.has_value()); + CHECK(std::abs( + changed->expected_tokens - low_confidence.expected_tokens()) < + 1e-6); + + for (std::uint64_t request = 99; request <= 103; ++request) { + ranker.observe_request_estimate(request, ddtree_confidence); + ranker.observe_request_yield(request, 5.0); } - CHECK(!ranker.has_stable_routing_prior_yield(-1.0)); - CHECK(ranker.forms_stable_routing_prior_cohort({ - {101, 5.0, 9.0, true, 1.0}, - {102, 5.0, 9.0, true, 1.0}, + CHECK(ranker.has_stable_confidence_yield( + SpeculatorKind::DDTree, 2.2)); + CHECK(!ranker.has_stable_confidence_yield( + SpeculatorKind::DSpark, 2.2)); + CHECK(ranker.forms_stable_confidence_cohort({ + {101, 5.0, 9.0, true, 2.2, 0, 0, false, + SpeculatorKind::DDTree}, + {102, 5.0, 9.0, true, 2.2, 0, 0, false, + SpeculatorKind::DDTree}, })); - CHECK(!ranker.forms_stable_routing_prior_cohort({ - {101, 5.0, 9.0, true, 1.0}, - {103, 1.25, 9.0, true, -1.0}, + CHECK(!ranker.forms_stable_confidence_cohort({ + {101, 5.0, 9.0, true, 2.2, 0, 0, false, + SpeculatorKind::DDTree}, + {104, 5.0, 9.0, true, 2.2, 0, 0, false, + SpeculatorKind::DSpark}, })); // Per-request evidence is local to one proposal-shape ranker and is // explicitly forgotten at request retirement. AdaptiveVerificationRanker compact; - ranker.observe_request_yield(99, 7.0); - CHECK(ranker.request_expected_tokens(99).value() == 7.0); - CHECK(ranker.request_yield_samples(99) == 1); - CHECK(!compact.request_expected_tokens(99).has_value()); - ranker.forget_request(99); - CHECK(!ranker.request_expected_tokens(99).has_value()); + ranker.observe_request_yield(199, 7.0); + CHECK(ranker.request_expected_tokens(199).value() == 7.0); + CHECK(ranker.request_yield_samples(199) == 1); + CHECK(!compact.request_expected_tokens(199).has_value()); + ranker.forget_request(199); + CHECK(!ranker.request_expected_tokens(199).has_value()); ranker.reset(); - CHECK(!ranker.routing_prior_expected_tokens(1.0).has_value()); + CHECK(!ranker.confidence_profile_expected_tokens( + SpeculatorKind::DDTree, 2.2).has_value()); } // Shared cohort evidence can rank a new request, but it cannot unlock the @@ -559,14 +629,14 @@ int main() { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(5, 100.0); ranker.observe_route(5, 1, 150.0); - std::vector prior_backed; + std::vector confidence_backed; for (int request = 1; request <= 5; ++request) { - prior_backed.push_back( - {request, 8.0, 8.0, true, 1.0, 0}); + confidence_backed.push_back( + {request, 8.0, 8.0, true, 8.0, 0}); } - CHECK(ranker.select(5, prior_backed, 1).requests.empty()); + CHECK(ranker.select(5, confidence_backed, 1).requests.empty()); - std::vector probe = prior_backed; + std::vector probe = confidence_backed; for (AdaptiveVerificationCandidate & candidate : probe) { candidate.calibrated = false; } @@ -576,12 +646,12 @@ int main() { CHECK(exploring.exploring); CHECK(exploring.requests.size() == 1); - for (AdaptiveVerificationCandidate & candidate : prior_backed) { + for (AdaptiveVerificationCandidate & candidate : confidence_backed) { candidate.evidence_samples = 4; } - CHECK(ranker.select(5, prior_backed, 1).requests.size() == 1); - prior_backed.back().expected_tokens = 1.0; - CHECK(ranker.select(5, prior_backed, 1).requests.empty()); + CHECK(ranker.select(5, confidence_backed, 1).requests.size() == 1); + confidence_backed.back().expected_tokens = 1.0; + CHECK(ranker.select(5, confidence_backed, 1).requests.empty()); } // With route costs already profiled, a no-confidence adapter calibrates @@ -593,8 +663,8 @@ int main() { ranker.observe_route(5, 2, 95.0); std::vector candidates = { {14, 4.0, 8.0, true, 1.0}, - {15, 1.0, 8.0, false, -1.0}, - {16, 1.0, 8.0, false, 0.0}, + {15, 1.0, 8.0, false}, + {16, 1.0, 8.0, false}, }; const AdaptiveVerificationDecision decision = ranker.select( 5, candidates, 2, @@ -602,7 +672,7 @@ int main() { CHECK(decision.exploring); CHECK(decision.requests.size() == 2); CHECK(decision.requests[0] == 14); - CHECK(decision.requests[1] == 16); + CHECK(decision.requests[1] == 15); } // A promising newcomer can replace one incumbent even when every executor @@ -627,22 +697,23 @@ int main() { } // Selection is tied to request value, not to a lane count. When the - // profitable request retires, speculation does not migrate to chat. + // profitable request retires, speculation does not migrate to a low-yield + // peer. { AdaptiveVerificationRanker ranker; ranker.observe_autoregressive(2, 100.0); ranker.observe_route(2, 1, 100.0); - std::vector code_and_chat = { + std::vector mixed_yield = { {21, 8.0, 8.0, true, 1.0}, {22, 1.0, 8.0, true, -1.0}, }; AdaptiveVerificationDecision decision = - ranker.select(2, code_and_chat, 1); + ranker.select(2, mixed_yield, 1); CHECK(decision.requests.size() == 1); CHECK(decision.requests[0] == 21); ranker.observe_autoregressive(1, 60.0); ranker.observe_route(1, 1, 100.0); - decision = ranker.select(1, {code_and_chat[1]}, 1); + decision = ranker.select(1, {mixed_yield[1]}, 1); CHECK(decision.requests.empty()); } @@ -765,6 +836,32 @@ int main() { } } + // With exact route costs known, request-granular selection has no C + // cutoff: one or two useful requests remain speculative at every + // occupancy through the supported C=16 while low-yield peers stay AR. + { + for (int concurrency = 1; concurrency <= 16; ++concurrency) { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(concurrency, 100.0); + for (int width = 1; width <= concurrency; ++width) { + ranker.observe_route(concurrency, width, 100.0); + } + std::vector one_useful; + std::vector two_useful; + for (int request = 0; request < concurrency; ++request) { + one_useful.push_back({ + request, request == 0 ? 4.0 : 1.0, 4.0, true}); + two_useful.push_back({ + request, request < 2 ? 4.0 : 1.0, 4.0, true}); + } + CHECK(ranker.select(concurrency, one_useful, concurrency) + .requests.size() == 1); + CHECK(ranker.select(concurrency, two_useful, concurrency) + .requests.size() == + static_cast(std::min(2, concurrency))); + } + } + // A promising calibrated prefix gets one bounded hardware-cost probe when // that subbatch shape has not been observed yet. { From cd6cf3016b13d5298c8c2e5b94698011e99cefa7 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 10:02:12 +0000 Subject: [PATCH 18/18] perf(concurrency): plan speculative verifier work adaptively --- .../concurrency/adaptive_verification.h | 599 ++++++++- .../common/concurrency/speculation_planning.h | 197 +++ .../qwen35/concurrency/qwen35_seq_engine.cpp | 1103 ++++++++++------- .../qwen35/concurrency/qwen35_seq_engine.h | 56 +- server/test/test_speculation_goodput.cpp | 652 ++++++++++ 5 files changed, 2108 insertions(+), 499 deletions(-) create mode 100644 server/src/common/concurrency/speculation_planning.h diff --git a/server/src/common/concurrency/adaptive_verification.h b/server/src/common/concurrency/adaptive_verification.h index ece531281..d2907f3d3 100644 --- a/server/src/common/concurrency/adaptive_verification.h +++ b/server/src/common/concurrency/adaptive_verification.h @@ -14,6 +14,7 @@ // keep separate rankers for distinct work buckets. #include "speculation_confidence.h" +#include "speculation_planning.h" #include #include @@ -66,6 +67,32 @@ struct AdaptiveVerificationDecision { int calibration_request = -1; bool exploring = false; double predicted_gain = 1.0; + // Absolute useful-token goodput for a measured steady decision. This is + // comparable across exact work profiles; zero means the selected action + // is cold/unmeasured. + double predicted_goodput = 0.0; +}; + +enum class AdaptiveVerificationWorkStatus : std::uint8_t { + Autoregressive, + Calibration, + Verification, + RequiredUnavailable, + InvalidMenu, +}; + +struct AdaptiveVerificationWorkDecision { + AdaptiveVerificationWorkStatus status = + AdaptiveVerificationWorkStatus::Autoregressive; + // A work key is returned for calibration and verification actions. + std::optional work; + AdaptiveVerificationDecision decision; + + bool has_work() const { + return work.has_value() && + (status == AdaptiveVerificationWorkStatus::Calibration || + status == AdaptiveVerificationWorkStatus::Verification); + } }; struct AdaptiveVerificationConfig { @@ -135,13 +162,22 @@ class AdaptiveVerificationRanker { std::uint64_t request, const SpeculationConfidenceEstimate & estimate) { if (!estimate.available()) return; + observe_request_estimate( + request, estimate.speculator, estimate.expected_tokens()); + } + + void observe_request_estimate( + std::uint64_t request, SpeculatorKind speculator, + double expected_tokens) { + if (!std::isfinite(expected_tokens) || expected_tokens < 1.0) return; const auto current = request_confidence_.find(request); - if (current != request_confidence_.end() && - current->second.available()) { + if (current != request_confidence_.end()) { const SpeculationConfidenceProfile old_profile = - speculation_confidence_profile(current->second); + speculation_confidence_profile( + current->second.speculator, + current->second.expected_tokens); const SpeculationConfidenceProfile new_profile = - speculation_confidence_profile(estimate); + speculation_confidence_profile(speculator, expected_tokens); if (old_profile.speculator != new_profile.speculator || std::abs(old_profile.expected_half_tokens - new_profile.expected_half_tokens) >= 2) { @@ -151,7 +187,7 @@ class AdaptiveVerificationRanker { request_yield_.erase(request); } } - request_confidence_[request] = estimate; + request_confidence_[request] = {speculator, expected_tokens}; } void observe_request_yield(std::uint64_t request, @@ -159,11 +195,11 @@ class AdaptiveVerificationRanker { if (!valid_yield(emitted_tokens)) return; update_yield(request_yield_[request], emitted_tokens); const auto confidence = request_confidence_.find(request); - if (confidence != request_confidence_.end() && - confidence->second.available()) { + if (confidence != request_confidence_.end()) { update_yield( confidence_yield_[speculation_confidence_profile( - confidence->second)], emitted_tokens); + confidence->second.speculator, + confidence->second.expected_tokens)], emitted_tokens); } } @@ -282,12 +318,13 @@ class AdaptiveVerificationRanker { const auto confidence_it = request_confidence_.find(request); const bool has_request = request_it != request_yield_.end(); const bool has_confidence = - confidence_it != request_confidence_.end() && - confidence_it->second.available(); + confidence_it != request_confidence_.end(); auto profile_it = confidence_yield_.end(); if (has_confidence) { profile_it = confidence_yield_.find( - speculation_confidence_profile(confidence_it->second)); + speculation_confidence_profile( + confidence_it->second.speculator, + confidence_it->second.expected_tokens)); } const bool has_profile = profile_it != confidence_yield_.end() && profile_it->second.samples >= @@ -299,11 +336,11 @@ class AdaptiveVerificationRanker { AdaptiveVerificationYieldEstimate out; if (has_confidence) { out.confidence_expected_tokens = - confidence_it->second.expected_tokens(); + confidence_it->second.expected_tokens; out.speculator = confidence_it->second.speculator; } double trusted_confidence = has_confidence - ? confidence_it->second.expected_tokens() : 1.0; + ? confidence_it->second.expected_tokens : 1.0; if (has_profile) { const bool stable_backlog_confidence = trust_stable_confidence && @@ -601,13 +638,17 @@ class AdaptiveVerificationRanker { fair_group_positions) { (void)yield_bucket; if (positions.size() < 2 || - !std::all_of( + !std::any_of( positions.begin(), positions.end(), [&](std::size_t position) { return has_stable_useful_yield(known[position]); })) { continue; } + // One proven-useful member is enough to rotate peers in the + // same narrow value bucket. Requiring every member to be + // stable permanently starves raw-confidence peers that never + // enter the bounded verifier prefix in the first place. std::vector members; members.reserve(positions.size()); for (std::size_t position : positions) { @@ -794,6 +835,9 @@ class AdaptiveVerificationRanker { if (admitted_prefix > required_count) { out.predicted_gain = admitted_goodput / baseline; } + out.predicted_goodput = admitted_goodput; + } else { + out.predicted_goodput = baseline; } return out; } @@ -856,9 +900,532 @@ class AdaptiveVerificationRanker { // Capped running means represent the complete request/profile while still // adapting gradually if a model or speculator changes behavior. std::map request_yield_; - std::map - request_confidence_; + struct StoredConfidenceEstimate { + SpeculatorKind speculator = SpeculatorKind::DDTree; + double expected_tokens = 1.0; + }; + std::map request_confidence_; std::map confidence_yield_; }; +// Owns one adaptive ranker per concrete verifier work shape. Speculative route +// timings and confidence calibration remain isolated by VerifierWorkKey, while +// the exact all-AR baseline is shared because it is independent of the +// speculator. Extra confidence-scout time has a separate profile and never +// contaminates either route cost. +class AdaptiveVerificationProfileBank { +public: + AdaptiveVerificationProfileBank() = default; + explicit AdaptiveVerificationProfileBank( + AdaptiveVerificationConfig config) + : config_(config), cost_ewma_alpha_(std::clamp( + config.cost_ewma_alpha, 0.0, 1.0)) {} + + AdaptiveVerificationRanker & profile( + const VerifierWorkKey & work) { + auto [it, inserted] = profiles_.try_emplace(work, config_); + if (inserted) { + // One replay of the current EWMA is sufficient: rankers use the AR + // sample as a baseline value, not as confidence evidence. + for (const auto & [active_requests, cost] : ar_cost_us_) { + it->second.observe_autoregressive( + active_requests, cost.elapsed_us); + } + } + return it->second; + } + + const AdaptiveVerificationRanker * find_profile( + const VerifierWorkKey & work) const { + const auto it = profiles_.find(work); + return it == profiles_.end() ? nullptr : &it->second; + } + + std::size_t profile_count() const { return profiles_.size(); } + + void observe_autoregressive(int active_requests, double elapsed_us) { + if (active_requests <= 0 || !valid_cost(elapsed_us)) return; + update_cost(ar_cost_us_[active_requests], elapsed_us); + for (auto & [work, ranker] : profiles_) { + (void)work; + ranker.observe_autoregressive(active_requests, elapsed_us); + } + } + + void observe_route(const VerifierWorkKey & work, + int active_requests, int speculative_requests, + double elapsed_us, + bool discard_first_sample = false) { + if (speculative_requests == 0) { + observe_autoregressive(active_requests, elapsed_us); + return; + } + profile(work).observe_route( + active_requests, speculative_requests, elapsed_us, + discard_first_sample); + } + + void observe_request_estimate( + const VerifierWorkKey & work, std::uint64_t request, + const SpeculationConfidenceEstimate & estimate) { + profile(work).observe_request_estimate(request, estimate); + } + + void observe_request_estimate( + const VerifierWorkKey & work, std::uint64_t request, + SpeculatorKind speculator, double expected_tokens) { + profile(work).observe_request_estimate( + request, speculator, expected_tokens); + } + + void observe_request_yield(const VerifierWorkKey & work, + std::uint64_t request, + double emitted_tokens) { + profile(work).observe_request_yield(request, emitted_tokens); + } + + void forget_request(std::uint64_t request) { + for (auto & [work, ranker] : profiles_) { + (void)work; + ranker.forget_request(request); + } + } + + // Select one exact verifier work shape from adapter-provided per-request + // menus for one configured speculator. DDTree, DSpark, and future adapters + // reuse this contract independently; simultaneous cross-speculator racing + // requires a separate portfolio policy. A measured profitable steady + // route wins over any new profiling. Otherwise only the lowest + // exploration_priority action is returned, so adapters can profile + // short/cheap shapes before wider alternatives. + AdaptiveVerificationWorkDecision select_work( + int active_requests, + const std::vector & menus, + bool probe_uncalibrated_with_verifier = false, + bool enforce_ar_peer_guard = true, + std::size_t minimum_speculative_route_samples = 1, + bool trust_stable_confidence = false, + bool allow_safe_peer_guard_relaxation = false) { + AdaptiveVerificationWorkDecision out; + if (active_requests <= 0 || + menus.size() != static_cast(active_requests)) { + out.status = AdaptiveVerificationWorkStatus::InvalidMenu; + return out; + } + + struct WorkOption { + const SpeculationRequestView * request = nullptr; + const VerifierWorkPlan * plan = nullptr; + }; + struct WorkGroup { + std::vector options; + bool traits_initialized = false; + int max_parallel_requests = std::numeric_limits::max(); + int adaptive_request_limit = std::numeric_limits::max(); + std::uint32_t exploration_priority = 0; + bool preferred_for_forced_mode = false; + }; + std::map groups; + std::map seen_slots; + std::map seen_requests; + std::optional configured_speculator; + std::size_t required_requests = 0; + bool required_without_plan = false; + + for (const RequestVerifierWorkMenu & menu : menus) { + if (menu.request.slot < 0 || + !seen_slots.emplace(menu.request.slot, true).second || + !seen_requests.emplace( + menu.request.request_id, true).second) { + out.status = AdaptiveVerificationWorkStatus::InvalidMenu; + return out; + } + if (menu.request.required) ++required_requests; + if (menu.speculative.empty()) { + if (menu.request.required) required_without_plan = true; + continue; + } + + std::map menu_work; + for (const VerifierWorkPlan & plan : menu.speculative) { + if (!plan.valid() || + (configured_speculator && + *configured_speculator != plan.work.speculator) || + !menu_work.emplace(plan.work, true).second) { + out.status = AdaptiveVerificationWorkStatus::InvalidMenu; + return out; + } + configured_speculator = plan.work.speculator; + WorkGroup & group = groups[plan.work]; + if (!group.traits_initialized) { + group.traits_initialized = true; + group.max_parallel_requests = + plan.max_parallel_requests; + group.adaptive_request_limit = + plan.adaptive_request_limit; + group.exploration_priority = + plan.exploration_priority; + group.preferred_for_forced_mode = + plan.preferred_for_forced_mode; + } else if ( + group.max_parallel_requests != + plan.max_parallel_requests || + group.adaptive_request_limit != + plan.adaptive_request_limit || + group.exploration_priority != + plan.exploration_priority || + group.preferred_for_forced_mode != + plan.preferred_for_forced_mode) { + out.status = AdaptiveVerificationWorkStatus::InvalidMenu; + return out; + } + group.options.push_back({&menu.request, &plan}); + } + } + if (required_without_plan) { + out.status = + AdaptiveVerificationWorkStatus::RequiredUnavailable; + return out; + } + if (groups.empty()) return out; + + struct Evaluation { + VerifierWorkKey work; + AdaptiveVerificationDecision decision; + std::uint32_t exploration_priority = 0; + std::size_t required_count = 0; + bool preferred_for_forced_mode = false; + bool exact_profile = false; + }; + std::vector evaluations; + evaluations.reserve(groups.size()); + + for (const auto & [work, group] : groups) { + Evaluation evaluation; + evaluation.work = work; + evaluation.exploration_priority = + group.exploration_priority; + evaluation.preferred_for_forced_mode = + group.preferred_for_forced_mode; + std::vector candidates; + candidates.reserve(group.options.size()); + for (const WorkOption & option : group.options) { + const SpeculationRequestView & request = *option.request; + const VerifierWorkPlan & plan = *option.plan; + + AdaptiveVerificationRanker & ranker = profile(work); + const double confidence = + plan.bounded_confidence_expected_tokens(); + if (plan.has_confidence()) { + ranker.observe_request_estimate( + request.request_id, work.speculator, + confidence); + } + + AdaptiveVerificationCandidate candidate; + candidate.request = request.slot; + candidate.maximum_tokens = plan.maximum_emitted_tokens; + candidate.confidence_expected_tokens = confidence; + candidate.progress_tokens = request.progress_tokens; + candidate.required = request.required; + candidate.speculator = work.speculator; + + if (plan.has_confidence()) { + const auto estimate = ranker.estimate_request_yield( + request.request_id, trust_stable_confidence); + candidate.expected_tokens = std::clamp( + estimate ? estimate->expected_tokens : confidence, + 1.0, plan.maximum_emitted_tokens); + candidate.evidence_samples = estimate + ? estimate->evidence_samples : 0; + candidate.calibrated = true; + } else { + const auto target = + ranker.request_expected_tokens(request.request_id); + if (target) { + candidate.expected_tokens = std::clamp( + *target, 1.0, + plan.maximum_emitted_tokens); + candidate.evidence_samples = + ranker.request_yield_samples(request.request_id); + candidate.calibrated = true; + } + } + candidates.push_back(candidate); + + if (request.required) { + ++evaluation.required_count; + } + } + // One executor shape must be able to carry every Always request. + if (evaluation.required_count != required_requests || + evaluation.required_count > + static_cast( + group.max_parallel_requests) || + candidates.empty()) { + continue; + } + + AdaptiveVerificationRanker & ranker = profile(work); + const int profile_limit = std::max( + static_cast(evaluation.required_count), + std::min({ + active_requests, + group.adaptive_request_limit, + static_cast(candidates.size())})); + evaluation.exact_profile = ranker.has_exact_profile( + active_requests, profile_limit, + minimum_speculative_route_samples, + static_cast(evaluation.required_count)); + const int executable_lanes = std::max( + static_cast(evaluation.required_count), + std::min( + group.adaptive_request_limit, + static_cast(candidates.size()))); + const bool broad_verifier_coverage = + adaptive_verification_can_relax_peer_guard( + active_requests, executable_lanes); + const bool stable_bounded_cohort = + evaluation.exact_profile && + candidates.size() == + static_cast(active_requests) && + adaptive_verification_can_extend_stable_cohort( + active_requests, executable_lanes) && + ranker.forms_stable_confidence_cohort(candidates); + const bool relax_peer_guard = + allow_safe_peer_guard_relaxation && + (broad_verifier_coverage || stable_bounded_cohort); + evaluation.decision = ranker.select( + active_requests, candidates, + group.adaptive_request_limit, + probe_uncalibrated_with_verifier, + /*probe_missing_routes=*/true, + enforce_ar_peer_guard && !relax_peer_guard, + minimum_speculative_route_samples); + evaluations.push_back(std::move(evaluation)); + } + if (evaluations.empty()) { + if (required_requests > 0) { + out.status = + AdaptiveVerificationWorkStatus::RequiredUnavailable; + } + return out; + } + + auto choose = [&](const Evaluation & evaluation) { + out.work = evaluation.work; + out.decision = evaluation.decision; + out.status = evaluation.decision.requests.empty() + ? AdaptiveVerificationWorkStatus::Calibration + : AdaptiveVerificationWorkStatus::Verification; + }; + + // Compare measured work alternatives by absolute useful-token + // goodput. Relative gain is profile-local and cannot rank two shapes + // whose required baselines have different costs. + const Evaluation * measured = nullptr; + for (const Evaluation & evaluation : evaluations) { + if (evaluation.decision.exploring || + evaluation.decision.requests.empty() || + evaluation.decision.predicted_goodput <= 0.0) { + continue; + } + if (!measured || + evaluation.decision.predicted_goodput > + measured->decision.predicted_goodput || + (evaluation.decision.predicted_goodput == + measured->decision.predicted_goodput && + evaluation.work < measured->work)) { + measured = &evaluation; + } + } + if (measured) { + choose(*measured); + return out; + } + + // Forced requests cannot fall back to AR. Adapter preference and + // exploration priority are used only while every executable shape is + // cold/unmeasured. + if (required_requests > 0) { + const Evaluation * forced = nullptr; + for (const Evaluation & evaluation : evaluations) { + if (evaluation.decision.requests.size() < required_requests) { + continue; + } + if (!forced || + (evaluation.preferred_for_forced_mode && + !forced->preferred_for_forced_mode) || + (evaluation.preferred_for_forced_mode == + forced->preferred_for_forced_mode && + evaluation.exploration_priority < + forced->exploration_priority) || + (evaluation.preferred_for_forced_mode == + forced->preferred_for_forced_mode && + evaluation.exploration_priority == + forced->exploration_priority && + evaluation.work < forced->work)) { + forced = &evaluation; + } + } + if (forced) choose(*forced); + return out; + } + + // Missing verifier widths and confidence needed to begin an incomplete + // profile are primary bounded exploration. A pure calibration request + // on a complete losing profile is deferred until every other work + // option has had its exploration opportunity. + const Evaluation * exploration = nullptr; + for (const Evaluation & evaluation : evaluations) { + const bool calibration = + evaluation.decision.requests.empty() && + evaluation.decision.calibration_request >= 0; + if (!evaluation.decision.exploring && + (!calibration || evaluation.exact_profile)) { + continue; + } + if (!exploration || + evaluation.exploration_priority < + exploration->exploration_priority || + (evaluation.exploration_priority == + exploration->exploration_priority && + evaluation.work < exploration->work)) { + exploration = &evaluation; + } + } + if (exploration) { + choose(*exploration); + return out; + } + + const Evaluation * fallback_calibration = nullptr; + for (const Evaluation & evaluation : evaluations) { + if (!evaluation.decision.requests.empty() || + evaluation.decision.calibration_request < 0) { + continue; + } + if (!fallback_calibration || + evaluation.exploration_priority < + fallback_calibration->exploration_priority || + (evaluation.exploration_priority == + fallback_calibration->exploration_priority && + evaluation.work < fallback_calibration->work)) { + fallback_calibration = &evaluation; + } + } + if (fallback_calibration) choose(*fallback_calibration); + return out; + } + + bool has_autoregressive_cost(int active_requests) const { + return ar_cost_us_.find(active_requests) != ar_cost_us_.end(); + } + + double autoregressive_cost_us(int active_requests) const { + const auto it = ar_cost_us_.find(active_requests); + return it == ar_cost_us_.end() + ? std::numeric_limits::infinity() + : it->second.elapsed_us; + } + + std::size_t autoregressive_cost_samples(int active_requests) const { + const auto it = ar_cost_us_.find(active_requests); + return it == ar_cost_us_.end() ? 0 : it->second.samples; + } + + void observe_scout(const ConfidenceScoutWorkKey & work, + int request_count, double elapsed_us) { + if (request_count <= 0 || !valid_cost(elapsed_us)) return; + update_cost(scout_cost_us_[{work, request_count}], elapsed_us); + } + + bool has_scout_cost(const ConfidenceScoutWorkKey & work, + int request_count) const { + return scout_cost_us_.find({work, request_count}) != + scout_cost_us_.end(); + } + + double scout_cost_us(const ConfidenceScoutWorkKey & work, + int request_count) const { + const auto it = scout_cost_us_.find({work, request_count}); + return it == scout_cost_us_.end() + ? std::numeric_limits::infinity() + : it->second.elapsed_us; + } + + std::size_t scout_cost_samples(const ConfidenceScoutWorkKey & work, + int request_count) const { + const auto it = scout_cost_us_.find({work, request_count}); + return it == scout_cost_us_.end() ? 0 : it->second.samples; + } + + // Compatibility for adapters not yet assigning stable scout shape IDs. + // New integrations should use ConfidenceScoutWorkKey so depths and + // executor paths cannot contaminate one another. + void observe_scout(SpeculatorKind speculator, int request_count, + double elapsed_us) { + observe_scout({speculator, 0, 0}, request_count, elapsed_us); + } + + bool has_scout_cost(SpeculatorKind speculator, + int request_count) const { + return has_scout_cost({speculator, 0, 0}, request_count); + } + + double scout_cost_us(SpeculatorKind speculator, + int request_count) const { + return scout_cost_us({speculator, 0, 0}, request_count); + } + + std::size_t scout_cost_samples(SpeculatorKind speculator, + int request_count) const { + return scout_cost_samples({speculator, 0, 0}, request_count); + } + + void reset() { + profiles_.clear(); + ar_cost_us_.clear(); + scout_cost_us_.clear(); + } + +private: + struct CostEstimate { + double elapsed_us = 0.0; + std::size_t samples = 0; + }; + + struct ScoutCostKey { + ConfidenceScoutWorkKey work; + int request_count = 0; + + bool operator<(const ScoutCostKey & other) const { + if (work != other.work) return work < other.work; + return request_count < other.request_count; + } + }; + + static bool valid_cost(double elapsed_us) { + return std::isfinite(elapsed_us) && elapsed_us > 0.0; + } + + void update_cost(CostEstimate & estimate, double elapsed_us) { + if (estimate.samples == 0) { + estimate.elapsed_us = elapsed_us; + } else { + estimate.elapsed_us = cost_ewma_alpha_ * elapsed_us + + (1.0 - cost_ewma_alpha_) * estimate.elapsed_us; + } + if (estimate.samples < std::numeric_limits::max()) { + ++estimate.samples; + } + } + + AdaptiveVerificationConfig config_; + double cost_ewma_alpha_ = 0.35; + std::map profiles_; + std::map ar_cost_us_; + std::map scout_cost_us_; +}; + } // namespace dflash::common diff --git a/server/src/common/concurrency/speculation_planning.h b/server/src/common/concurrency/speculation_planning.h new file mode 100644 index 000000000..d631d01ce --- /dev/null +++ b/server/src/common/concurrency/speculation_planning.h @@ -0,0 +1,197 @@ +#pragma once + +// Model-neutral value types for adaptive speculative planning. +// +// A concrete speculator owns proposal generation and durable state. These +// types only describe request identity, optional confidence scouting, and the +// verifier work shapes that the concrete executor can offer for each request. +// Ordinary autoregressive decoding is deliberately implicit in every menu. + +#include "speculation_confidence.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct SpeculationRequestView { + // Stable across slot reuse and therefore suitable for confidence caches. + std::uint64_t request_id = 0; + // Engine-local execution handle. It must not be used as a cache key. + int slot = -1; + std::int32_t seed_token = -1; + int progress_tokens = 0; + // Always/forced speculation is represented per request. A work selector + // must not choose a shape that omits any required request. + bool required = false; +}; + +// Adapter-stable identity for one confidence-scout execution shape. The +// adapter owns shape_id (for example DDTree top-k projection or a DSpark +// confidence-head executor); candidate_tokens keeps different decode depths +// from contaminating the same timing profile. +struct ConfidenceScoutWorkKey { + SpeculatorKind speculator = SpeculatorKind::DDTree; + std::uint32_t shape_id = 0; + int candidate_tokens = 0; + + bool valid() const { return candidate_tokens > 0; } + + bool operator==(const ConfidenceScoutWorkKey & other) const { + return speculator == other.speculator && + shape_id == other.shape_id && + candidate_tokens == other.candidate_tokens; + } + + bool operator!=(const ConfidenceScoutWorkKey & other) const { + return !(*this == other); + } + + bool operator<(const ConfidenceScoutWorkKey & other) const { + if (speculator != other.speculator) { + return static_cast(speculator) < + static_cast(other.speculator); + } + if (shape_id != other.shape_id) return shape_id < other.shape_id; + return candidate_tokens < other.candidate_tokens; + } +}; + +enum class ConfidenceScoutStatus : std::uint8_t { + Ready, + // The speculator can expose confidence only after generating a proposal. + ProposalRequired, + // A transient setup/compute failure. The caller may preserve service with + // AR and retry according to its normal calibration policy. + RetryableFailure, + Ineligible, +}; + +struct ConfidenceScoutRequest { + SpeculationRequestView request; + ConfidenceScoutWorkKey work; + // Conditional candidate positions requested, excluding the mandatory + // root/AR token. A concrete adapter may return a shorter estimate. + int max_candidate_tokens = 0; + + bool valid() const { + return work.valid() && max_candidate_tokens > 0 && + max_candidate_tokens <= work.candidate_tokens; + } +}; + +struct ConfidenceScoutResult { + std::uint64_t request_id = 0; + ConfidenceScoutStatus status = ConfidenceScoutStatus::RetryableFailure; + SpeculationConfidenceEstimate confidence; + // Exact request-local time when the adapter executes scouts separately. + // A fused adapter may leave this zero and report only batch elapsed_us. + double elapsed_us = 0.0; + + bool ready() const { + return status == ConfidenceScoutStatus::Ready && + confidence.available(); + } +}; + +struct ConfidenceScoutBatchResult { + // One entry per request, in input order. request_id makes association + // explicit even when a concrete implementation partially fails. + std::vector requests; + // Extra scouting time is intentionally not part of a verifier route cost. + // It may instead inform calibration cadence and cold-start accounting. + double elapsed_us = 0.0; +}; + +// Stable structural identity for one executor/profile shape. shape_id is +// scoped by speculator; proposal_nodes and verifier_rows stay separate because +// a branching tree can verify many rows while exposing a much shorter path. +struct VerifierWorkKey { + SpeculatorKind speculator = SpeculatorKind::DDTree; + std::uint32_t shape_id = 0; + int proposal_nodes = 0; + int verifier_rows = 0; + + bool valid() const { + return proposal_nodes > 0 && verifier_rows > 0; + } + + bool operator==(const VerifierWorkKey & other) const { + return speculator == other.speculator && + shape_id == other.shape_id && + proposal_nodes == other.proposal_nodes && + verifier_rows == other.verifier_rows; + } + + bool operator!=(const VerifierWorkKey & other) const { + return !(*this == other); + } + + bool operator<(const VerifierWorkKey & other) const { + if (speculator != other.speculator) { + return static_cast(speculator) < + static_cast(other.speculator); + } + if (shape_id != other.shape_id) return shape_id < other.shape_id; + if (proposal_nodes != other.proposal_nodes) { + return proposal_nodes < other.proposal_nodes; + } + return verifier_rows < other.verifier_rows; + } +}; + +struct VerifierWorkPlan { + VerifierWorkKey work; + // Maximum useful output including the ordinary root token. This cannot be + // inferred from verifier_rows for branching proposals. + double maximum_emitted_tokens = 1.0; + // Raw request-local confidence for this exact work budget. NaN means the + // plan is executable but currently uncalibrated. + double confidence_expected_tokens = + std::numeric_limits::quiet_NaN(); + // Forced speculation uses the adapter's preferred executable plan when no + // measured adaptive decision is available. + bool preferred_for_forced_mode = false; + // Hard executor capacity. No route, including Always/required requests, + // may exceed this bound. + int max_parallel_requests = std::numeric_limits::max(); + // Efficiency/policy bound for adaptive peers. Required requests may exceed + // this value when the hard executor capacity still permits them. + int adaptive_request_limit = std::numeric_limits::max(); + // Lower values are explored first. Adapters should place cheap/short work + // before wider work so wider shapes are reached only after cheaper shapes + // have a complete losing profile. + std::uint32_t exploration_priority = 0; + + bool has_confidence() const { + return std::isfinite(confidence_expected_tokens) && + confidence_expected_tokens >= 1.0; + } + + double bounded_confidence_expected_tokens() const { + return has_confidence() + ? std::clamp( + confidence_expected_tokens, 1.0, + maximum_emitted_tokens) + : std::numeric_limits::quiet_NaN(); + } + + bool valid() const { + return work.valid() && std::isfinite(maximum_emitted_tokens) && + maximum_emitted_tokens >= 1.0 && max_parallel_requests > 0 && + adaptive_request_limit > 0 && + adaptive_request_limit <= max_parallel_requests; + } +}; + +struct RequestVerifierWorkMenu { + SpeculationRequestView request; + // AR is implicit. Empty means this request has no executable speculative + // plan at the current position/capacity. + std::vector speculative; +}; + +} // namespace dflash::common diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index af02236f5..613fc4dae 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -71,8 +71,48 @@ int adaptive_calibration_interval(int active_requests) { int adaptive_confidence_refresh_interval(int active_requests) { // Revisit AR-routed requests without turning confidence estimation into a - // per-token tax. Wider batches amortize a rejected probe for longer. - return 64 * std::max(1, (active_requests + 3) / 4); + // per-token tax. One request-local estimate covers an ordinary 96-token + // completion; longer continuations are revisited as their regime changes. + return 128 * std::max(1, (active_requests + 3) / 4); +} + +// Adapter-stable topology identifiers. Work depth remains a separate key so +// measurements for a four-token chain, an eight-token chain, and the full +// branching tree can never contaminate one another. +constexpr std::uint32_t kDDTreeChainWorkShape = 1; +constexpr std::uint32_t kDDTreeBranchingWorkShape = 2; +constexpr std::uint32_t kDDTreeTop1ScoutShape = 1; +constexpr std::uint32_t kDDTreeTop8ScoutShape = 2; + +int ddtree_top_k_for_budget(int budget, int draft_block_size) { + return budget > draft_block_size - 1 ? 8 : 1; +} + +VerifierWorkKey ddtree_work_key(int budget, int draft_block_size) { + return { + SpeculatorKind::DDTree, + ddtree_top_k_for_budget(budget, draft_block_size) == 1 + ? kDDTreeChainWorkShape : kDDTreeBranchingWorkShape, + budget, + budget + 1, + }; +} + +ConfidenceScoutWorkKey ddtree_scout_work_key( + int budget, int draft_block_size) { + const int top_k = ddtree_top_k_for_budget(budget, draft_block_size); + return { + SpeculatorKind::DDTree, + top_k == 1 ? kDDTreeTop1ScoutShape : kDDTreeTop8ScoutShape, + std::max(1, draft_block_size - 1), + }; +} + +int ddtree_scout_top_k(const ConfidenceScoutWorkKey & work) { + if (work.speculator != SpeculatorKind::DDTree) return 0; + if (work.shape_id == kDDTreeTop1ScoutShape) return 1; + if (work.shape_id == kDDTreeTop8ScoutShape) return 8; + return 0; } } // namespace @@ -97,7 +137,6 @@ Qwen35SeqEngine::Qwen35SeqEngine( tree_scratch_stride_(tree_scratch_stride) { const int n_slots = slots_.slot_count(); slot_draft_kv_.resize((size_t)n_slots); - compact_tree_cohort_.resize((size_t)n_slots, 0); // The concurrent DDTree stack is gated to a local same-device drafter. // Build metadata-only BF16 views over each slot's disjoint target feature @@ -204,11 +243,23 @@ bool Qwen35SeqEngine::ddtree_input_eligible(const StepInput & in) const { return slots_.slot(in.slot).generated_tokens() >= min_floor; } -std::optional -Qwen35SeqEngine::estimate_ddtree_confidence( - const StepInput & in) { +std::optional +Qwen35SeqEngine::scout_ddtree_confidence( + const ConfidenceScoutRequest & request) { const int q_len = b_.dw_.block_size; const int hidden = b_.w_.n_embd; + const int K = ddtree_scout_top_k(request.work); + const SpeculationRequestView & view = request.request; + if (K <= 0 || !request.valid() || view.slot < 0 || + view.slot >= slots_.slot_count() || !slots_.is_active(view.slot) || + slots_.slot(view.slot).request_id != view.request_id) { + return std::nullopt; + } + StepInput in; + in.slot = view.slot; + in.token = view.seed_token; + in.allow_speculation = true; + if (!ddtree_input_eligible(in)) return std::nullopt; if (q_len <= 1 || !build_lm_head_projection_step( b_.proj_sg_, b_.w_, b_.target_backend_, q_len)) { return std::nullopt; @@ -218,8 +269,7 @@ Qwen35SeqEngine::estimate_ddtree_confidence( DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); if (!draft || !mirror) return std::nullopt; - auto fail = [&]() -> - std::optional { + auto fail = [&]() -> std::optional { // begin_step advances only host cache bookkeeping, but a failed graph // can leave the appended rows incomplete. Rebuild from committed // target features on the next proposal. @@ -251,12 +301,12 @@ Qwen35SeqEngine::estimate_ddtree_confidence( return fail(); } - std::vector top_lp((size_t)q_len); - std::vector top_ids((size_t)q_len); + std::vector top_lp((size_t)q_len * K); + std::vector top_ids((size_t)q_len * K); bool topk_ready = false; #ifdef DFLASH27B_HAVE_DRAFT_TOPK topk_ready = geometric_extract_draft_topk_cuda( - b_.proj_sg_.logits->data, q_len, b_.w_.n_vocab, 1, + b_.proj_sg_.logits->data, q_len, b_.w_.n_vocab, K, top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); #endif if (!topk_ready) { @@ -266,7 +316,7 @@ Qwen35SeqEngine::estimate_ddtree_confidence( b_.proj_sg_.logits, logits.data(), 0, sizeof(float) * logits.size()); extract_draft_topk( - logits.data(), q_len, b_.w_.n_vocab, 1, + logits.data(), q_len, b_.w_.n_vocab, K, top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); } @@ -279,14 +329,81 @@ Qwen35SeqEngine::estimate_ddtree_confidence( for (int pos = 1; pos <= confidence_tokens; ++pos) { confidence[(size_t)pos - 1] = static_cast( std::clamp(std::exp(static_cast( - top_lp[(size_t)pos])), + top_lp[(size_t)pos * K])), 0.0, 1.0)); } - return make_speculation_confidence_estimate( + DDTreeConfidenceScout out; + out.request_id = view.request_id; + out.slot = in.slot; + out.seed_token = in.token; + out.committed_tokens = slots_.slot(in.slot).cur_pos; + out.top_k = K; + out.top_log_probs = std::move(top_lp); + out.top_token_ids = std::move(top_ids); + out.confidence = make_speculation_confidence_estimate( SpeculatorKind::DDTree, confidence.data(), static_cast(confidence.size()), SpeculationConfidenceCost::ExtraDraftPass, /*posthoc_calibrated=*/false); + return out; +} + +ConfidenceScoutBatchResult +Qwen35SeqEngine::scout_ddtree_confidence_batch( + const std::vector & requests, + std::vector & prepared) { + using Clock = std::chrono::steady_clock; + const auto started = Clock::now(); + ConfidenceScoutBatchResult batch; + batch.requests.reserve(requests.size()); + prepared.reserve(prepared.size() + requests.size()); + // This adapter batches association and accounting, but proposal execution + // remains serial because each slot owns a distinct persistent draft-KV + // ring. A future DSpark adapter may implement the same contract with one + // fused confidence-head graph. + for (const ConfidenceScoutRequest & request : requests) { + const auto request_started = Clock::now(); + ConfidenceScoutResult result; + result.request_id = request.request.request_id; + StepInput current; + current.slot = request.request.slot; + current.token = request.request.seed_token; + current.allow_speculation = true; + const bool associated = current.slot >= 0 && + current.slot < slots_.slot_count() && + slots_.is_active(current.slot) && + slots_.slot(current.slot).request_id == result.request_id; + if (!associated || !request.valid() || + ddtree_scout_top_k(request.work) <= 0 || + !ddtree_input_eligible(current)) { + result.status = ConfidenceScoutStatus::Ineligible; + result.elapsed_us = std::max( + 1.0, std::chrono::duration( + Clock::now() - request_started).count()); + batch.requests.push_back(std::move(result)); + continue; + } + std::optional scout = + scout_ddtree_confidence(request); + result.elapsed_us = std::max( + 1.0, std::chrono::duration( + Clock::now() - request_started).count()); + if (!scout || !scout->confidence.available()) { + result.status = ConfidenceScoutStatus::RetryableFailure; + batch.requests.push_back(std::move(result)); + continue; + } + result.status = ConfidenceScoutStatus::Ready; + result.confidence = scout->confidence.limited_to( + request.max_candidate_tokens); + scout->elapsed_us = result.elapsed_us; + prepared.push_back(std::move(*scout)); + batch.requests.push_back(std::move(result)); + } + batch.elapsed_us = std::max( + 1.0, std::chrono::duration( + Clock::now() - started).count()); + return batch; } void Qwen35SeqEngine::remember_ddtree_confidence( @@ -294,26 +411,16 @@ void Qwen35SeqEngine::remember_ddtree_confidence( const SpeculationConfidenceEstimate & estimate, int progress_tokens) { if (!estimate.available()) return; - const auto current = ddtree_confidence_.find(request_id); - if (current != ddtree_confidence_.end() && - current->second.available()) { - const SpeculationConfidenceProfile old_profile = - speculation_confidence_profile(current->second); - const SpeculationConfidenceProfile new_profile = - speculation_confidence_profile(estimate); - if (old_profile.speculator != new_profile.speculator || - std::abs(old_profile.expected_half_tokens - - new_profile.expected_half_tokens) >= 2) { - ddtree_target_yield_.erase(request_id); - } - } ddtree_confidence_[request_id] = estimate; ddtree_confidence_progress_[request_id] = progress_tokens; } std::optional Qwen35SeqEngine::step_ddtree( const StepPlan & speculative_plan, const StepPlan & ar_plan, - int tree_budget) { + int tree_budget, + const std::vector * prepared, + int * reused_scouts) { + if (reused_scouts) *reused_scouts = 0; StepResult result; const int active = (int)speculative_plan.decode.size(); const int total_active = active + (int)ar_plan.decode.size(); @@ -364,50 +471,72 @@ std::optional Qwen35SeqEngine::step_ddtree( // Proposal is sequential by slot: immutable draft weights are shared, // while each slot owns an independent persistent context-KV ring. for (const StepInput & in : speculative_plan.decode) { - DraftKvState * draft = ensure_slot_draft_kv(in.slot); - DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); - if (!draft || !mirror) return proposal_fallback(); - noise[0] = in.token; - std::fill(noise.begin() + 1, noise.end(), b_.w_.mask_token_id); - if (!b_.w_.embedder.embed( - noise.data(), q_len, noise_embed.data()) || - !draft_kv_begin_step(*draft, b_.dw_, b_.draft_backend_, - *mirror, slots_.slot(in.slot).cur_pos)) { - return proposal_fallback(); + const DDTreeConfidenceScout * scouted = nullptr; + if (prepared) { + const auto found = std::find_if( + prepared->begin(), prepared->end(), + [&](const DDTreeConfidenceScout & candidate) { + return candidate.request_id == + slots_.slot(in.slot).request_id && + candidate.slot == in.slot && + candidate.seed_token == in.token && + candidate.committed_tokens == + slots_.slot(in.slot).cur_pos && + candidate.top_k == K && + candidate.top_log_probs.size() == + static_cast(q_len * K) && + candidate.top_token_ids.size() == + static_cast(q_len * K); + }); + if (found != prepared->end()) scouted = &*found; } - ggml_backend_tensor_set( - draft->inp_embed, noise_embed.data(), 0, - sizeof(float) * noise_embed.size()); - if (ggml_backend_graph_compute(b_.draft_backend_, draft->gf) != - GGML_STATUS_SUCCESS) { - return proposal_fallback(); - } - // The draft and target backends own separate HIP streams even when - // both are placed on hip:0. Projection consumes the draft hidden - // state on the target stream, so establish the producer/consumer - // ordering explicitly before the cross-backend tensor copy. Without - // this barrier the first tree proposal can race stale hidden rows and - // collapse acceptance to the one-token fallback. - ggml_backend_synchronize(b_.draft_backend_); - ggml_backend_tensor_copy( - draft->hidden_states, b_.proj_sg_.hidden_input); - if (ggml_backend_graph_compute( - b_.target_backend_, b_.proj_sg_.gf) != GGML_STATUS_SUCCESS) { - return proposal_fallback(); - } - bool topk_ready = false; + if (scouted) { + top_lp = scouted->top_log_probs; + top_ids = scouted->top_token_ids; + if (reused_scouts) ++*reused_scouts; + } else { + DraftKvState * draft = ensure_slot_draft_kv(in.slot); + DraftFeatureMirror * mirror = slot_feature_mirror(in.slot); + if (!draft || !mirror) return proposal_fallback(); + noise[0] = in.token; + std::fill(noise.begin() + 1, noise.end(), b_.w_.mask_token_id); + if (!b_.w_.embedder.embed( + noise.data(), q_len, noise_embed.data()) || + !draft_kv_begin_step(*draft, b_.dw_, b_.draft_backend_, + *mirror, slots_.slot(in.slot).cur_pos)) { + return proposal_fallback(); + } + ggml_backend_tensor_set( + draft->inp_embed, noise_embed.data(), 0, + sizeof(float) * noise_embed.size()); + if (ggml_backend_graph_compute(b_.draft_backend_, draft->gf) != + GGML_STATUS_SUCCESS) { + return proposal_fallback(); + } + // Projection consumes the draft hidden state on another backend + // stream. Establish producer/consumer ordering before the copy. + ggml_backend_synchronize(b_.draft_backend_); + ggml_backend_tensor_copy( + draft->hidden_states, b_.proj_sg_.hidden_input); + if (ggml_backend_graph_compute( + b_.target_backend_, b_.proj_sg_.gf) != + GGML_STATUS_SUCCESS) { + return proposal_fallback(); + } + bool topk_ready = false; #ifdef DFLASH27B_HAVE_DRAFT_TOPK - topk_ready = geometric_extract_draft_topk_cuda( - b_.proj_sg_.logits->data, q_len, b_.w_.n_vocab, K, - top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); -#endif - if (!topk_ready) { - ggml_backend_tensor_get( - b_.proj_sg_.logits, logits.data(), 0, - sizeof(float) * logits.size()); - extract_draft_topk( - logits.data(), q_len, b_.w_.n_vocab, K, + topk_ready = geometric_extract_draft_topk_cuda( + b_.proj_sg_.logits->data, q_len, b_.w_.n_vocab, K, top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); +#endif + if (!topk_ready) { + ggml_backend_tensor_get( + b_.proj_sg_.logits, logits.data(), 0, + sizeof(float) * logits.size()); + extract_draft_topk( + logits.data(), q_len, b_.w_.n_vocab, K, + top_lp.data(), top_ids.data(), b_.cfg_.ddtree_temp); + } } Proposal p; @@ -1135,32 +1264,25 @@ SeqEngine::AdmitResult Qwen35SeqEngine::admit( uint64_t request_id, const std::vector & prompt, const SamplerCfg & sampler) { + bool was_idle = true; + for (int slot = 0; slot < slots_.slot_count(); ++slot) { + if (slots_.is_active(slot)) { + was_idle = false; + break; + } + } AdmitResult result = slots_.admit(request_id, prompt, sampler); if (result.status == AdmitResult::Status::admitted) { - std::uint8_t inherited_compact_budget = 0; - bool has_active_peer = false; - bool compact_peers_agree = true; - for (int slot = 0; slot < slots_.slot_count(); ++slot) { - if (slot == result.slot || !slots_.is_active(slot)) continue; - const std::uint8_t peer_budget = - compact_tree_cohort_[(size_t)slot]; - if (!has_active_peer) { - inherited_compact_budget = peer_budget; - has_active_peer = true; - } else if (peer_budget != inherited_compact_budget) { - compact_peers_agree = false; - } - if (peer_budget == 0) compact_peers_agree = false; - } - adaptive_verification_.forget_request(request_id); - compact_short_adaptive_verification_.forget_request(request_id); - compact_adaptive_verification_.forget_request(request_id); + if (was_idle) { + previous_decode_requests_.clear(); + stable_decode_cohort_steps_ = 0; + draining_decode_cohort_ = false; + adaptive_calibration_cooldowns_.clear(); + adaptive_scout_failure_cooldowns_.clear(); + } + adaptive_verification_profiles_.forget_request(request_id); ddtree_confidence_.erase(request_id); ddtree_confidence_progress_.erase(request_id); - ddtree_target_yield_.erase(request_id); - compact_tree_cohort_[(size_t)result.slot] = - has_active_peer && compact_peers_agree - ? inherited_compact_budget : 0; reset_recurrent_slot(b_.cache_, result.slot); if (slots_.residency_active()) { slots_.slot(result.slot).kvflash_last_reselect_generated = @@ -1401,6 +1523,39 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { } if (inputs.empty() && plan.prefills.empty()) return result; + const int active_requests = static_cast(inputs.size()); + std::vector current_decode_requests; + current_decode_requests.reserve(inputs.size()); + for (const StepInput & in : inputs) { + current_decode_requests.push_back( + slots_.slot(in.slot).request_id); + } + std::sort( + current_decode_requests.begin(), current_decode_requests.end()); + if (current_decode_requests == previous_decode_requests_) { + if (stable_decode_cohort_steps_ < + std::numeric_limits::max()) { + ++stable_decode_cohort_steps_; + } + } else { + draining_decode_cohort_ = + !previous_decode_requests_.empty() && + current_decode_requests.size() < + previous_decode_requests_.size() && + std::includes( + previous_decode_requests_.begin(), + previous_decode_requests_.end(), + current_decode_requests.begin(), + current_decode_requests.end()); + if (!draining_decode_cohort_) { + // Hardware profiles survive cohort turnover, but a new/refilled + // request set must get one bounded chance to expose confidence. + adaptive_calibration_cooldowns_.clear(); + } + previous_decode_requests_ = std::move(current_decode_requests); + stable_decode_cohort_steps_ = 0; + } + if (!ddtree_available(plan)) return step_regular(plan); const char * adaptive = std::getenv("DFLASH_DDTREE_ADAPTIVE"); @@ -1423,170 +1578,166 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { // Confidence is a decode-regime signal, not a permanent prompt // label. Periodically forget it so an AR-routed request can be // re-probed after its continuation changes character. - adaptive_verification_.forget_request(seq.request_id); - compact_short_adaptive_verification_.forget_request( - seq.request_id); - compact_adaptive_verification_.forget_request(seq.request_id); + adaptive_verification_profiles_.forget_request(seq.request_id); ddtree_confidence_.erase(seq.request_id); ddtree_confidence_progress_.erase(observed); - ddtree_target_yield_.erase(seq.request_id); } } - const int inherited_compact_budget = inputs.empty() - ? 0 : compact_tree_cohort_[(size_t)inputs.front().slot]; - const bool inherited_compact_tree = - inherited_compact_budget > 0 && - std::all_of( - inputs.begin(), inputs.end(), [&](const StepInput & in) { - return compact_tree_cohort_[(size_t)in.slot] == - inherited_compact_budget; - }); - int confidence_candidates = 0; - int useful_confidence_candidates = 0; - int eligible_candidates = 0; - for (const StepInput & in : inputs) { - if (!ddtree_input_eligible(in)) continue; - ++eligible_candidates; - const Qwen35Slot & seq = slots_.slot(in.slot); - const auto estimate = ddtree_confidence_.find(seq.request_id); - if (estimate == ddtree_confidence_.end() || - !estimate->second.available()) { - continue; - } - ++confidence_candidates; - double expected_tokens = estimate->second.expected_tokens(); - const auto target = ddtree_target_yield_.find(seq.request_id); - if (target != ddtree_target_yield_.end() && target->second.samples > 0) { - // Raw confidence orders cold requests. Target-accepted output - // progressively corrects that estimate, independent of which - // proposal shape produced the observation. - const double target_weight = std::min( - 1.0, static_cast(target->second.samples) / 4.0); - expected_tokens += target_weight * - (target->second.expected_tokens - expected_tokens); - } - if (expected_tokens >= 1.5) { - ++useful_confidence_candidates; - } - } - const bool confidence_cohort_ready = eligible_candidates > 0 && - confidence_candidates == eligible_candidates; - const bool mixed_confidence_cohort = confidence_cohort_ready && - useful_confidence_candidates > 0 && - useful_confidence_candidates < static_cast(inputs.size()); - // Full DDTree depth remains the neutral low-occupancy shape. Once the - // concrete speculator has estimated the whole eligible cohort, a mixed C=4 - // backlog can use a compact proposal for its small useful subset. No prompt - // category or model-specific task table participates in this decision. constexpr int kFullTreeMaxConcurrency = 4; constexpr int kCompactShortTreeBudget = 4; constexpr int kCompactTreeBudget = 8; - const bool starts_short_backlog_shape = - plan.has_refill_backlog && inputs.size() == kFullTreeMaxConcurrency && - confidence_cohort_ready && useful_confidence_candidates > 0 && - useful_confidence_candidates <= 2 && mixed_confidence_cohort; - const bool retains_compact_mixed_cohort = - inherited_compact_tree && - (!confidence_cohort_ready || mixed_confidence_cohort); - const bool compact_tree_route = - inputs.size() > kFullTreeMaxConcurrency || - retains_compact_mixed_cohort || starts_short_backlog_shape; - const int selected_compact_budget = - inherited_compact_tree && - (!plan.has_refill_backlog || !confidence_cohort_ready) - ? inherited_compact_budget - : (confidence_cohort_ready && - useful_confidence_candidates > 0 && - useful_confidence_candidates <= 2 - ? kCompactShortTreeBudget - : kCompactTreeBudget); - const int adaptive_tree_budget = - compact_tree_route - ? std::min(b_.cfg_.ddtree_budget, selected_compact_budget) - : b_.cfg_.ddtree_budget; - const int tree_budget = std::clamp( - oracle.tree_budget > 0 ? oracle.tree_budget - : adaptive_tree_budget, - 1, b_.cfg_.ddtree_budget); - const bool compact_short_shape = - compact_tree_route && tree_budget <= kCompactShortTreeBudget; - AdaptiveVerificationRanker & route_ranker = compact_tree_route - ? (compact_short_shape - ? compact_short_adaptive_verification_ - : compact_adaptive_verification_) - : adaptive_verification_; - - // Keep the scheduler independent of the concrete speculation algorithm. - // DDTree contributes either a cheap draft-confidence estimate or a - // profile-local mean of target-accepted useful tokens. A future DSpark - // adapter can contribute calibrated prefix survival from its confidence - // head while reusing the same selection policy. - std::vector candidates; - candidates.reserve(inputs.size()); - auto collect_candidates = [&](bool trust_stable_confidence = false) { - candidates.clear(); + constexpr int kDrainingTailExplorationDelay = 8; + const bool suppress_draining_tail_exploration = + adaptive_enabled && !oracle.forces_selection() && + draining_decode_cohort_ && + stable_decode_cohort_steps_ < kDrainingTailExplorationDelay; + const int draft_block_size = b_.dw_.block_size; + const int direct_speculation_limit = std::max( + 1, detail::target_paged_tree_direct_request_limit( + b_.cache_.tree_capture_lanes)); + + struct DDTreeWorkSpec { + VerifierWorkKey work; + int adaptive_request_limit = 1; + std::uint32_t priority = 0; + bool preferred_for_forced_mode = false; + }; + std::vector offered_work; + auto offer_work = [&](int requested_budget, int adaptive_limit, + std::uint32_t priority, bool preferred) { + const int budget = std::clamp( + requested_budget, 1, b_.cfg_.ddtree_budget); + const VerifierWorkKey work = + ddtree_work_key(budget, draft_block_size); + const auto existing = std::find_if( + offered_work.begin(), offered_work.end(), + [&](const DDTreeWorkSpec & candidate) { + return candidate.work == work; + }); + if (existing != offered_work.end()) { + existing->preferred_for_forced_mode |= preferred; + existing->priority = std::min(existing->priority, priority); + return; + } + offered_work.push_back({ + work, + std::max(1, std::min(active_requests, adaptive_limit)), + priority, + preferred, + }); + }; + + if (oracle.tree_budget > 0) { + offer_work(oracle.tree_budget, active_requests, 0, true); + } else if (!adaptive_enabled || active_requests <= 3) { + // Keep the proven C<=3 executor shape intact; adaptation changes only + // which requests use it. Disabling adaptation restores fixed full + // DDTree at every occupancy. + offer_work(b_.cfg_.ddtree_budget, active_requests, 0, true); + } else if (active_requests == kFullTreeMaxConcurrency) { + // Profile the established full tree first. Only a measured losing + // full shape can advance C=4 to the cheaper compact verifier. + offer_work(b_.cfg_.ddtree_budget, active_requests, 0, true); + offer_work( + kCompactShortTreeBudget, + std::min(2, direct_speculation_limit), 1, false); + } else { + // Above C=4, start with bounded verifier work. The shared selector + // advances to eight nodes only after the exact short profile loses; + // raw confidence orders requests but never chooses the work shape. + offer_work( + kCompactShortTreeBudget, + std::min(2, direct_speculation_limit), 0, false); + offer_work( + kCompactTreeBudget, direct_speculation_limit, 1, true); + } + + auto make_work_menus = [&]() { + std::vector menus; + menus.reserve(inputs.size()); for (const StepInput & in : inputs) { - if (!ddtree_input_eligible(in)) continue; const Qwen35Slot & seq = slots_.slot(in.slot); - const auto confidence = ddtree_confidence_.find(seq.request_id); - if (confidence != ddtree_confidence_.end() && - confidence->second.available()) { - route_ranker.observe_request_estimate( - seq.request_id, - confidence->second.limited_to(tree_budget)); + RequestVerifierWorkMenu menu; + menu.request.request_id = seq.request_id; + menu.request.slot = in.slot; + menu.request.seed_token = in.token; + menu.request.progress_tokens = seq.generated_tokens(); + const bool eligible = ddtree_input_eligible(in); + menu.request.required = eligible && + in.speculation_policy == SpeculationPolicy::Always; + if (eligible) { + const auto confidence = + ddtree_confidence_.find(seq.request_id); + for (const DDTreeWorkSpec & offered : offered_work) { + VerifierWorkPlan work_plan; + work_plan.work = offered.work; + const bool chain = offered.work.shape_id == + kDDTreeChainWorkShape; + work_plan.maximum_emitted_tokens = chain + ? static_cast( + offered.work.proposal_nodes + 1) + : static_cast(draft_block_size); + if (confidence != ddtree_confidence_.end() && + confidence->second.available()) { + work_plan.confidence_expected_tokens = + confidence->second.limited_to(std::min( + offered.work.proposal_nodes, + draft_block_size - 1)).expected_tokens(); + } + work_plan.preferred_for_forced_mode = + offered.preferred_for_forced_mode; + work_plan.max_parallel_requests = active_requests; + work_plan.adaptive_request_limit = + offered.adaptive_request_limit; + work_plan.exploration_priority = offered.priority; + menu.speculative.push_back(std::move(work_plan)); + } } - const std::optional estimate = - route_ranker.estimate_request_yield( - seq.request_id, trust_stable_confidence); - candidates.push_back({ - in.slot, - estimate ? estimate->expected_tokens : 1.0, - static_cast(compact_tree_route - ? tree_budget + 1 : b_.dw_.block_size), - estimate.has_value(), - estimate ? estimate->confidence_expected_tokens - : std::numeric_limits::quiet_NaN(), - estimate ? estimate->evidence_samples : 0, - seq.generated_tokens(), - in.speculation_policy == SpeculationPolicy::Always, - SpeculatorKind::DDTree, - }); + menus.push_back(std::move(menu)); } + return menus; }; - collect_candidates(); + std::vector work_menus = make_work_menus(); + std::vector prepared_scouts; AdaptiveVerificationDecision decision; - const int direct_speculation_limit = - detail::target_paged_tree_direct_request_limit( - b_.cache_.tree_capture_lanes); - const int adaptive_speculation_limit = compact_tree_route - ? std::min( - direct_speculation_limit, - compact_short_shape ? 2 : direct_speculation_limit) - : static_cast(inputs.size()); - const int active_requests = static_cast(inputs.size()); - const bool broad_verifier_coverage = - adaptive_verification_can_relax_peer_guard( - active_requests, adaptive_speculation_limit); - const bool bounded_stable_cohort_extension = - adaptive_verification_can_extend_stable_cohort( - active_requests, adaptive_speculation_limit); - const bool use_stable_cohort_extension = - plan.has_refill_backlog && compact_tree_route && - !broad_verifier_coverage && bounded_stable_cohort_extension; - if (use_stable_cohort_extension) { - collect_candidates(/*trust_stable_confidence=*/true); - } - bool enforce_ar_peer_guard = true; - std::size_t minimum_route_samples = 1; + std::optional route_work; + AdaptiveVerificationRanker * route_ranker = nullptr; + int tree_budget = offered_work.front().work.proposal_nodes; + // Cold graph/allocation work affects short-request goodput too. Keep the + // established one-sample policy rather than doubling every C<=4 width; + // selected steady routes continue updating their EWMA after the probe. + const std::size_t minimum_route_samples = 1; + // Closed cohorts always protect AR peers. A continuous refill may relax + // that guard only inside the generic selector when this exact work shape + // has broad verifier coverage or a stable target-verified cohort. + const bool allow_safe_peer_guard_relaxation = + plan.has_refill_backlog; + AdaptiveVerificationWorkDecision work_selection; + + auto bind_selected_work = [&]() { + route_work = work_selection.work; + route_ranker = nullptr; + if (!route_work) return; + tree_budget = route_work->proposal_nodes; + route_ranker = &adaptive_verification_profiles_.profile(*route_work); + }; + if (oracle.forces_selection()) { + route_work = offered_work.front().work; + tree_budget = route_work->proposal_nodes; const int force_limit = std::min( - adaptive_speculation_limit, - oracle.speculative_requests); + oracle.speculative_requests, + static_cast(std::count_if( + inputs.begin(), inputs.end(), + [&](const StepInput & in) { + return ddtree_input_eligible(in); + }))); decision.requests.reserve((size_t)force_limit); - for (const AdaptiveVerificationCandidate & candidate : candidates) { + for (const StepInput & in : inputs) { + if (!ddtree_input_eligible(in)) continue; if ((int)decision.requests.size() >= force_limit) break; - decision.requests.push_back(candidate.request); + decision.requests.push_back(in.slot); } static bool logged = false; if (!logged) { @@ -1594,191 +1745,215 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { "[parallel-ddtree] oracle force_speculative=%d " "tree_budget=%d direct_limit=%d\n", oracle.speculative_requests, tree_budget, - adaptive_speculation_limit); + direct_speculation_limit); logged = true; } - } else if (adaptive_enabled) { - const int mandatory_speculative_requests = - static_cast(std::count_if( - candidates.begin(), candidates.end(), - [](const AdaptiveVerificationCandidate & candidate) { - return candidate.required; - })); - const int exact_profile_width = std::max( - mandatory_speculative_requests, - std::min(adaptive_speculation_limit, - static_cast(candidates.size()))); - // Replay a cold graph shape only in the narrow capacity extension that - // can plausibly become steady. This avoids doubling rejected probes at - // occupancies such as Strix C=16 and leaves established low-C profiles - // unchanged. - minimum_route_samples = - use_stable_cohort_extension && - confidence_cohort_ready && - useful_confidence_candidates == active_requests - ? 2 : 1; - int & calibration_cooldown = compact_tree_route - ? (compact_short_shape - ? compact_short_adaptive_calibration_cooldown_ - : compact_adaptive_calibration_cooldown_) - : adaptive_calibration_cooldown_; - const bool exact_profile_ready = route_ranker.has_exact_profile( - static_cast(inputs.size()), exact_profile_width, - minimum_route_samples, mandatory_speculative_requests); - // A closed cohort keeps request-local AR-peer protection. Continuous - // backlog may instead optimize aggregate output once either verifier - // coverage is broad enough or target-verified outcomes have made a - // useful confidence bucket stable. Exact route costs still decide - // whether speculation wins, so the same evidence can unlock C=8 while - // an unprofitable C=16 profile remains AR. A future speculator supplies - // its own work-bucket ranker and inherits the same rule. - const bool has_stable_backlog_cohort = - plan.has_refill_backlog && exact_profile_ready && - route_ranker.forms_stable_confidence_cohort(candidates); - const bool has_stable_backlog_candidate = - plan.has_refill_backlog && exact_profile_ready && - std::any_of( - candidates.begin(), candidates.end(), - [&](const AdaptiveVerificationCandidate & candidate) { - return route_ranker.has_stable_confidence_yield( - candidate.speculator, - candidate.confidence_expected_tokens); - }); - const bool relax_ar_peer_guard = - plan.has_refill_backlog && compact_tree_route && - (broad_verifier_coverage || - (bounded_stable_cohort_extension && - has_stable_backlog_cohort)); - enforce_ar_peer_guard = !relax_ar_peer_guard; - const bool exact_profile_started = - route_ranker.has_speculative_profile_sample( - static_cast(inputs.size()), exact_profile_width); - // Finish a profile already started at this occupancy, but do not make - // a shrinking compact cohort pay for brand-new C-tail exploration. - // A fresh batch at that same C still profiles normally, and a prior - // exact profile remains reusable. - const bool probe_missing_routes = - !compact_tree_route || !inherited_compact_tree || - exact_profile_started; - // Full DDTree uses its drafter-side confidence adapter. Compact DDTree - // avoids a duplicate draft pass and calibrates one unknown request in - // useful verification instead. DSpark can take the former path with - // its confidence head; target outcomes calibrate both adapters. - const bool drafter_side_calibration = - !compact_tree_route && probe_missing_routes; - const bool suppress_uncovered_mixed_calibration = - exact_profile_ready && mixed_confidence_cohort && - !broad_verifier_coverage && has_stable_backlog_candidate; - if (!exact_profile_ready || inputs.size() <= 3) { - calibration_cooldown = 0; - } else if (calibration_cooldown > 0) { - --calibration_cooldown; - } - // Once a closed mixed cohort has measured every exact width, probing - // each remaining low-priority request cannot remove its AR critical - // path. A refill backlog changes that objective, so keep calibrating - // newcomers there; homogeneous/unknown cohorts can still prove the - // existing all-speculative exception. - const bool verifier_side_calibration = - compact_tree_route && probe_missing_routes && - !suppress_uncovered_mixed_calibration && - (!exact_profile_ready || - (calibration_cooldown == 0 && - (plan.has_refill_backlog || !mixed_confidence_cohort))); - std::vector verifier_candidates; - const std::vector * decision_candidates = - &candidates; - const bool potentially_homogeneous = - route_ranker.forms_homogeneous_cohort( - candidates, /*require_stable_local_evidence=*/false); - if (verifier_side_calibration) { - verifier_candidates = candidates; - for (AdaptiveVerificationCandidate & candidate : - verifier_candidates) { - if (candidate.calibrated && - potentially_homogeneous && - !route_ranker.has_stable_evidence( - candidate.evidence_samples) && - route_ranker.has_useful_yield( - candidate.expected_tokens)) { - // A slow route may earn the homogeneous exception only - // when every active request still looks useful. Stable - // low-yield peers make further local proof pure overhead. - candidate.calibrated = false; - } + } else if (!adaptive_enabled) { + route_work = offered_work.front().work; + tree_budget = route_work->proposal_nodes; + for (const StepInput & in : inputs) { + if (ddtree_input_eligible(in)) decision.requests.push_back(in.slot); + } + } else { + work_selection = adaptive_verification_profiles_.select_work( + active_requests, work_menus, + /*probe_uncalibrated_with_verifier=*/false, + /*enforce_ar_peer_guard=*/true, minimum_route_samples, + /*trust_stable_confidence=*/plan.has_refill_backlog, + allow_safe_peer_guard_relaxation); + if (work_selection.status == + AdaptiveVerificationWorkStatus::InvalidMenu) { + return fail_step("adaptive speculation produced an invalid work menu"); + } + if (work_selection.status == + AdaptiveVerificationWorkStatus::RequiredUnavailable) { + return fail_step( + "required speculation has no common executable verifier work"); + } + const bool has_required_request = std::any_of( + work_menus.begin(), work_menus.end(), + [](const RequestVerifierWorkMenu & menu) { + return menu.request.required; + }); + if (suppress_draining_tail_exploration && + !has_required_request && + (work_selection.status == + AdaptiveVerificationWorkStatus::Calibration || + work_selection.decision.exploring)) { + work_selection = AdaptiveVerificationWorkDecision{}; + } + decision = work_selection.decision; + bind_selected_work(); + + if (route_work) { + int & cooldown = adaptive_calibration_cooldowns_[ + {*route_work, active_requests}]; + if (cooldown > 0) --cooldown; + const ConfidenceScoutWorkKey scout_work = + ddtree_scout_work_key(tree_budget, draft_block_size); + int & scout_failure_cooldown = + adaptive_scout_failure_cooldowns_[ + {scout_work, active_requests}]; + if (scout_failure_cooldown > 0) { + --scout_failure_cooldown; } - decision_candidates = &verifier_candidates; - } - decision = route_ranker.select( - static_cast(inputs.size()), *decision_candidates, - adaptive_speculation_limit, verifier_side_calibration, - probe_missing_routes, - enforce_ar_peer_guard, minimum_route_samples); - if (compact_tree_route && exact_profile_ready && - verifier_side_calibration && decision.exploring) { - calibration_cooldown = - adaptive_calibration_interval( - static_cast(inputs.size())); - } - if (drafter_side_calibration && - decision.calibration_request >= 0 && - calibration_cooldown == 0) { - const auto input = std::find_if( - inputs.begin(), inputs.end(), - [&](const StepInput & in) { - return in.slot == decision.calibration_request; - }); - if (input != inputs.end()) { - const std::optional - confidence = - estimate_ddtree_confidence(*input); - if (confidence) { - const std::uint64_t request_id = - slots_.slot(input->slot).request_id; - remember_ddtree_confidence( - request_id, *confidence, - slots_.slot(input->slot).generated_tokens()); - route_ranker.observe_request_estimate( - request_id, confidence->limited_to(tree_budget)); + const bool may_scout = decision.calibration_request >= 0 && + scout_failure_cooldown == 0 && + (active_requests <= kFullTreeMaxConcurrency || + cooldown == 0); + if (may_scout) { + int scout_limit = + ddtree_scout_top_k(scout_work) == 1 ? 2 : 1; + if (scout_limit > 1 && + adaptive_verification_profiles_.has_autoregressive_cost( + active_requests) && + adaptive_verification_profiles_.has_scout_cost( + scout_work, 2) && + adaptive_verification_profiles_.scout_cost_us( + scout_work, 2) > + 0.5 * adaptive_verification_profiles_. + autoregressive_cost_us(active_requests)) { + scout_limit = 1; + } + + const int required_count = static_cast(std::count_if( + work_menus.begin(), work_menus.end(), + [](const RequestVerifierWorkMenu & menu) { + return menu.request.required; + })); + int adaptive_limit = 1; + for (const DDTreeWorkSpec & offered : offered_work) { + if (offered.work == *route_work) { + adaptive_limit = offered.adaptive_request_limit; + break; + } + } + const int exact_width = std::max( + required_count, + std::min(adaptive_limit, static_cast(std::count_if( + work_menus.begin(), work_menus.end(), + [](const RequestVerifierWorkMenu & menu) { + return !menu.speculative.empty(); + })))); + const bool exact_profile_before = + route_ranker && route_ranker->has_exact_profile( + active_requests, exact_width, + minimum_route_samples, required_count); + + std::vector scout_requests; + scout_requests.reserve((size_t)scout_limit); + auto add_scout = [&](int slot) { + if ((int)scout_requests.size() >= scout_limit || + std::any_of( + scout_requests.begin(), scout_requests.end(), + [&](const ConfidenceScoutRequest & request) { + return request.request.slot == slot; + })) { + return; + } + const auto menu = std::find_if( + work_menus.begin(), work_menus.end(), + [&](const RequestVerifierWorkMenu & candidate) { + return candidate.request.slot == slot && + !candidate.speculative.empty(); + }); + if (menu == work_menus.end()) return; + ConfidenceScoutRequest request; + request.request = menu->request; + request.work = scout_work; + request.max_candidate_tokens = std::min( + tree_budget, scout_work.candidate_tokens); + if (request.valid()) { + scout_requests.push_back(std::move(request)); + } + }; + add_scout(decision.calibration_request); + for (const RequestVerifierWorkMenu & menu : work_menus) { + if ((int)scout_requests.size() >= scout_limit) break; + const auto known = + ddtree_confidence_.find(menu.request.request_id); + if (!menu.speculative.empty() && + (known == ddtree_confidence_.end() || + !known->second.available())) { + add_scout(menu.request.slot); + } + } + + if (!scout_requests.empty()) { + const std::size_t prepared_before = + prepared_scouts.size(); + ConfidenceScoutBatchResult scout = + scout_ddtree_confidence_batch( + scout_requests, prepared_scouts); + adaptive_verification_profiles_.observe_scout( + scout_work, + static_cast(scout_requests.size()), + scout.elapsed_us); + const std::size_t ready = + prepared_scouts.size() - prepared_before; + for (std::size_t i = prepared_before; + i < prepared_scouts.size(); ++i) { + const DDTreeConfidenceScout & prepared = + prepared_scouts[i]; + remember_ddtree_confidence( + prepared.request_id, prepared.confidence, + slots_.slot(prepared.slot).generated_tokens()); + std::fprintf(stderr, + "[parallel-ddtree] confidence request=%llu " + "slot=%d expected_tokens=%.3f work_nodes=%d\n", + (unsigned long long)prepared.request_id, + prepared.slot, + prepared.confidence.expected_tokens(), + tree_budget); + } std::fprintf(stderr, - "[parallel-ddtree] confidence request=%llu slot=%d " - "expected_tokens=%.3f\n", - (unsigned long long)request_id, - input->slot, confidence->expected_tokens()); - collect_candidates(); - const AdaptiveVerificationDecision steady_after = - route_ranker.select( - static_cast(inputs.size()), candidates, - adaptive_speculation_limit, + "[parallel-ddtree] confidence scout requested=%zu " + "ready=%zu elapsed_us=%.1f\n", + scout_requests.size(), ready, scout.elapsed_us); + scout_failure_cooldown = ready == 0 + ? adaptive_calibration_interval(active_requests) + : 0; + if (active_requests > kFullTreeMaxConcurrency && + (exact_profile_before || ready == 0)) { + cooldown = adaptive_calibration_interval( + active_requests); + } + + work_menus = make_work_menus(); + work_selection = + adaptive_verification_profiles_.select_work( + active_requests, work_menus, /*probe_uncalibrated_with_verifier=*/false, - /*probe_missing_routes=*/false, - enforce_ar_peer_guard, - minimum_route_samples); - if (inputs.size() > 3 && exact_profile_ready && - static_cast(steady_after.requests.size()) >= - exact_profile_width) { - // Drafter confidence is not free on the sequential - // per-slot DDTree path. Fill the executor frontier - // immediately, then amortize lower-priority newcomers. - calibration_cooldown = - adaptive_calibration_interval( - static_cast(inputs.size())); + /*enforce_ar_peer_guard=*/true, + minimum_route_samples, + /*trust_stable_confidence=*/ + plan.has_refill_backlog, + allow_safe_peer_guard_relaxation); + if (work_selection.status == + AdaptiveVerificationWorkStatus::InvalidMenu) { + return fail_step( + "adaptive speculation produced an invalid work menu"); + } + if (work_selection.status == + AdaptiveVerificationWorkStatus:: + RequiredUnavailable) { + return fail_step( + "required speculation has no common executable " + "verifier work"); + } + if (suppress_draining_tail_exploration && + !has_required_request && + (work_selection.status == + AdaptiveVerificationWorkStatus::Calibration || + work_selection.decision.exploring)) { + work_selection = + AdaptiveVerificationWorkDecision{}; } - decision = route_ranker.select( - static_cast(inputs.size()), candidates, - adaptive_speculation_limit, - verifier_side_calibration, - probe_missing_routes, - enforce_ar_peer_guard, - minimum_route_samples); + decision = work_selection.decision; + bind_selected_work(); } } } - } else { - decision.requests.reserve(candidates.size()); - for (const AdaptiveVerificationCandidate & candidate : candidates) { - decision.requests.push_back(candidate.request); - } } std::vector selected((size_t)n_slots, 0); @@ -1798,25 +1973,42 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { ar_plan.decode.push_back(in); } } - if (adaptive_enabled && !oracle.forces_selection() && - compact_tree_route) { - // This is work-shape metadata, not a decision to speculate. Propagate - // it on AR steps too so newly refilled slots cannot make a draining - // cohort rediscover C-tail graph costs. - for (const StepInput & in : inputs) { - compact_tree_cohort_[(size_t)in.slot] = - static_cast(tree_budget); - } - } using Clock = std::chrono::steady_clock; StepResult routed_result; const int speculative_count = static_cast(speculative_plan.decode.size()); + if (speculative_count > 0 && !route_work) { + return fail_step("adaptive speculation selected no verifier work"); + } + const int proposal_top_k = + ddtree_top_k_for_budget(tree_budget, draft_block_size); + int expected_reused_scouts = 0; + double reused_scout_us = 0.0; + for (const StepInput & in : speculative_plan.decode) { + const auto prepared = std::find_if( + prepared_scouts.begin(), prepared_scouts.end(), + [&](const DDTreeConfidenceScout & candidate) { + return candidate.request_id == + slots_.slot(in.slot).request_id && + candidate.slot == in.slot && + candidate.seed_token == in.token && + candidate.committed_tokens == + slots_.slot(in.slot).cur_pos && + candidate.top_k == proposal_top_k; + }); + if (prepared != prepared_scouts.end()) { + ++expected_reused_scouts; + reused_scout_us += prepared->elapsed_us; + } + } + int reused_scouts = 0; const auto started = Clock::now(); if (speculative_count > 0) { std::optional mixed = - step_ddtree(speculative_plan, ar_plan, tree_budget); + step_ddtree( + speculative_plan, ar_plan, tree_budget, + &prepared_scouts, &reused_scouts); if (!mixed) { // Proposal setup failed before target/cache mutation. Preserve // service with one ordinary packed step and retry speculation on a @@ -1829,65 +2021,52 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { routed_result = step_regular(ar_plan); if (!routed_result.ok()) return routed_result; } - const double route_us = std::max( + const double execution_us = std::max( 1.0, std::chrono::duration( Clock::now() - started).count()); + // Only request-local artifacts actually consumed by this route contribute + // scouting time. Failed and unselected sequential scouts remain separate + // calibration cost and cannot poison the verifier profile. + const bool route_timing_valid = + expected_reused_scouts == reused_scouts; + const double route_us = execution_us + reused_scout_us; if (adaptive_enabled && !oracle.forces_selection()) { - route_ranker.observe_route( - static_cast(inputs.size()), speculative_count, route_us, - /*discard_first_sample=*/ - speculative_count > 0 && minimum_route_samples > 1); if (speculative_count == 0) { - // AR has the same work shape for every proposal profile. Sharing - // this exact (C,0) timing avoids re-running the baseline merely - // because confidence moves a cohort from full to compact DDTree. - for (AdaptiveVerificationRanker * ranker : { - &adaptive_verification_, - &compact_short_adaptive_verification_, - &compact_adaptive_verification_}) { - if (ranker != &route_ranker) { - ranker->observe_route( - static_cast(inputs.size()), 0, route_us); - } - } + adaptive_verification_profiles_.observe_autoregressive( + static_cast(inputs.size()), route_us); + } else if (route_work && route_timing_valid) { + adaptive_verification_profiles_.observe_route( + *route_work, static_cast(inputs.size()), + speculative_count, route_us, + /*discard_first_sample=*/minimum_route_samples > 1); } } if (decision.exploring && !speculative_plan.decode.empty()) { std::fprintf(stderr, "[parallel-ddtree] adaptive verification probe active=%zu " - "speculative=%zu ar=%zu\n", + "speculative=%zu ar=%zu work_nodes=%d reused_scouts=%d\n", inputs.size(), speculative_plan.decode.size(), - ar_plan.decode.size()); + ar_plan.decode.size(), tree_budget, reused_scouts); } for (DecodeOutput & out : routed_result.decode) { - if (oracle.forces_selection()) break; + if (!adaptive_enabled || oracle.forces_selection()) break; if (out.failed || out.slot < 0 || out.slot >= n_slots) continue; - if (selected[(size_t)out.slot]) { + if (selected[(size_t)out.slot] && route_work) { const double emitted = static_cast(out.ddtree_accepted_tokens + 1); Qwen35Slot & observed_seq = slots_.slot(out.slot); const auto confidence = ddtree_confidence_.find(observed_seq.request_id); if (confidence != ddtree_confidence_.end()) { - route_ranker.observe_request_estimate( + adaptive_verification_profiles_.observe_request_estimate( + *route_work, observed_seq.request_id, confidence->second.limited_to(tree_budget)); } - route_ranker.observe_request_yield( - observed_seq.request_id, emitted); - RequestTargetYield & target = - ddtree_target_yield_[observed_seq.request_id]; - if (target.samples == 0) { - target.expected_tokens = emitted; - } else { - constexpr double kTargetYieldAlpha = 0.35; - target.expected_tokens += kTargetYieldAlpha * - (emitted - target.expected_tokens); - } - target.samples = std::min( - target.samples + 1, 64); + adaptive_verification_profiles_.observe_request_yield( + *route_work, observed_seq.request_id, emitted); } } @@ -2312,14 +2491,10 @@ SeqEngine::StepResult Qwen35SeqEngine::step_regular(const StepPlan & plan) { void Qwen35SeqEngine::retire(int slot) { if (!slots_.is_active(slot)) return; const std::uint64_t request_id = slots_.slot(slot).request_id; - adaptive_verification_.forget_request(request_id); - compact_short_adaptive_verification_.forget_request(request_id); - compact_adaptive_verification_.forget_request(request_id); + adaptive_verification_profiles_.forget_request(request_id); ddtree_confidence_.erase(request_id); ddtree_confidence_progress_.erase(request_id); - ddtree_target_yield_.erase(request_id); slots_.retire(slot); - compact_tree_cohort_[(size_t)slot] = 0; } } // namespace dflash::common diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 3149c6e65..5036c14c2 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -33,6 +33,7 @@ #include #include #include +#include #include namespace dflash::common { @@ -122,8 +123,22 @@ class Qwen35SeqEngine final : public SeqEngine { DraftKvState * ensure_slot_draft_kv(int slot); bool ddtree_available(const StepPlan & plan) const; bool ddtree_input_eligible(const StepInput & input) const; - std::optional estimate_ddtree_confidence( - const StepInput & input); + struct DDTreeConfidenceScout { + std::uint64_t request_id = 0; + int slot = -1; + int32_t seed_token = -1; + int committed_tokens = 0; + int top_k = 0; + double elapsed_us = 0.0; + std::vector top_log_probs; + std::vector top_token_ids; + SpeculationConfidenceEstimate confidence; + }; + std::optional scout_ddtree_confidence( + const ConfidenceScoutRequest & request); + ConfidenceScoutBatchResult scout_ddtree_confidence_batch( + const std::vector & requests, + std::vector & prepared); void remember_ddtree_confidence( std::uint64_t request_id, const SpeculationConfidenceEstimate & estimate, @@ -133,7 +148,9 @@ class Qwen35SeqEngine final : public SeqEngine { // caller may safely use the ordinary packed AR path for this iteration. std::optional step_ddtree( const StepPlan & speculative_plan, const StepPlan & ar_plan, - int tree_budget); + int tree_budget, + const std::vector * prepared = nullptr, + int * reused_scouts = nullptr); Qwen35Backend & b_; Qwen35SlotManager slots_; @@ -142,27 +159,28 @@ class Qwen35SeqEngine final : public SeqEngine { int tree_scratch_base_ = 0; int tree_scratch_stride_ = 0; bool capture_features_ = false; - AdaptiveVerificationRanker adaptive_verification_; - AdaptiveVerificationRanker compact_short_adaptive_verification_; - AdaptiveVerificationRanker compact_adaptive_verification_; + AdaptiveVerificationProfileBank adaptive_verification_profiles_; // Latest drafter confidence is request-owned and can be projected onto // the current full/compact work budget without looking at prompt text. std::map ddtree_confidence_; std::map ddtree_confidence_progress_; - struct RequestTargetYield { - double expected_tokens = 1.0; - std::size_t samples = 0; - }; - // Shape choice uses target outcomes when available; this map deliberately - // spans full/compact rankers while their hardware route costs stay split. - std::map ddtree_target_yield_; - int adaptive_calibration_cooldown_ = 0; - int compact_short_adaptive_calibration_cooldown_ = 0; - int compact_adaptive_calibration_cooldown_ = 0; - // Zero means no sticky compact cohort. Non-zero entries store the exact - // DDTree work-shape budget so short and wide observations never mix. - std::vector compact_tree_cohort_; + // Calibration cadence is both work- and occupancy-local. A losing C=8 + // DDTree profile must not delay a C=4 request or a later DSpark adapter. + std::map, int> + adaptive_calibration_cooldowns_; + // Transient scout failures are adapter-shape-local. Backoff prevents a + // broken confidence path from taxing every AR token, including at low C. + std::map, int> + adaptive_scout_failure_cooldowns_; + // A draining closed batch should reuse measured work, but it should not + // pay immediately to discover every transient lower-C tail. If occupancy + // and request identity remain stable, exploration is re-enabled after a + // short bounded delay. Identity prevents a new cohort from inheriting an + // unrelated predecessor's draining state. + std::vector previous_decode_requests_; + int stable_decode_cohort_steps_ = 0; + bool draining_decode_cohort_ = false; ggml_context * feature_view_ctx_ = nullptr; std::vector slot_feature_mirrors_; std::vector> slot_draft_kv_; diff --git a/server/test/test_speculation_goodput.cpp b/server/test/test_speculation_goodput.cpp index 2a77b2547..5d0f2a07c 100644 --- a/server/test/test_speculation_goodput.cpp +++ b/server/test/test_speculation_goodput.cpp @@ -622,6 +622,658 @@ int main() { SpeculatorKind::DDTree, 2.2).has_value()); } + // Equal-value peers outside a bounded verifier prefix receive service once + // one member has proven the bucket useful. A materially lower-value peer + // remains in its own bucket and cannot displace them merely for fairness. + { + AdaptiveVerificationRanker ranker; + ranker.observe_autoregressive(8, 100.0); + ranker.observe_route(8, 1, 80.0); + ranker.observe_route(8, 2, 80.0); + std::vector candidates = { + {1, 4.0, 5.0, true, + std::numeric_limits::quiet_NaN(), 4, 12}, + {2, 4.0, 5.0, true, + std::numeric_limits::quiet_NaN(), 4, 12}, + {3, 4.0, 5.0, true, + std::numeric_limits::quiet_NaN(), 0, 4}, + {4, 4.0, 5.0, true, + std::numeric_limits::quiet_NaN(), 0, 4}, + {5, 2.0, 5.0, true, + std::numeric_limits::quiet_NaN(), 0, 0}, + }; + const AdaptiveVerificationDecision rotated = ranker.select( + 8, candidates, /*max_speculative_requests=*/2); + CHECK(rotated.requests.size() == 2); + CHECK(rotated.requests[0] == 3); + CHECK(rotated.requests[1] == 4); + CHECK(std::find(rotated.requests.begin(), rotated.requests.end(), 5) == + rotated.requests.end()); + } + + // Scouting and verifier menus carry only model-neutral request, confidence, + // and structural work metadata. AR remains implicit, and a branching + // verifier's row count does not pretend to be its maximum emitted path. + { + SpeculationRequestView request; + request.request_id = 700; + request.slot = 3; + request.seed_token = 42; + request.progress_tokens = 9; + + ConfidenceScoutRequest scout; + scout.request = request; + scout.work = {SpeculatorKind::DDTree, 7, 4}; + scout.max_candidate_tokens = 4; + CHECK(scout.request.request_id == 700); + CHECK(scout.work.valid()); + CHECK(scout.valid()); + CHECK(scout.max_candidate_tokens == 4); + scout.max_candidate_tokens = 5; + CHECK(!scout.valid()); + scout.max_candidate_tokens = 0; + CHECK(!scout.valid()); + scout.max_candidate_tokens = 4; + + const float confidence_values[] = {0.8f, 0.5f}; + ConfidenceScoutResult ready; + ready.request_id = request.request_id; + ready.status = ConfidenceScoutStatus::Ready; + ready.confidence = make_speculation_confidence_estimate( + SpeculatorKind::DDTree, confidence_values, 2, + SpeculationConfidenceCost::ExtraDraftPass); + ready.elapsed_us = 12.0; + CHECK(ready.ready()); + CHECK(ready.elapsed_us == 12.0); + ready.status = ConfidenceScoutStatus::ProposalRequired; + CHECK(!ready.ready()); + + ConfidenceScoutBatchResult batch; + batch.requests.push_back(ready); + batch.elapsed_us = 25.0; + CHECK(batch.requests.size() == 1); + CHECK(batch.elapsed_us == 25.0); + + const VerifierWorkKey tree_work{ + SpeculatorKind::DDTree, 1, 32, 33}; + const VerifierWorkKey compact_work{ + SpeculatorKind::DDTree, 2, 4, 5}; + const VerifierWorkKey dspark_work{ + SpeculatorKind::DSpark, 1, 3, 4}; + CHECK(tree_work.valid()); + CHECK(tree_work != compact_work); + CHECK(compact_work < tree_work || tree_work < compact_work); + CHECK(tree_work < dspark_work); + + VerifierWorkPlan plan; + plan.work = tree_work; + plan.maximum_emitted_tokens = 16.0; + plan.confidence_expected_tokens = 2.2; + plan.preferred_for_forced_mode = true; + plan.max_parallel_requests = 3; + plan.adaptive_request_limit = 2; + plan.exploration_priority = 2; + CHECK(plan.valid()); + CHECK(plan.has_confidence()); + CHECK(plan.work.verifier_rows == 33); + CHECK(plan.maximum_emitted_tokens == 16.0); + CHECK(plan.max_parallel_requests == 3); + CHECK(plan.adaptive_request_limit == 2); + CHECK(plan.exploration_priority == 2); + + RequestVerifierWorkMenu menu; + menu.request = request; + menu.speculative.push_back(plan); + CHECK(menu.request.slot == 3); + CHECK(menu.speculative.size() == 1); + CHECK(menu.speculative[0].preferred_for_forced_mode); + } + + // Work profiles isolate speculative shapes while sharing the exact AR + // baseline. Extra scouting cost has its own EWMA and cannot create or + // mutate a verifier profile. + { + AdaptiveVerificationConfig config; + config.cost_ewma_alpha = 0.5; + AdaptiveVerificationProfileBank bank(config); + const VerifierWorkKey short_tree{ + SpeculatorKind::DDTree, 1, 4, 5}; + const VerifierWorkKey wide_tree{ + SpeculatorKind::DDTree, 2, 8, 9}; + const VerifierWorkKey dspark_linear{ + SpeculatorKind::DSpark, 1, 3, 4}; + + // An AR observation made before a profile exists is replayed when that + // work key is first requested. + bank.observe_autoregressive(5, 100.0); + AdaptiveVerificationRanker & short_ranker = bank.profile(short_tree); + CHECK(short_ranker.has_autoregressive_cost(5)); + CHECK(short_ranker.autoregressive_cost_us(5) == 100.0); + bank.observe_route(short_tree, 5, 1, 80.0); + CHECK(short_ranker.has_route_cost(5, 1)); + + AdaptiveVerificationRanker & wide_ranker = bank.profile(wide_tree); + CHECK(wide_ranker.has_autoregressive_cost(5)); + CHECK(!wide_ranker.has_route_cost(5, 1)); + CHECK(bank.profile_count() == 2); + + bank.observe_autoregressive(5, 120.0); + CHECK(bank.autoregressive_cost_us(5) == 110.0); + CHECK(bank.autoregressive_cost_samples(5) == 2); + CHECK(short_ranker.autoregressive_cost_us(5) == 110.0); + CHECK(wide_ranker.autoregressive_cost_us(5) == 110.0); + + const ConfidenceScoutWorkKey short_scout{ + SpeculatorKind::DDTree, 1, 4}; + const ConfidenceScoutWorkKey wide_scout{ + SpeculatorKind::DDTree, 2, 8}; + const ConfidenceScoutWorkKey dspark_scout{ + SpeculatorKind::DSpark, 1, 4}; + bank.observe_scout(short_scout, 2, 50.0); + bank.observe_scout(short_scout, 2, 70.0); + bank.observe_scout(wide_scout, 2, 90.0); + bank.observe_scout(dspark_scout, 2, 110.0); + bank.observe_scout(short_scout, 1, 20.0); + CHECK(bank.profile_count() == 2); + CHECK(bank.scout_cost_us(short_scout, 2) == 60.0); + CHECK(bank.scout_cost_samples(short_scout, 2) == 2); + CHECK(bank.scout_cost_us(wide_scout, 2) == 90.0); + CHECK(bank.scout_cost_us(dspark_scout, 2) == 110.0); + CHECK(bank.scout_cost_us(short_scout, 1) == 20.0); + CHECK(short_ranker.autoregressive_cost_us(5) == 110.0); + CHECK(short_ranker.route_cost_us(5, 1) == 80.0); + + bank.observe_request_yield(short_tree, 900, 4.0); + bank.observe_request_yield(wide_tree, 900, 3.0); + CHECK(short_ranker.request_expected_tokens(900).has_value()); + CHECK(wide_ranker.request_expected_tokens(900).has_value()); + bank.forget_request(900); + CHECK(!short_ranker.request_expected_tokens(900).has_value()); + CHECK(!wide_ranker.request_expected_tokens(900).has_value()); + + // A k=0 route is shared AR work and does not manufacture the referenced + // speculative profile. It is inherited if that profile appears later. + bank.observe_route(dspark_linear, 4, 0, 90.0); + CHECK(bank.profile_count() == 2); + CHECK(bank.profile(dspark_linear).autoregressive_cost_us(4) == 90.0); + CHECK(bank.profile_count() == 3); + + bank.reset(); + CHECK(bank.profile_count() == 0); + CHECK(!bank.has_autoregressive_cost(5)); + CHECK(!bank.has_scout_cost(short_scout, 2)); + } + + // Work-menu exploration is ordered by adapter priority. A measured + // profitable short route suppresses a fresh wide probe, while a complete + // losing short profile advances to the wider shape. + { + const VerifierWorkKey short_work{ + SpeculatorKind::DDTree, 101, 4, 5}; + const VerifierWorkKey wide_work{ + SpeculatorKind::DDTree, 102, 8, 9}; + auto menus = [&]() { + std::vector out; + for (int slot = 0; slot < 2; ++slot) { + RequestVerifierWorkMenu menu; + menu.request.request_id = + static_cast(1000 + slot); + menu.request.slot = slot; + menu.request.progress_tokens = slot; + VerifierWorkPlan short_plan; + short_plan.work = short_work; + short_plan.maximum_emitted_tokens = 5.0; + short_plan.confidence_expected_tokens = + slot == 0 ? 4.0 : 2.0; + short_plan.max_parallel_requests = 1; + short_plan.adaptive_request_limit = 1; + short_plan.exploration_priority = 0; + VerifierWorkPlan wide_plan; + wide_plan.work = wide_work; + wide_plan.maximum_emitted_tokens = 9.0; + wide_plan.confidence_expected_tokens = + short_plan.confidence_expected_tokens; + wide_plan.max_parallel_requests = 1; + wide_plan.adaptive_request_limit = 1; + wide_plan.exploration_priority = 1; + menu.speculative = {short_plan, wide_plan}; + out.push_back(std::move(menu)); + } + return out; + }(); + + AdaptiveVerificationProfileBank profitable; + profitable.observe_autoregressive(2, 100.0); + AdaptiveVerificationWorkDecision selected = + profitable.select_work(2, menus); + CHECK(selected.has_work()); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Verification); + CHECK(selected.work.value() == short_work); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Verification); + CHECK(selected.decision.exploring); + CHECK(selected.decision.requests.size() == 1); + + profitable.observe_route(short_work, 2, 1, 70.0); + selected = profitable.select_work(2, menus); + CHECK(selected.work.value() == short_work); + CHECK(!selected.decision.exploring); + CHECK(selected.decision.requests.size() == 1); + CHECK(!profitable.profile(wide_work).has_route_cost(2, 1)); + + AdaptiveVerificationProfileBank losing; + losing.observe_autoregressive(2, 100.0); + losing.observe_route(short_work, 2, 1, 300.0); + selected = losing.select_work(2, menus); + CHECK(selected.has_work()); + CHECK(selected.work.value() == wide_work); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Verification); + CHECK(selected.decision.exploring); + CHECK(selected.decision.requests.size() == 1); + } + + // Refill relaxation is scoped to each exact work shape's executable + // coverage. Two lanes may optimize aggregate C=5 goodput, while a + // one-lane C=16 route still protects its fifteen AR peers even when its + // aggregate token estimate looks attractive. + { + auto make_menus = []( + int active, int eligible, const VerifierWorkKey & work, + int adaptive_limit, double confidence) { + std::vector menus; + menus.reserve((size_t)active); + for (int slot = 0; slot < active; ++slot) { + RequestVerifierWorkMenu menu; + menu.request.request_id = + static_cast(9000 + slot); + menu.request.slot = slot; + if (slot < eligible) { + VerifierWorkPlan plan; + plan.work = work; + plan.maximum_emitted_tokens = confidence; + plan.confidence_expected_tokens = confidence; + plan.max_parallel_requests = active; + plan.adaptive_request_limit = adaptive_limit; + menu.speculative = {plan}; + } + menus.push_back(std::move(menu)); + } + return menus; + }; + + const VerifierWorkKey broad_work{ + SpeculatorKind::DDTree, 901, 4, 5}; + AdaptiveVerificationProfileBank broad; + broad.observe_autoregressive(5, 100.0); + broad.observe_route(broad_work, 5, 1, 130.0); + broad.observe_route(broad_work, 5, 2, 120.0); + const auto broad_menus = make_menus( + 5, 2, broad_work, 2, 4.0); + CHECK(broad.select_work(5, broad_menus).decision.requests.empty()); + const AdaptiveVerificationWorkDecision refill_broad = + broad.select_work( + 5, broad_menus, + /*probe_uncalibrated_with_verifier=*/false, + /*enforce_ar_peer_guard=*/true, + /*minimum_speculative_route_samples=*/1, + /*trust_stable_confidence=*/true, + /*allow_safe_peer_guard_relaxation=*/true); + CHECK(refill_broad.decision.requests.size() == 2); + + const VerifierWorkKey narrow_work{ + SpeculatorKind::DDTree, 902, 4, 5}; + AdaptiveVerificationProfileBank narrow; + narrow.observe_autoregressive(16, 100.0); + narrow.observe_route(narrow_work, 16, 1, 120.0); + const auto narrow_menus = make_menus( + 16, 1, narrow_work, 1, 21.0); + const AdaptiveVerificationWorkDecision refill_narrow = + narrow.select_work( + 16, narrow_menus, + /*probe_uncalibrated_with_verifier=*/false, + /*enforce_ar_peer_guard=*/true, + /*minimum_speculative_route_samples=*/1, + /*trust_stable_confidence=*/true, + /*allow_safe_peer_guard_relaxation=*/true); + CHECK(refill_narrow.status == + AdaptiveVerificationWorkStatus::Autoregressive); + CHECK(refill_narrow.decision.requests.empty()); + CHECK(narrow.select_work( + 16, narrow_menus, + /*probe_uncalibrated_with_verifier=*/false, + /*enforce_ar_peer_guard=*/false) + .decision.requests.size() == 1); + } + + // Target-verified yield is local to each exact work shape and replaces raw + // confidence once stable. The wider route wins despite its lower raw score + // because only its target outcomes are useful. + { + const VerifierWorkKey short_work{ + SpeculatorKind::DDTree, 201, 4, 5}; + const VerifierWorkKey wide_work{ + SpeculatorKind::DDTree, 202, 8, 9}; + AdaptiveVerificationProfileBank bank; + bank.observe_autoregressive(1, 100.0); + bank.observe_route(short_work, 1, 1, 100.0); + bank.observe_route(wide_work, 1, 1, 100.0); + bank.observe_request_estimate( + short_work, 2000, SpeculatorKind::DDTree, 8.0); + bank.observe_request_estimate( + wide_work, 2000, SpeculatorKind::DDTree, 2.0); + for (int sample = 0; sample < 4; ++sample) { + bank.observe_request_yield(short_work, 2000, 1.0); + bank.observe_request_yield(wide_work, 2000, 4.0); + } + + RequestVerifierWorkMenu menu; + menu.request.request_id = 2000; + menu.request.slot = 7; + VerifierWorkPlan short_plan; + short_plan.work = short_work; + short_plan.maximum_emitted_tokens = 8.0; + short_plan.confidence_expected_tokens = 8.0; + short_plan.max_parallel_requests = 1; + short_plan.adaptive_request_limit = 1; + short_plan.exploration_priority = 0; + VerifierWorkPlan wide_plan; + wide_plan.work = wide_work; + wide_plan.maximum_emitted_tokens = 8.0; + wide_plan.confidence_expected_tokens = 2.0; + wide_plan.max_parallel_requests = 1; + wide_plan.adaptive_request_limit = 1; + wide_plan.exploration_priority = 1; + menu.speculative = {short_plan, wide_plan}; + + const AdaptiveVerificationWorkDecision selected = + bank.select_work(1, {menu}); + CHECK(selected.has_work()); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Verification); + CHECK(selected.work.value() == wide_work); + CHECK(selected.decision.requests == std::vector{7}); + CHECK(!selected.decision.exploring); + } + + // Always requests cannot disappear into AR. The adapter's preferred plan + // wins at cold start; without one, lower exploration priority is the + // deterministic forced-mode fallback. + { + const VerifierWorkKey short_work{ + SpeculatorKind::DDTree, 301, 4, 5}; + const VerifierWorkKey wide_work{ + SpeculatorKind::DDTree, 302, 8, 9}; + RequestVerifierWorkMenu menu; + menu.request.request_id = 3000; + menu.request.slot = 5; + menu.request.required = true; + VerifierWorkPlan short_plan; + short_plan.work = short_work; + short_plan.maximum_emitted_tokens = 5.0; + short_plan.max_parallel_requests = 1; + short_plan.adaptive_request_limit = 1; + short_plan.exploration_priority = 0; + VerifierWorkPlan wide_plan; + wide_plan.work = wide_work; + wide_plan.maximum_emitted_tokens = 9.0; + wide_plan.preferred_for_forced_mode = true; + wide_plan.max_parallel_requests = 1; + wide_plan.adaptive_request_limit = 1; + wide_plan.exploration_priority = 4; + menu.speculative = {short_plan, wide_plan}; + + AdaptiveVerificationProfileBank preferred_bank; + AdaptiveVerificationWorkDecision selected = + preferred_bank.select_work(1, {menu}); + CHECK(selected.work.value() == wide_work); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Verification); + CHECK(selected.decision.requests == std::vector{5}); + + menu.speculative[1].preferred_for_forced_mode = false; + AdaptiveVerificationProfileBank priority_bank; + selected = priority_bank.select_work(1, {menu}); + CHECK(selected.work.value() == short_work); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Verification); + CHECK(selected.decision.requests == std::vector{5}); + } + + // Work decisions distinguish normal AR, confidence calibration, malformed + // adapter input, and a forced request that has no executable common shape. + { + const VerifierWorkKey work{ + SpeculatorKind::DDTree, 401, 4, 5}; + VerifierWorkPlan plan; + plan.work = work; + plan.maximum_emitted_tokens = 5.0; + plan.max_parallel_requests = 2; + plan.adaptive_request_limit = 1; + + RequestVerifierWorkMenu empty; + empty.request.request_id = 4000; + empty.request.slot = 0; + AdaptiveVerificationProfileBank ar_bank; + AdaptiveVerificationWorkDecision selected = + ar_bank.select_work(1, {empty}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Autoregressive); + CHECK(!selected.has_work()); + + RequestVerifierWorkMenu unknown = empty; + unknown.speculative = {plan}; + AdaptiveVerificationProfileBank calibration_bank; + calibration_bank.observe_autoregressive(1, 100.0); + selected = calibration_bank.select_work(1, {unknown}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Calibration); + CHECK(selected.work.value() == work); + CHECK(selected.decision.calibration_request == 0); + + selected = calibration_bank.select_work(2, {unknown}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::InvalidMenu); + CHECK(!selected.has_work()); + + RequestVerifierWorkMenu duplicate_slot = empty; + duplicate_slot.request.request_id = 4001; + selected = ar_bank.select_work(2, {empty, duplicate_slot}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::InvalidMenu); + + RequestVerifierWorkMenu duplicate_request = empty; + duplicate_request.request.slot = 1; + selected = ar_bank.select_work(2, {empty, duplicate_request}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::InvalidMenu); + + RequestVerifierWorkMenu inconsistent_a = unknown; + RequestVerifierWorkMenu inconsistent_b = unknown; + inconsistent_b.request.request_id = 4001; + inconsistent_b.request.slot = 1; + inconsistent_b.speculative[0].exploration_priority = 1; + selected = ar_bank.select_work( + 2, {inconsistent_a, inconsistent_b}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::InvalidMenu); + + RequestVerifierWorkMenu duplicate_work = unknown; + duplicate_work.speculative.push_back(plan); + selected = ar_bank.select_work(1, {duplicate_work}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::InvalidMenu); + + RequestVerifierWorkMenu mixed_speculators = unknown; + VerifierWorkPlan dspark_plan = plan; + dspark_plan.work = { + SpeculatorKind::DSpark, 403, 4, 5}; + mixed_speculators.speculative.push_back(dspark_plan); + selected = ar_bank.select_work(1, {mixed_speculators}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::InvalidMenu); + + RequestVerifierWorkMenu required_empty = empty; + required_empty.request.required = true; + selected = ar_bank.select_work(1, {required_empty}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::RequiredUnavailable); + + const VerifierWorkKey other_work{ + SpeculatorKind::DDTree, 402, 8, 9}; + RequestVerifierWorkMenu required_a = unknown; + required_a.request.required = true; + RequestVerifierWorkMenu required_b = unknown; + required_b.request.request_id = 4001; + required_b.request.slot = 1; + required_b.request.required = true; + required_b.speculative[0].work = other_work; + selected = ar_bank.select_work(2, {required_a, required_b}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::RequiredUnavailable); + } + + // The adaptive limit bounds optional profiling, while required/Always + // requests may cross it only when the hard executor capacity permits all. + { + const VerifierWorkKey work{ + SpeculatorKind::DDTree, 501, 4, 5}; + VerifierWorkPlan plan; + plan.work = work; + plan.maximum_emitted_tokens = 5.0; + plan.preferred_for_forced_mode = true; + plan.max_parallel_requests = 2; + plan.adaptive_request_limit = 1; + std::vector required; + for (int slot = 0; slot < 2; ++slot) { + RequestVerifierWorkMenu menu; + menu.request.request_id = + static_cast(5000 + slot); + menu.request.slot = slot; + menu.request.required = true; + menu.speculative = {plan}; + required.push_back(std::move(menu)); + } + AdaptiveVerificationProfileBank bank; + AdaptiveVerificationWorkDecision selected = + bank.select_work(2, required); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Verification); + CHECK(selected.decision.requests.size() == 2); + + for (RequestVerifierWorkMenu & menu : required) { + menu.speculative[0].max_parallel_requests = 1; + } + selected = bank.select_work(2, required); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::RequiredUnavailable); + } + + // Once required routes are measured, absolute goodput overrides the cold + // forced-mode preference. Relative gain is one for both required baselines. + { + const VerifierWorkKey preferred_work{ + SpeculatorKind::DDTree, 601, 4, 5}; + const VerifierWorkKey faster_work{ + SpeculatorKind::DDTree, 602, 8, 9}; + RequestVerifierWorkMenu menu; + menu.request.request_id = 6000; + menu.request.slot = 0; + menu.request.required = true; + VerifierWorkPlan preferred; + preferred.work = preferred_work; + preferred.maximum_emitted_tokens = 4.0; + preferred.confidence_expected_tokens = 2.0; + preferred.preferred_for_forced_mode = true; + preferred.max_parallel_requests = 1; + preferred.adaptive_request_limit = 1; + VerifierWorkPlan faster = preferred; + faster.work = faster_work; + faster.preferred_for_forced_mode = false; + menu.speculative = {preferred, faster}; + + AdaptiveVerificationProfileBank bank; + bank.observe_route(preferred_work, 1, 1, 200.0); + bank.observe_route(faster_work, 1, 1, 100.0); + const AdaptiveVerificationWorkDecision selected = + bank.select_work(1, {menu}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Verification); + CHECK(selected.work.value() == faster_work); + CHECK(std::abs( + selected.decision.predicted_goodput - 0.02) < 1e-12); + CHECK(selected.decision.predicted_gain == 1.0); + } + + // Adapter confidence cannot claim more useful output than the selected + // work shape can emit; both calibration storage and goodput use the clamp. + { + const VerifierWorkKey work{ + SpeculatorKind::DSpark, 701, 1, 2}; + RequestVerifierWorkMenu menu; + menu.request.request_id = 7000; + menu.request.slot = 0; + VerifierWorkPlan plan; + plan.work = work; + plan.maximum_emitted_tokens = 2.0; + plan.confidence_expected_tokens = 100.0; + plan.max_parallel_requests = 1; + plan.adaptive_request_limit = 1; + CHECK(plan.bounded_confidence_expected_tokens() == 2.0); + menu.speculative = {plan}; + + AdaptiveVerificationProfileBank bank; + bank.observe_autoregressive(1, 100.0); + bank.observe_route(work, 1, 1, 100.0); + const AdaptiveVerificationWorkDecision selected = + bank.select_work(1, {menu}); + CHECK(selected.work.value() == work); + CHECK(std::abs( + selected.decision.predicted_goodput - 0.02) < 1e-12); + const auto estimate = + bank.profile(work).estimate_request_yield(7000); + CHECK(estimate.has_value()); + CHECK(estimate->expected_tokens == 2.0); + CHECK(estimate->confidence_expected_tokens == 2.0); + } + + // An exact losing short profile's fallback calibration cannot indefinitely + // block bounded exploration of the next wider work shape. + { + const VerifierWorkKey short_work{ + SpeculatorKind::DDTree, 801, 4, 5}; + const VerifierWorkKey wide_work{ + SpeculatorKind::DDTree, 802, 8, 9}; + RequestVerifierWorkMenu menu; + menu.request.request_id = 8000; + menu.request.slot = 0; + VerifierWorkPlan short_plan; + short_plan.work = short_work; + short_plan.maximum_emitted_tokens = 5.0; + short_plan.max_parallel_requests = 1; + short_plan.adaptive_request_limit = 1; + short_plan.exploration_priority = 0; + VerifierWorkPlan wide_plan; + wide_plan.work = wide_work; + wide_plan.maximum_emitted_tokens = 9.0; + wide_plan.confidence_expected_tokens = 4.0; + wide_plan.max_parallel_requests = 1; + wide_plan.adaptive_request_limit = 1; + wide_plan.exploration_priority = 1; + menu.speculative = {short_plan, wide_plan}; + + AdaptiveVerificationProfileBank bank; + bank.observe_autoregressive(1, 100.0); + bank.observe_route(short_work, 1, 1, 300.0); + const AdaptiveVerificationWorkDecision selected = + bank.select_work(1, {menu}); + CHECK(selected.status == + AdaptiveVerificationWorkStatus::Verification); + CHECK(selected.work.value() == wide_work); + CHECK(selected.decision.exploring); + } + // Shared cohort evidence can rank a new request, but it cannot unlock the // slow-route homogeneous exception. A bounded verifier probe may cross the // steady peer guard to gather the missing request-local evidence.