From 8902901e61fcc514244bd5cd2ec0109bc56f6571 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Wed, 12 Aug 2026 22:55:17 +0000 Subject: [PATCH 1/8] 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 2/8] 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 3/8] 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 4/8] 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 5/8] 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 6/8] 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 32bc12b7379ae5acb428a4575153d34a6bc72b43 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 13:41:20 +0000 Subject: [PATCH 7/8] feat(qwen35): gate speculation by measured acceptance Add the packed mixed DDTree/AR execution substrate, a model-neutral acceptance-and-cost gate, and per-request decode modes. Replace cohort thresholds and scout policy with bounded in-band probation, measured break-even routing, hysteresis, and re-probing. --- .../benchmarks/concurrency/FEATURE_MATRIX.md | 2 +- server/CMakeLists.txt | 9 + server/README.md | 23 + server/src/common/concurrency/seq_engine.h | 9 +- .../src/common/concurrency/speculation_gate.h | 370 +++++++++ server/src/common/model_backend.h | 4 + server/src/common/speculation_policy.h | 34 + server/src/common/step_graph.h | 8 + server/src/internal.h | 33 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 722 +++++++++++++++--- .../qwen35/concurrency/qwen35_seq_engine.h | 9 +- .../concurrency/qwen35_slot_manager.cpp | 30 - .../qwen35/concurrency/qwen35_slot_manager.h | 17 - server/src/qwen35/graph_builders.cpp | 60 +- server/src/qwen35/graph_builders.h | 26 +- server/src/qwen35/qwen35_backend.cpp | 35 +- server/src/qwen35/qwen35_target_graph.cpp | 227 ++++-- 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 | 31 +- server/test/test_recurrent_snapshot.cpp | 6 + server/test/test_seq_slot_manager.cpp | 56 -- server/test/test_server_unit.cpp | 21 + server/test/test_speculation_gate.cpp | 276 +++++++ 25 files changed, 1758 insertions(+), 296 deletions(-) create mode 100644 server/src/common/concurrency/speculation_gate.h create mode 100644 server/src/common/speculation_policy.h create mode 100644 server/test/test_speculation_gate.cpp diff --git a/harness/benchmarks/concurrency/FEATURE_MATRIX.md b/harness/benchmarks/concurrency/FEATURE_MATRIX.md index 62fce6a46..1915ae7f1 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-aligned continuous DDTree policy; leave it at the default `1` when measuring acceptance-gated speculation. The default policy gives cold requests bounded probation through ordinary speculative rounds, then ranks each request from its measured accepted-token EMA and the measured cost of the exact `(active requests, speculative requests)` route shape. No scout or prompt classifier runs. Selected DDTree paths and unselected AR roots share one packed target pass, with accepted paths committed directly when capture capacity permits and through fused replay otherwise. 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/CMakeLists.txt b/server/CMakeLists.txt index ad6ef874f..20d97beb5 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_gate.cpp") + # Model-neutral acceptance/cost gate: no model or GPU. + add_executable(test_speculation_gate + test/test_speculation_gate.cpp) + target_include_directories(test_speculation_gate PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/test) + list(APPEND _raw_unit_test_targets test_speculation_gate) + 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/README.md b/server/README.md index 1e96b697d..f78b9ac42 100644 --- a/server/README.md +++ b/server/README.md @@ -170,6 +170,29 @@ Run it directly: --model-name luce-dflash ``` +With a decode drafter configured, the default is measured per-request +routing. Operators can change the default without unloading the drafter: + +```bash +--ddtree --decode-mode adaptive # acceptance/cost-gated 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" +} +``` + +Accepted values are `adaptive`, `ar`, and `speculation`. Forced speculation +requires a supported drafter. The control is speculator-neutral: DDTree uses +it today, and later speculators can consume the same request policy. + ### Compression proxy mode `dflash_server` can run as a **PFlash compression proxy** in front of any diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index 6371e2552..a0aa35316 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,10 @@ 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 measured engine policy. + 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; @@ -306,7 +311,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/concurrency/speculation_gate.h b/server/src/common/concurrency/speculation_gate.h new file mode 100644 index 000000000..65d8d2aca --- /dev/null +++ b/server/src/common/concurrency/speculation_gate.h @@ -0,0 +1,370 @@ +#pragma once + +// Model- and executor-neutral admission policy for mixed speculative/AR +// decode. The gate consumes only measured wall times, accepted-token yield, +// executor capacity, and an optional speculator-provided request prior. + +#include "common/speculation_policy.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +struct SpecGateConfig { + double ema_alpha = 0.4; + double cost_ewma_alpha = 0.35; + int probe_rounds = 2; + int max_probers = 2; + int bad_rounds = 2; + int reprobe_tokens = 64; + double margin = 1.05; + double slack = 1.10; +}; + +struct SpecCandidate { + int slot = -1; + std::uint64_t request_id = 0; + SpeculationPolicy policy = SpeculationPolicy::Adaptive; + bool eligible = false; + double prior_accept = std::numeric_limits::quiet_NaN(); + int generated_tokens = 0; +}; + +class SpeculationGate { +public: + explicit SpeculationGate(SpecGateConfig config, double max_accept) + : cfg_(sanitize(config)), + max_accept_(valid_positive(max_accept) + ? std::max(1.0, max_accept) : 1.0) {} + + // The project is C++17, so these public range-taking methods provide the + // same contiguous/range call shape as the spec's std::span API without + // raising the language level of every CUDA/HIP translation unit. + template + std::vector plan(int C, const CandidateRange & candidates, + int k_cap) { + if (C <= 0) return {}; + + struct Ranked { + const SpecCandidate * candidate = nullptr; + double score = 1.0; + bool probing = false; + bool reprobe = false; + }; + std::vector forced; + std::vector steady; + std::vector probers; + + for (const SpecCandidate & candidate : candidates) { + if (!candidate.eligible || + candidate.policy == SpeculationPolicy::Never) { + continue; + } + RequestState & state = requests_[candidate.request_id]; + state.request_id = candidate.request_id; + if (candidate.policy == SpeculationPolicy::Always) { + forced.push_back({&candidate, + score_for(state, candidate), false, false}); + continue; + } + if (state.mode == Mode::spec) { + steady.push_back( + {&candidate, state.ema_yield, false, false}); + continue; + } + if (state.mode == Mode::probing) { + probers.push_back( + {&candidate, probe_score(candidate), true, false}); + continue; + } + const long long generated = candidate.generated_tokens; + const long long last = state.tokens_at_last_spec; + if (generated - last >= cfg_.reprobe_tokens) { + probers.push_back( + {&candidate, probe_score(candidate), true, true}); + } + } + + // In particular, an all-Never/ineligible cohort must not allocate + // either request or per-concurrency state. + if (forced.empty() && steady.empty() && probers.empty()) return {}; + ensure_cost(C); + + const auto rank = [](const Ranked & a, const Ranked & b) { + if (a.score != b.score) return a.score > b.score; + return a.candidate->request_id < b.candidate->request_id; + }; + std::sort(forced.begin(), forced.end(), rank); + std::sort(probers.begin(), probers.end(), rank); + if ((int)probers.size() > cfg_.max_probers) { + probers.resize((size_t)cfg_.max_probers); + } + steady.insert(steady.end(), probers.begin(), probers.end()); + std::sort(steady.begin(), steady.end(), rank); + + std::vector selected = forced; + const int capacity = std::max(0, std::min(C, k_cap)); + if (costs_[(size_t)C].ar.samples > 0 && + (int)forced.size() <= capacity) { + double surplus = 0.0; + for (const Ranked & item : forced) surplus += item.score - 1.0; + const double ar_us = costs_[(size_t)C].ar.value; + for (const Ranked & item : steady) { + const int next_k = (int)selected.size() + 1; + if (next_k > capacity) break; + const double next_surplus = surplus + item.score - 1.0; + const double spec_us = estimated_spec_us(C, next_k); + const bool goodput = valid_positive(spec_us) && + (C + next_surplus) / spec_us >= + cfg_.margin * C / ar_us; + const bool peers_ok = next_k == C || + spec_us <= cfg_.slack * ar_us; + if (!goodput || !peers_ok) break; + selected.push_back(item); + surplus = next_surplus; + } + } + + std::vector slots; + slots.reserve(selected.size()); + std::vector reprobed; + for (const Ranked & item : selected) { + const SpecCandidate & candidate = *item.candidate; + RequestState & state = requests_[candidate.request_id]; + state.pending_generated_tokens = candidate.generated_tokens; + state.pending_forced = + candidate.policy == SpeculationPolicy::Always; + if (item.reprobe && state.mode == Mode::ar) { + state.mode = Mode::probing; + state.bad_rounds = 0; + reprobed.push_back(&state); + } + slots.push_back(candidate.slot); + } + for (RequestState * state : reprobed) { + log_transition(state->request_id, Mode::ar, Mode::probing, + state->ema_yield, C, (int)selected.size()); + } + return slots; + } + + void observe_ar(int C, double step_us) { + if (C <= 0 || !valid_positive(step_us)) return; + ensure_cost(C); + costs_[(size_t)C].ar.observe(step_us, cfg_.cost_ewma_alpha); + trace(C, 0); + } + + template + void observe_spec(int C, int k, double step_us, + const AcceptedRange & accepted) { + if (C <= 0 || k < 1 || k > C || !valid_positive(step_us)) return; + ensure_cost(C); + CostState & cost = costs_[(size_t)C]; + cost.ensure_k(k); + cost.spec[(size_t)k].observe(step_us, cfg_.cost_ewma_alpha); + if (cost.ar.samples > 0) { + cost.delta.observe( + (step_us - cost.ar.value) / k, cfg_.cost_ewma_alpha); + } + + for (const std::pair & sample : accepted) { + if (sample.second < 1) continue; + RequestState & state = requests_[sample.first]; + state.request_id = sample.first; + const double emitted = std::min( + max_accept_, static_cast(sample.second)); + if (state.rounds == 0) state.ema_yield = emitted; + else state.ema_yield = cfg_.ema_alpha * emitted + + (1.0 - cfg_.ema_alpha) * state.ema_yield; + ++state.rounds; + if (state.pending_generated_tokens >= 0) { + state.tokens_at_last_spec = + state.pending_generated_tokens; + } + state.pending_generated_tokens = -1; + const bool forced = state.pending_forced; + state.pending_forced = false; + if (forced || cost.ar.samples == 0) continue; + + const double marginal_us = + estimated_spec_us(C, k) - estimated_spec_us(C, k - 1); + const bool pays = state.ema_yield - 1.0 >= + C * marginal_us / cost.ar.value; + const Mode before = state.mode; + if (state.mode == Mode::probing && + state.rounds >= cfg_.probe_rounds) { + state.mode = pays ? Mode::spec : Mode::ar; + state.bad_rounds = 0; + } else if (state.mode == Mode::spec) { + state.bad_rounds = pays ? 0 : state.bad_rounds + 1; + if (state.bad_rounds >= cfg_.bad_rounds) { + state.mode = Mode::ar; + state.bad_rounds = 0; + } + } + if (before != state.mode) { + log_transition(sample.first, before, state.mode, + state.ema_yield, C, k); + } + } + trace(C, k); + } + + void forget(std::uint64_t request_id) { requests_.erase(request_id); } + +private: + enum class Mode { probing, spec, ar }; + + struct Ewma { + double value = 0.0; + int samples = 0; + + void observe(double sample, double alpha) { + value = samples == 0 ? sample + : alpha * sample + (1.0 - alpha) * value; + ++samples; + } + }; + + struct RequestState { + std::uint64_t request_id = 0; + Mode mode = Mode::probing; + double ema_yield = 1.0; + int rounds = 0; + int bad_rounds = 0; + int tokens_at_last_spec = 0; + int pending_generated_tokens = -1; + bool pending_forced = false; + }; + + struct CostState { + Ewma ar; + std::vector spec{1}; + Ewma delta; + + void ensure_k(int k) { + if ((int)spec.size() <= k) spec.resize((size_t)k + 1); + } + }; + + static bool valid_positive(double value) { + return std::isfinite(value) && value > 0.0; + } + + static SpecGateConfig sanitize(SpecGateConfig cfg) { + cfg.ema_alpha = std::isfinite(cfg.ema_alpha) + ? std::clamp(cfg.ema_alpha, 0.0, 1.0) : 0.4; + cfg.cost_ewma_alpha = std::isfinite(cfg.cost_ewma_alpha) + ? std::clamp(cfg.cost_ewma_alpha, 0.0, 1.0) : 0.35; + cfg.probe_rounds = std::max(1, cfg.probe_rounds); + cfg.max_probers = std::max(0, cfg.max_probers); + cfg.bad_rounds = std::max(1, cfg.bad_rounds); + cfg.reprobe_tokens = std::max(0, cfg.reprobe_tokens); + if (!valid_positive(cfg.margin)) cfg.margin = 1.05; + if (!valid_positive(cfg.slack)) cfg.slack = 1.10; + return cfg; + } + + void ensure_cost(int C) { + if ((int)costs_.size() <= C) costs_.resize((size_t)C + 1); + } + + double probe_score(const SpecCandidate & candidate) const { + return std::isfinite(candidate.prior_accept) + ? std::clamp(candidate.prior_accept, 1.0, max_accept_) + : max_accept_; + } + + double score_for(const RequestState & state, + const SpecCandidate & candidate) const { + return state.rounds > 0 + ? std::clamp(state.ema_yield, 1.0, max_accept_) + : probe_score(candidate); + } + + double nearest_delta(int C) const { + int best_distance = std::numeric_limits::max(); + double best = 0.0; + for (int other = 1; other < (int)costs_.size(); ++other) { + if (costs_[(size_t)other].delta.samples == 0) continue; + const int distance = std::abs(other - C); + if (distance < best_distance) { + best_distance = distance; + best = costs_[(size_t)other].delta.value; + } + } + return best; + } + + double estimated_spec_us(int C, int k) const { + const CostState & cost = costs_[(size_t)C]; + if (k == 0) return cost.ar.value; + if (k < (int)cost.spec.size() && + cost.spec[(size_t)k].samples > 0) { + return cost.spec[(size_t)k].value; + } + const double delta = cost.delta.samples > 0 + ? cost.delta.value : nearest_delta(C); + return cost.ar.value + k * delta; + } + + static const char * mode_name(Mode mode) { + switch (mode) { + case Mode::probing: return "probing"; + case Mode::spec: return "spec"; + case Mode::ar: return "ar"; + } + return "unknown"; + } + + static void log_transition(std::uint64_t request_id, Mode before, + Mode after, double ema_yield, int C, int k) { + std::fprintf(stderr, + "[speculation-gate] request=%llu mode=%s->%s " + "ema_yield=%.3f C=%d k=%d\n", + (unsigned long long)request_id, mode_name(before), + mode_name(after), ema_yield, C, k); + } + + void trace(int C, int k) const { + if (std::getenv("DFLASH_SPECULATION_GATE_TRACE") == nullptr) return; + const CostState & cost = costs_[(size_t)C]; + const double spec = k > 0 && k < (int)cost.spec.size() && + cost.spec[(size_t)k].samples > 0 + ? cost.spec[(size_t)k].value + : std::numeric_limits::quiet_NaN(); + const double delta = cost.delta.samples > 0 + ? cost.delta.value + : std::numeric_limits::quiet_NaN(); + std::fprintf(stderr, + "[speculation-gate] C=%d k=%d T_ar=%.3f T_spec=%.3f " + "delta=%.3f requests=%zu\n", + C, k, cost.ar.samples > 0 ? cost.ar.value + : std::numeric_limits::quiet_NaN(), + spec, delta, requests_.size()); + for (const auto & entry : requests_) { + const RequestState & state = entry.second; + std::fprintf(stderr, + "[speculation-gate] request=%llu mode=%s ema_yield=%.3f " + "rounds=%d bad_rounds=%d\n", + (unsigned long long)entry.first, mode_name(state.mode), + state.ema_yield, state.rounds, state.bad_rounds); + } + } + + SpecGateConfig cfg_; + double max_accept_ = 1.0; + std::map requests_; + std::vector costs_{1}; +}; + +} // namespace dflash::common diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 12445e6b2..01c46968d 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 request policy. Concurrent engines consume it directly; + // sequential paths map Never to their established AR fallback. + 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..e46bff503 --- /dev/null +++ b/server/src/common/speculation_policy.h @@ -0,0 +1,34 @@ +#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/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 9cb2a03b6..0a4d0cc17 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); @@ -684,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. @@ -724,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 52949acfc..89916c748 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -15,12 +15,19 @@ #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 #include #include +#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 { @@ -35,6 +42,15 @@ int decode_bucket_width(int live_count) { return 64; } +int forced_speculative_requests() { + static const int forced = []() { + const char * value = std::getenv( + "DFLASH_ADAPTIVE_VERIFY_FORCE_SPECULATIVE_REQUESTS"); + return value ? std::max(0, std::atoi(value)) : -1; + }(); + return forced; +} + } // namespace Qwen35SeqEngine::Qwen35SeqEngine( @@ -52,6 +68,7 @@ Qwen35SeqEngine::Qwen35SeqEngine( prefill_quantum_(std::max(1, prefill_quantum)), b_(backend), slots_(pool, max_ctx, std::max(1, tree_width), backend.paged_kv_residency_.get()), + speculation_gate_(SpecGateConfig{}, std::max(1, tree_width)), scratch_row_(scratch_row), tree_width_(tree_width), tree_scratch_base_(tree_scratch_base), tree_scratch_stride_(tree_scratch_stride) { @@ -140,37 +157,33 @@ 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 true; + return slots_.slot(in.slot).generated_tokens() >= min_floor; } 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; @@ -200,7 +213,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]); @@ -216,7 +229,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(); @@ -278,21 +291,72 @@ 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; + 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)) { + 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; - std::vector flat_tokens((size_t)total_tree, 0); + const int total_packed = total_tree + n_ar; + std::vector tree_feature_rows; + if (direct_commit) { + 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_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 @@ -301,11 +365,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) { @@ -324,18 +391,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; } @@ -357,10 +446,26 @@ 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, 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) != @@ -368,9 +473,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) { @@ -387,6 +504,262 @@ 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)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]; + 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; + } + + + 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 (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=one-pass-tree-ar speculative=%d ar=%zu\n", + active, ar_plan.decode.size()); + } + return result; + } + replay_total += (int)ar_plan.decode.size(); std::vector replay_segments; std::vector replay_tokens; @@ -394,12 +767,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; @@ -437,6 +810,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"; @@ -450,7 +853,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) { @@ -500,18 +903,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()); @@ -522,34 +927,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; } } - 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); + result.decode.reserve((size_t)total_active); for (Proposal & p : proposals) { DecodeOutput out; out.slot = p.slot; @@ -563,20 +962,14 @@ 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)); + } + 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)); } @@ -787,8 +1180,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 +1226,171 @@ 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); + const int C = static_cast(inputs.size()); + const int k_cap = std::min( + C, std::max(0, b_.cache_.tree_capture_lanes)); + + // DDTree's adapter has no cold prior: normal speculative rounds provide + // both accepted yield and the hardware cost. DSpark can later populate + // prior_accept from its calibrated confidence head without gate changes. + std::vector candidates; + candidates.reserve(inputs.size()); + for (const StepInput & in : inputs) { + const Qwen35Slot & seq = slots_.slot(in.slot); + candidates.push_back({ + in.slot, + seq.request_id, + in.speculation_policy, + ddtree_input_eligible(in), + std::numeric_limits::quiet_NaN(), + seq.generated_tokens(), + }); + } + + const int oracle_limit = forced_speculative_requests(); + const bool gate_active = adaptive_enabled && oracle_limit < 0; + std::vector spec_slots; + if (gate_active) { + spec_slots = speculation_gate_.plan(C, candidates, k_cap); + } else if (oracle_limit >= 0) { + // Explicit Always requests remain mandatory even when the benchmark + // oracle asks for a smaller synthetic speculative subbatch. + for (const SpecCandidate & candidate : candidates) { + if (candidate.eligible && + candidate.policy == SpeculationPolicy::Always) { + spec_slots.push_back(candidate.slot); + } + } + const int target = std::max( + oracle_limit, static_cast(spec_slots.size())); + for (const SpecCandidate & candidate : candidates) { + if ((int)spec_slots.size() >= target) break; + if (!candidate.eligible || + candidate.policy != SpeculationPolicy::Adaptive) { + continue; + } + spec_slots.push_back(candidate.slot); + } + } else { + // Burn-in/parity mode: preserve fixed speculation for every eligible + // request except an explicit per-request Never override. + for (const SpecCandidate & candidate : candidates) { + if (candidate.eligible && + candidate.policy != SpeculationPolicy::Never) { + spec_slots.push_back(candidate.slot); + } + } + } + if (gate_active && (int)spec_slots.size() > k_cap) { + return fail_step( + "forced speculation requests exceed DDTree executor capacity"); + } + + std::vector selected((size_t)n_slots, 0); + for (int slot : spec_slots) { + if (slot >= 0 && slot < n_slots) selected[(size_t)slot] = 1; + } + + StepPlan speculative_plan; + StepPlan ar_plan; + speculative_plan.decode.reserve(spec_slots.size()); + ar_plan.decode.reserve(inputs.size() - spec_slots.size()); + 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); + } + } + + using Clock = std::chrono::steady_clock; + 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. Do not learn from this contaminated timing. + return step_regular(plan); + } + 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 (gate_active) { + if (speculative_count == 0) { + speculation_gate_.observe_ar(C, route_us); + } else { + std::vector> accepted_yields; + accepted_yields.reserve((size_t)speculative_count); + for (const DecodeOutput & out : routed_result.decode) { + if (out.failed || out.slot < 0 || out.slot >= n_slots || + !selected[(size_t)out.slot]) { + continue; + } + accepted_yields.emplace_back( + slots_.slot(out.slot).request_id, + static_cast(out.ddtree_accepted_tokens + 1)); + } + speculation_gate_.observe_spec( + C, speculative_count, route_us, accepted_yields); + } } + 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(routed_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; @@ -1222,6 +1771,7 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { void Qwen35SeqEngine::retire(int slot) { if (!slots_.is_active(slot)) return; + speculation_gate_.forget(slots_.slot(slot).request_id); slots_.retire(slot); } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 75275df39..74e96dcb1 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -22,6 +22,7 @@ #pragma once #include "common/concurrency/seq_engine.h" +#include "common/concurrency/speculation_gate.h" #include "common/dflash_draft_kv.h" #include "common/dflash_feature_ring.h" #include "common/ddtree.h" @@ -118,13 +119,17 @@ 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; + 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_; + SpeculationGate speculation_gate_; int64_t scratch_row_ = 0; int tree_width_ = 0; int tree_scratch_base_ = 0; diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index e559a224e..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 { @@ -205,8 +204,6 @@ 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; @@ -225,33 +222,6 @@ 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; diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index 9a84fb3be..a2c13f968 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -66,12 +66,6 @@ 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; - uint64_t ddtree_sampled_steps = 0; - bool active() const { return phase == Qwen35SlotPhase::prefill || phase == Qwen35SlotPhase::decode; @@ -83,10 +77,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. @@ -156,13 +146,6 @@ 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/graph_builders.cpp b/server/src/qwen35/graph_builders.cpp index 1becd302c..dc2796770 100644 --- a/server/src/qwen35/graph_builders.cpp +++ b/server/src/qwen35/graph_builders.cpp @@ -740,21 +740,36 @@ 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, + 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 && + (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)) { + 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; @@ -767,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); } @@ -783,10 +798,22 @@ 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( 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 +832,19 @@ 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); + } sg.gf = ggml_new_graph_custom(sg.ctx, graph_capacity, false); QwenGraphInputs gi{}; @@ -812,8 +852,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; @@ -821,8 +861,13 @@ 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; gi.tree_width = tree_width; gi.tree_scratch_base = tree_scratch_base; @@ -832,6 +877,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..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, @@ -220,7 +224,9 @@ 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, + 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 efc7e617b..4df7c373c 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; } @@ -551,7 +572,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=one-pass-tree-ar+bounded-replay adaptive=%s " + "policy=acceptance-gate 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..5065ea225 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; } @@ -839,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. @@ -907,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. @@ -931,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)); } @@ -941,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 @@ -965,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 @@ -987,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); @@ -1077,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; @@ -1086,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) { @@ -1095,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; - GGML_ASSERT(!active_slot_ids || !cap); + 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] @@ -1135,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 && @@ -1151,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(); @@ -1181,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 @@ -1188,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, @@ -1210,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 { @@ -1233,18 +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]) { - 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], cap->conv_input->ne[2], - 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)); @@ -1264,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)); } @@ -1274,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); @@ -1328,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 { @@ -1357,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 @@ -1372,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); } @@ -1391,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 @@ -1410,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); } @@ -1451,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 @@ -1460,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); } } @@ -1741,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/src/server/http_server.cpp b/server/src/server/http_server.cpp index 52a4a845c..c64abb9aa 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 0fa0ef718..04983d2fb 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 server policy and may be overridden by a 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 045786cb2..437b2fd1d 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -702,8 +702,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..f9dac611a 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] != '-') { @@ -681,6 +692,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 +1131,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 +1184,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_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_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index 7a2d7a26c..51b5aed6d 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -492,62 +492,6 @@ 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. { diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 2ff471a3d..7d7ce35cd 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -2557,6 +2557,27 @@ 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); + TEST_ASSERT(parse_decode_mode("ar") == SpeculationPolicy::Never); + TEST_ASSERT(parse_decode_mode("speculation") == SpeculationPolicy::Always); + TEST_ASSERT(!parse_decode_mode("sometimes")); + + 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_gate.cpp b/server/test/test_speculation_gate.cpp new file mode 100644 index 000000000..3fbb37e7d --- /dev/null +++ b/server/test/test_speculation_gate.cpp @@ -0,0 +1,276 @@ +#include "common/concurrency/speculation_gate.h" +#include "host_check.h" + +#include +#include +#include +#include +#include + +using namespace dflash::common; + +static int g_checks = 0; + +namespace { + +SpecCandidate candidate( + std::uint64_t request_id, int slot, + SpeculationPolicy policy = SpeculationPolicy::Adaptive, + bool eligible = true, + double prior = std::numeric_limits::quiet_NaN(), + int generated = 0) { + return {slot, request_id, policy, eligible, prior, generated}; +} + +std::vector> accepted( + std::initializer_list> values) { + return values; +} + +bool equal_slots(const std::vector & actual, + std::initializer_list expected) { + return actual == std::vector(expected); +} + +SpecGateConfig permissive_config() { + SpecGateConfig cfg; + cfg.ema_alpha = 1.0; + cfg.cost_ewma_alpha = 1.0; + cfg.margin = 1.0; + cfg.slack = 10.0; + return cfg; +} + +} // namespace + +int main() { + // Cold start is one general all-AR baseline step. The following step + // admits at most max_probers, ordered by score then stable request id. + { + SpecGateConfig cfg = permissive_config(); + cfg.max_probers = 2; + SpeculationGate gate(cfg, 8); + std::vector candidates = { + candidate(30, 3), candidate(10, 1), candidate(20, 2), + }; + CHECK(gate.plan(3, candidates, 3).empty()); + gate.observe_ar(3, 100.0); + CHECK(equal_slots(gate.plan(3, candidates, 3), {1, 2})); + } + + // Priors rank cold requests but probation still performs the real + // measurement. Equal priors use request_id rather than reusable slot id. + { + SpecGateConfig cfg = permissive_config(); + cfg.max_probers = 2; + SpeculationGate gate(cfg, 8); + gate.observe_ar(4, 100.0); + std::vector candidates = { + candidate(9, 0, SpeculationPolicy::Adaptive, true, 3.0), + candidate(7, 3, SpeculationPolicy::Adaptive, true, 6.0), + candidate(5, 2, SpeculationPolicy::Adaptive, true, 6.0), + }; + CHECK(equal_slots(gate.plan(4, candidates, 4), {2, 3})); + } + + // Never and ineligible requests do not create request state. Always is + // returned before a baseline and remains exempt from profitability. + { + SpeculationGate gate(permissive_config(), 8); + std::vector never = { + candidate(1, 0, SpeculationPolicy::Never), + candidate(2, 1, SpeculationPolicy::Never), + }; + CHECK(gate.plan(2, never, 2).empty()); + gate.observe_ar(2, 100.0); + std::vector policies = { + candidate(1, 0), + candidate(3, 1, SpeculationPolicy::Always), + candidate(4, 2, SpeculationPolicy::Always, false), + }; + CHECK(equal_slots(gate.plan(2, policies, 2), {1, 0})); + gate.observe_spec(2, 2, 1000.0, accepted({{3, 1}, {1, 1}})); + std::vector forced = { + candidate(3, 1, SpeculationPolicy::Always), + }; + CHECK(equal_slots(gate.plan(2, forced, 1), {1})); + } + + // Forced requests beyond executor capacity are returned intact so the + // engine can surface the configuration error instead of silently routing + // a user-forced request through AR. + { + SpeculationGate gate(permissive_config(), 8); + std::vector forced = { + candidate(1, 0, SpeculationPolicy::Always), + candidate(2, 1, SpeculationPolicy::Always), + }; + CHECK(equal_slots(gate.plan(2, forced, 1), {0, 1})); + } + + // C=1 covers probation, the break-even transition, two-round bad-yield + // hysteresis, AR token-cadence re-probing, and one-round re-admission. + { + SpecGateConfig cfg = permissive_config(); + cfg.probe_rounds = 2; + cfg.bad_rounds = 2; + cfg.reprobe_tokens = 64; + SpeculationGate gate(cfg, 8); + gate.observe_ar(1, 100.0); + + std::vector one = {candidate(11, 0)}; + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + gate.observe_spec(1, 1, 100.0, accepted({{11, 4}})); + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + gate.observe_spec(1, 1, 100.0, accepted({{11, 4}})); + + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + gate.observe_spec(1, 1, 400.0, accepted({{11, 1}})); + CHECK(gate.plan(1, one, 1).empty()); + gate.observe_ar(1, 500.0); + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + gate.observe_spec(1, 1, 1000.0, accepted({{11, 1}})); + + one[0].generated_tokens = 63; + CHECK(gate.plan(1, one, 1).empty()); + one[0].generated_tokens = 64; + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + gate.observe_spec(1, 1, 100.0, accepted({{11, 8}})); + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + } + + // Synthetic marginal timings, rather than a universal acceptance + // threshold, decide whether probation converges to speculation or AR. + { + SpecGateConfig cfg = permissive_config(); + cfg.probe_rounds = 1; + SpeculationGate gate(cfg, 8); + gate.observe_ar(2, 100.0); + std::vector one = {candidate(21, 0)}; + CHECK(equal_slots(gate.plan(2, one, 1), {0})); + gate.observe_spec(2, 1, 200.0, accepted({{21, 2}})); + CHECK(gate.plan(2, one, 1).empty()); + } + + // Once an observed shape is hopeless even under max_accept, new cold + // requests stay AR without running more speculative probes. + { + SpecGateConfig cfg = permissive_config(); + cfg.probe_rounds = 1; + cfg.max_probers = 1; + cfg.margin = 1.05; + SpeculationGate gate(cfg, 4); + gate.observe_ar(8, 100.0); + std::vector first = {candidate(31, 0)}; + CHECK(equal_slots(gate.plan(8, first, 8), {0})); + gate.observe_spec(8, 1, 1000.0, accepted({{31, 1}})); + std::vector next = {candidate(32, 1)}; + CHECK(gate.plan(8, next, 8).empty()); + CHECK(gate.plan(8, next, 8).empty()); + } + + // A missing (C,k) cost uses the nearest concurrency's per-lane affine + // increment, then the first real measurement corrects that estimate. + { + SpecGateConfig cfg = permissive_config(); + cfg.probe_rounds = 1; + SpeculationGate gate(cfg, 8); + gate.observe_ar(2, 100.0); + std::vector at_two = {candidate(41, 0)}; + CHECK(equal_slots(gate.plan(2, at_two, 1), {0})); + gate.observe_spec(2, 1, 120.0, accepted({{41, 8}})); + + gate.observe_ar(3, 150.0); + std::vector at_three = {candidate(42, 1)}; + CHECK(equal_slots(gate.plan(3, at_three, 1), {1})); + gate.observe_spec(3, 1, 1000.0, accepted({{42, 1}})); + std::vector corrected = {candidate(43, 2)}; + CHECK(gate.plan(3, corrected, 1).empty()); + } + + // Cost EWMAs accept valid samples and invalid observations never poison + // either the hardware profile or request state. + { + SpecGateConfig cfg = permissive_config(); + cfg.cost_ewma_alpha = 0.5; + cfg.probe_rounds = 1; + SpeculationGate gate(cfg, 2); + gate.observe_ar(1, 100.0); + gate.observe_ar(1, 200.0); // T_ar = 150 + gate.observe_ar(1, 0.0); + gate.observe_ar(1, std::numeric_limits::quiet_NaN()); + std::vector one = {candidate(51, 0)}; + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + gate.observe_spec(1, 1, + std::numeric_limits::quiet_NaN(), accepted({{51, 1}})); + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + gate.observe_spec(1, 1, 300.0, accepted({{51, 2}})); + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + gate.observe_ar(1, 100.0); // T_ar = 125 + CHECK(gate.plan(1, one, 1).empty()); + } + + // Eligibility may disappear for a step without discarding learned state. + // forget() removes it on finish, and a new request reusing the same slot + // receives independent cold probation. + { + SpecGateConfig cfg = permissive_config(); + cfg.probe_rounds = 1; + SpeculationGate gate(cfg, 8); + gate.observe_ar(1, 100.0); + std::vector one = {candidate(61, 0)}; + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + gate.observe_spec(1, 1, 100.0, accepted({{61, 8}})); + one[0].eligible = false; + CHECK(gate.plan(1, one, 1).empty()); + one[0].eligible = true; + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + + gate.forget(61); + std::vector reused = {candidate(62, 0)}; + CHECK(equal_slots(gate.plan(1, reused, 1), {0})); + gate.forget(62); // finishing during probation leaves no residue + std::vector reused_again = {candidate(63, 0)}; + CHECK(equal_slots(gate.plan(1, reused_again, 1), {0})); + } + + // A profitable cold C=2 cohort converges to all-spec within the bounded + // probation window, using the same mechanism as every other occupancy. + { + SpecGateConfig cfg = permissive_config(); + cfg.probe_rounds = 2; + SpeculationGate gate(cfg, 8); + std::vector cohort = { + candidate(71, 0), candidate(72, 1), + }; + CHECK(gate.plan(2, cohort, 2).empty()); + gate.observe_ar(2, 100.0); + for (int round = 0; round < cfg.probe_rounds; ++round) { + CHECK(equal_slots(gate.plan(2, cohort, 2), {0, 1})); + gate.observe_spec( + 2, 2, 100.0, accepted({{71, 4}, {72, 4}})); + } + CHECK(equal_slots(gate.plan(2, cohort, 2), {0, 1})); + } + + // At C=8 one expensive probation measurement makes the optimistic + // hopeless check reject all later probes; steady state is pure AR. + { + SpecGateConfig cfg = permissive_config(); + cfg.probe_rounds = 2; + cfg.max_probers = 2; + cfg.margin = 1.05; + SpeculationGate gate(cfg, 8); + std::vector cohort; + for (int i = 0; i < 8; ++i) cohort.push_back(candidate(80 + i, i)); + CHECK(gate.plan(8, cohort, 8).empty()); + gate.observe_ar(8, 100.0); + CHECK(equal_slots(gate.plan(8, cohort, 8), {0, 1})); + gate.observe_spec(8, 2, 1000.0, accepted({{80, 1}, {81, 1}})); + CHECK(gate.plan(8, cohort, 8).empty()); + CHECK(gate.plan(8, cohort, 8).empty()); + } + + std::printf("speculation gate: %d checks passed\n", g_checks); + return 0; +} From 616864640a598771da6ef9081adfd67ab58857d1 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Tue, 18 Aug 2026 13:57:20 +0000 Subject: [PATCH 8/8] refactor(concurrency): use a marginal speculation cut Derive every request score from measured EMA freshness, remove the mode and demotion state machine, and admit only lanes that improve prefix goodput. Carry successful decode progress in observations and cover probation, staleness, EMA hysteresis, noisy costs, and freeloader rejection. --- .../src/common/concurrency/speculation_gate.h | 182 ++++++-------- .../qwen35/concurrency/qwen35_seq_engine.cpp | 7 +- server/test/test_speculation_gate.cpp | 224 +++++++++++------- 3 files changed, 204 insertions(+), 209 deletions(-) diff --git a/server/src/common/concurrency/speculation_gate.h b/server/src/common/concurrency/speculation_gate.h index 65d8d2aca..20c58da94 100644 --- a/server/src/common/concurrency/speculation_gate.h +++ b/server/src/common/concurrency/speculation_gate.h @@ -13,7 +13,6 @@ #include #include #include -#include #include namespace dflash::common { @@ -23,7 +22,6 @@ struct SpecGateConfig { double cost_ewma_alpha = 0.35; int probe_rounds = 2; int max_probers = 2; - int bad_rounds = 2; int reprobe_tokens = 64; double margin = 1.05; double slack = 1.10; @@ -38,6 +36,12 @@ struct SpecCandidate { int generated_tokens = 0; }; +struct SpecObservation { + std::uint64_t request_id = 0; + int emitted_tokens = 0; + int generated_tokens = 0; +}; + class SpeculationGate { public: explicit SpeculationGate(SpecGateConfig config, double max_accept) @@ -56,8 +60,6 @@ class SpeculationGate { struct Ranked { const SpecCandidate * candidate = nullptr; double score = 1.0; - bool probing = false; - bool reprobe = false; }; std::vector forced; std::vector steady; @@ -69,28 +71,22 @@ class SpeculationGate { continue; } RequestState & state = requests_[candidate.request_id]; - state.request_id = candidate.request_id; if (candidate.policy == SpeculationPolicy::Always) { - forced.push_back({&candidate, - score_for(state, candidate), false, false}); - continue; - } - if (state.mode == Mode::spec) { - steady.push_back( - {&candidate, state.ema_yield, false, false}); - continue; - } - if (state.mode == Mode::probing) { - probers.push_back( - {&candidate, probe_score(candidate), true, false}); + forced.push_back({&candidate, state.rounds > 0 + ? measured_score(state) : optimistic_score(candidate)}); continue; } const long long generated = candidate.generated_tokens; const long long last = state.tokens_at_last_spec; - if (generated - last >= cfg_.reprobe_tokens) { - probers.push_back( - {&candidate, probe_score(candidate), true, true}); - } + const bool stale = + generated - last >= cfg_.reprobe_tokens; + const bool optimistic = + state.rounds < cfg_.probe_rounds || stale; + (optimistic ? probers : steady).push_back({ + &candidate, + optimistic ? optimistic_score(candidate) + : measured_score(state), + }); } // In particular, an all-Never/ineligible cohort must not allocate @@ -114,45 +110,46 @@ class SpeculationGate { const int capacity = std::max(0, std::min(C, k_cap)); if (costs_[(size_t)C].ar.samples > 0 && (int)forced.size() <= capacity) { - double surplus = 0.0; - for (const Ranked & item : forced) surplus += item.score - 1.0; const double ar_us = costs_[(size_t)C].ar.value; + std::vector estimated_us((size_t)capacity + 1, ar_us); + for (int k = 1; k <= capacity; ++k) { + const double raw = estimated_spec_us(C, k); + estimated_us[(size_t)k] = valid_positive(raw) + ? std::max(estimated_us[(size_t)k - 1], raw) + : estimated_us[(size_t)k - 1]; + } + + double emitted = C; + for (const Ranked & item : forced) emitted += item.score - 1.0; + int current_k = (int)forced.size(); + double current_goodput = + emitted / estimated_us[(size_t)current_k]; + const double ar_goodput = C / ar_us; for (const Ranked & item : steady) { - const int next_k = (int)selected.size() + 1; + const int next_k = current_k + 1; if (next_k > capacity) break; - const double next_surplus = surplus + item.score - 1.0; - const double spec_us = estimated_spec_us(C, next_k); - const bool goodput = valid_positive(spec_us) && - (C + next_surplus) / spec_us >= - cfg_.margin * C / ar_us; + const double next_emitted = emitted + item.score - 1.0; + const double next_us = estimated_us[(size_t)next_k]; + const double next_goodput = next_emitted / next_us; + const bool marginal = next_goodput > current_goodput; + const bool overall = + next_goodput >= cfg_.margin * ar_goodput; const bool peers_ok = next_k == C || - spec_us <= cfg_.slack * ar_us; - if (!goodput || !peers_ok) break; + next_us <= cfg_.slack * ar_us; + if (!marginal || !overall || !peers_ok) break; selected.push_back(item); - surplus = next_surplus; + current_k = next_k; + emitted = next_emitted; + current_goodput = next_goodput; } } std::vector slots; slots.reserve(selected.size()); - std::vector reprobed; for (const Ranked & item : selected) { - const SpecCandidate & candidate = *item.candidate; - RequestState & state = requests_[candidate.request_id]; - state.pending_generated_tokens = candidate.generated_tokens; - state.pending_forced = - candidate.policy == SpeculationPolicy::Always; - if (item.reprobe && state.mode == Mode::ar) { - state.mode = Mode::probing; - state.bad_rounds = 0; - reprobed.push_back(&state); - } - slots.push_back(candidate.slot); - } - for (RequestState * state : reprobed) { - log_transition(state->request_id, Mode::ar, Mode::probing, - state->ema_yield, C, (int)selected.size()); + slots.push_back(item.candidate->slot); } + trace_plan(C, capacity, slots); return slots; } @@ -176,45 +173,17 @@ class SpeculationGate { (step_us - cost.ar.value) / k, cfg_.cost_ewma_alpha); } - for (const std::pair & sample : accepted) { - if (sample.second < 1) continue; - RequestState & state = requests_[sample.first]; - state.request_id = sample.first; + for (const SpecObservation & sample : accepted) { + if (sample.emitted_tokens < 1) continue; + RequestState & state = requests_[sample.request_id]; const double emitted = std::min( - max_accept_, static_cast(sample.second)); + max_accept_, static_cast(sample.emitted_tokens)); if (state.rounds == 0) state.ema_yield = emitted; else state.ema_yield = cfg_.ema_alpha * emitted + (1.0 - cfg_.ema_alpha) * state.ema_yield; ++state.rounds; - if (state.pending_generated_tokens >= 0) { - state.tokens_at_last_spec = - state.pending_generated_tokens; - } - state.pending_generated_tokens = -1; - const bool forced = state.pending_forced; - state.pending_forced = false; - if (forced || cost.ar.samples == 0) continue; - - const double marginal_us = - estimated_spec_us(C, k) - estimated_spec_us(C, k - 1); - const bool pays = state.ema_yield - 1.0 >= - C * marginal_us / cost.ar.value; - const Mode before = state.mode; - if (state.mode == Mode::probing && - state.rounds >= cfg_.probe_rounds) { - state.mode = pays ? Mode::spec : Mode::ar; - state.bad_rounds = 0; - } else if (state.mode == Mode::spec) { - state.bad_rounds = pays ? 0 : state.bad_rounds + 1; - if (state.bad_rounds >= cfg_.bad_rounds) { - state.mode = Mode::ar; - state.bad_rounds = 0; - } - } - if (before != state.mode) { - log_transition(sample.first, before, state.mode, - state.ema_yield, C, k); - } + state.tokens_at_last_spec = + std::max(0, sample.generated_tokens); } trace(C, k); } @@ -222,8 +191,6 @@ class SpeculationGate { void forget(std::uint64_t request_id) { requests_.erase(request_id); } private: - enum class Mode { probing, spec, ar }; - struct Ewma { double value = 0.0; int samples = 0; @@ -236,14 +203,9 @@ class SpeculationGate { }; struct RequestState { - std::uint64_t request_id = 0; - Mode mode = Mode::probing; double ema_yield = 1.0; int rounds = 0; - int bad_rounds = 0; int tokens_at_last_spec = 0; - int pending_generated_tokens = -1; - bool pending_forced = false; }; struct CostState { @@ -267,7 +229,6 @@ class SpeculationGate { ? std::clamp(cfg.cost_ewma_alpha, 0.0, 1.0) : 0.35; cfg.probe_rounds = std::max(1, cfg.probe_rounds); cfg.max_probers = std::max(0, cfg.max_probers); - cfg.bad_rounds = std::max(1, cfg.bad_rounds); cfg.reprobe_tokens = std::max(0, cfg.reprobe_tokens); if (!valid_positive(cfg.margin)) cfg.margin = 1.05; if (!valid_positive(cfg.slack)) cfg.slack = 1.10; @@ -278,17 +239,14 @@ class SpeculationGate { if ((int)costs_.size() <= C) costs_.resize((size_t)C + 1); } - double probe_score(const SpecCandidate & candidate) const { + double optimistic_score(const SpecCandidate & candidate) const { return std::isfinite(candidate.prior_accept) ? std::clamp(candidate.prior_accept, 1.0, max_accept_) : max_accept_; } - double score_for(const RequestState & state, - const SpecCandidate & candidate) const { - return state.rounds > 0 - ? std::clamp(state.ema_yield, 1.0, max_accept_) - : probe_score(candidate); + double measured_score(const RequestState & state) const { + return std::clamp(state.ema_yield, 1.0, max_accept_); } double nearest_delta(int C) const { @@ -317,22 +275,16 @@ class SpeculationGate { return cost.ar.value + k * delta; } - static const char * mode_name(Mode mode) { - switch (mode) { - case Mode::probing: return "probing"; - case Mode::spec: return "spec"; - case Mode::ar: return "ar"; - } - return "unknown"; - } - - static void log_transition(std::uint64_t request_id, Mode before, - Mode after, double ema_yield, int C, int k) { + static void trace_plan(int C, int capacity, + const std::vector & slots) { + if (std::getenv("DFLASH_SPECULATION_GATE_TRACE") == nullptr) return; std::fprintf(stderr, - "[speculation-gate] request=%llu mode=%s->%s " - "ema_yield=%.3f C=%d k=%d\n", - (unsigned long long)request_id, mode_name(before), - mode_name(after), ema_yield, C, k); + "[speculation-gate] plan C=%d capacity=%d selected=%zu slots=", + C, capacity, slots.size()); + for (size_t i = 0; i < slots.size(); ++i) { + std::fprintf(stderr, "%s%d", i == 0 ? "" : ",", slots[i]); + } + std::fputc('\n', stderr); } void trace(int C, int k) const { @@ -354,10 +306,10 @@ class SpeculationGate { for (const auto & entry : requests_) { const RequestState & state = entry.second; std::fprintf(stderr, - "[speculation-gate] request=%llu mode=%s ema_yield=%.3f " - "rounds=%d bad_rounds=%d\n", - (unsigned long long)entry.first, mode_name(state.mode), - state.ema_yield, state.rounds, state.bad_rounds); + "[speculation-gate] request=%llu ema_yield=%.3f " + "rounds=%d tokens_at_last_spec=%d\n", + (unsigned long long)entry.first, state.ema_yield, + state.rounds, state.tokens_at_last_spec); } } diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index 89916c748..cc08007a3 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -1336,16 +1336,17 @@ SeqEngine::StepResult Qwen35SeqEngine::step(const StepPlan & plan) { if (speculative_count == 0) { speculation_gate_.observe_ar(C, route_us); } else { - std::vector> accepted_yields; + std::vector accepted_yields; accepted_yields.reserve((size_t)speculative_count); for (const DecodeOutput & out : routed_result.decode) { if (out.failed || out.slot < 0 || out.slot >= n_slots || !selected[(size_t)out.slot]) { continue; } - accepted_yields.emplace_back( + accepted_yields.push_back({ slots_.slot(out.slot).request_id, - static_cast(out.ddtree_accepted_tokens + 1)); + static_cast(out.ddtree_accepted_tokens + 1), + slots_.slot(out.slot).generated_tokens()}); } speculation_gate_.observe_spec( C, speculative_count, route_us, accepted_yields); diff --git a/server/test/test_speculation_gate.cpp b/server/test/test_speculation_gate.cpp index 3fbb37e7d..a5ac66b26 100644 --- a/server/test/test_speculation_gate.cpp +++ b/server/test/test_speculation_gate.cpp @@ -1,10 +1,8 @@ #include "common/concurrency/speculation_gate.h" #include "host_check.h" -#include #include #include -#include #include using namespace dflash::common; @@ -22,8 +20,13 @@ SpecCandidate candidate( return {slot, request_id, policy, eligible, prior, generated}; } -std::vector> accepted( - std::initializer_list> values) { +std::vector observed( + std::initializer_list values) { + return values; +} + +std::vector batch( + std::initializer_list values) { return values; } @@ -44,8 +47,8 @@ SpecGateConfig permissive_config() { } // namespace int main() { - // Cold start is one general all-AR baseline step. The following step - // admits at most max_probers, ordered by score then stable request id. + // Cold start seeds one AR baseline, then admits only the best bounded + // optimistic cohort. Equal scores are ordered by stable request id. { SpecGateConfig cfg = permissive_config(); cfg.max_probers = 2; @@ -58,8 +61,8 @@ int main() { CHECK(equal_slots(gate.plan(3, candidates, 3), {1, 2})); } - // Priors rank cold requests but probation still performs the real - // measurement. Equal priors use request_id rather than reusable slot id. + // Adapter priors rank cold requests; measured acceptance remains the + // production signal after probation. { SpecGateConfig cfg = permissive_config(); cfg.max_probers = 2; @@ -73,8 +76,9 @@ int main() { CHECK(equal_slots(gate.plan(4, candidates, 4), {2, 3})); } - // Never and ineligible requests do not create request state. Always is - // returned before a baseline and remains exempt from profitability. + // Never and ineligible requests stay AR. Always is returned before a + // baseline and remains exempt from profitability. Capacity overflow is + // returned intact so the engine can report the configuration error. { SpeculationGate gate(permissive_config(), 8); std::vector never = { @@ -89,71 +93,95 @@ int main() { candidate(4, 2, SpeculationPolicy::Always, false), }; CHECK(equal_slots(gate.plan(2, policies, 2), {1, 0})); - gate.observe_spec(2, 2, 1000.0, accepted({{3, 1}, {1, 1}})); - std::vector forced = { - candidate(3, 1, SpeculationPolicy::Always), - }; - CHECK(equal_slots(gate.plan(2, forced, 1), {1})); - } + gate.observe_spec(2, 1, 1000.0, observed({{3, 1, 1}})); + CHECK(equal_slots(gate.plan(2, batch({ + candidate(3, 1, SpeculationPolicy::Always)}), 1), {1})); - // Forced requests beyond executor capacity are returned intact so the - // engine can surface the configuration error instead of silently routing - // a user-forced request through AR. - { - SpeculationGate gate(permissive_config(), 8); std::vector forced = { - candidate(1, 0, SpeculationPolicy::Always), - candidate(2, 1, SpeculationPolicy::Always), + candidate(5, 0, SpeculationPolicy::Always), + candidate(6, 1, SpeculationPolicy::Always), }; CHECK(equal_slots(gate.plan(2, forced, 1), {0, 1})); } - // C=1 covers probation, the break-even transition, two-round bad-yield - // hysteresis, AR token-cadence re-probing, and one-round re-admission. + // C=1: one bad first sample cannot end probation. After two low samples + // the measured score loses the cut; 64 AR tokens make it optimistic for + // one re-probe. EMA smoothing alone keeps one later rejection from + // immediately removing a productive request. { SpecGateConfig cfg = permissive_config(); + cfg.ema_alpha = 0.4; cfg.probe_rounds = 2; - cfg.bad_rounds = 2; cfg.reprobe_tokens = 64; SpeculationGate gate(cfg, 8); gate.observe_ar(1, 100.0); std::vector one = {candidate(11, 0)}; CHECK(equal_slots(gate.plan(1, one, 1), {0})); - gate.observe_spec(1, 1, 100.0, accepted({{11, 4}})); + gate.observe_spec(1, 1, 200.0, observed({{11, 1, 1}})); CHECK(equal_slots(gate.plan(1, one, 1), {0})); - gate.observe_spec(1, 1, 100.0, accepted({{11, 4}})); - - CHECK(equal_slots(gate.plan(1, one, 1), {0})); - gate.observe_spec(1, 1, 400.0, accepted({{11, 1}})); + gate.observe_spec(1, 1, 200.0, observed({{11, 1, 2}})); CHECK(gate.plan(1, one, 1).empty()); - gate.observe_ar(1, 500.0); - CHECK(equal_slots(gate.plan(1, one, 1), {0})); - gate.observe_spec(1, 1, 1000.0, accepted({{11, 1}})); - one[0].generated_tokens = 63; + one[0].generated_tokens = 65; CHECK(gate.plan(1, one, 1).empty()); - one[0].generated_tokens = 64; + one[0].generated_tokens = 66; + CHECK(equal_slots(gate.plan(1, one, 1), {0})); + gate.observe_spec(1, 1, 120.0, observed({{11, 8, 67}})); + one[0].generated_tokens = 67; CHECK(equal_slots(gate.plan(1, one, 1), {0})); - gate.observe_spec(1, 1, 100.0, accepted({{11, 8}})); + gate.observe_spec(1, 1, 120.0, observed({{11, 1, 68}})); + one[0].generated_tokens = 68; CHECK(equal_slots(gate.plan(1, one, 1), {0})); } - // Synthetic marginal timings, rather than a universal acceptance - // threshold, decide whether probation converges to speculation or AR. + // The marginal cut rejects a zero-surplus freeloader even though two + // high-yield requests keep the aggregate route far above the AR baseline. + { + SpeculationGate gate(permissive_config(), 8); + gate.observe_ar(6, 100.0); + gate.observe_spec(6, 2, 106.0, + observed({{21, 4, 1}, {22, 4, 1}})); + gate.observe_spec(6, 2, 106.0, + observed({{21, 4, 2}, {22, 4, 2}})); + gate.observe_spec(6, 1, 103.0, observed({{23, 1, 1}})); + gate.observe_spec(6, 1, 103.0, observed({{23, 1, 2}})); + std::vector cohort = { + candidate(21, 0, SpeculationPolicy::Adaptive, true, + std::numeric_limits::quiet_NaN(), 2), + candidate(22, 1, SpeculationPolicy::Adaptive, true, + std::numeric_limits::quiet_NaN(), 2), + candidate(23, 2, SpeculationPolicy::Adaptive, true, + std::numeric_limits::quiet_NaN(), 2), + }; + CHECK(equal_slots(gate.plan(6, cohort, 6), {0, 1})); + } + + // Marginal improvement alone is insufficient: the route must also clear + // the global 5% margin and the AR-peer latency bound. { SpecGateConfig cfg = permissive_config(); + cfg.margin = 1.05; cfg.probe_rounds = 1; - SpeculationGate gate(cfg, 8); - gate.observe_ar(2, 100.0); - std::vector one = {candidate(21, 0)}; - CHECK(equal_slots(gate.plan(2, one, 1), {0})); - gate.observe_spec(2, 1, 200.0, accepted({{21, 2}})); - CHECK(gate.plan(2, one, 1).empty()); + SpeculationGate margin_gate(cfg, 8); + margin_gate.observe_ar(1, 100.0); + margin_gate.observe_spec(1, 1, 195.0, observed({{31, 2, 1}})); + CHECK(margin_gate.plan(1, batch({candidate( + 31, 0, SpeculationPolicy::Adaptive, true, + std::numeric_limits::quiet_NaN(), 1)}), 1).empty()); + + cfg.margin = 1.0; + cfg.slack = 1.10; + SpeculationGate slack_gate(cfg, 8); + slack_gate.observe_ar(2, 100.0); + slack_gate.observe_spec(2, 1, 111.0, observed({{32, 8, 1}})); + CHECK(slack_gate.plan(2, batch({candidate( + 32, 0, SpeculationPolicy::Adaptive, true, + std::numeric_limits::quiet_NaN(), 1)}), 1).empty()); } - // Once an observed shape is hopeless even under max_accept, new cold - // requests stay AR without running more speculative probes. + // Once an observed shape is hopeless even with max_accept, later cold + // requests cost arithmetic only and run no speculative probe. { SpecGateConfig cfg = permissive_config(); cfg.probe_rounds = 1; @@ -161,100 +189,113 @@ int main() { cfg.margin = 1.05; SpeculationGate gate(cfg, 4); gate.observe_ar(8, 100.0); - std::vector first = {candidate(31, 0)}; - CHECK(equal_slots(gate.plan(8, first, 8), {0})); - gate.observe_spec(8, 1, 1000.0, accepted({{31, 1}})); - std::vector next = {candidate(32, 1)}; - CHECK(gate.plan(8, next, 8).empty()); - CHECK(gate.plan(8, next, 8).empty()); + CHECK(equal_slots(gate.plan(8, batch({candidate(41, 0)}), 8), {0})); + gate.observe_spec(8, 1, 1000.0, observed({{41, 1, 1}})); + CHECK(gate.plan(8, batch({candidate(42, 1)}), 8).empty()); + CHECK(gate.plan(8, batch({candidate(42, 1)}), 8).empty()); } - // A missing (C,k) cost uses the nearest concurrency's per-lane affine - // increment, then the first real measurement corrects that estimate. + // Missing shapes use the nearest concurrency's affine increment until a + // real measurement corrects them. { SpecGateConfig cfg = permissive_config(); cfg.probe_rounds = 1; SpeculationGate gate(cfg, 8); gate.observe_ar(2, 100.0); - std::vector at_two = {candidate(41, 0)}; - CHECK(equal_slots(gate.plan(2, at_two, 1), {0})); - gate.observe_spec(2, 1, 120.0, accepted({{41, 8}})); + gate.observe_spec(2, 1, 120.0, observed({{51, 8, 1}})); gate.observe_ar(3, 150.0); - std::vector at_three = {candidate(42, 1)}; - CHECK(equal_slots(gate.plan(3, at_three, 1), {1})); - gate.observe_spec(3, 1, 1000.0, accepted({{42, 1}})); - std::vector corrected = {candidate(43, 2)}; - CHECK(gate.plan(3, corrected, 1).empty()); + CHECK(equal_slots(gate.plan(3, batch({candidate(52, 1)}), 1), {1})); + gate.observe_spec(3, 1, 1000.0, observed({{52, 1, 1}})); + CHECK(gate.plan(3, batch({candidate(53, 2)}), 1).empty()); } - // Cost EWMAs accept valid samples and invalid observations never poison - // either the hardware profile or request state. + // Noisy exact width samples are made nondecreasing before the cut. A + // lower measured T(2) therefore cannot make a zero-yield second lane look + // profitable after T(1). + { + SpecGateConfig cfg = permissive_config(); + cfg.probe_rounds = 1; + SpeculationGate gate(cfg, 8); + gate.observe_ar(2, 100.0); + gate.observe_spec(2, 1, 130.0, observed({{61, 4, 1}})); + gate.observe_spec(2, 2, 120.0, + observed({{61, 4, 2}, {62, 1, 1}})); + std::vector candidates = { + candidate(61, 0, SpeculationPolicy::Adaptive, true, + std::numeric_limits::quiet_NaN(), 2), + candidate(62, 1, SpeculationPolicy::Adaptive, true, + std::numeric_limits::quiet_NaN(), 1), + }; + CHECK(equal_slots(gate.plan(2, candidates, 2), {0})); + } + + // Invalid observations never poison costs or consume probation. { SpecGateConfig cfg = permissive_config(); cfg.cost_ewma_alpha = 0.5; cfg.probe_rounds = 1; - SpeculationGate gate(cfg, 2); + SpeculationGate gate(cfg, 8); gate.observe_ar(1, 100.0); gate.observe_ar(1, 200.0); // T_ar = 150 gate.observe_ar(1, 0.0); gate.observe_ar(1, std::numeric_limits::quiet_NaN()); - std::vector one = {candidate(51, 0)}; + std::vector one = {candidate(71, 0)}; CHECK(equal_slots(gate.plan(1, one, 1), {0})); gate.observe_spec(1, 1, - std::numeric_limits::quiet_NaN(), accepted({{51, 1}})); - CHECK(equal_slots(gate.plan(1, one, 1), {0})); - gate.observe_spec(1, 1, 300.0, accepted({{51, 2}})); + std::numeric_limits::quiet_NaN(), + observed({{71, 1, 1}})); CHECK(equal_slots(gate.plan(1, one, 1), {0})); - gate.observe_ar(1, 100.0); // T_ar = 125 + gate.observe_spec(1, 1, 300.0, observed({{71, 2, 1}})); + one[0].generated_tokens = 1; CHECK(gate.plan(1, one, 1).empty()); } - // Eligibility may disappear for a step without discarding learned state. - // forget() removes it on finish, and a new request reusing the same slot - // receives independent cold probation. + // Eligibility can disappear without discarding measurements. forget() + // removes state, so a new request reusing the slot starts cold. { SpecGateConfig cfg = permissive_config(); cfg.probe_rounds = 1; SpeculationGate gate(cfg, 8); gate.observe_ar(1, 100.0); - std::vector one = {candidate(61, 0)}; + std::vector one = {candidate(81, 0)}; CHECK(equal_slots(gate.plan(1, one, 1), {0})); - gate.observe_spec(1, 1, 100.0, accepted({{61, 8}})); + gate.observe_spec(1, 1, 100.0, observed({{81, 8, 1}})); + one[0].generated_tokens = 1; one[0].eligible = false; CHECK(gate.plan(1, one, 1).empty()); one[0].eligible = true; CHECK(equal_slots(gate.plan(1, one, 1), {0})); - gate.forget(61); - std::vector reused = {candidate(62, 0)}; - CHECK(equal_slots(gate.plan(1, reused, 1), {0})); - gate.forget(62); // finishing during probation leaves no residue - std::vector reused_again = {candidate(63, 0)}; - CHECK(equal_slots(gate.plan(1, reused_again, 1), {0})); + gate.forget(81); + CHECK(equal_slots(gate.plan(1, batch({candidate(82, 0)}), 1), {0})); + gate.forget(82); + CHECK(equal_slots(gate.plan(1, batch({candidate(83, 0)}), 1), {0})); } - // A profitable cold C=2 cohort converges to all-spec within the bounded - // probation window, using the same mechanism as every other occupancy. + // A profitable C=2 cohort converges through the same probation rule used + // at every occupancy. { SpecGateConfig cfg = permissive_config(); cfg.probe_rounds = 2; SpeculationGate gate(cfg, 8); std::vector cohort = { - candidate(71, 0), candidate(72, 1), + candidate(91, 0), candidate(92, 1), }; CHECK(gate.plan(2, cohort, 2).empty()); gate.observe_ar(2, 100.0); for (int round = 0; round < cfg.probe_rounds; ++round) { CHECK(equal_slots(gate.plan(2, cohort, 2), {0, 1})); - gate.observe_spec( - 2, 2, 100.0, accepted({{71, 4}, {72, 4}})); + gate.observe_spec(2, 2, 100.0, + observed({{91, 4, round + 1}, {92, 4, round + 1}})); + cohort[0].generated_tokens = round + 1; + cohort[1].generated_tokens = round + 1; } CHECK(equal_slots(gate.plan(2, cohort, 2), {0, 1})); } - // At C=8 one expensive probation measurement makes the optimistic - // hopeless check reject all later probes; steady state is pure AR. + // At C=8 one expensive probation round makes even optimistic requests + // fail the cut, so steady state is pure AR. { SpecGateConfig cfg = permissive_config(); cfg.probe_rounds = 2; @@ -262,11 +303,12 @@ int main() { cfg.margin = 1.05; SpeculationGate gate(cfg, 8); std::vector cohort; - for (int i = 0; i < 8; ++i) cohort.push_back(candidate(80 + i, i)); + for (int i = 0; i < 8; ++i) cohort.push_back(candidate(100 + i, i)); CHECK(gate.plan(8, cohort, 8).empty()); gate.observe_ar(8, 100.0); CHECK(equal_slots(gate.plan(8, cohort, 8), {0, 1})); - gate.observe_spec(8, 2, 1000.0, accepted({{80, 1}, {81, 1}})); + gate.observe_spec(8, 2, 1000.0, + observed({{100, 1, 1}, {101, 1, 1}})); CHECK(gate.plan(8, cohort, 8).empty()); CHECK(gate.plan(8, cohort, 8).empty()); }